from rsf.proj import *
from math import *

#2D Example------------------------------------------------------------------------------------
xmax = 6.0
zmax = 3.0

# Forel
#layers = ((0.00,0.00,0.00,0.00,0.00),
#               (0.30,0.50,0.375,0.20,0.30),
#        (0.55,0.75,0.6,0.45,0.55),
#        (0.65,0.85,0.7,0.55,0.65),
#               (1.30,1.30,1.45,1.60,1.20),
#               (2.0,2.0,2.0,2.0,2.0))
layers = ((0.00,0.00,0.00,0.00,0.00,0.00,0.00,0.00),
                (0.3,0.3,0.3,0.3,0.3,0.3,0.3,0.3),
        (0.5,0.5,0.5,0.5,0.65,0.9,0.65,0.5),
        (0.85,0.85,0.85,0.95,1.2,1.3,1.3,1.3),
                (1.55,1.55,1.55,1.55,1.55,1.55,1.55,1.55),
                (2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0))

# Flat
#flatlayers = ((0.00,0.00,0.00,0.00,0.00),
#               (0.375,0.375,0.375,0.375,0.375),
#        (0.6,0.6,0.6,0.6,0.6),
#        (0.7,0.7,0.7,0.7,0.7),
#               (1.45,1.45,1.45,1.45,1.45),
#               (2.0,2.0,2.0,2.0,2.0))
flatlayers = ((0.00,0.00,0.00,0.00,0.00,0.00,0.00,0.00),
                (1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0),
        (2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0),
        (3.0,3.0,3.0,3.0,3.0,3.0,3.0,3.0),
                (4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0),
                (5.0,5.0,5.0,5.0,5.0,5.0,5.0,5.0))

def arr2str(array,sep=' '):
    return sep.join(map(str,array))

n1 = len(layers[0])
n2 = len(layers)

# Generate reflectors ####################################################################
Flow('layers.asc',None,
     '''
        echo %s
     n1=%d n2=%d o1=0 d1=%g
     data_format=ascii_float in=$TARGET     
     ''' % (' '.join(map(arr2str,layers)),
            n1,n2,xmax/(n1-1)))
Flow('layers','layers.asc','dd form=native')

d = 0.0101 # non-round for reproducibility

Flow('refl1','layers',''' window n2=1 f2=0| spline o1=0 d1=%g n1=%d '''% (d,int(1.5+xmax/d)))
Flow('refl2','layers',''' window n2=1 f2=1| spline o1=0 d1=%g n1=%d '''% (d,int(1.5+xmax/d)))
Flow('refl3','layers',''' window n2=1 f2=2| spline o1=0 d1=%g n1=%d '''% (d,int(1.5+xmax/d)))
Flow('refl4','layers',''' window n2=1 f2=3| spline o1=0 d1=%g n1=%d '''% (d,int(1.5+xmax/d)))
Flow('refl5','layers',''' window n2=1 f2=4| spline o1=0 d1=%g n1=%d '''% (d,int(1.5+xmax/d)))
Flow('refl6','layers',''' window n2=1 f2=5| spline o1=0 d1=%g n1=%d '''% (d,int(1.5+xmax/d)))

Flow('ref','refl1 refl2 refl3 refl4 refl5 refl6', '''cat axis=2 ${SOURCES[1:6]} | math output="input*3/2"''')

Plot('ref','''graph yreverse=y wanttitle=n label1=Distance unit1=km plotfat=6
                     label2=Depth unit2=km min2=0 max2=3 min1=0 max1=6
                                         screenht=5.0 screenratio=0.333 yll=3.5 xll=1.5
                                         axisfat=3 titlefat=3 titlesz=10 labelfat=3 labelsz=6''')

Flow('flat.asc',None,
     '''
        echo %s
     n1=%d n2=%d o1=0 d1=%g
     data_format=ascii_float in=$TARGET     
     ''' % (' '.join(map(arr2str,flatlayers)),
            n1,n2,xmax/(n1-1)))
Flow('flats','flat.asc','dd form=native')

