from rsf.proj import *

sys.path.append('..')
import sigsbee

sigsbee.getvel('vel','migvel')
Result('vel','grey color=j allpos=y bias=1.5 scalebar=y barreverse=y wanttitle=n')

sigsbee.getzo('zodata')
Result('zodata','grey title="Zero Offset" ')

# From velocity to slowness
Flow('slowness','vel','transp | transp plane=23 | math output=1/input')

# Fourier transform and transpose

fmax = 60 # maximum frequency
nf = 721  # number of frequencies

Flow('fft','zodata',
     'fft1 | window max1=%g | transp plane=12 | transp plane=23' % fmax)

# Extended split-step migration
Flow('mig','fft slowness',
     '''
     zomig3 ompnth=1 mode=m --readwrite=y verb=y
     nrmax=3 dtmax=0.00005 slo=${SOURCES[1]} |
     window | transp
     ''',split=[3,nf,[0]],reduce='add')

Result('mig',
       '''
       grey title=Migration 
       label2=Depth    unit2=km 
       label1=Distance unit1=km
       ''')


nw=600
jw=1
ow=1
nt=1500
dt=0.004
ot=0
nmx=500
dmx=0.04572
omx=3.32994

sigsbee.getshots('shot')


Result('shot','byte | grey3 flat=n frame1=250 frame2=80 frame3=100 title=Shots')

Result('zero','shot',
       '''
       window  min2=0 max2=0 size2=1 |
       grey  pclip=98 color=I screenratio=1.5 gainpanel=a
       label2=Position label1=Time title= label3=  unit2=km unit1=s
       labelsz=3
       ''')
Result('shot499','shot',
       '''
       window n3=1 f3=499 |
       grey  pclip=99 color=I gainpanel=a wantframenum=y  unit1=s label1=Time
       label2=Offset unit2=km label3=Shot unit3=km title=
       screenratio=1.35 labelsz=3
       ''')
# Receiver shot gather wavefield FFT
Flow('rfft','shot',
     '''
      fft1 | window squeeze=n n1=%d min1=%g j1=%d |
      spray axis=3 n=1 o=0 d=1 label=hy | spray axis=5 n=1 o=0 d=1 label=sy
       ''' % (nw,ow,jw), local=1)

# Source wavefield FFT
Flow('sfft',None,
     '''
      spike k1=1 n1=%d d1=%g |
      fft1 | window squeeze=n n1=%d min1=%g j1=%d | put label1=w
      ''' % (nt,dt,nw,ow,jw),local=1)

# Interpolate wavefields on surface grid
Flow('rwav swav','rfft sfft','srsyn nx=%d dx=%g ox=%g wav=${SOURCES[1]} swf=${TARGETS[1]}' % (nmx,dmx,omx),local=1)
# w, x, y, s

# Transpose and setting coordinates for 3-D migration
Flow('rtra','rwav','transp plane=12 | transp plane=23',local=1)
Flow('stra','swav','transp plane=12 | transp plane=23',local=1)
# x, y, w, s
# Prepare slowness on 3-D grid
Flow('slow','vel',
     '''
     transp | window j1=1 | math output=1/input |
      spray axis=2 n=1 d=1 o=0 |
      window j3=1 squeeze=n
      ''',local=1)

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

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