from rsf.proj import *
from rsf.recipes.beg import server
import math



#memsize for transpose
memsize=5000

#frequency for modeling?
freq = 20
#t = 1
dt = 0.00025
nt=1351
ot = 0
#z = 1
dz = 2
nz = 416
oz = 0
#x = 2
nx = 600
dx = 25
ox = 0
#y = 3
ny = 700
dy = 25
oy = 0

#Plane Wave Destruction Parameters
rect1d = nz/10
rect1t = nt/10
rect2 = nx/10
rect3 = rect2
order1 = 4
#smoothing parameters
srect1 = nz / 20
srect2 = nx / 20
srect3 = srect2 #make 2nd and 3rd axis have same smoothing length

######################
#Plotting Information#
######################
frame1d=250
frame1t=800
frame2=300
frame3=350

#Plotting Ratio
volpt2 = .7

def grey3(byte,color,title,label,frame1,frame2,frame3,point2):
         return '''
         byte gainpanel=all %s mean=y |
         grey3 flat=n frame1=%i frame2=%i frame3=%i
         color=%s title="%s" label1="%s" point2=%g
         ''' % (byte,frame1,frame2,frame3,color,title,label,point2)


#Get data
Fetch('ordovician-vp.HH','xavier',server)
Flow('vp','ordovician-vp.HH',
     '''
     dd form=native | 
     transp plane=12 memsize=%i |
     put o1=0 o2=0 o3=0  
     label1="Z Depth" label2="X Inline" label3="Y Crossline"
     unit1=m unit2=m unit3=m
     '''%(memsize))
Flow('vpbar','vp','bar gainpannel=a allpos=y min=3500 max=6500')
Plot('vp','vp vpbar',
     '''
     byte allpos=y min=3500 max=6500|
     grey3 title=Velocity color=a barlabel=Velocity barunit="m/s"
     frame1=%i frame2=%i frame3=%i flat=n
     scalebar=y bar=${SOURCES[1]} point2=%g 
     '''%(frame1d,frame2,frame3,volpt2))
Result('vp','vp vpbar',
     '''
     put label2=Inline label3=Crossline unit2=km unit3=km
     d2=.025 d3=.025   label1=Depth|
     byte allpos=y min=3500 max=6500|
     grey3 title="Ordovician Velocity" color=a barlabel=Velocity barunit="m/s"
     frame1=%i frame2=%i frame3=%i flat=n
     scalebar=y bar=${SOURCES[1]} point2=%g  
     '''%(frame1d,frame2,frame3,volpt2))
#smooth slowness
Flow('vp_s','vp',
     '''
     math output="1/input" |
     smooth rect1=%i rect2=%i rect3=%i|
     math output="1/input"
     '''%(srect1,srect2,srect3))

#Calculate Reflectivity
#No density data with model
Flow('refl','vp','ai2refl')
Flow('refl_t','refl vp',
    '''
    depth2time velocity=${SOURCES[1]} dt=%g nt=%i
    '''%(dt,nt))

#Transpose vp for modeling use
Flow('wvp','vp',
     '''
     transp plane=12 memsize=%i |
     transp plane=23 memsize=%i
     '''%(memsize,memsize))
Flow('wvp_s','vp_s',
     '''
     transp plane=12 memsize=%i |
     transp plane=23 memsize=%i
     '''%(memsize,memsize))
#Low Rank
Flow('fft2','wvp','fft1 | fft3 axis=2 pad=1 | fft3 axis=3 pad=1')
Flow('right2 left2','wvp fft2',
     '''
     scale dscale=0.5 |
     isolr3 seed=2012 dt=%g fft=${SOURCES[1]} 
     left=${TARGETS[1]} seed=2014 npk=10 eps=0.00001
     '''%(dt))