Flow('reflf1','flats',''' window n2=1 f2=0| spline o1=0 d1=%g n1=%d '''% (d,int(1.5+xmax/d)))
Flow('reflf2','flats',''' window n2=1 f2=1| spline o1=0 d1=%g n1=%d '''% (d,int(1.5+xmax/d)))
Flow('reflf3','flats',''' window n2=1 f2=2| spline o1=0 d1=%g n1=%d '''% (d,int(1.5+xmax/d)))
Flow('reflf4','flats',''' window n2=1 f2=3| spline o1=0 d1=%g n1=%d '''% (d,int(1.5+xmax/d)))
Flow('reflf5','flats',''' window n2=1 f2=4| spline o1=0 d1=%g n1=%d '''% (d,int(1.5+xmax/d)))
Flow('reflf6','flats',''' window n2=1 f2=5| spline o1=0 d1=%g n1=%d '''% (d,int(1.5+xmax/d)))

Flow('reff','reflf1 reflf2 reflf3 reflf4 reflf5 reflf6', '''cat axis=2 ${SOURCES[1:6]} ''')

# Model specifications ########################################################################
v1 = 2.5
v2 = 3.2
v3 = 3.4
v4 = 3.7
v5 = 3.86

updown=[1,2,3,4,5,4,3,2,1,0]
N = len(updown)-1
vstat = 1 # velocity gradient
gx = [-0.06,0.15,0.2,0.3,0.35]
gz = [0.0,0.0,0.0,0.0,0.0]
v = [v1,v2,v3,v4,v5]
vcol = [1.508,1.581,1.69,1.826,2,2.2]
xref = [3.0,3.0,3.0,3.0,3.0]
zref = [0.0,0.375,0.6,0.7,1.45]

xrefstr = ','.join(map(str,xref)) # convert xref to a string
zrefstr = ','.join(map(str,zref)) # convert zref to a string
gxstr = ','.join(map(str,gx)) # convert gx to a string
gzstr = ','.join(map(str,gz)) # convert gz to a string
updownstr = ','.join(map(str,updown)) # convert updown to a string
vstr = ','.join(map(str,v)) # convert v to a string
vcolstr = ','.join(map(str,vcol)) # convert vcol to a string

Flow('refcut','ref','window n2=6 | put o2=0')
Flow('xrefl','refcut',
        '''isaac2 niter=500 number=%d vstatus=%d debug=n xs=2 xr=4 velocity=%s layer=%s
           xref=%s zref=%s xgradient=%s zgradient=%s tol=0.0001''' % (N,vstat,vstr,updownstr,xrefstr,zrefstr,gxstr,gzstr))
Plot('xrefl',
         '''
         dd type=complex | window | 
         graph wanttitle=n wantaxis=n yreverse=y min2=0 max2=3 min1=0 max1=6 plotcol=7 plotfat=6 screenht=5.0 screenratio=0.333 yll=3.5 xll=1.5
         axisfat=3 titlefat=3 titlesz=10 labelfat=3 labelsz=6
         ''')
Plot('xrefl-points','xrefl',
         '''
         dd type=complex | window | 
         graph wanttitle=n wantaxis=n yreverse=y min2=0 max2=3 min1=0 max1=6 plotcol=7 plotfat=6  symbol=* symbolsz=10 screenht=5.0 screenratio=0.333 yll=3.5 xll=1.5
         axisfat=3 titlefat=3 titlesz=10 labelfat=3 labelsz=6
         ''')
Result('v','xrefl xrefl-points ref','Overlay')# Plot of 2D ray 


# Plot of model 1 ########################################################
Flow('mod1','ref',
     '''
     window n2=5 f2=1 |
     unif2 d1=%g n1=%d v00=%s dvdx=%s x0=%s
     ''' % (d,int(1.5+zmax/d),vstr,gxstr,xrefstr))
#     ''' % (d,int(1.5+zmax/d),vstr+',0.0',gxstr+',0.0',xrefstr+',0.0'))

Flow('modcol','ref',
     '''
     window n2=5 f2=1 |
     unif2 d1=%g n1=%d v00=%s
     ''' % (d,int(1.5+zmax/d),vcolstr))

Flow('velocity','reff',
     '''
     window n2=5 f2=1 |
     unif2 d1=%g n1=%d v00=%s dvdx=%s x0=%s |
     transp
     ''' % (1,5,vstr,gxstr,xrefstr))


