from rsf.proj import*
from rsf.prog import RSFROOT
from math import *
import math

## Plot font and screen ratio, width, height for uniform ppt figs.
p1=0.7
sr=1.0
sw=7.0
sh=10.0
xll=2.0
fat=2

#####
##Synthetic CMP Parameters
nx=128
delx=0.0125*2
ox=-1.6
ny=nx
dely=delx
oy=ox
nt=1001
dt=0.004
fw=10.0
####

################################################################################################################
########  4 Reflection Events ################################################

#111111111 Make a layer: z = z0 + x*dx + y*dy

dx=0.4
dy=0.1
dt=.004
z0=0.8
##z0=0.75

cg = 1/math.sqrt(1+dx*dx+dy*dy)
ca = -dx*cg
cb = -dy*cg
d = z0*cg

mx = 0
my = 0

D = d - mx*ca - my*cb

v0 = 2.5
t0 = 2*D/v0
it0=int(t0/dt)
it1=it0
#print it1
Flow('spike1',None,
     '''
     spike n1=%d k1=%d nsp=1 | ricker1 frequency=%g |
     spray axis=2 n=%d o=%g d=%g |
     spray axis=3 n=%d o=%g d=%g
     ''' % (nt,it1,fw,ny,oy,dely,nx,ox,delx))

wx = 1-ca*ca
wy = 1-cb*cb
wxy = - ca*cb

Flow('vel1.asc',None,'echo %g %g %g n1=3 data_format=ascii_float in=$TARGET' % (wx,wy,wxy))
Flow('vel1','vel1.asc','dd form=native | scale dscale=%g | spray axis=1 n=%d' % (1/(v0*v0),nt),local=1)


Flow('cmp1','spike1 vel1','inmo3 velocity=${SOURCES[1]}',local=1)

#222222222222 Make another layer: z = z0 + x*dx + y*dy

dx=0.7
dy=0.41
dt=.004
z0=0.85

cg = 1/math.sqrt(1+dx*dx+dy*dy)
ca = -dx*cg
cb = -dy*cg
d = z0*cg

mx = 0
my = 0

D = d - mx*ca - my*cb

v0 = 1.7
t0 = 2*D/v0
it0=int(t0/dt)
it2=it0
#print it2
Flow('spike2',None,
     '''
     spike n1=%d k1=%d nsp=1 | ricker1 frequency=%g |
     spray axis=2 n=%d o=%g d=%g |
     spray axis=3 n=%d o=%g d=%g
     ''' % (nt,it2,fw,ny,oy,dely,nx,ox,delx))

wx = 1-ca*ca
wy = 1-cb*cb
wxy = - ca*cb

Flow('vel2.asc',None,'echo %g %g %g n1=3 data_format=ascii_float in=$TARGET' % (wx,wy,wxy))
Flow('vel2','vel2.asc','dd form=native | scale dscale=%g | spray axis=1 n=%d' % (1/(v0*v0),nt),local=1)


Flow('cmp2','spike2 vel2','inmo3 velocity=${SOURCES[1]}',local=1)


#3333333333333 Make yet another layer: z = z0 + x*dx + y*dy

dx=0.1
dy=0.9
dt=.004
z0=1.1

cg = 1/math.sqrt(1+dx*dx+dy*dy)
ca = -dx*cg
cb = -dy*cg
d = z0*cg

mx = 0
my = 0

D = d - mx*ca - my*cb

v0 = 1.75
t0 = 2*D/v0
it0=int(t0/dt)
it3=it0
#print it3
Flow('spike3',None,
     '''
     spike n1=%d k1=%d nsp=1 | ricker1 frequency=%g |
     spray axis=2 n=%d o=%g d=%g |
     spray axis=3 n=%d o=%g d=%g
     ''' % (nt,it3,fw,ny,oy,dely,nx,ox,delx))

wx = 1-ca*ca
wy = 1-cb*cb
wxy = - ca*cb

Flow('vel3.asc',None,'echo %g %g %g n1=3 data_format=ascii_float in=$TARGET' % (wx,wy,wxy))
Flow('vel3','vel3.asc','dd form=native | scale dscale=%g | spray axis=1 n=%d' % (1/(v0*v0),nt),local=1)

Flow('cmp3','spike3 vel3','inmo3 velocity=${SOURCES[1]}',local=1)

#444444444444 Make a 4th layer: z = z0 + x*dx + y*dy

dx=0.2
dy=0.1
dt=.004
z0=1.3

cg = 1/math.sqrt(1+dx*dx+dy*dy)
ca = -dx*cg
cb = -dy*cg
d = z0*cg