Flow('fft2_s','wvp_s','fft1 | fft3 axis=2 pad=1 | fft3 axis=3 pad=1')
Flow('right2_s left2_s','wvp_s fft2_s',
     '''
     scale dscale=0.5 |
     isolr3 seed=2012 dt=%g fft=${SOURCES[1]} 
     left=${TARGETS[1]} seed=2014 npk=10 eps=0.00001
     '''%(dt))

Flow('model','refl_t vp',
    '''
    ricker1 frequency=%g |
    time2depth velocity=${SOURCES[1]}
    dz=%g nz=%i
    ''' %(freq,dz,nz))
Plot('model',
     '''
     put label2=Inline d2=.025 unit2=km
     label3=Crossline d3=.025 unit3=km|
     byte|
     grey3 title="20 Hz Convolution"
     frame1=%i frame2=%i frame3=%i flat=n
     point2=%g
     '''%(frame1d,frame2,frame3,volpt2))
Flow('exp','model left2 right2',
     '''
     fftexp3 ompchunk=1
     left=${SOURCES[1]} right=${SOURCES[2]} 
     nt=%i dt=%g 
     ''' %(nt,dt))
Flow('exp0','exp',
     '''
     transp plane=13 memsize=%i|
     transp plane=23 memsize=%i|
     put n4=1 d4=25 label4=offset unit4=m
     '''%(memsize,memsize))
Plot('exp0',
     '''
     put label1=Time unit1=ms d1=.25
     label2=Inline d2=.025 unit2=km
     label3=Crossline d3=.025 unit3=km|
     byte|
     grey3 title="Modeled Data"
     frame1=%i frame2=%i frame3=%i flat=n
     point2=%g
     '''%(frame1t,frame2,frame3,volpt2))
Result('exp0',
     '''
     put label1=Time unit1=s d1=.00025
     label2="Inline" d2=.025 unit2=km
     label3="Crossline" d3=.025 unit3=km|
     byte|
     grey3 title="Modeled Data"
     frame1=%i frame2=%i frame3=%i flat=n
     point2=%g  
     '''%(frame1t,frame2,frame3,volpt2))

#Plane Wave Construction
#USED IN
#        tccs/strfilter/hongliu
#spray radius
ns2=2
ns3=2

#Inline and Crossline Slope
Flow('dip','exp0',
     '''
     dip rect1=%i rect2=%i rect3=%i order=%i verb=y
     '''%(rect1t,rect2,rect3,order1))
#split dip into inline and crossline
Flow('dipI','dip','window n4=1')
Flow('dipX','dip','window f4=1')
#Plot Dips
Flow('dip_bar','dip','bar')
Plot('dipI','dipI dip_bar',
     '''
     byte|
     grey3 scalebar=y bar=${SOURCES[1]}
     title="Inline Dip" color=j
     frame1=%i frame2=%i frame3=%i
     point2=%g
     '''%(frame1t,frame2,frame3,volpt2))
Plot('dipX','dipX dip_bar',
     '''
     byte|
     grey3 scalebar=y bar=${SOURCES[1]}
     title="Crossline Dip" color=j
     frame1=%i frame2=%i frame3=%i
     point2=%g
     '''%(frame1t,frame2,frame3,volpt2))
Plot('dips','dipI dipX','SideBySideIso')
Result('dips','dipI dipX','SideBySideIso')
#Spray for enhanced reflections
Flow('reflections','exp0 dip',
     '''
     pwspray2 ns2=%d ns3=%d dip=${SOURCES[1]} order=%i verb=y eps=0.1|
     stack axis=2
     '''% (ns2,ns3,order1))
#2nd Axis stack provides reflection Volume

#Subtract enhanced reflections from model data to get diffractions
Flow('diffractions','exp0 reflections',
     '''
     add scale=1,-1 ${SOURCES[1]} 
     ''')
Plot('reflections',
     '''
     put label1=Time unit1=ms d1=.25
     label2=Inline d2=.025 unit2=km
     label3=Crossline d3=.025 unit3=km|
     byte|
     grey3 title="Enhanced Reflections"
     frame1=%i frame2=%i frame3=%i flat=n
     point2=%g
     '''%(frame1t,frame2,frame3,volpt2))
