from rsf.proj import *
import math

s02 = 1.
qx  = 0.052/2

min2 = 0.5
max2 = 6.5

# Problem setup ###############################################################################
# velocity model
Flow('slow2',None,
     '''
     math n1=101 n2=361 d1=0.02 d2=0.02 output="%g-2*%g*x2"
     label1=Depth unit1=km label2=Position unit2=km 
     label=Slowness-squared unit="s\^2\_/km\^2"
     ''' % (s02,qx))
Flow('v2','slow2','math output="1/input" ')
Plot('v2',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="True w(x,z)" 
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 allpos=y minval=1 maxval=1.5 bias=1 clip=0.5
     screenratio=0.75 screenht=9 allpos=y barlabel="v\^2\_" barunit="km\^2\_/s\^2"
     ''' % (min2,max2))

# analytical sigma in (z,x)
Flow('depsigma','slow2',
     '''
     math output="sqrt(((%g-2*%g*x2)-sqrt((%g-2*%g*x2)^2-4*%g*%g*x1*x1))/(2*%g*%g))" |
     put label=Sigma unit=
     ''' % (s02,qx,s02,qx,qx,qx,qx,qx))

# analytical x0
Flow('ax0','depsigma',
     '''
     math output="x2+%g*input*input/2" |
     put label=Position unit=km
     ''' % qx)

Plot('ax0',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y allpos=y barreverse=y
     labelsz=15 titlesz=16 titlefat=8 labelfat=6
     title="Analytical x\_0\^" screenratio=0.75 screenht=9
     minval=0 maxval=7.5 clip=7.5
     ''' % (min2,max2))
Plot('cax0','ax0',
     '''
     window min2=%g max2=%g |
     contour nc=100 scalebar=y plotcol=7 plotfat=7
     wantaxis=n wanttitle=n screenratio=0.75 screenht=9
     ''' % (min2,max2))
Plot('x0','ax0 cax0','Overlay')

# analytical t0
Flow('at0','ax0 depsigma',
     '''
     math sigma=${SOURCES[1]}
     output="(%g-2*%g*input)*sigma+%g*%g*sigma*sigma*sigma/3" |
     put label=Time unit=s
     ''' % (s02,qx,qx,qx))

Plot('at0',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y allpos=y barreverse=y
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 minval=0 maxval=2 clip=2
     title="Analytical t\_0\^" screenratio=0.75 screenht=9
     ''' % (min2,max2))
Plot('cat0','at0',
     '''
     window min2=%g max2=%g |
     contour nc=40 scalebar=y plotcol=7 plotfat=7
     wantaxis=n wanttitle=n screenratio=0.75 screenht=9
     ''' % (min2,max2))
Plot('t0','at0 cat0','Overlay')

# analytical geometrical spreading
Flow('ageo','slow2',
     '''
     math output="2*((%g-2*%g*x2)^2-4*%g*%g*x1*x1)/(%g-2*%g*x2)/(%g-2*%g*x2+sqrt((%g-2*%g*x2)^2-4*%g*%g*x1*x1))"
     ''' % (s02,qx,qx,qx,s02,qx,s02,qx,s02,qx,qx,qx))

Plot('ageo',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y allpos=y 
     labelsz=15 titlesz=16 titlefat=8 labelfat=6
     barlabel="Q\^2\_" barunit= minval=0.65 maxval=1 bias=0.65 clip=0.35
     title="Geometrical Spreading" screenratio=0.65 screenht=9
     ''' % (min2,max2))

# analytical sigma in (t0,x0)
Flow('timsigma',None,
     '''
     math n1=626 n2=361 d1=0.004 d2=0.02
     output="((3*x1+sqrt(9*x1*x1+4*(%g-2*%g*x2)^3/(%g*%g)))/(2*%g*%g))^(1./3)-(%g-2*%g*x2)/%g*(2/(3*%g*x1+sqrt(9*%g*%g*x1*x1+4*(%g-2*%g*x2)^3)))^(1./3)"
     label1=Time unit1=s label2=Position unit2=km 
     label=Sigma unit= |
     cut n1=1
     ''' % (s02,qx,qx,qx,qx,qx,s02,qx,qx,qx,qx,qx,s02,qx))

