from rsf.proj import *

#############################
# Data fetching, SEG-Y -> RSF
#############################

# Fetch the dataset and convert to multiple rsf files
rawsegy=['L23535','L23536','L23537']
for name in rawsegy:
    Fetch(name+'.SGY',
          server='http://certmapper.cr.usgs.gov',
          top='nersl/NPRA/SEISMIC/1981/31_81',
          dir='DMUX')
    Flow([name,'t'+name,name+'.bin',name+'.asc'],name+'.SGY',
         'segyread tfile=${TARGETS[1]} bfile=${TARGETS[2]} hfile=${TARGETS[3]}')

# Raw data and headers combined
Flow('rawline',rawsegy,'cat ${SOURCES[1:3]} axis=2')
Flow('trawline',map(lambda x: 't'+x, rawsegy),'cat ${SOURCES[1:3]} axis=2')

###############
# Trace editing
###############

# Form mask to window out misfires
Flow('shotmask','trawline',
     '''
     window n1=1 f1=2 | mask min=157 max=157 |
     add add=-1 | add scale=-1
     ''')

# Window out misfires, test shots, and auxiliary traces;
# define shot and offset dimensions in local survey coordinates
Flow('shotline',['rawline','shotmask'],
     '''
     headerwindow mask=${SOURCES[1]} |
     put n2=101 n3=67 | window n2=96 f3=10 |
     put o2=-5225 d2=110 label2=Offset unit2=ft
         o3=0 d3=440 label3=Shot unit3=ft
     ''')

#############################
# Elevation profile, datuming
#############################

# Shot elevation information
Flow('spelev','../line31-81/spnElev.txt',
     '''
     echo in=$SOURCE data_format=ascii_float n1=2 n2=33 |
     dd form=native
     ''',stdin=0)
# Elevation header (make shotpoint 100 the origin of the survey coordinate system)
Flow('helev','spelev','window n1=1 f1=0 | math output="(input - 100)*440"')
# Elevation data
Flow('delev','spelev','window n1=1 f1=1')
# Interpolate to the regular survey coordinate system
Flow('elev',['delev','helev'],'shapebin1 nx=321 x0=-5280 dx=110 head=${SOURCES[1]} eps=0.1')
# Datum shifts, weathering layer velocity - v0
Flow('datum',['shotline','elev'],'datshift v0=8000 elev=${SOURCES[1]}')
# Shift traces
Flow('dshotline',['shotline','datum'],'datstretch datum=${SOURCES[1]}')

Plot('selev','spelev','''
     dd type=complex | transp | math output="(real(input) - 100)*440 + imag(input)" |
     graph screenratio=0.3 min2=0 max2=140 min1=-3000 max1=30000
           title="Elevation profile" symbol="x"
     ''')
Plot('ielev','elev','''
     graph screenratio=0.3 min2=0 max2=140 min1=-3000 max1=30000
           title="Elevation profile"
     ''')
Result('elev',['selev','ielev'],'Overlay')
Result('datum','''
       put label3= unit3=s |
       grey color=j allpos=y pclip=100 bias=0.02 scalebar=y title="Datum shift"
       ''')

################################
# Filtration, signal enhancement
################################

# Gain correction, T-X -> F-K
Flow('fkshots','dshotline','tpow tpow=2.0 | bandpass flo=3 fhi=35 | fft1 | fft3 axis=2')
# F-K filter for one shot
fkslope = 10000
fshift = -2.0
Flow('fkfilter','fkshots','''
     window n3=1 | real |
     math output="-%g*x2+x1-%g" | clip2 upper=1 lower=0.99 | math output="input^1e+6" |
     math output="(%g*x2+x1-%g)*input" | clip2 upper=1 lower=0.99 | math output="input^1e+6" |
     smooth rect1=5 rect2=5 repeat=2
     ''' % (fkslope,fshift,fkslope,fshift))
# Apply filter to all shots, convert back to T-X
Flow('fshotline',['fkfilter','fkshots'],'''
     rtoc | spray axis=3 n=57 |
     math s=${SOURCES[1]} output="real(input)*real(s) + real(input)*imag(s)" |
     fft3 axis=2 inv=y | fft1 inv=y
     ''')

# Shot for visualization
fksv = 27
Plot('fkshotb','dshotline','''
     window n3=1 f3=%d | tpow tpow=2.0 |
     grey title="Shot at %g ft (before F-K)"
     ''' % (fksv,fksv*440))
Plot('fkshota','fshotline','''
     window n3=1 f3=%d |
     grey title="Shot at %g ft (after F-K)"
     ''' % (fksv,fksv*440))
