from rsf.proj import *

#################### 
# Fetch the dataset
####################
tgz = '2D_Land_data_2ms.tgz'
Fetch(tgz,
      server='http://www.freeusp.org',
      top='RaceCarWebsite/TechTransfer/Tutorials/Processing_2D',
      dir='Data')
files = map(lambda x: 'Line_001.'+x,Split('TXT SPS RPS XPS sgy'))
Flow(files,tgz,
     'gunzip -c $SOURCE | tar -xvf -',stdin=0,stdout=-1)

######################## 
# Convert to RSF format
########################
Flow('line tline','Line_001.sgy','segyread tfile=${TARGETS[1]}')
Result('first','line',
       '''
       window n2=1000 |
       agc rect1=250 rect2=100 |
       grey title="First 1000 traces" labelsz=6 
       ''')

###################
# Get true geometry
###################
lines = {'S':251,'R':782}
color = {'S':4, 'R':2}
for case in 'SR':
    # X-Y geometry
    Flow(case+'.asc','Line_001.%cPS' % case,
         '''awk 'NR > 20 {print $8, " ", $9}' ''')
    Flow(case,case+'.asc',
         '''
         echo in=$SOURCE data_format=ascii_float n1=2 n2=%d |
         dd form=native 
         ''' % lines[case],stdin=0)
    Plot(case,
         '''
         scale dscale=0.001 | dd type=complex |
         graph symbol=* title=%c plotcol=%d
         min1=684 max1=705 min2=3837 max2=3842
         ''' % (case,color[case]))

Plot('s118','S',
     '''
     window n2=1 f2=118 | scale dscale=0.001 | dd type=complex |
     graph symbol=O wanttitle=n plotcol=3 symbolsz=4 plotfat=5
     min1=684 max1=705 min2=3837 max2=3842
     ''')
Result('SRO','R S s118','Overlay')

# Arrange receiver coordinates 
shots = []
for shot in range(lines['S']):
    line = 'line%d' % shot
    Flow(line,'R','window f2=%d n2=282' % (2*shot))
    shots.append(line)
Flow('rece',shots,'cat axis=3 ${SOURCES[1:%d]}' % len(shots))

Plot('rec','rece',
     '''
     window n1=1 j2=5 j3=2 | scale dscale=0.001 | 
     rtoc | math output="input+I*x2" |
     graph symbol=* plotcol=%d title="Stacking Diagram"
     label1=Distance unit1=km label2="Shot Number"
     min1=684 max1=705
     ''' % color['R'])
Plot('sou','S',
     '''
     window n1=1 j2=2 | scale dscale=0.001 | 
     rtoc | math output="input+I*x1" |
     graph symbol=* plotcol=%d wanttitle=n wantaxis=n min1=684 max1=705
     ''' % color['S'])
Result('diagram','rec sou','Overlay')

# Display one shot using true geometry
Flow('sour','S','spray axis=2 n=282 o=0 d=1')
Flow('sx','sour','window n1=1 | scale dscale=0.001')
Flow('sy','sour','window n1=1 f1=1 | scale dscale=0.001')
Flow('rx','rece','window n1=1 | scale dscale=0.001')
Flow('ry','rece','window n1=1 f1=1 | scale dscale=0.001')

Flow('offset','sx sy rx ry',
     '''
     math SX=${SOURCES[0]} SY=${SOURCES[1]}
     RX=${SOURCES[2]} RY=${SOURCES[3]}
     output="sqrt((RX-SX)^2+(RY-SY)^2)"
     ''')

##############################
# Convert to regular geometry
##############################
Flow('lines','line',
     '''
     intbin xk=cdpt yk=fldr | window f2=2 |
     put
     label3=Source d3=0.05  o3=688  unit3=km
     label2=Offset d2=0.025 o2=-3.5 unit2=km
     label1=Time unit1=s
     ''')

Result('lines',
       '''
       transp memsize=1000 plane=23 |
       byte gainpanel=each |
       grey3 frame1=500 frame2=100 frame3=120 flat=n
       title="Raw Data"
       ''')

Flow('shot118','lines','window n3=1 f3=118')
Flow('offset118','offset','window n2=1 f2=118')
Flow('boff1','offset118','window n1=141 | math output="-input"')
Flow('boff2','offset118','window f1=141')
Flow('boff118','boff1 boff2','cat axis=1 ${SOURCES[1]}')
Result('shot118','shot118 boff118',
       '''
       agc rect1=50 rect2=50 |
       wiggle xpos=${SOURCES[1]} transp=y yreverse=y poly=y
       wherexlabel=t wheretitle=b title="Shot 118"
       ''')

