from rsf.proj import *
from rsf.recipes.beg import server
import math

# Fetch and read the data
Fetch('filt_mig.sgy','teapot',server)
Flow('mig tmig mig.asc mig.bin', 'filt_mig.sgy',
     '''
     segyread tfile=${TARGETS[1]} hfile=${TARGETS[2]}
     bfile=${TARGETS[3]}
     ''')

# Convert from 2D to 3D
Flow('cube mask','mig','intbin xk=ep yk=tracf mask=${TARGETS[1]}')

Flow('cubew','cube','window n1=751')
Flow('slice','cubew','window n1=1 f1=420')

# Compute crossline and inline
az = 70*math.pi/180
Flow('x', 'slice',
     'math output="%g+(%g)*x1+(%g)*x2" ' % (345.0/2,math.cos(az),math.sin(az)))
Flow('y', 'slice',
     'math output="%g+(%g)*x1+(%g)*x2" ' % (188.0/2,-math.sin(az),math.cos(az)))
Flow('xy','x y','cat axis=3 ${SOURCES[1]} | put n1=64860 n2=1 | window | transp')


# Window a cube according the crossline and inline we get
Flow('cuber','cubew xy',
     '''
     window min1=0.45 n1=512 | put n2=64860 n3=1 | transp memsize=1000 |
     bin xkey=0 ykey=1 head=${SOURCES[1]}
     ny=256 nx=64 x0=310 dx=1 y0=-150 dy=1 interp=2 |
     transp plane=13 memsize=1000 | put o2=1 o3=1 
     label2="Inline number" label3="Crossline number" unit2= unit3=
     ''')

# Define the plot function for cube
def cubeplot(frame=[420,199,99],title='',clip='',extra=''):
    return '''
    byte gainpanel=all bar=bar.rsf %s |
    grey3 frame1=%d frame2=%d frame3=%d flat=y
    label1=Time unit1=s label2="Inline number" label3="Crossline number"
    title="%s" %s point1=0.7 point2=0.7 flat=n 
    ''' % (clip,frame[0],frame[1],frame[2],title,extra)

# Display the cube
Result('cuber',cubeplot([195,140,36]))
Result('cube',cubeplot([195,140,36]))


# Estimate dips
Flow('patch','cuber',
     'pad beg1=25 end1=25 | patch p=1,4,1')
Flow('maskdip','cuber',
     '''
     math output=1 | pad beg1=25 end1=25 |
     patch p=1,4,1
     ''')

Flow('dip','patch maskdip',
     'dip rect1=5 rect2=5 rect3=5 order=3 mask=${SOURCES[1]}',
     split=[5,4,[0,1]])

# Inline dip
Flow('dip1','dip','window n4=1 squeeze=n | patch inv=y weight=y dim=3')
# Crossline dip
Flow('dip2','dip','window f4=1 squeeze=n | patch inv=y weight=y dim=3')

# Display dips
Result('dip1',
       'window n1=512 f1=25 |' +
       cubeplot([195,140,36],'Inline Dip','','color=j scalebar=y barlabel="Slope"'))
Result('dip2',
       'window n1=512 f1=25 |' +
       cubeplot([195,140,36],'Crossline Dip','','color=j scalebar=y barlabel="Slope"'))

Flow('dips','dip1 dip2','cat axis=4 ${SOURCES[1]}')
Flow('seed','dips','window n2=1 n3=1 n4=1 | math output=x1')

# Compute cost time
Flow('shift1','cuber','window f2=1')
Flow('shift2','cuber','window f3=1')
Flow('last1','cuber','window f2=255 squeeze=n')
Flow('last2','cuber','window f3=63 squeeze=n')
Flow('ref1','shift1 last1','cat axis=2 ${SOURCES[1]}')
Flow('ref2','shift2 last2','cat axis=3 ${SOURCES[1]}')
Flow('ref1s','ref1','add mode=p $SOURCE | stack axis=1 norm=n')
Flow('ref2s','ref2','add mode=p $SOURCE | stack axis=1 norm=n')
Flow('corr1', 'ref1 cuber',
     'add mode=p ${SOURCES[1]} | stack axis=1 norm=n | clip2 lower=1e-3')