# analytical Dix
Flow('dix','timsigma',
     '''
     math output="sqrt(%g-2*%g*x2)/(%g-2*%g*x2-%g*%g*input*input)" |
     put label=Velocity unit=km/s
     ''' % (s02,qx,s02,qx,qx,qx))

Plot('dix',
     '''
     window min2=%g max2=%g | math output="input^2"|
     grey color=j scalebar=y allpos=y bias=1
     labelsz=15 titlesz=16 titlefat=8 labelfat=6
     title="w\_d\^(x\_0\^,t\_0\^)" barlabel="v\^2\_" barunit="km\^2\_/s\^2"
     screenratio=0.65 screenht=9 minval=1 maxval=1.35 bias=1 clip=0.35
     ''' % (min2,max2))

# convert Dix to depth
Flow('dixdepth','dix',
     '''
     time2depth velocity=$SOURCE intime=y twoway=n nz=101 dz=0.02 |
     put label1=Depth unit1=km
     ''')

Result('model-slow','v2 x0 t0','OverUnderAniso')

## mask
#Flow('mask','ax0',
#     '''
#     math output=-1 | cut n2=25 | cut n2=35 f2=326 | dd type=int
#     ''')

## inversion
#Flow('t2d tt2d xt2d ft2d gt2d ct2d','init dix mask',
#     '''
#     tdconvert niter=3 verb=n cgiter=2000 eps=4 shape=y rect1=4 rect2=10 dix=${SOURCES[1]} mask=${SOURCES[2]}
#     t0=${TARGETS[1]} x0=${TARGETS[2]} f0=${TARGETS[3]} grad=${TARGETS[4]} cost=${TARGETS[5]}
#     ''')

## cost
#Plot('cost0','ct2d',
#     '''
#     window min2=%g max2=%g | window n3=1 f3=0 | 
#     grey color=j scalebar=y barreverse=y barlabel=Cost barunit=
#     title="Initial f" minval=-0.45 maxval=0.63 clip=0.63
#     labelsz=15 titlesz=16 titlefat=8 labelfat=6
#     screenratio=0.65 screenht=9
#     ''' % (min2,max2))
#Plot('cost3','ct2d',
#     '''
#     window min2=%g max2=%g | window n3=1 f3=3 | 
#     grey color=j scalebar=y barreverse=y barlabel=Cost barunit=
#     title="Final f" minval=-0.45 maxval=0.63 clip=0.63
#     labelsz=15 titlesz=16 titlefat=8 labelfat=6
#     screenratio=0.65 screenht=9
#     ''' % (min2,max2))

#Result('hs2cost','cost0 cost3','OverUnderAniso')

## error
#Plot('error0','vel init',
#     '''
#     add scale=-1,1 ${SOURCES[1]} | window min2=%g max2=%g |
#     grey color=j scalebar=y barreverse=y
#     title="Initial Model Error" minval=-0.004 maxval=0.176 clip=0.176
#     labelsz=15 titlesz=16 titlefat=8 labelfat=6
#     screenratio=0.65 screenht=9
#     ''' % (min2,max2))
#Plot('error3','vel t2d',
#     '''
#     add scale=-1,1 ${SOURCES[1]} | window min2=%g max2=%g |
#     grey color=j scalebar=y barreverse=y
#     title="Final Model Error" minval=-0.004 maxval=0.176 clip=0.176
#     labelsz=15 titlesz=16 titlefat=8 labelfat=6
#     screenratio=0.65 screenht=9
#     ''' % (min2,max2))

#Result('hs2error','error0 error3','OverUnderAniso')

############################################################################################################
# Weak assumption 
############################################################################################################

# Reference 1D model in (z,x) from the central trace of Dix velocity 
# velocity model
Flow('slowz','dixdepth','window n2=1 f2=180 | math output="1/input" | spray axis=2 n=361 o=0 d=0.02')
Flow('dslow','slowz slow2','math true=${SOURCES[1]} output="sqrt(input^2-true)"')