mx = 0
my = 0

D = d - mx*ca - my*cb

v0 = 2.0
t0 = 2*D/v0
it0=int(t0/dt)
it4=it0
#print it4
Flow('spike4',None,
     '''
     spike n1=%d k1=%d nsp=1 | ricker1 frequency=%g |
     spray axis=2 n=%d o=%g d=%g |
     spray axis=3 n=%d o=%g d=%g
     ''' % (nt,it4,fw,ny,oy,dely,nx,ox,delx))

wx = 1-ca*ca
wy = 1-cb*cb
wxy = - ca*cb

Flow('vel4.asc',None,'echo %g %g %g n1=3 data_format=ascii_float in=$TARGET' % (wx,wy,wxy))
Flow('vel4','vel4.asc','dd form=native | scale dscale=%g | spray axis=1 n=%d' % (1/(v0*v0),nt),local=1)

Flow('cmp4','spike4 vel4','inmo3 velocity=${SOURCES[1]}',local=1)

## Add events to create CMP.
Flow('cmp','cmp1 cmp2 cmp3 cmp4',
     '''
     math cmp2=${SOURCES[1]} cmp3=${SOURCES[2]} cmp4=${SOURCES[3]} output="input+cmp2+cmp3+cmp4" |
     window max1=1.5 min1=0.5 j1=2 | put o1=0 | put label2=X label3=Y unit2=km unit3=km
     ''',local=1)


def Grey(data,other):
      Result(data,
      '''
      put d1=0.004 d2=1 o1=0 o2=0 | 
      grey clip=0.35 transp=y yreverse=y  
      label2=Trace  unit2="" label1=Time unit1="s" title=""
      wherexlabel=t scalebar=n wheretitle=b screenratio=1.2  %s
      '''
      %other)

def Greyfk(data,data0,other):
        Result(data,data0,
                '''put d1=0.004 d2=1 d3=1| fft1 | fft3 axis=2 pad=1| fft3 axis=3 pad=1| window max1=60| cabs |byte clip=300 allpos=y | grey3 labelfat=4 font=2 titlefat=4 flat=n label2="Inline wavenumber" unit2= label3="Xline wavenumber" unit1=Hz label1=Frequency unit2= frame1=10 frame2=16 frame3=16 point1=0.8 point2=0.8 title= screenratio=1.2 unit2= unit3= %s color=j'''%other)


def Grey3(data,other):
      Result(data,
       '''
       put d1=0.004 o3=0 d2=50 d3=25 | byte  clip=0.2 |
       grey3 wanttitle=y flat=n labelfat=4 font=2 titlefat=4
       label1="Time" unit1=s label2=Inline label3=Xline unit3=
       frame1=28 frame2=16 frame3=10 point1=0.85 point2=0.75 screenratio=1.2
       clip= title= wheretitle=t 
       scalebar=n label1=Time unit1=s title= unit3=m unit2=m %s
       '''%other)

def Grey3n(data,other):
      Result(data,
       '''
       put d1=0.004 o3=0 | byte  clip=0.35 |
       grey3 wanttitle=y flat=n labelfat=4 font=2 titlefat=4
       label1="Time" unit1=s label2=Inline label3=Xline unit3=
       frame1=28 frame2=16 frame3=10 point1=0.85 point2=0.75 screenratio=1.2
       clip= title= wheretitle=t 
       scalebar=n label1=Time unit1=s title= unit3=m unit2=m %s
       '''%other)

def Grey3nn(data,other):
      Result(data,
       '''
       put d1=0.004 o3=0 | byte  clip=0.35 |
       grey3 wanttitle=y flat=n labelfat=4 font=2 titlefat=4
       label1="Time" unit1=s label2=Inline label3=Xline unit3=
       frame1=28 frame2=16 frame3=10 point1=0.85 point2=0.75 screenratio=1.2
       clip= title= wheretitle=t 
       scalebar=n %s
       '''%other)

def Greys(data,other):
        Result(data,
                '''put d1=0.004 o3=0 d2=50 d3=25 | byte clip=0.3 bar=bar.rsf| grey3 labelfat=4 font=2 titlefat=4 label2=Inline unit2=km label3=Xline unit2=km flat=n frame1=28 frame2=15 frame3=10 point1=0.85 point2=0.75 screenratio=1.2 title= label1=Time unit1=s title= unit3=m unit2=m %s'''%(other))

