#############################################################
#
# Adjustable parameters (modify for the assignment)
#
#############################################################

velw=1  # water velocity (km/s)
vels=2  # sediment velocity (km/s)
water=2 # water depth (km)
bsr=1   # bsr depth (km)

#############################################################
#
# You don't need to change anything below 
#
#############################################################

from rsfproj import *

#############################################################
#
# Get data
#
#############################################################

Fetch('cmps-tp.HH','blake')

#############################################################
#
# Plot CMP Gather
#
#############################################################

Flow('cmp','cmps-tp.HH',
     '''
     window f3=950 n3=1 max1=6 |
     dd form=native |
     reverse which=2 |
     put o2=0.0 d2=1
     ''')
Plot('cmp',
     '''
     grey title="CMP Gather" wheretitle=t label1=Time unit1=s
     label2="Trace Number" wherexlabel=b max1=6
     ''')
Plot('floorbox',None,
     'box x0=6.153 y0=6.565 label="Seafloor" xt=1 yt=1')
Plot('bsrbox',None,
     'box x0=4.722 y0=4.778 label="BSR" xt=-1 yt=-1')
Plot('gcmp','cmp floorbox bsrbox','Overlay')

Flow('off1',None,'math n1=24 o1=0.4 d1=0.1  output=x1')
Flow('off2',None,'math n1=24 o1=2.8 d1=0.05 output=x1')
Flow('off','off1 off2','cat axis=1 ${SOURCES[1]}')
Plot('wcmp','cmp off',
     '''
     wiggle xpos=${SOURCES[1]} transp=y yreverse=y 
     label1=Time unit1=s label2=Offset unit2=km
     poly=y title="CMP Gather"
     ''')

Result('cmp','wcmp gcmp','SideBySideAniso')

#############################################################
#
# Create a model
#
#############################################################

depth = {'floor': water,
         'bsr': water+bsr}

for ref in depth.keys():
    Flow(ref,None,
         'spike n1=501 d1=0.005 o1=-0.25 mag=%g' % depth[ref])
    Plot(ref,
         '''
         graph min2=0 max2=5 yreverse=y plotcol=0 plotfat=10 
         wanttitle=n wantaxis=n pad=n scalebar=y
         ''')

Flow('vel','floor',
     '''
     unif2 d1=0.005 n1=1001 v00=%g,%g dvdz=0,0.1
     ''' % (velw,vels))
Plot('vel',
     '''
     grey color=j bias=%g scalebar=y wanttitle=n
     label1=Depth unit1=km label2=Lateral unit2=km
     barlabel=Velocity barunit="km/s" barreverse=y
     wherexlabel=b
     ''' % (0.5*(velw+vels)))

#############################################################
#
# Ray tracing
#
#############################################################

rays=Split('vel floor bsr')
times=['cmp',]
for ref in depth.keys():
    ray=ref+'ray'
    for case in range(2):
        Flow(ray+str(case),'vel',
             '''
             shoot2 yshot=0 tol=0.001
             zshot=%g nr=24 r0=%g dr=%g
             ''' % ((depth[ref],0.2,0.05),
                    (depth[ref],1.4,0.025))[case])
    Flow(ray,[ray+'0',ray+'1'],'cat axis=1 ${SOURCES[1]}')
    Flow(['t'+ray,ray+'.ray'],['vel',ray],
         '''
         cell2 yshot=0 zshot=%g
         anglefile=${SOURCES[1]} traj=${TARGETS[1]}
         ''' % depth[ref])
    Plot(ray,[ray+'.ray','vel'],
         '''
         plotrays frame=${SOURCES[1]} plotcol=7 scalebar=y
         ''')
    rays.append(ray)
    Plot('t'+ray,
         '''
         scale dscale=2 |
         graph yreverse=y min2=4 max2=6 plotfat=5
         wantaxis=n wanttitle=n
         ''')
    times.append('t'+ray)

Plot('ray',rays,'Overlay')
Plot('time',times,'Overlay')

Result('ray','ray time','SideBySideAniso')

End()