Plot('diffractions',
     '''
     put label1=Time unit1=s d1=.00025
     label2=Inline d2=.025 unit2=km
     label3=Crossline d3=.025 unit3=km|
     byte|
     grey3 title="Remaining Diffractions"
     frame1=%i frame2=%i frame3=%i flat=n
     point2=%g
     '''%(frame1t,frame2,frame3,volpt2))
Result('diffractions',
     '''
     put label1=Time unit1=s d1=.00025
     label2="Inline" d2=.025 unit2=km
     label3="Crossline" d3=.025 unit3=km|
     byte|
     grey3 title="Separated Diffractions"
     frame1=%i frame2=%i frame3=%i flat=n
     point2=%g  
     '''%(frame1t,frame2,frame3,volpt2))
#smooth slowness migrations

mig = ['exp0','reflections','diffractions']
for item in mig:
     Flow(item+'-mig',[item,'left2_s','right2_s'],
          '''
          reverse which=1 |
          transp plane=12 memsize=%i|
          transp plane=23 memsize=%i|
          fftexp3 mig=y ompchunk=1
          left=${SOURCES[1]} right=${SOURCES[2]} 
          nz=%d dz=%g
          ''' %(memsize,memsize,nz,dz))
     Plot(item+'-mig',
          '''
          put label1=Time unit1=ms d1=.25
          label2=Inline d2=.025 unit2=km
          label3=Crossline d3=.025 unit3=km|
          byte|
          grey3 title="%s"
          frame1=%i frame2=%i frame3=%i flat=n
          point2=%g  
          '''%(item,frame1d,frame2*2/5,frame3*2/5,volpt2))
Result('exp0-mig',
          '''
          put 
          label2=Inline d2=.025 unit2=km
          label3=Crossline d3=.025 unit3=km|
          byte|
          grey3 title="Conventional Image"
          frame1=%i frame2=%i frame3=%i flat=n
          point2=%g  
          '''%(frame1d,frame2,frame3,volpt2))
Result('diffractions-mig',
          '''
          put 
          label2=Inline d2=.025 unit2=km
          label3=Crossline d3=.025 unit3=km|
          byte gainpanel=a|
          grey3 title="Diffraction Image"
          frame1=%i frame2=%i frame3=%i flat=n
          point2=%g  
          '''%(frame1d,frame2,frame3,volpt2))
##############################################
# create prediction volume of migrated image #
##############################################

# method via /parvaneh/pre-coherency/three-d

Flow('image_slope','exp0-mig',
     '''
     dip rect1=%i rect2=%i rect3=%i order=4
     '''%(rect1d,rect2,rect3))

Flow('image_slope_bar','image_slope','bar gainpannel=a')
Plot('image_slope_I','image_slope image_slope_bar',
     '''
     window n4=1|
     put label1=Time unit1=ms d1=.25
     label2=Inline d2=.05 unit2=km
     label3=Crossline d3=.05 unit3=km|
     byte gainpannel=a|
     grey3 title="Inline Image Slope" color=j 
     frame1=%i frame2=%i frame3=%i flat=n
     scalebar=y bar=${SOURCES[1]} point2=%g
     barlabel=Slope 
     '''%(frame1t,frame2/5,frame3/5,volpt2))
Plot('image_slope_X','image_slope image_slope_bar',
     '''
     window n4=1 f4=1|
     put label1=Time unit1=ms d1=.25
     label2=Inline d2=.05 unit2=km
     label3=Crossline d3=.05 unit3=km|
     byte gainpannel=a|
     grey3 title="Crossline Image Slope" color=j 
     frame1=%i frame2=%i frame3=%i flat=n
     scalebar=y bar=${SOURCES[1]} point2=%g
     barlabel=Slope
     '''%(frame1t,frame2/5,frame3/5,volpt2))