def Greys2d(data,other):
        Result(data,
                '''grey clip=0.3 label2=Inline unit2=km label3=Xline unit2=km flat=n screenratio=1.2 title=                             wanttitle=y flat=n labelfat=4 font=2 titlefat=4                         
                                label2=Trace unit2="" label1="Time"  unit1=s
                                title="" wherexlabel=b wheretitle=t poly=y 
                                wheretitle=t wherexlabel=b labelsz=10 %s'''%(other))

def Graph(data,other):
        Result(data,'graph label1="Iter #no" label2="SNR" unit2=dB unit1="" title="" wherexlabel=b wheretitle=t %s' %other)

def Greyz(data,other):
        Result(data,
       '''
       grey label2=Trace label1="Time" unit1= clip=0.35 title= screenratio=0.6
       color=g %s'''%other)

"""
def Wig(data,other): 
        Result(data,'''
                                put d1=0.004 |
                                wiggle transp=y yreverse=y screenratio=1.2
                                wanttitle=y flat=n labelfat=4 font=2 titlefat=4                         
                                label2=Trace unit2="" label1="Time"  unit1=s
                                title="" wherexlabel=b wheretitle=t poly=y 
                                wheretitle=t wherexlabel=b clip=0.05 labelsz=10 %s'''%other)
"""
def Wig(data,other):
        Result(data,'''
                                put d1=0.004 |
                                grey color=g color=b transp=y yreverse=y screenratio=1.2
                                wanttitle=y flat=n labelfat=4 font=2 titlefat=4                         
                                label2=Trace unit2="" label1="Time"  unit1=s
                                title="" wherexlabel=b wheretitle=t poly=y 
                                wheretitle=t wherexlabel=b clip=0.2 labelsz=10 %s'''%other)

########################################################################
# INITIALIZATION
########################################################################
matlab         = WhereIs('matlab')
matROOT = '../matfun'
matfun = 'Hyper_N20'
matlabpath = os.environ.get('MATLABPATH',os.path.join(RSFROOT,'lib'))

if not matlab:
   sys.stderr.write('\nCannot find Matlab.\n')
   sys.exit(1)

############################################################
## generate and process synthetic data
############################################################
Flow('syn-c syn-n syn-fk syn-lr syn-dlr syn-olr',[os.path.join(matROOT,matfun+'.m'),'cmp'],
    '''MATLABPATH=%(matlabpath)s %(matlab)s 
    -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun)s('${SOURCES[1]}','${TARGETS[0]}','${TARGETS[1]}','${TARGETS[2]}','${TARGETS[3]}','${TARGETS[4]}','${TARGETS[5]}');quit"
    '''%vars(),stdin=0,stdout=-1)

Flow('hyp3d-20-c','syn-c','put n1=126 n2=32 n3=32')
Flow('hyp3d-20-n','syn-n','put n1=126 n2=32 n3=32')
Flow('hyp3d-20-fk','syn-fk','put n1=126 n2=32 n3=32')
Flow('hyp3d-20-lr','syn-lr','put n1=126 n2=32 n3=32')
Flow('hyp3d-20-dlr','syn-dlr','put n1=126 n2=32 n3=32')
Flow('hyp3d-20-olr','syn-olr','put n1=126 n2=32 n3=32')

Flow('hyp3d-20-n-fk','hyp3d-20-n hyp3d-20-fk',' add scale=1,-1 ${SOURCES[1]}')
Flow('hyp3d-20-n-lr','hyp3d-20-n hyp3d-20-lr',' add scale=1,-1 ${SOURCES[1]}')
Flow('hyp3d-20-n-dlr','hyp3d-20-n hyp3d-20-dlr',' add scale=1,-1 ${SOURCES[1]}')
Flow('hyp3d-20-n-olr','hyp3d-20-n hyp3d-20-olr',' add scale=1,-1 ${SOURCES[1]}')

Flow('hyp3d-20-simi1','hyp3d-20-n-fk hyp3d-20-fk','similarity other=${SOURCES[1]} niter=20 rect1=3 rect2=3 rect3=3')
Flow('hyp3d-20-simi2','hyp3d-20-n-lr hyp3d-20-lr','similarity other=${SOURCES[1]} niter=20 rect1=3 rect2=3 rect3=3')
Flow('hyp3d-20-simi3','hyp3d-20-n-dlr hyp3d-20-dlr','similarity other=${SOURCES[1]} niter=20 rect1=3 rect2=3 rect3=3 ')
Flow('hyp3d-20-simi4','hyp3d-20-n-olr hyp3d-20-olr','similarity other=${SOURCES[1]} niter=20 rect1=3 rect2=3 rect3=3 ')
Greys('hyp3d-20-simi1','color=j scalebar=y clip=0.8 minval=0 maxval=0.3 title="FK" barlabel="Similarity" ')
Greys('hyp3d-20-simi2','color=j scalebar=y clip=0.8 minval=0 maxval=0.3 title="RR" barlabel="Similarity"')
Greys('hyp3d-20-simi3','color=j scalebar=y clip=0.8 minval=0 maxval=0.3 title="DRR" barlabel="Similarity"')
Greys('hyp3d-20-simi4','color=j scalebar=y clip=0.8 minval=0 maxval=0.3 title="ODRR" barlabel="Similarity"')


