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))
# 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.5,0.5,0.5,0.5),
#         (0.85,0.85,0.85,0.85,0.85,0.85,0.85,0.85),
#               (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]} ''')


#anisotropy (Q must not equal to 1 in the case of anisotropy)
# c11,c33, Q1, Q3

aniso = ((9.0,9.01,1.0001,1.001), # constant velocity 3 km/s
          (12.716,10.233,1.128,1.137), # Africa shale 
          (13.0,13.01,1.0001,1.001), #constant velocity sqrt(13) km/s
          (14.47,9.57,1.58,1.68), # Greenhorn Jones & Wang
                  (22.05,14.89,1.344,1.423))  # North sea Vernik & Liu

n3 = len(aniso[0])
n4 = len(aniso)

Flow('aniso.asc',None,
     '''
        echo %s
     n1=%d n2=%d o1=0 d1=1
     data_format=ascii_float in=$TARGET     
     ''' % (' '.join(map(arr2str,aniso)),
            n3,n4))
Flow('aniso','aniso.asc','dd form=native')
# Model specifications ########################################################################

updown=[1,2,3,4,5,4,3,2,1,0]
N = len(updown)-1
vstat = 2
gx = [0.0,0.0,0.0,0.0,0.0]
gz = [0.0,0.0,0.0,0.0,0.0]
v = [1,1,1,1,1]
vcol = [1.508,1.581,1.69,1.826,2,2.2]
xref = [0.1,0.1,0.1,0.1,0.1]
zref = [0.2,0.45,0.6,1.2,1.5]

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('xrefl','ref aniso',
        '''isaac2 niter=500 number=%d vstatus=%d debug=n xs=2.3 xr=4.3 velocity=%s layer=%s
           xref=%s zref=%s xgradient=%s zgradient=%s tol=0.0001 aniso=${SOURCES[1]}''' % (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=6 labelfat=1 labelsz=4
         ''')
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=6 labelfat=1 labelsz=4
         ''')
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
     ''' % (d,int(1.5+zmax/d),vstr))

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

Result('mod1',
       '''
       grey color=j
       screenratio=%g screenht=4 wanttitle=n
       mean=y labelfat=1 labelsz=4
       label1="Depth (km)"
       label2="Distance (km)"
       ''' % (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')

v1 = 3
v2 = 3.2
v3 = 3.6
v4 = 3.1
v5 = 3.86

vnmo1 = sqrt(v1*v1)
vnmo2 = sqrt(12.716/1.137)
vnmo3 = sqrt(v3*v3)
vnmo4 = sqrt(14.47/1.68)
vnmo5 = sqrt(22.05/1.423)

vver = (v1,v2,v3,v4,v5)
Flow('vver.asc',None,
     '''
        echo %s
     n1=%d o1=0 d1=1
     data_format=ascii_float in=$TARGET     
     ''' % (' '.join(map(str,vver)),5 ))
Flow('vver','vver.asc','dd form=native')

vnmo = (vnmo1,vnmo2,vnmo3,vnmo4,vnmo5)
Flow('vnmo.asc',None,
     '''
        echo %s
     n1=%d o1=0 d1=1
     data_format=ascii_float in=$TARGET     
     ''' % (' '.join(map(str,vnmo)),5))
Flow('vnmo','vnmo.asc','dd form=native')

Flow('velocityver','vver','spray d=%f n=%d | transp' % (d,int(1.5+xmax/d)))
Flow('velocitynmo','vnmo','spray d=%f n=%d | transp' % (d,int(1.5+xmax/d)))

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


Flow('diffrac1','refs','math output="1.0" | cut n2=4 | math output="1.0 - input" ')
Flow('diffrac2','refs','math output="1.0" | cut f1=420 n1=11 | math output="1.0 - input" ')
Flow('diffractivity','diffrac1 diffrac2','add ${SOURCES[1]} | mask min=1 | dd type=float  ')

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

Flow('zo','refs dips diffractivity aniso',
        '''
        kirmod_newton nt=751 dt=0.004 freq=15 cmp=y
        ns=201 s0=3.2825 ds=0.0101 nh=1 dh=0.04 h0=0 verb=y 
        debug=n fwdxini=y
        vstatus=%d debug=n
        dip=${SOURCES[1]} refl=${SOURCES[2]} aniso=${SOURCES[3]}
        ''' %(vstat))
Flow('zotaper','zo','window | pow pow1=1 | costaper nw1=25 | costaper nw2=25')
Result('zotaper',
        '''
        window | grey transp=y yreverse=y poly=y
        title="Zero-offset section" label2=Midpoint unit2=km screenratio=1.3
        axisfat=3 titlefat=3 titlesz=18 labelfat=3 labelsz=14 max1=2
        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 velocityver','math v=${SOURCES[1]} output="input/v"')
Flow('v2t0','t0 velocitynmo','math  v=${SOURCES[1]} output="input*v^2" ')

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

Flow('offset',None,'spike n1=201 d1=0.0101 o1=-1.01 | math output=x1')

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

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

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

# Plotting comparison ###################################################################################################
s3=426

for j in [s3]:
        sh = 3
        shot = str(sh)
        for i in range(4,5):
                num = str(i+1)
                Flow('zo'+shot,'zotaper',' window ')
                Plot('zo'+shot,
                        '''
                        grey transp=y yreverse=y poly=y pclip=95
                        title="Zero-offset diffraction at %s km" min1=0 max1=2 label2=Midpoint unit2=km screenratio=1.5
                        axisfat=3 titlefat=3 titlesz=6 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.5
                        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.5 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',       ''' zo2 
#                                                       hyperhetto4-2 hyperhetto5-2
#                                               ''','Overlay')
Result('hypercompareanisodiff3',        ''' zo3 hyperto5-3 
                                                        hyperhetto5-3
                                                ''','Overlay')



# Plotting flatness comparison ###################################################################################
# For last reflectors
for j in [s3]:
        sh = 3
        shot = str(sh)
        Flow('hypertimeshift'+shot,'vnmosqext t0sumext',
                ''' math t0=${SOURCES[1]} output="2*sqrt(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="2*sqrt(t0^2 + (x3)^2/input) - 2*t0"| 
                window n1=1 f1=%d n2=1 f2=4
                '''  % j)
        Flow('warped'+shot,'zo'+shot+' hypertimeshift'+shot,'datstretch datum=${SOURCES[1]}')
        Flow('warpedhet'+shot,'zo'+shot+' hyperhettimeshift'+shot,'datstretch datum=${SOURCES[1]}')
        Plot('warped'+shot,
                '''
               grey transp=y yreverse=y poly=y
                title="Flattened at %s km (1D)" min1=1.65 max1=1.95 label2=Midpoint 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="Flattened at %s km (Proposed)" min1=1.65 max1=1.95 label2=Midpoint 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.65 max2=1.95 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('warpedhypercompareaniso3','warpedn3 warpedhetn3','TwoRows')




# Plotting migrated results comparison ###################################################################################
Flow('vold','vnmosq','math output="sqrt(input)" | window n1=1 f1=426')
vold = '3,3.229,3.336,3.272,3.413'

Flow('vnew','vnmosqhet','math output="sqrt(input)" | window n1=1 f1=426')
vnew = '3,3.229,3.423,3.442,3.92'

# Shift the pulse a bit to ensure the correct velocity is applied
Flow('refdiff','t0sum','window n1=1 f1=425 | add scale=2 | spray axis=2 n=201 o=3.2825 d=0.0101 | transp')
Flow('vdiffmodelold','refdiff','unif2 d1=0.004 n1=751 v00=%s | smooth rect1=10' %vold)
Flow('vdiffmodelnew','refdiff','unif2 d1=0.004 n1=751 v00=%s | smooth rect1=10' %vnew)

Flow('oldmig','zotaper vdiffmodelold',' window | mig2 antialias=1.0 vel=${SOURCES[1]}')
Flow('newmig','zotaper vdiffmodelnew',' window | mig2 antialias=1.0 vel=${SOURCES[1]}')


Plot('oldmig',
       '''
       grey transp=y yreverse=y poly=y pclip=99
       title="Migrated diffraction (1D)" min1=0 max1=2 label2=Midpoint unit2=km screenratio=0.75
       axisfat=3 titlefat=6 titlesz=14 labelfat=6 labelsz=12
       wherexlabel=top
       ''' )
Plot('newmig',
       '''
       grey transp=y yreverse=y poly=y pclip=99
       title="Migrated diffraction (Proposed)" min1=0 max1=2 label2=Midpoint unit2=km screenratio=0.75
       axisfat=3 titlefat=6 titlesz=14 labelfat=6 labelsz=12
       wherexlabel=top
       ''' )

Result('migcompareaniso','oldmig newmig','SideBySideAniso')


# Focusing measurement
Flow('oldmigfocus','oldmig','window min1=1.65 max1=1.9 | focus dim=2 rect1=10 rect2=3')
Flow('newmigfocus','newmig','window min1=1.65 max1=1.9 | focus dim=2 rect1=10 rect2=3')

Plot('oldmigfocus',
       '''
       grey transp=y yreverse=y poly=y pclip=99
       title="Focus of migrated diffraction (1D)" min1=1.65 max1=1.9 label2=Midpoint unit2=km screenratio=0.55
       axisfat=3 titlefat=6 titlesz=14 labelfat=6 labelsz=10
       wherexlabel=top color=j allpos=y
       ''' )
Plot('newmigfocus',
       '''
       grey transp=y yreverse=y poly=y pclip=99
       title="Focus of migrated diffraction (Proposed)" min1=1.65 max1=1.9 label2=Midpoint unit2=km screenratio=0.55
       axisfat=3 titlefat=6 titlesz=14 labelfat=6 labelsz=10
       wherexlabel=top color=j allpos=y
       ''' )

Result('focuscompareaniso','oldmigfocus newmigfocus','TwoRows')


# Focusing measurement with envelope
Flow('oldmigenv','oldmig','window min1=1.65 max1=1.9 | envelope')
Flow('newmigenv','newmig','window min1=1.65 max1=1.9 | envelope')
Plot('oldmigenv',
       '''
       grey transp=y yreverse=y poly=y
       title="Focus of migrated diffraction (1D)" min1=1.65 max1=1.9 label2=Midpoint unit2=km screenratio=0.55
       axisfat=3 titlefat=6 titlesz=14 labelfat=6 labelsz=10
       wherexlabel=top color=j allpos=y scalebar=y maxval=0.082 clip=0.082
       ''' )
Plot('newmigenv',
       '''
       grey transp=y yreverse=y poly=y 
       title="Focus of migrated diffraction (Proposed)" min1=1.65 max1=1.9 label2=Midpoint unit2=km screenratio=0.55
       axisfat=3 titlefat=6 titlesz=14 labelfat=6 labelsz=10
       wherexlabel=top color=j allpos=y scalebar=y maxval=0.082 clip=0.082
       ''' )

Result('envcompareaniso','oldmigenv newmigenv','TwoRows')

End()