from rsf.proj import *

# Download pre-processed CMP gathers
# from the Viking Graben dataset
Fetch('paracdp.segy','viking')

# Convert to RSF
Flow('paracdp tparacdp','paracdp.segy',
     'segyread tfile=${TARGETS[1]}')

# Convert to CDP gathers, time-power gain and high-pass filter
Flow('cmps','paracdp',
     '''
     intbin xk=cdpt yk=cdp | window max1=4 | 
     pow pow1=2 | bandpass flo=5 |
     put label3=Midpoint unit3=km o3=1.619 d3=0.0125
     ''')

# Extract offsets
Flow('offsets mask','tparacdp',
     '''
     headermath output=offset | 
     intbin head=$SOURCE xk=cdpt yk=cdp mask=${TARGETS[1]} | 
     dd type=float |
     scale dscale=0.001
     ''')

# Window bad traces
Flow('maskbad','cmps',
     'mul $SOURCE | stack axis=1 | mask min=1e-20')

Flow('mask2','maskbad mask','spray axis=1 n=1 | mul ${SOURCES[1]}')

# NMO stack with an ensemble of constant velocities
Flow('stacks','cmps offsets mask2',
     '''
     stacks half=n v0=1.4 nv=121 dv=0.02 
     offset=${SOURCES[1]} mask=${SOURCES[2]}
     ''',split=[3,'omp'])

# Taper midpoint
Flow('stackst','stacks','costaper nw3=100')

Result('stacks','stackst',
       '''
       byte gainpanel=all | transp plane=23 memsize=5000 |
       grey3 frame1=500 frame2=100 frame3=30 point1=0.8 point2=0.8
       title="Constant-Velocity Stacks" label3=Velocity unit3=km/s
       ''')

# Apply double Fourier transform (cosine transform)
Flow('cosft','stackst','pad n3=2401 | cosft sign1=1 sign3=1')

# Transpose f-v-k to v-f-k
Flow('transp','cosft','transp',split=[3,'omp'])

# Fowler DMO: mapping velocities
Flow('map','transp',
     '''
     math output="x1/sqrt(1+0.25*x3*x3*x1*x1/(x2*x2))" | 
     cut n2=1
     ''')

Result('map',
       '''
       byte gainpanel=all allpos=y bar=bar.rsf | 
       grey3 title="Fowler Map" label1=Velocity 
       unit1=km/s label3=Wavenumber barlabel=Velocity barunit=km/s
       frame1=50 frame2=500 frame3=1000 color=x scalebar=y 
       ''')

Flow('fowler','transp map','iwarp warp=${SOURCES[1]} | transp',
     split=[3,'omp'])

# Inverse Fourier transform
Flow('dmo','fowler','cosft sign1=-1 sign3=-1 | window n3=2142')

Result('dmo',
       '''
       byte gainpanel=all | transp plane=23 memsize=5000 |
       grey3 frame1=500 frame2=100 frame3=30 point1=0.8 point2=0.8
       title="Constant-Velocity DMO Stacks" 
       label3=Velocity unit3=km/s
       ''')

# Compute envelope for picking
Flow('envelope','dmo','envelope | scale axis=2',split=[3,'omp'])

# Pick velocity
Flow('vpick','envelope','pick rect1=25 rect2=50 vel0=1.45')

Result('vpick',
       '''
       grey mean=y color=j scalebar=y barreverse=y barunit=km/s
       title="Picked Migration Velocity" label2="CMP Number" unit2= 
       ''')

# Take a slice
Flow('slice','dmo vpick','slice pick=${SOURCES[1]}')

Result('slice','grey title="Viking Graben DMO Stack" ')

# Check one CMP location

p = 1500 # !!! MODIFY ME !!!

Flow('before','stackst','window n3=1 f3=%d | envelope' % p)
Flow('after','envelope','window n3=1 f3=%d' % p)

for case in ('before','after'):
    Plot(case,
         '''
         window max1=3.5 |
         grey color=j allpos=y title="%s DMO" 
         label2=Velocity unit2=km/s
         ''' % case.capitalize())

Flow('vpick1','vpick','window n2=1 f2=%d' % p)
Plot('vpick1',
     '''
     graph yreverse=y transp=y plotcol=7 plotfat=7 
     pad=n min2=1.4 max2=3.8 wantaxis=n wanttitle=n
     ''')
Plot('after2','after vpick1','Overlay')

Result('envelope','before after2','SideBySideAniso')


End()