#Prediction Volume Parameters
ns2=2
ns3=2

#prediction volume
Flow('slope_spray','exp0-mig image_slope',
     'pwspray2 ns2=%d ns3=%d dip=${SOURCES[1]}'% (ns2,ns3))

# Prediction error
Flow('copy','exp0-mig','spray axis=2 n=25 d=25 o=0')
Flow('sweight','copy slope_spray',
     '''
     similarity other=${SOURCES[1]} rect1=10 rect4=1 | math output=1/input
     ''')

Flow('diff','exp0-mig slope_spray',
     '''
     spray axis=2 n=25 d=25 o=0|
     add scale=1,-1 ${SOURCES[1]} 
     ''')
Flow('diff2','diff','mul $SOURCE')

# coherence
for box in range(2,6):
    coh = 'cohP%d' % box
    Flow(coh,'diff2',
        '''
        transp plane=12 | spray axis=2 n=1 | put n1=%d n2=%d | 
        boxsmooth rect1=%d | boxsmooth rect2=%d | 
        window f1=%d f2=%d | 
        stack axis=2 min=y | stack axis=1 min=y 
        ''' % (2*ns2+1,2*ns3+1,box,box,box-1,box-1))

#Make figures for Feb 4 Meeting
frame1lst=[208,275,350]
frame2lst=[300,300,300]
frame3lst=[350,350,350]
dlst = [0.415,0.55,0.7]
regularize='put label1="Depth" d1=0.002 unit1="km" d2=0.025 label2="Inline" unit2="km" d3=0.025 label3="Crossline" unit3="km" |'

#Predictive Coherence Clip
clip=.01

for frame in xrange(0,len(frame1lst)):
     #Velocity model
     Plot('vp_%i'%frame,'vp vpbar',
          regularize+
          '''
          byte allpos=y min=3500 max=6500|
          grey3 title="Velocity Model" color=a barlabel=Velocity
          frame1=%i frame2=%i frame3=%i flat=n barunit="m/s"
          scalebar=y bar=${SOURCES[1]} point2=%g
          '''%(frame1lst[frame],frame2lst[frame],frame3lst[frame],volpt2))
     Plot('vpa_%i'%frame,'vp vpbar',
          regularize+
          '''
          byte allpos=y min=3500 max=6500|
          grey3 title="Velocity Model" color=a 
          frame1=%i frame2=%i frame3=%i flat=n
          scalebar=n bar=${SOURCES[1]} point2=%g
          '''%(frame1lst[frame],frame2lst[frame],frame3lst[frame],volpt2))
     #Smooth Slowness 
     Plot('vps_%i'%frame,'vp_s vpbar',
          regularize+
          '''
          byte allpos=y min=3500 max=6500|
          grey3 title="Smooth Slowness Velocity" color=a 
          frame1=%i frame2=%i frame3=%i flat=n
          scalebar=y bar=${SOURCES[1]} point2=%g
          '''%(frame1lst[frame],frame2lst[frame],frame3lst[frame],volpt2))
     #Conventional Image With Smooth Slowness
     Plot('image_%i'%frame,'exp0-mig',
          regularize+
          '''
          byte|
          grey3 title="Smooth Slowness RTM Image"
          frame1=%i frame2=%i frame3=%i flat=n
          point2=%g
          '''%(frame1lst[frame],frame2lst[frame],frame3lst[frame],volpt2))
     #Diffraction Image With Smooth Slowness
     Plot('diffimage_%i'%frame,'diffractions-mig',
          regularize+
          '''
          byte|
          grey3 title="Diffraction Image"
          frame1=%i frame2=%i frame3=%i flat=n
          point2=%g
          '''%(frame1lst[frame],frame2lst[frame],frame3lst[frame],volpt2))
     #Predictive Coherence (length 2)
     Plot('coherence_%i'%frame,'cohP2',
         regularize+
         '''
         byte clip=%g |
         grey3 frame1=%i frame2=%i frame3=%i point2=%g
         flat=n color=e title="Prediction length 2 Coherence"
         '''%(clip,frame1lst[frame],frame2lst[frame],frame3lst[frame],volpt2))#e!,h
     #Some 2D stuff
     Plot('coh2d_%i'%frame,'cohP2','window n1=1 f1=%i | grey color=e title="Prediction length 2 Coherence"'%(frame1lst[frame]))
     Plot('dif2d_%i'%frame,'diffractions-mig','window n1=1 f1=%i | grey title="Diffraction Image"'%(frame1lst[frame]))
     Plot('img2d_%i'%frame,'exp0-mig','window n1=1 f1=%i | grey title="Conventional Image from %g km"'%(frame1lst[frame],dlst[frame]))
     Plot('vp2d_%i'%frame,'vp','window n1=1 f1=%i | grey title="Velocity Model" color=j allpos=y bias=3000'%(frame1lst[frame]))
     #Some 2D Stack Attributes
     stk= 'window n1=10 f1=%i | stack axis=1 | transp|put unit2=km unit1=km  d2=.025 d1=.025 |'%(frame1lst[frame]-5)
     stka= 'window n1=10 f1=%i  | stack axis=1 | transp|put unit2=km unit1=km  d2=.025 d1=.025 | window min2=8.5 max2=11.5 min1=3.5 max1=8.5  | '%(frame1lst[frame]-5)
     stkb= 'window n1=10 f1=%i  | stack axis=1 | transp|put unit2=km unit1=km  d2=.025 d1=.025 | window min2=.5 max2=2.5 min1=11.75 max1=13.75  | '%(frame1lst[frame]-5)
     stkc= 'window n1=10 f1=%i  | stack axis=1 | transp|put unit2=km unit1=km  d2=.025 d1=.025 | window min2=.5 max2=2 min1=7 max1=8.5  | '%(frame1lst[frame]-5)
     Plot('coh2dstk_%i'%frame,'cohP2',stk+'grey color=e title="Prediction length 2 Coherence"')