##################
# Convert to CMPs
##################
Flow('rcmps mask','lines',
     '''
     shot2cmp mask=${TARGETS[1]} half=n |
     put o2=-1.75 d2=0.05 label2="Half-offset"
     ''')
Result('rcmps',
       '''
       byte gainpanel=each | window j3=3 |
       grey3 frame1=500 frame2=35 frame3=214 flat=n
       point1=0.8 point2=0.4 title="CMPs"
       ''')

###############
# Raw stacking
##############
Flow('fold','mask','dd type=float | stack axis=1 norm=n')
Flow('rstack','rcmps','stack')
Result('rstack','fold rstack',
       '''
       spray axis=1 n=1501 d=0.002 o=0 label=Time unit=s |
       add scale=1,1000 ${SOURCES[1]} |
       grey color=j title="Raw Stack (with Fold)"
       ''')

###################
# Statics corrections
###################
for case in 'SR':
    stat = case+'-statics'
    Flow(stat+'.asc','Line_001.%cPS' % case,'''awk 'NR > 20 {print $10}' ''')
    Flow(stat,stat+'.asc',
         '''
         echo in=$SOURCE data_format=ascii_float n1=%d | dd form=native |
         scale dscale=%g | put label=Time unit=s
         ''' % (lines[case],1./1900),stdin=0)

# Compute statics corrections
shots = []
for shot in range(lines['S']):
    line = 'statics-line%d' % shot
    Flow(line,'R-statics','window f1=%d n1=282' % (2*shot))
    shots.append(line)
Flow('rstat',shots,'cat axis=2 ${SOURCES[1:%d]}' % len(shots))
Flow('statics','S-statics rstat',
     '''
     spray axis=1 n=282 o=0 d=1 |
     add ${SOURCES[1]}
     ''')
Result('statics',
       '''
       grey color=j scalebar=y title="Static Correction"
       transp=n label1=Trace label2=Shot bias=0.055
       ''')
# Apply statics correction
Flow('lines2','lines statics','datstretch datum=${SOURCES[1]}')

###################
# First break mute
##################
# Select 4 shots every tenth sequential shot
Flow('inpmute','lines2',
     '''
     window f3=198 j3=10 n3=4 
     ''')
Result('inpmute',
       '''
       put n2=1128 n3=1 |
       agc rect1=50 rect2=20 | grey wanttitle=n
       ''')

# Select muting parameter for background noise
Flow('outmute','inpmute',
     '''
     mutter t0=0. v0=5.2
     ''')
Result('outmute',
       '''
       put n2=1128 n3=1 |
       agc rect1=50 rect2=20 | grey wanttitle=n
       ''')

# First break muting for all shots
Flow('mutes','lines2',
     '''
     mutter t0=0. v0=5.2 
     ''' )

###########################
# Time-frequency ananlysis
###########################
# Shot 198
Flow('shot198','mutes','window n3=1 f3=198')
Plot('shot198','grey title="Shot 198" ')

# Spectra
Flow('spec198','shot198','spectra2')
Plot('spec198',
     '''
     grey color=j yreverse=n title="Spectra 198" bias=0.08
     label1=Frequency unit1=Hz label2=Wavenumber unit2=1/km
     ''')
Result('input198','shot198 spec198','SideBySideAniso')

# Test antialiasing filter
Flow('spike',None,'spike n1=1501 d1=0.002 o1=0 k1=750 mag=1 nsp=1')
Flow('bandp','spike','bandpass flo=3 fhi=125 nphi=8')
Result('sbandp','bandp',
       '''
       spectra |
       graph title="Transfer Function"
       labelsz=8 plotfat=5 grid=y
       ''')

# Subsampling all shots to 4ms
Flow('subsample','mutes',
     'bandpass flo=3 fhi=125 nphi=8 | window j1=2')
Flow('subshot198','subsample','window n3=1 f3=198')
Plot('subshot198','grey title="Subsampled 198" ')
# Spectra
Flow('subspec198','subshot198','spectra2')
Plot('subspec198',
     '''
     grey color=j yreverse=n title="Spectra 198" bias=0.08
     label1=Frequency unit1=Hz label2=Wavenumber unit2=1/km
     ''')
Result('output198','subshot198 subspec198','SideBySideAniso')

# Calculate Time-frequency spectra using LTFT
Flow('ltft198','subshot198',
     '''
     ltft rect=20 verb=n nw=50 dw=2 niter=50
     ''')
