from rsf.proj import *
import math
from rsf.recipes.lines import vlines,hlines
from rsf.recipes.itime import itimec,itimef
from rsf.recipes.fmean import Freqmean,Freqplots

PI=math.pi
def sz(n):
   return(' labelsz=%d titlesz=%d '%(n,n))

sz1 = sz(12)
sz2 = sz(10)



##### Fourier ######

Flow('spik',None,'spike d1=0.002 n1=401 k1=201 mag=1 | put label2="u(t)"')
Flow('rick','spik','ricker1 frequency=40')
Flow('chrp','spik',
     '''math output="sin(2*%g*(50*x1+(5-50)/2/0.8*x1^2))" |
        cut max1=0.3 | cut min1=0.5 ''' %PI)
Flow('noiz','spik','math output=0 | noise seed=2012 var=1')
itimef('taus','spik',div='rect1=1')
itimef('taur','rick',div='rect1=1')
itimef('tauc','chrp',div='rect1=1')
itimef('taun','noiz',div='rect1=1')

taxis=  'min1=0 max1=0.8 wanttitle=n crowd2=0.6 '+sz(16)
faxis='''min1=0 max1=125 wanttitle=n crowd2=0.6 
         min2=-0.2 max2=1. o2num=0 d2num=0.4 n2tic=3 '''+sz(16)
plot='plotfat=5 labelfat=5 axisfat=5'
Plot('rick','graph %s' %(taxis+plot));
Plot('chrp','graph %s' %(taxis+plot));
Plot('noiz','graph %s' %(taxis+plot));

Plot('taur','graph %s' %(faxis+plot));
Plot('tauc','graph %s' %(faxis+plot));
Plot('taun','graph %s' %(faxis+plot));

Result('rick','rick taur','SideBySideAniso')
Result('noiz','noiz taun','SideBySideAniso')
Result('chrp','chrp tauc','SideBySideAniso')



##### Synthetic with 3 spikes #####
fhi=125; niter=100;
tfpar=' grey color=j scalebar=y'

Flow('syn',None,
     '''
     spike d1=0.004 n1=512 nsp=3 k1=100,300,350 mag=1,-0.5,0.5 |
     ricker1 frequency=30 |
     noise seed=2011 var=1.e-6 | put label1="Time" unit1="s" label2="u(t)"
     ''')
trace = 'syn';
Plot(trace,trace,'graph wanttitle=n crowd2=0.35 plotfat=5 max1=2 '+sz2)
Result(trace,trace,'graph wanttitle=n crowd2=0.5 plotfat=5 max1=2 '+sz2)
Result('syn0',trace,'graph wanttitle=n crowd2=0.5 plotfat=5 max1=2 '+sz2)

Flow('time11',  trace,  'math output="x1"')
#Plot('time12', 'time11','graph wanttitle=n')


#### LTFT synthetic

rc=7
axes  = 'min1=0. max1=2. min2=0. max2=100'
ticks = '' #'o1num=0 d1num=30 n1tic=5 o2num=.4 d2num=0.4 n2tic=5'
grid  = '' #'gridcol=5 grid1=y g1num0=10 g1num=50' 

# inst. traveltime (t,w)
itimec('tt0', trace, cltf='rect=%d'%rc,
       div='rect1=%d rect2=%d niter=%d verb=n'%(5,rc,niter))  # generates: syn_cltf.rsf, syn_dcltf.rsf, tt0.rsf
Plot('ltf0','syn_cltf',
       '''math output="abs(input)" | real | 
          %s allpos=y title="u(t,f)" scalebar=y wherexlabel=bottom
          wheretitle=top barlabel="Amplitude" 
          %s %s %s %s''' %(tfpar,axes,ticks,grid,sz2))
Plot('tt0',
     '''%s allpos=y title="tau(t,f)" scalebar=y 
        wherexlabel=bottom wheretitle=top barlabel="Traveltime"
        %s %s %s %s''' %(tfpar,axes,ticks,grid,sz2))
Plot('zc0', 'tt0',
     '''math output="input-x1" | zerocross levels=3 |  
        contour plotfat=10 allpos=n nc=1 c0=-1 scalebar=y wanttitle=n wantaxis=n %s %s %s 
     ''' %(axes,ticks,grid))

# mean and std of frequency
Freqmean('syn_cltf',r=2,std='y')  # generates syn_cltf_s0,  syn_cltf_s1, syn_cltf_s2,
                                  #           syn_cltf_fmn, syn_cltf_fstd
Flow('syn_flo','syn_cltf_fmn syn_cltf_fstd','add ${SOURCES[1]} scale=1,-1')
Flow('syn_fhi','syn_cltf_fmn syn_cltf_fstd','add ${SOURCES[1]} ')
Freqplots('syn_flo', 'syn_flo',axes=axes, scalebar='y')
Freqplots('syn_fhi', 'syn_fhi',axes=axes, scalebar='y')
Plot('syn_freq','syn_flo syn_fhi','Overlay')
#              plotfat=10,plotcol=1, dash=1, transp='y', yreverse='y')

# Figures
Result('tt0', 'tt0 zc0 syn_freq', 'Overlay')
Result('ltf0','ltf0 syn_freq','Overlay')


#### inst. traveltime (t)
Flow('tau0', 'tt0 syn_flo syn_fhi',
     'stackn min=${SOURCES[1]} max=${SOURCES[2]} | put label2="tau(t)" unit2=s')
#     'window min2=10 max2=60 | stack | put label2="tau(t)" unit2=s')

axes  = 'min1=0. max1=2. min2=0. max2=2.'
ticks = 'o1num=0. d1num=0.4 n1tic=6 o2num=0.4 d2num=0.8 n2tic=3'
#grid1  = 'gridcol=5  grid1=y g1num0=0.4 g1num=0.8'
#grid2  = 'gridcol=5  grid2=y g2num0=0.4 g2num=0.8'


Plot('time12', 'time11',
     '''graph wanttitle=n wantaxis=n 
        dash=1 %s %s crowd2=0.5''' %(sz1,axes))
Plot('tau0',
       '''graph wanttitle=n 
          %s %s %s crowd2=0.5 plotfat=10''' %(sz1,axes,ticks))
Plot('tau11', 'tau0 time12', 'Overlay')
Flow('timesyn',None,'spike n1=3 o1=0 d1=1 nsp=3 k1=1,2,3 mag=0.4,1.2,1.4')


vlines('vtimesyn', 'timesyn', max=2, axes=axes, plot='plotcol=5 crowd2=0.5')
hlines('htimesyn', 'timesyn', max=2, axes=axes, plot='plotcol=5 crowd2=0.5')

Result('taut0', 'tau0 time12 vtimesyn htimesyn', 'Overlay') # plot of tau and t vs. t

Flow('tau_t0', 'tau0',
     ''' math output="input-x1" | put label2="tau(t)-t"''')
axes = 'min1=0. max1=2. min2=-1 max2=1'
#ticks= 'o1num=0. d1num=0.4 n1tic=6 o2num=-0.4 d2num=0.4 n2tic=3' 
grid1 = 'gridcol=5  grid1=y g1num0=0.4 g1num=0.8'
grid2 = 'gridcol=5  grid2=y g2num0=0.  g2num=1'

Plot('tau_t0',
     '''graph %s crowd2=0.5  
        wanttitle=n plotfat=10 %s %s %s''' %(sz1,axes,grid1,grid2))
Result('tau-t0','tau_t0 vtimesyn','Overlay')




End()