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')

# 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''')
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')

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

#############################################################
# Improved Automatic Picking from Radii Interceptor
#############################################################
mute = Program('mute.c')
Flow('envmute','envelope %s'%(mute[0]),
     './${SOURCES[1]} t1=1.4 v1=3.2')
Flow('vpick','envmute','pick vel0=1.45 rect1=25 rect2=50')

Result('vpick',
       '''
       grey mean=y color=x scalebar=y title="DMO Velocity"
       barreverse=y barlabel=Velocity barunit=km/s
       ''')

# Take a slice from mutting envelope
Flow('slice','dmo vpick','slice pick=${SOURCES[1]}')
Result('slice','grey title="DMO Stack"')

############################################################
# Dix conversion to interval velocity
############################################################

Flow('weight','envmute vpick','slice pick=${SOURCES[1]}')
Flow('vdix','vpick weight',
     'dix rect1=25 rect2=50 weight=${SOURCES[1]}')

Result('vdix',
       '''
       grey allpos=y bias=1.3 clip=3.2 
       color=x scalebar=y title="Dix Velocity in Time"
       barreverse=y barlabel=Velocity barunit=km/s
       ''')

Flow('vofz','vdix',
     '''
     time2depth velocity=$SOURCE intime=y nz=1001 z0=0 dz=0.005 | 
     put label1=Depth unit1=km
     ''')

Result('vofz',
       '''
       grey allpos=y bias=1.3 clip=3.2 
       color=x scalebar=y title="Dix Velocity in Depth"
       barreverse=y barlabel=Velocity barunit=km/s
       ''')

###########################################################
# Zero-offset reverse-time migration
###########################################################

Flow('fft','vofz','transp | fft1 | fft3 axis=2 pad=1')
Flow('right left','vofz fft',
     '''
     transp | scale dscale=0.5 |
     isolr2 seed=2016 dt=0.002 npk=50
     fft=${SOURCES[1]} left=${TARGETS[1]} 
     ''')

Flow('rtm snaps','slice left right',
     '''
     spline n1=2000 o1=0 d1=0.002 |
     reverse which=1 |
     transp |
     fftexp0 mig=y snap=10 snaps=${TARGETS[1]}
     left=${SOURCES[1]} right=${SOURCES[2]}
     nz=1001 dz=0.005  
     ''')

Result('rtm',
       '''
       window max1=3.125 | 
       grey title="Post-Stack Depth Migration" unit2=km
       ''')

Plot('snaps','grey title=Snapshots gainpanel=all unit2=km',view=1)

End()