Result('ltft198',
       '''
       math output="abs(input)" | real |
       byte allpos=y gainpanel=100 pclip=99 |
       grey3 color=j frame1=120 frame2=7 frame3=71 label1=Time flat=y 
       unit1=s label3=Offset label2="\F5 f \F-1" unit3=km
       screenht=10 screenratio=0.7 parallel2=n format2=%3.1f
       point1=0.8 point2=0.3 wanttitle=n 
       ''')
# Thresholding
Flow('thr198','ltft198',
     '''
     transp plane=23 memsize=1000 |
     threshold2 pclip=25 verb=y |
     transp plane=23 memsize=1000
     ''')
Result('thr198',
       '''
       math output="abs(input)" | real |
       byte allpos=y gainpanel=100 pclip=99 |
       grey3 color=j  frame1=120 frame2=7 frame3=71 label1=Time flat=y 
       unit1=s label3=Offset label2="\F5 f \F-1" unit3=km
       screenht=10 screenratio=0.7 parallel2=n format2=%3.1f
       point1=0.8 point2=0.3 wanttitle=n 
       ''')
# Denoise
Flow('noise198','thr198','ltft inv=y | mutter t0=-0.5 v0=0.7')
Plot('noise198','grey title="Ground-roll 198" unit2=km')

Flow('signal198','subshot198 noise198',
     'add scale=1,-1 ${SOURCES[1]}')
Plot('signal198','grey title="Ground-roll Removal" ')
Result('sn198','signal198 noise198','SideBySideAniso')

# Cheating here
Flow('noise','subsample',
     '''
     threshold2 pclip=20 verb=y
     ''')
Flow('signals','subsample noise','add scale=1,-1 ${SOURCES[1]}')

## # Apply LTFT on entire 2-D line
## Flow('ltfts','subsample',
##      '''
##      ltft rect=20 verb=y nw=50 dw=2 niter=50
##      ''',split=[3,282],reduce="cat axis=4")
## Flow('thresholds','ltfts',
##      '''
##      transp plane=23 memsize=1000 | threshold2 pclip=25 verb=y  |
##      transp plane=23 memsize=1000
##      ''',split=[3,251])
## Flow('noise','thresholds',
##      '''
##       ltft inv=y | 
##       mutter t0=-0.5 v0=0.7
##       ''')
## Flow('signals','subsample noise','add scale=1,-1 ${SOURCES[1]}')

############################
# Initial velocity analysis
############################
# Convert shots to CMPs
Flow('cmps','signals',
     '''
     mutter v0=3. |
     shot2cmp half=n | put o2=-1.75 d2=0.05 label2="Half-offset"
     ''')
Result('cmps',
       '''
       byte gainpanel=each |
       grey3 frame1=500 frame2=36 frame3=642 flat=n
       title="CMP Gathers" point1=0.7 label2=Offset label3=Midpoint
       ''')

# Set up velocity scan parameters
v0 = 1.0
dv = 0.05
nv = 75

# Velocity scanning for all CMP gathers
Flow('scn','cmps',
     '''
     vscan semblance=y v0=%g nv=%d dv=%g half=y str=0 |
     mutter v0=0.9 t0=-4.5 inner=y
     ''' % (v0,nv,dv),split=[3,1285])

Flow('vel','scn','pick rect1=15 rect2=25 gate=100 an=10 | window')
Result('vel',
       '''
       grey title="NMO Velocity" label1="Time" label2="Lateral"
       color=j scalebar=y allpos=y bias=2.1 barlabel="Velocity"
       barreverse=y o2num=1 d2num=1 n2tic=3 
       ''')

#################
# Brute stacking
#################
# NMO
Flow('nmo','cmps vel',
     '''
     nmo velocity=${SOURCES[1]} half=y
     ''')

Result('nmo',
       '''
       byte gainpanel=each | 
       grey3 frame1=500 frame2=36 frame3=642 flat=n
       title="NMOed Data" point1=0.7
       label2=Offset label3=Midpoint
       ''')

# NMO stacking
Flow('bstack','nmo','stack')
Result('bstack',
       'shapeagc rect1=50 rect2=100 | grey title="NMO Stack" ')

################
# Toy migration
################
# Prestack Kirchhoff time migration
Flow('tcmps','cmps','transp memsize=1000 plane=23')
Flow('pstm','tcmps vel',
     '''
     mig2 vel=${SOURCES[1]} apt=5 antialias=1
     ''',split=[3,71,[0]],reduce='add')
Result('pstm',
       '''
       grey title="Prestack Kirchhoff Time Migration"
       ''')

End()