Plot('fkspec','fkshots','''
     window n3=1 f3=%d max1=50 |
     math output="sqrt(real(input)*real(input)-imag(input)*imag(input))" | 
     real | grey allpos=y pclip=100 color=j title="F-K spectrum"
     ''')
Plot('fkfilt','fkfilter','''
     window max1=50 |
     grey title="F-K filter" allpos=y color=j scalebar=y pclip=100
     ''')
Result('fkshot',['fkshotb','fkspec','fkfilt','fkshota'],'SideBySideAniso')

# Shot gathers -> CMP gathers
Flow('cmp','fshotline','shot2cmp half=n')

##################
# Velocity picking
##################

# Try two different approaches here

# Semblance for velocity picking from every 10th CMP
Flow('semb1','cmp',
     '''
     window n3=45 f3=50 j3=10 |
     sfvscan half=n semblance=y v0=5000 nv=251 dv=50 nb=80
     ''')
# Semblance for velocity picking from supergathers (20 CMPs each)
Flow('semb2','cmp',
     '''
     window n3=500 f3=25 | put n4=25 o4=-660 d4=1100 n3=20 o3=-577.5 |
     stack axis=3 |
     sfvscan half=n semblance=y v0=5000 nv=251 dv=50 nb=80
     ''')

Result('semb1','''
       grey color=j allpos=y pclip=100 gainpanel=e title="Semblance (every 10th CMP)"
       ''')
Result('semb2','''
       grey color=j allpos=y pclip=100 gainpanel=e title="Semblance (supergathers)"
       ''')

# Velocity picking
Flow('pick1','semb1','''
     mutter x0=10000 abs=n half=n inner=n slope0=0.00075 | 
     mutter x0=7500  abs=n half=n inner=y slope0=0.0015 |
     pick rect1=80 rect2=3 vel0=8000
     ''')
Flow('pick2','semb2','''
     mutter x0=10000 abs=n half=n inner=n slope0=0.00075 | 
     mutter x0=7500  abs=n half=n inner=y slope0=0.0015 |
     pick rect1=80 rect2=3 vel0=8000
     ''')

# Interpolate onto full CMP grid
Flow('vel1','pick1',
     '''
     transp | spline n1=551 d1=55 o1=-2612.5 |
     transp
     ''')
Flow('vel2','pick2',
     '''
     transp | spline n1=551 d1=55 o1=-2612.5 |
     transp
     ''')

Result('vel1','''
       window n2=450 f2=50 |
       grey allpos=y bias=8000 color=j scalebar=y pclip=100
            title="Stacking velocities (every 10th CMP)"
       ''')
Result('vel2','''
       window n2=450 f2=50 |
       grey allpos=y bias=8000 color=j scalebar=y pclip=100
            title="Stacking velocities (supergathers, 20 CMPs each)"
       ''')

###################
# Karl's velocities
###################

# Take picked velocities and covert to ascii RSF file
Flow(['vpick','vpick.dat'],['./vpick2rsf.sh','../line31-81/vpickorig.txt'],'''
     ./${SOURCES[0]} ${SOURCES[1]} ${TARGETS[1]} |
     dd form=native
     ''',stdin=0)
# Interpolate onto a regular grid
Flow('kvel','vpick','''
     window n1=1 f1=2 | 
     shapebin head=${SOURCES[0]} xkey=0 ykey=1
              nx=3000 x0=0 dx=0.002 ny=551 y0=-2612.5 dy=55
              niter=100 filt1=50 filt2=30 eps=0.1 |
     put label1=Time unit1=s label2=Midpoint unit2=ft
     ''')

Result('kvel','''
       window n2=450 f2=50 |
       grey allpos=y bias=8000 color=j scalebar=y pclip=100
            title="Karl's velocities"
       ''')

##################
# NMO and stacking
##################

# NMO correction on CMP gathers
Flow('nmo',['cmp','vel1'],'nmo velocity=${SOURCES[1]} half=n')

# Stacking
Flow('stack','nmo','stack axis=2 norm=n')
Result('stack','agc rect1=250 | grey title="Stack"')

Flow('mstack','nmo','transp | despike2 | median')
Result('mstack','agc rect1=250 | grey title="Despike + Median"')

# Stack from USGS
Fetch('31_81_PR.SGY',
      server='http://certmapper.cr.usgs.gov',
      top='nersl/NPRA/SEISMIC/1981/31_81',
      dir='PROCESSED')
Flow(['31_81_PR','t31_81_PR','31_81_PR.bin','31_81_PR.asc'],'31_81_PR.SGY',
     'segyread tfile=${TARGETS[1]} bfile=${TARGETS[2]} hfile=${TARGETS[3]}')

End()