#     Plot('dif2dstk_%i'%frame,'diffractions-mig',stk+' grey title="Diffraction Image from %g km"'%(frame1lst[frame],dlst[frame]))
 #    Plot('img2dstk_%i'%frame,'exp0-mig',stk+'grey title="Conventional Image from %g km"'%(frame1lst[frame],dlst[frame]))
     Result('dif2dstk-%i'%frame,'diffractions-mig',stk+' grey title="Diffraction Image from %g km"  '%(dlst[frame]))
     Result('img2dstk-%i'%frame,'exp0-mig',stk+'grey title="Conventional Image from %g km"  '%(dlst[frame]))
     Result('dif2dstka-%i'%frame,'diffractions-mig',stka+' grey title="Zoomed Diffraction Image from %g km"  '%(dlst[frame]))
     Result('img2dstka-%i'%frame,'exp0-mig',stka+'grey title="Zoomed Conventional Image from %g km"  '%(dlst[frame]))
     Result('dif2dstkb-%i'%frame,'diffractions-mig',stkb+' grey title="Zoomed Diffraction Image from %g km"  '%(dlst[frame]))
     Result('img2dstkb-%i'%frame,'exp0-mig',stkb+'grey title="Zoomed Conventional Image from %g km"  '%(dlst[frame]))
     Result('dif2dstkc_%i'%frame,'diffractions-mig',stkc+' grey title="Zoomed Diffraction Image from %g km"  '%(dlst[frame]))
     Result('img2dstkc_%i'%frame,'exp0-mig',stkc+'grey title="Zoomed Conventional Image from %g km"  '%(dlst[frame]))
     Plot('vp2dstk_%i'%frame,'vp',stk+'grey title="Velocity Model" color=j allpos=y bias=5000')


End()