Result('mod1',
       '''
       grey color=j
       screenratio=%g screenht=4 wanttitle=n
       labelfat=3 labelsz=6 scalebar=y
       maxval=5.2 minval=2.2 allpos=y bias=2.2 clip=3
       label1="Depth (km)"
       label2="Distance (km)"
       barlabel="Velocity (km/s)"
       ''' % (zmax/xmax))

Flow('lineshot','modcol','window n2=1 f2=426 | math output=x1')
Flow('lineshotc','lineshot','math output="426*0.0101" | cat axis=2 $SOURCE | transp | dd type=complex | window ')
Plot('lineshotc','graph screenratio=%g yreverse=y screenht=4 min1=0 max1=6 wanttitle=n wantaxis=n plotfat=4 plotcol=7' % (zmax/xmax))
Plot('modcol',
       '''
       grey color=j
       screenratio=%g screenht=4 wanttitle=n
       mean=y labelfat=2 labelsz=4
       label1="Depth (km)"
       label2="Distance (km)"
       ''' % (zmax/xmax))
Result('modcol','modcol lineshotc','Overlay')

Flow('refs','ref','window n2=5 f2=1 | put o2=0')
Flow('dips','refs','deriv scale=y')

# Kirchhoff modeling ############################################################
shotscmp = (2.5,3.0,3.5)
plotscmp = []

Flow('shotcmp','refs dips',
        '''
        kirmod_newton nt=751 dt=0.004 freq=15 cmp=y
        ns=61 s0=0.0 ds=0.1 nh=101 dh=0.04 h0=-2 verb=y 
        debug=n fwdxini=y
        vstatus=%d velocity=%s debug=n
        xref=%s zref=%s xgradient=%s zgradient=%s dip=${SOURCES[1]}
        ''' %(vstat,vstr,xrefstr,zrefstr,gxstr,gzstr))

Flow('shotcmp15','shotcmp','window n3=1 f3=15')
Flow('shotcmp30','shotcmp','window n3=1 f3=30')
Flow('shotcmp45','shotcmp','window n3=1 f3=45')
Plot('shotcmp15',
        '''
        grey transp=y yreverse=y poly=y
        title="CMP at 1.5 km" max1=1.5 label2=Offset unit2=km screenratio=0.75
        axisfat=3 titlefat=3 titlesz=18 labelfat=3 labelsz=14
        wherexlabel=top
        ''')
Plot('shotcmp30',
        '''
        grey transp=y yreverse=y poly=y
        title="CMP at 3 km" max1=1.5 label2=Offset unit2=km screenratio=0.75
        axisfat=3 titlefat=3 titlesz=18 labelfat=3 labelsz=14
        wherexlabel=top
        ''')
Plot('shotcmp45',
        '''
        grey transp=y yreverse=y poly=y
        title="CMP at 4.5 km" max1=1.5 label2=Offset unit2=km screenratio=0.75
        axisfat=3 titlefat=3 titlesz=18 labelfat=3 labelsz=14
        wherexlabel=top
        ''')

# NMO velocity for the model ##########################################################################################
# Compute depth
Flow('depth','ref refs','window n2=5 | math s=${SOURCES[1]} output="s-input" | put o2=0')
Flow('t0','depth velocity','math v=${SOURCES[1]} output="input/v"')
Flow('v2t0','t0 velocity','math  v=${SOURCES[1]} output="input*v^2" ')

Flow('t0sum','t0','transp | causint | transp')
Flow('t0sumext','t0sum','spray axis=3 n=101 d=0.04 o=-2')
Flow('vnmosq','v2t0 t0sum','transp | causint | transp | math t0=${SOURCES[1]} output="input/t0" ')
Flow('vnmosqext','vnmosq','spray axis=3 n=101 d=0.04 o=-2')
Flow('hypertime','vnmosqext t0sumext','math t0=${SOURCES[1]} output="sqrt(4*t0^2 + x3^2/input)"')

Flow('offset',None,'spike n1=101 d1=0.04 o1=-2 | math output=x1')

# Taking into account heterogeneities ##################################################################################

Flow('slow','velocity',' math output="1/input" ')
Flow('vnmosqbylayer','velocity',' math output="input^2" ')

Flow('vnmosqhet','refs vnmosqbylayer slow t0sum',
        '''
        sffermatrecursion vnmosq=${SOURCES[1]} slow=${SOURCES[2]} t0sum=${SOURCES[3]}
        ''')
