from rsf.proj import*

padno=1024 #padding for seislet tranform
r1=20           #smoothing radius
r2=20           #smoothing radius
n1=1001         #number of temporal samples
n2=751          #trace numbers

## module defining
def Grey(data,other):
        Result(data,
        '''
        grey label2=Trace unit2=""  labelsz=11 label1=Time unit1="s" 
        title="" wherexlabel=b wheretitle=t %s'''%other)

##########################################
#    Make synthetic test
##########################################
Flow('hyper',None,
     '''
     spike n1=1001 |
     noise seed=2011 rep=y |
     math output="input^3" |
     cut n1=100 |
     ricker1 frequency=10 |
     spray axis=2 n=751 d=0.01 o=-3.5 label=Offset unit=km |
     nmostretch inv=y v0=4 half=n
     ''')
Flow('noise1',None,
     '''
     spike n1=1001 |
     noise seed=2011 rep=y |
     math output="input^3" |
     cut n1=100 |
     ricker1 frequency=10 |
     spray axis=2 n=751 d=0.01 o=-3.5 label=Offset unit=km |
     nmostretch inv=y v0=0.5 half=n|
     mutter hyper=y tp=20 v0=1
     ''')
Flow('noise2',None,
     '''
     spike n1=1001 |
     noise seed=2011 rep=y |
     math output="input^3" |
     cut n1=100 |
     ricker1 frequency=10 |
     spray axis=2 n=2000 d=0.01 o=-10 label=Offset unit=km |
     nmostretch inv=y v0=4 half=y | window f2=1249   
     ''')
Flow('test1','hyper noise1 noise2',
        '''
        add scale=1,1,1 ${SOURCES[1]} ${SOURCES[2]} | 
        scale aixs=2 | put o2=0 d2=1''')

Flow('dip','test1',
        '''
        dip rect1=%d rect2=%d | pad n2=%d'''%(r1,r2,padno))

Flow('zero',None,'spike n1=1 mag=0')
Flow('freq',None,'spike n1=1 mag=10')
Flow('slet00','test1 dip',
        '''
        pad n2=%d | seislet dip=${SOURCES[1]} 
        eps=0.1 adj=y inv=y unit=y type=b'''%padno)
# hilbert transform
Flow('slet00hilb','slet00','envelope hilb=y ' )


# transformed to 2-D seislet domain
slets=[]
for i in range(padno):
        slet='slet-'+str(i)
        Flow(slet,'slet00 slet00hilb freq',
        '''
        cmplx ${SOURCES[1]} | window n2=1 f2=%d | 
        freqlet freq=${SOURCES[2]} type=b'''%i)
        slets.append(slet)

Flow('slet1',slets,'cat axis=2 ${SOURCES[1:%d]}'%len(slets))
Flow('slet0','slet1','math output="abs(input)" | real')

# transformed to 2-D wavelet domain
Flow('wlet','test1',
        '''
        pad n2=%d | transp | dwt unit=y type=l | 
        transp | dwt unit=y type=l'''%padno)

# transformed to 2-D Fourier domain
Flow('ft','test1',
        '''
        pad n2=%d | rtoc|fft3 axis=1 padno=1 |
        fft3 axis=2 padno=1|cabs'''%padno)

# transform ed to 2-D curvelet domain
Flow('dct','test1',
        '''
        pad n2=%d | fdct ac=y adj=n nba=8 nbs=8''')

# Ploting
Grey('test1','')

Grey('slet0',
        '''
        wanttitle=n label1="Time Scale" unit1=s 
        label2="Distance Scale"  labelsz=11 unit2=""''')
Grey('wlet',
        '''
        wanttitle=n label1="Time Scale" unit1=s
        label2="Distance Scale"  labelsz=11 unit2=""''')
Grey('ft',
        '''
        wanttitle=n label2="Normalized Wavenumber"  labelsz=11 unit2= 
        label1=Frequency unit1=Hz color=j''')

# sorting the coefficients
Flow('sletcoef','slet0',
        '''
        put n1=1025024 o1=1 d1=1 n2=1 
        unit1= unit2= | sort ''')
Flow('ftcoef','ft',
        '''
        put n1=1025024 o1=1 d1=1 n2=1 
        unit1= unit2= | sort''')
Flow('wletcoef','wlet',
        '''put n1=1025024 o1=1 d1=1 n2=1 
        unit1= unit2= | sort''')
#Flow('dctcoef','dct','put d1=1 | sort | window n1=1025024')


# transformed domain coefficients decaying diagram
Plot('sigcoef','sletcoef wletcoef ftcoef',
'''
cat axis=2 ${SOURCES[1:3]} |
window n1=350000 | scale axis=1 | 
math output="20*log(input)/log(10)"|
graph dash=1,0,2 label1=n label2="a\_n\^" 
unit2="dB" wanttitle=n  labelsz=11''')

#Plot('sigcoef-dct','sletcoef wletcoef ftcoef dctcoef',
#'''
#cat axis=2 ${SOURCES[1:4]} |
#window n1=350000 | scale axis=1 | 
#math output="20*log(input)/log(10)"|
#graph dash=1,0,2 label1=n label2="a\_n\^" 
#unit2="dB" wanttitle=n''')

# Making frames
Plot('label0',None,
        '''
        box x0=10 y0=7.5 label="Fourier" xt=0.5 yt=0.5
        ''')
Plot('label1',None,
        '''
        box x0=7.5 y0=6 label="Wavelet" xt=0.5 yt=0.5 
        ''') # xt,yt relative position 0.5
Plot('label2',None,
        '''
        box x0=4.4 y0=4 label="Seislet" xt=0.5 yt=0.5
        ''')
Plot('label3',None,
        '''
        box x0=4.7 y0=7.3 label="Curvelet" xt=0.5 yt=0.5
        ''')
Result('sigcoef','sigcoef label0 label1 label2','Overlay')
#Result('sigcoef-dct','sigcoef-dct label0 label1 label2 label3',
#       '''Overlay''')

End()