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.65,1.0,0.65,0.5),
#        (1.0,1.1,1.5,1.5,1.1,1.2,1.1,1.1),
#               (1.75,1.75,1.75,1.75,1.75,1.55,1.55,1.55),
#               (2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0))
# Flat
#layers = ((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),
                (1.0,1.0,1.0,1.0,1.0),
        (2.0,2.0,2.0,2.0,2.0),
        (3.0,3.0,3.0,3.0,3.0),
                (4.0,4.0,4.0,4.0,4.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)


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=6 labelfat=1 labelsz=4''')


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

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

# Kirchhoff modeling examples############################################################################
Flow('shotcmp','refs dips aniso',
        '''
        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 
        vstatus=%d debug=n fwdxini=n niter=100
        xref=0 zref=0 dip=${SOURCES[1]} aniso=${SOURCES[2]}
        ''' %vstat)

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=2 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=2 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=2 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 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=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','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=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=200
s2=426
s3=435

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="Anisotropic" 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('hypercompareaniso2',    ''' 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 hypercompare2 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.65 max1=1.95 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.65 max1=1.95 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.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('warpedcompareaniso2','warpedn2 warpedhetn2','TwoRows')

# Velocity scan and picked ##################################################################################################################################
#Flow('vscan','shotcmp',' window max1=1.5 | vscan semblance=y v0=2.5 nb=10 dv=0.02 nv=101 half=n')
#Flow('pick','vscan','pick rect1=30 rect2=10 vel0=3 gate=10 an=3')

#Plot('vscan1','vscan',
#       '''
#       window n3=1 f3=30 |
#       grey transp=y yreverse=y color=j allpos=y
#       title="Semblance" max1=1.5 label2=Offset unit2=km screenratio=1.5 
#       axisfat=3 titlefat=3 titlesz=14 labelfat=1 labelsz=10
#       wherexlabel=top 
#       ''')
#Plot('pick1','pick',
#       '''
#       window n2=1 f2=30 |
#       graph transp=y yreverse=y plotcol=0 plotfat=10 min1=0 max1=1.5 min2=2.5 max2=4.5 screenratio=1.5
#       wanttitle=n wantaxis=n
#       ''')
#Result('pickcmp1','vscan1 pick1','Overlay')

## NMO
#Flow('nmo','shotcmp pick','nmo half=n velocity=${SOURCES[1]}')

#Result('pick',
#       '''     
#       grey color=j bias=3.375 maxval=4 minval=2.75 clip=0.625
#       screenratio=%g screenht=4 title="Picked velocity"
#       labelfat=1 labelsz=4 wherexlabel=top scalebar=y barlabel="Stacking velocity (km/s)"
#       label1="Time"
#       label2="Distance (km)"
#       ''' % (zmax/xmax))

## RMS velocity for stratified medium
## Create a dense model
#Flow('tlayer1',None,'spike n1=61 d1=0.1 o1=0 mag=%g' %t0to1)
#Flow('tlayer2',None,'spike n1=61 d1=0.1 o1=0 mag=%g' %t0to2)
#Flow('tlayer3',None,'spike n1=61 d1=0.1 o1=0 mag=%g' %t0to3)
#Flow('tlayer4',None,'spike n1=61 d1=0.1 o1=0 mag=%g' %t0to4)
#Flow('tlayer5',None,'spike n1=61 d1=0.1 o1=0 mag=%g' %t0to5)

#Flow('tlayer','tlayer1 tlayer2 tlayer3 tlayer4 tlayer5', '''cat axis=2 ${SOURCES[1:5]}''')

#Flow('rms','tlayer',
#     '''
#     unif2 d1=0.004 n1=376 o1=0.0 v00=%g,%g,%g,%g,%g |
#     math output="sqrt(input)"
#     ''' % (vnmoto1,vnmoto2,vnmoto3,vnmoto4,vnmoto5))
#     
#Result('rms',
#       '''     
#       grey color=j  bias=3.375 maxval=4 minval=2.75 clip=0.625
#       screenratio=%g screenht=4 title="RMS velocity"
#       labelfat=1 labelsz=4 wherexlabel=top scalebar=y barlabel="RMS velocity (km/s)"
#       label1="Time"
#       label2="Distance (km)"
#       ''' % (zmax/xmax))

End()