from rsf.proj import *
from math import *
import srpmig


par = {
    'nx':2133,  'dx':0.01143,  'ox':3.05562,   # horizontal (km)
    'nz':1201,  'dz':0.00762,  'oz':0.,        # depth (km)
    'nt':1500,  'dt':0.008,    'ot':0.,        # trace (s)
    'ns':500,   'ds':150.,     'os':10925.,    # shots position (ft)
    'nh':348,   'dh':75.,      'oh':0.         # offset positions (ft)
    }

# 1 foot = 0.3048 m
par['ft2km'] = 0.0003048

# Shots and offsets in km
par['dsm'] = par['ds']*par['ft2km']
par['osm'] = par['os']*par['ft2km']
par['dhm'] = par['dh']*par['ft2km']
par['ohm'] = par['oh']*par['ft2km']


# ----------------------------------------------------------------------
# MIGRATION VELOCITY
# ----------------------------------------------------------------------


migvfile = 'sigsbee2a_migvel.sgy'
Fetch(migvfile,'sigsbee')

# Convert SEG-Y data
Flow('dmvel tmvel dm.head dm.bhead',migvfile,
     '''
     segyread
     tape=$SOURCE
     tfile=${TARGETS[1]}
     hfile=${TARGETS[2]}
     bfile=${TARGETS[3]}
     ''',stdin=0,local=1)

Flow('mvel','dmvel',
     '''
     scale rscale=%(ft2km)g |
     put o1=%(oz)g d1=%(dz)g label1=Depth unit1=km
     o2=%(ox)g d2=%(dx)g label2=Distance unit2=km
     ''' % par,local=1 )

# Plot velocity
Result('mvel',
     '''
     grey color=j allpos=y title= scalebar=y
     barlabel=Velocity barunit=km/s
     label1=Depth unit1=km label2=Lateral unit2=km
     barreverse=y pclip=100 screenratio=0.375 screenht=6.5
     ''',local=1)


# ----------------------------------------------------------------------
# FINITE DIFFERENCE DATA
# ----------------------------------------------------------------------


data = 'sigsbee2a_nfs.sgy'
Fetch(data ,'sigsbee')

# Convert SEG-Y data
Flow('ddata tdata dd.head dd.bhead',data,
     '''
     segyread
     tape=$SOURCE
     tfile=${TARGETS[1]}
     hfile=${TARGETS[2]}
     bfile=${TARGETS[3]}
     ''',stdin=0,local=1)


# ----------------------------------------------------------------------
# SHOT GATHERS
# ----------------------------------------------------------------------


# Shot positions
Flow('tss','tdata','dd type=float | headermath output="%(os)g + fldr*(%(ds)g)" | window' % par,local=1)
Flow('tsi','tss','math output=input/%(ds)g' % par,local=1)

# Offset positions
Flow('too','tdata','dd type=float | headermath output="offset" | window',local=1)
Flow('toi','too','math output=input/%(dh)g' % par,local=1)

# Create sraw by binning
Flow('tos','toi tsi','cat axis=2 space=n ${SOURCES[1]} | transp | dd type=int',local=1)
Flow('sraw','ddata tos','intbin head=${SOURCES[1]} xkey=0 ykey=1',local=1)
# t,o,s

# Prepare shots by mutting
Flow('shots','sraw',
     '''
     put label1=Time unit1=s
     d2=0.075 o2=0.0 label2=Offset
     d3=0.150 o3=10.925 label3=Shot |
     mutter half=false t0=1.0 v0=6.0 |
     put d2=%(dhm)g o2=%(ohm)g unit2=km d3=%(dsm)g o3=%(osm)g unit3=km
     ''' % par,local=1)
# t,h,s

# Shot positions for plot
par.update({'s1':6., 's2':12., 's3':18., 'zero':0})

Plot('shot1','shots','window n3=1 min3=%g | grey wantframenum=y title=Shot' % par['s1'],local=1)
Plot('shot2','shots','window n3=1 min3=%g | grey wantframenum=y title=Shot' % par['s2'],local=1)
Plot('shot3','shots','window n3=1 min3=%g | grey wantframenum=y title=Shot' % par['s3'],local=1)

Result('shotimg','shot1 shot2 shot3','SideBySideAniso',local=1)

Result('zero','shots','window n2=1 min2=0 | grey title=Zero',local=1)


# SELECT SHOT NUMBERS
# par.update({'nss':500, 'jss':1, 'fss':0})
# par.update({'nss':98,  'jss':5, 'fss':5})
# par.update({'nss':128, 'jss':3, 'fss':50})

par.update({'nss':500, 'jss':1, 'fss':0})

Flow('shot','shots','window n3=%(nss)d f3=%(fss)d j3=%(jss)d' % par,local=1)


# ----------------------------------------------------------------------
# MIGRATION
# ----------------------------------------------------------------------


# Frequency sampling
par.update({'nwt':751, 'dwt':0.0833333})

# Frequency resolution
# 50 Hz = par.update({'nw':600, 'jw':1, 'ow':1})

# par.update({'nw':200, 'jw':1, 'ow':1})
# par.update({'nw':300, 'jw':2, 'ow':1})

par.update({'nw':600, 'jw':1, 'ow':1})


# Migration horizontal resolution grid (km)
# par.update({'nmx':1067, 'dmx':0.02286, 'omx':3.05562})

par.update({'nmx':2133, 'dmx':0.01143, 'omx':3.05562})


# Migration slowness resolution
# par.update({'jsx':2, 'jsz':1})

par.update({'jsx':1, 'jsz':1})

srpmig.fwaveslow('shot','mvel',par)

Plot('rwav',
     '''
     window | math output="abs(input)" | real |
     grey title=Wavefield
     label1=Frequency unit1=Hz label2=Distance unit2=km
     ''',local=1,view=1)

# PARALLEL
# --------
# PSPI migration
Flow('img cig','stra rtra slow',
     '''
     srmig3 nrmax=20 dtmax=5e-05 eps=0.01 --readwrite=y verb=y ompnth=1
     tmx=16 rwf=${SOURCES[1]} slo=${SOURCES[2]} cig=${TARGETS[1]}
     itype=o
     ''',split=[4,500,[0,1]],reduce='add')

# IMAGE
# x, y, z
Flow('img2','img','window | transp | put label1=Depth unit1=km label2=Lateral unit2=km',local=1)
Result('image','img2','grey title="PSPI wave equation image"',local=1)

Result('agc','img2','agc rect1=600 rect2=400 | grey title="PSPI wave equation image"',local=1)

# STORE BY PACKING DATA AND HEADER

Flow('imgsigsbee','img','window --out=stdout',local=1)

End()