Flow('vnmosqhetext','vnmosqhet','spray axis=3 n=101 d=0.04 o=-2')
Flow('hyperhettime','vnmosqhetext t0sumext','math t0=${SOURCES[1]} output="sqrt(4*t0^2 + x3^2/input)"')

# Plotting comparison ###################################################################################################
s1=390
s2=426
s3=530

for j in [s1,s2,s3]:
        if j == s1:
                sh = 1
        elif j == s2:
                sh = 2
        elif j == s3:
                sh = 3
        shot = str(sh)
        for i in range(5):
                num = str(i+1)
                Flow('shotcmp'+shot,'shotcmp','window n3=1 f3=%d '  %floor(j*0.1))
                Plot('shotcmp'+shot,
                        '''
                        grey transp=y yreverse=y poly=y
                        title="CMP at %s km" max1=2 label2=Offset unit2=km screenratio=1.15
                        axisfat=3 titlefat=3 titlesz=8 labelfat=3 labelsz=6
                        wherexlabel=top
                        ''' % (str(j*0.0101)) )
                Plot('hyperto'+num+'-'+shot,'hypertime offset',
                        '''
                        window n1=1 f1=%d n2=1 f2=%d | cat axis=2 ${SOURCES[1]} order=2,1 | transp | dd type=complex | window |
                        graph yreverse=y plotcol=%d plotfat=4 min2=0 max2=2 screenratio=1.15
                        axisfat=3 wanttitle=n wantaxis=n
                        ''' % (j,i,i+1) )
                Plot('hyperhetto'+num+'-'+shot,'hyperhettime offset',
                        '''
                        window n1=1 f1=%d n2=1 f2=%d | cat axis=2 ${SOURCES[1]} order=2,1 | transp | dd type=complex | window |
                        graph yreverse=y plotcol=%d plotfat=4 min2=0 max2=2 screenratio=1.15 dash=1
                        axisfat=3 wanttitle=n wantaxis=n
                        ''' % (j,i,i+1) )

Plot('box',None,'box font=2 x0=4.8 y0=8.6 label="Isotropic with velocity gradient" xt=0.000000 yt=0.000000')
Plot('hypercompare1',   ''' shotcmp1 hyperto1-1 hyperto2-1 hyperto3-1 hyperto4-1 hyperto5-1 
                                                        hyperhetto1-1 hyperhetto2-1 hyperhetto3-1 hyperhetto4-1 hyperhetto5-1
                                                ''','Overlay')
Result('hypercompare2', ''' shotcmp2 hyperto1-2 hyperto2-2 hyperto3-2 hyperto4-2 hyperto5-2 
                                                        hyperhetto1-2 hyperhetto2-2 hyperhetto3-2 hyperhetto4-2 hyperhetto5-2 box
                                                ''','Overlay')
Plot('hypercompare3',   ''' shotcmp3 hyperto1-3 hyperto2-3 hyperto3-3 hyperto4-3 hyperto5-3 
                                                        hyperhetto1-3 hyperhetto2-3 hyperhetto3-3 hyperhetto4-3 hyperhetto5-3
                                                ''','Overlay')

#Result('hypercompare','hypercompare1 hypercompare3','SideBySideAniso')

# Plotting flatness comparison ###################################################################################