Plot('slowz',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y allpos=y title="Reference V(z)" pclip=100
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 barlabel="Slowness" barunit="s/km"
     minval=0.12 maxval=0.95 clip=0.95 screenratio=0.65 screenht=9
     ''' % (min2,max2))
#Plot('dslow',
#     '''
#     window min2=%g max2=%g |
#     grey color=j scalebar=y title=Model pclip=100
#     labelsz=15 titlesz=16 titlefat=8 labelfat=6
#     screenratio=0.65 screenht=9 barlabel="Slowness" barunit="s/km"
#     ''' % (min2,max2))
# analytical x0
Flow('ax0z','slowz',
     '''
     math output="x2" |
     put label=Position unit=km
     ''')

# analytical t0
Flow('at0z','slowz',
     '''
     math output="x1*input" |
     put label=Time unit=s
     ''')
#Plot('at0z',
#     '''
#     window min2=%g max2=%g |
#     grey color=j scalebar=y title=Model pclip=100
#     labelsz=15 titlesz=16 titlefat=8 labelfat=6
#     screenratio=0.65 screenht=9
#     ''' % (min2,max2))
# 
# Derivatives of Dix velocity squared
Flow('dv2dt0','dix','math output="input^2" | smoothder')
Flow('dv2dx0','dix','math output="input^2" | transp | smoothder | transp')
Flow('beta','dv2dt0 dix',
     '''
     time2depth velocity=${SOURCES[1]} intime=y twoway=n nz=101 dz=0.02 |
     put label1=Depth unit1=km
     ''')
Flow('alpha','dv2dx0 dix',
     '''
     time2depth velocity=${SOURCES[1]} intime=y twoway=n nz=101 dz=0.02 |
     put label1=Depth unit1=km
     ''')
Plot('alpha',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="dw\_d\^/dx\_0" pclip=100
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 minval=0 allpos=y
     screenratio=0.75 screenht=9 allpos=y barlabel="dw\_d\^/dx\_0\^" barunit="km/s\^2"
     ''' % (min2,max2))
Plot('beta',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="dw\_d\^/dt\_0" pclip=100
     labelsz=15 titlesz=16 titlefat=8 labelfat=6
     screenratio=0.75 screenht=9 allpos=y barlabel="dw\_d\^/dt\_0\^" barunit="km\^2\_/s\^3"
     ''' % (min2,max2))

# Reference velocity squared
Flow('refdix','dixdepth','math output="input^2" ')
Flow('refvz','slowz','math output="1/input^2" ')

Plot('refdix',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="Ref w\_dr\^(x,z)" 
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 allpos=y minval=1 maxval=1.5 bias=1 clip=0.5
     screenratio=0.75 screenht=9 allpos=y barlabel="v\^2" barunit="km\^2\_/s\^2"
     ''' % (min2,max2))
Plot('refvz',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="Ref w\_r\^(z)" 
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 allpos=y  minval=1 maxval=1.35 bias=1 clip=0.35
     screenratio=0.75 screenht=9 allpos=y barlabel="v\^2" barunit="km\^2\_/s\^2"
     ''' % (min2,max2))

Result('input-slow','refdix alpha beta','OverUnderAniso')

 # Remap models for FD stability
for i in ['refdix','refvz','alpha','beta']:
        Flow(i+'_remap',i,'window')

# Find differential ########################################################################################
Flow('depth dx0 dt0 dv','refdix_remap refvz_remap alpha_remap beta_remap',
        '''
        time2depthweak zsubsample=10 nsmooth=3
        velocity=$SOURCE refvelocity=${SOURCES[1]} dvdx0=${SOURCES[2]} dvdt0=${SOURCES[3]}
        outdx0=${TARGETS[1]} outdt0=${TARGETS[2]} outdv=${TARGETS[3]}
        ''')

# t0
Flow('dt0true','at0 at0z','math est=${SOURCES[1]} output="input-est" ')
Plot('dt0true',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="Analytical \F10 D\F3 t\_0"  
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 minval=-0.15 maxval=0.15 clip=0.15 bias=0
     screenratio=0.75 screenht=9 barlabel="Time (s)" barunit=
     ''' % (min2,max2))
Plot('dt0',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="Estimated \F10 D\F3 t\_0"
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 minval=-0.15 maxval=0.15 clip=0.15 bias=0
     screenratio=0.75 screenht=9 barlabel="Time (s)" barunit=
     ''' % (min2,max2))
Result('dt0compare','dt0true dt0','OverUnderAniso')

Flow('dt0err','dt0true dt0','math est=${SOURCES[1]} output="input-est"')
Plot('dt0err',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="Error from estimated \F10 D\F3 t\_0"
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 minval=-0.15 maxval=0.15 clip=0.15 bias=0 
     screenratio=0.75 screenht=9 barlabel="Time (s)" barunit=
     ''' % (min2,max2))



# x0
Flow('dx0true','ax0 ax0z','math est=${SOURCES[1]} output="input-est" ')
Plot('dx0true ',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="Analytical \F10 D\F3 x\_0"
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 minval=0 maxval=0.1 clip=0.1 allpos=y
     screenratio=0.75 screenht=9 barlabel="Position (km)" barunit=
     ''' % (min2,max2))

Plot('dx0',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="Estimated \F10 D\F3 x\_0"
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 minval=0 maxval=0.1 clip=0.1 allpos=y
     screenratio=0.75 screenht=9 barlabel="Position (km)" barunit=
     ''' % (min2,max2))

Result('dx0compare','dx0true dx0','OverUnderAniso')

Flow('dx0err','dx0true dx0','math est=${SOURCES[1]} output="input-est"')
Plot('dx0err',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="Error from estimated \F10 D\F3 x\_0"
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 allpos=y minval=0 maxval=0.1 clip=0.1 
     screenratio=0.75 screenht=9 barlabel="Position (km)" barunit= 
     ''' % (min2,max2))



# v
Flow('dv2true','slow2 refvz','math est=${SOURCES[1]} output="1/input-est" ')
Plot('dv2true ',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="Analytical \F10 D\F3 w"
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 minval=-0.25 maxval=0.25 clip=0.25 bias=0
     screenratio=0.75 screenht=9 barlabel="v\^2" barunit="km\^2\_/s\^2"
     ''' % (min2,max2))

Plot('dv',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="Estimated \F10 D\F3 w "
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 minval=-0.25 maxval=0.25 clip=0.25 bias=0
     screenratio=0.75 screenht=9 barlabel="v\^2" barunit="km\^2\_/s\^2"
     ''' % (min2,max2))

Result('dv2compare','dv2true dv','OverUnderAniso')

Flow('dv2err','dv2true dv','math est=${SOURCES[1]} output="(input-est)"')
Plot('dv2err',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="Error from estimated \F10 D\F3 w"
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 minval=-0.25 maxval=0.25 clip=0.25 bias=0
     screenratio=0.75 screenht=9 barlabel="v\^2" barunit="km\^2\_/s\^2"
     ''' % (min2,max2))


Result('truecompare','dt0true dx0true dv2true','OverUnderAniso')
Result('estcompare-slow','dt0 dx0 dv','OverUnderAniso')
Result('errcompare-slow','dt0err dx0err dv2err','OverUnderAniso')

# Velocity difference
Flow('errvdix','slow2 dixdepth','math est=${SOURCES[1]} output="abs(1/sqrt(input)-est)" ')
Flow('errvest','slow2 refvz dv','math b=${SOURCES[1]} c=${SOURCES[2]} output="abs(1/sqrt(input)-sqrt(b+c))" ')
Plot('errvdix',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="Velocity error from conventional Dix inversion"
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 allpos=y minval=0 maxval=0.01 clip=0.01 bias=0 
     screenratio=0.75 screenht=9 barlabel="v" barunit="km/s"
     ''' % (min2,max2))
Plot('errvest',
     '''
     window min2=%g max2=%g |
     grey color=j scalebar=y title="Velocity error from the proposed method"
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 allpos=y minval=0 maxval=0.01 clip=0.01 bias=0 
     screenratio=0.75 screenht=9 barlabel="v" barunit="km/s"
     ''' % (min2,max2))

Result('errvcompare','errvdix errvest','OverUnderAniso')
End()