Grey3('hyp3d-20-c','color=g color=b title=Clean')
Grey3('hyp3d-20-n','color=g color=b title=Noisy')
Grey3('hyp3d-20-lr','color=g color=b title=RR')
Grey3('hyp3d-20-dlr','color=g color=b title=DRR')
Grey3('hyp3d-20-olr','color=g color=b title=ODRR')
Grey3('hyp3d-20-fk','color=g color=b title=FK')

Greyfk('hyp3d-20-c-fkk','hyp3d-20-c','title=Clean')
Greyfk('hyp3d-20-n-fkk','hyp3d-20-n','title=Noisy')
Greyfk('hyp3d-20-lr-fkk','hyp3d-20-lr','title=RR')
Greyfk('hyp3d-20-dlr-fkk','hyp3d-20-dlr','title=DRR')
Greyfk('hyp3d-20-olr-fkk','hyp3d-20-olr','title=ODRR')
Greyfk('hyp3d-20-fk-fkk','hyp3d-20-fk','title=FK')

Grey3('hyp3d-20-n-lr','color=g color=b title=RR')
Grey3('hyp3d-20-n-dlr','color=g color=b title=DRR')
Grey3('hyp3d-20-n-olr','color=g color=b title=ODRR')
Grey3('hyp3d-20-n-fk','color=g color=b title=FK')

Flow('hyp3d-20-c-s','hyp3d-20-c','window n3=1 f3=5')
Flow('hyp3d-20-n-s','hyp3d-20-n','window n3=1 f3=5')
Flow('hyp3d-20-lr-s','hyp3d-20-lr','window n3=1 f3=5')
Flow('hyp3d-20-dlr-s','hyp3d-20-dlr','window n3=1 f3=5')
Flow('hyp3d-20-olr-s','hyp3d-20-olr','window n3=1 f3=5')
Flow('hyp3d-20-fk-s','hyp3d-20-fk','window n3=1 f3=5')

Flow('hyp3d-20-n-lr-s','hyp3d-20-n-lr','window n3=1 f3=5')
Flow('hyp3d-20-n-dlr-s','hyp3d-20-n-dlr','window n3=1 f3=5')
Flow('hyp3d-20-n-olr-s','hyp3d-20-n-olr','window n3=1 f3=5')
Flow('hyp3d-20-n-fk-s','hyp3d-20-n-fk','window n3=1 f3=5')

Wig('hyp3d-20-c-s','title=Clean')
Wig('hyp3d-20-n-s','title=Noisy')
Wig('hyp3d-20-lr-s','title=RR')
Wig('hyp3d-20-dlr-s','title=DRR')
Wig('hyp3d-20-olr-s','title=ODRR')
Wig('hyp3d-20-fk-s','title=FK')
Wig('hyp3d-20-n-lr-s','title=RR')
Wig('hyp3d-20-n-dlr-s','title=DRR')
Wig('hyp3d-20-n-olr-s','title=ODRR')
Wig('hyp3d-20-n-fk-s','title=FK')


Flow('hyp3d-20-simi1-s','hyp3d-20-simi1','window n3=1 f3=5')
Flow('hyp3d-20-simi2-s','hyp3d-20-simi2','window n3=1 f3=5')
Flow('hyp3d-20-simi3-s','hyp3d-20-simi3','window n3=1 f3=5')
Flow('hyp3d-20-simi4-s','hyp3d-20-simi4','window n3=1 f3=5')


Greys2d('hyp3d-20-simi1-s','color=j scalebar=y minval=0 maxval=0.3 title="FK" barlabel="Similarity"')
Greys2d('hyp3d-20-simi2-s','color=j scalebar=y minval=0 maxval=0.3 title="RR" barlabel="Similarity"')
Greys2d('hyp3d-20-simi3-s','color=j scalebar=y minval=0 maxval=0.3 title="DRR" barlabel="Similarity"')
Greys2d('hyp3d-20-simi4-s','color=j scalebar=y minval=0 maxval=0.3 title="ODRR"  barlabel="Similarity"')



End()