# For last reflectors
for j in [s1,s2,s3]:
        if j == s1:
                sh = 1
        elif j == s2:
                sh = 2
        elif j == s3:
                sh = 3
        shot = str(sh)
        Flow('hypertimeshift'+shot,'vnmosqext t0sumext',
                ''' math t0=${SOURCES[1]} output="sqrt(4*t0^2 + x3^2/input) - 2*t0"| 
                window n1=1 f1=%d n2=1 f2=4
                ''' % j)
        Flow('hyperhettimeshift'+shot,'vnmosqhetext t0sumext',
                ''' math t0=${SOURCES[1]} output="sqrt(4*t0^2 + x3^2/input) - 2*t0"| 
                window n1=1 f1=%d n2=1 f2=4
                '''  % j)
        Flow('warped'+shot,'shotcmp'+shot+' hypertimeshift'+shot,'datstretch datum=${SOURCES[1]}')
        Flow('warpedhet'+shot,'shotcmp'+shot+' hyperhettimeshift'+shot,'datstretch datum=${SOURCES[1]}')
        Plot('warped'+shot,
                '''
               grey transp=y yreverse=y poly=y
                title=" Corrected CMP at %s km (1D)" min1=1.55 max1=1.85 label2=Offset unit2=km screenratio=0.5
                axisfat=3 titlefat=6 titlesz=14 labelfat=6 labelsz=10
                wherexlabel=top
                ''' % str(j*0.0101) )
        Plot('warpedhet'+shot,
                '''
               grey transp=y yreverse=y poly=y
                title="Corrected CMP at %s km (Proposed)" min1=1.55 max1=1.85 label2=Offset unit2=km screenratio=0.5
                axisfat=3 titlefat=6 titlesz=14 labelfat=6 labelsz=10
                wherexlabel=top
                ''' % str(j*0.0101) )
        Plot('t0to-'+shot,'t0sumext offset',
                '''
                window n1=1 f1=%d n2=1 f2=4 | scale dscale=2 | cat axis=2 ${SOURCES[1]} order=2,1 | transp | dd type=complex | window |
                graph yreverse=y plotcol=5 plotfat=16 min2=1.55 max2=1.85 screenratio=0.5 dash=1
                axisfat=3 wanttitle=n wantaxis=n
                ''' % j )
        Plot('warpedn'+shot,'warped'+shot+' t0to-'+shot,'Overlay')
        Plot('warpedhetn'+shot,'warpedhet'+shot+' t0to-'+shot,'Overlay')

Result('warpedcompare2','warpedn2 warpedhetn2','TwoRows')


# Ray tracing for numerical calculation of the second traveltime derivative ##########
updown2=[1,2,3,4,5]
N2 = len(updown2)-1
xinitial = [4.,4.,4.,4.]
updown2str = ','.join(map(str,updown2)) # convert updown to a string
xinistr = ','.join(map(str,xinitial)) # convert updown to a string

Flow('ray-2','refcut',
        '''isaac2 niter=10 number=%d vstatus=%d debug=n xs=4.2626 xr=4.3026 velocity=%s layer=%s
           xref=%s zref=%s xgradient=%s zgradient=%s xinitial=%s tol=1e-5''' % (N2,vstat,vstr,updown2str,xrefstr,zrefstr,gxstr,gzstr,xinistr))
Flow('ray-1','refcut',
        '''isaac2 niter=10 number=%d vstatus=%d debug=n xs=4.2826 xr=4.3026 velocity=%s layer=%s
           xref=%s zref=%s xgradient=%s zgradient=%s xinitial=%s tol=1e-5''' % (N2,vstat,vstr,updown2str,xrefstr,zrefstr,gxstr,gzstr,xinistr))
Flow('ray0','refcut',
        '''isaac2 niter=10 number=%d vstatus=%d debug=n xs=4.3026 xr=4.5652 velocity=%s layer=%s
           xref=%s zref=%s xgradient=%s zgradient=%s xinitial=%s tol=1e-5''' % (N2,vstat,vstr,updown2str,xrefstr,zrefstr,gxstr,gzstr,xinistr))
Flow('ray+1','refcut',
        '''isaac2 niter=10 number=%d vstatus=%d debug=n xs=4.3226 xr=4.3026 velocity=%s layer=%s
           xref=%s zref=%s xgradient=%s zgradient=%s xinitial=%s tol=1e-5''' % (N2,vstat,vstr,updown2str,xrefstr,zrefstr,gxstr,gzstr,xinistr))
Flow('ray+2','refcut',
        '''isaac2 niter=10 number=%d vstatus=%d debug=n xs=4.4226 xr=4.3026 velocity=%s layer=%s
           xref=%s zref=%s xgradient=%s zgradient=%s xinitial=%s tol=1e-5''' % (N2,vstat,vstr,updown2str,xrefstr,zrefstr,gxstr,gzstr,xinistr))

for k in ['-2','-1','0','+1','+2']:
        Plot('ray'+k,
                 '''
                 dd type=complex | window | 
                 graph wanttitle=n wantaxis=n yreverse=y min2=0 max2=3 min1=0 max1=6 plotcol=7 plotfat=6 screenht=5.0 screenratio=0.333 yll=3.5 xll=1.5
                 axisfat=3 titlefat=3 titlesz=10 labelfat=3 labelsz=6
                 ''')
Result('dt2dx2','ray0 ref','Overlay')


End()