Flow('corr2', 'ref2 cuber',
     'add mode=p ${SOURCES[1]} | stack axis=1 norm=n | clip2 lower=1e-3')
Flow('cuber2','cuber','add mode=p $SOURCE | stack axis=1 norm=n')
Flow('cos1','corr1 cuber2 ref1s',
     'math s1=${SOURCES[1]} s2=${SOURCES[2]} output="(s1*s2)/(input*input)"')
Flow('cos2','corr2 cuber2 ref2s',
     'math s1=${SOURCES[1]} s2=${SOURCES[2]} output="(s1*s2)/(input*input)"')

Flow('cos','cos1 cos2',
     '''
     cat axis=3 ${SOURCES[1]} |
     smooth rect1=40 rect2=40
     ''')

# Compute RT with different reference traces
picks = []
for z, y in [(2, 25), (3, 25), (5, 25), (15, 20), (40, 25), (60, 23), (70, 25),
             (75, 40), (80, 25), (80, 40), (85, 24), (100, 28), (130, 22),
             (135, 28), (140, 25), (140, 36), (190, 39), (200, 25), (230, 25),
             (235, 25), (235, 26), (235, 24), (240, 25), (245, 25), (250, 25),
             (250, 26), (251,25)]:
    time = 'time-%d-%d' % (z,y)
    pick = 'pick-%d-%d' % (z,y)
    Flow(time,'cos',
         '''
         mul $SOURCE | stack axis=3 norm=n |
         put o1=0 d1=1 o2=0 d2=1 o3=0 d3=1 |
         eikonal vel=n zshot=%d yshot=%d
         ''' % (z,y))
    Flow(pick,['dips','seed',time],
         'pwpaint2 seed=${SOURCES[1]} cost=${SOURCES[2]} order=3 eps=1')
    picks.append(pick)

# Compute RT
np = len(picks)
Flow('pick',picks,
     'add ${SOURCES[1:%d]} | scale dscale=%g' % (np,1.0/np))
Flow('pickt','pick','transp plane=23')
Result('pick',
       'window n1=512 f1=25 |' +
        cubeplot([195,140,36],'Relative Time','allpos=y',
       'color=j allpos=y scalebar=y barlabel="RT"'))

# Faltten the data
Flow('cuberpad','cuber','pad beg1=25 end1=25')
Flow('flat','cuberpad pick','iwarp  warp=${SOURCES[1]} eps=1')
Result('flat',cubeplot([195,140,36]))


# RT-seislet transform
Flow('rtseis-inline','cuberpad pick pickt',
     '''
     rtseislet rt=${SOURCES[1]} eps=1 adj=y inv=y unit=y type=b
     ''')
Flow('rtseis','rtseis-inline pickt',
    '''
    transp plane=23 |
    rtseislet rt=${SOURCES[1]} eps=1 adj=y inv=y unit=y type=b |
    transp plane=23 
    ''')
Result('rtseis-inline','window n1=512 f1=25 |' + cubeplot([195,140,36]))
Result('rtseis','window n1=512 f1=25 |' + cubeplot([195,140,36]))

# Inverse RT-seislet transform using 1% most significant coefficients
Flow('teapot-rtseisrec5','rtseis pickt pick',
     '''
     threshold pclip=5 |
     transp plane=23 |
     rtseislet rt=${SOURCES[1]} eps=1 inv=y unit=y type=b |
     transp plane=23 |
     rtseislet rt=${SOURCES[2]} eps=1 inv=y unit=y type=b |
     window n1=512 f1=25 
     ''')

Result('teapot-rtseisrec5',cubeplot([195,140,36]))

Flow('teapot-diff','cuber teapot-rtseisrec5','add ${SOURCES[1]} scale=1,-1')
Result('teapot-diff',cubeplot([195,140,36]))

End()