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

def Wig(data,other):
        Result(data,'put d1=0.004 d2=1 o1=0 o2=1 |window j2=2 | wiggle labelsz=8 titlesz=8 titlefat=2 labelfat=2 font=2 color=g grid=n poly=y transp=y yreverse=y clip=0.55  label2=linear  unit2="" label1=Time unit1="s" title="" wherexlabel=t scalebar=n wheretitle=t   %s'%other)

def Wigzoom(data,other):
        Result(data,'put d1=0.004 d2=1 o1=1.5 o2=0 | window j2=2 | wiggle grid=n poly=y transp=y yreverse=y clip=0.55 label2=linear unit2="" label1=Time unit1="s" title="" wherexlabel=t scalebar=n  wheretitle=b  %s'%other)

def Grey(data,other):
        Result(data,'grey grid=n labelsz=8 titlesz=8 titlefat=2 labelfat=2 font=2 label2=linear  unit2="" label1=Time unit1="s" title="" wherexlabel=b scalebar=n wheretitle=t   %s'%other)

def Graph(data,other):
        Result(data,'graph grid=n label2=Amplitude  unit2="" labelsz=8 titlesz=8 titlefat=2 labelfat=2 font=2 label1=Time unit1="s" title="" wherexlabel=b scalebar=n wheretitle=t %s'%other)

def Grey3(data,other):
        Result(data,'transp plane=12|byte allpos=y | grey3 flat=n color=j frame2=60 frame1=120 frame3=15 grid=n label1=Time  unit1="s" label3=Trace label2=Frequency unit2="Hz" title="" wherexlabel=b scalebar=n wheretitle=t labelsz=8 titlesz=8 titlefat=2 labelfat=2 font=2 point1=0.8 point1=0.85 %s'%other)

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

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

pi=math.pi
t1=0.5
t2=0.8
dt=0.004
k1=t1/dt
k2=t2/dt
nt=252
f0=50
Q1=41
Q2=80
nq=40
fhi=1/dt/2
flo=0

############################################################
## generate and process synthetic data
############################################################
Flow('data-t',[os.path.join(matROOT,matfun+'.m')],
     '''MATLABPATH=%(matlabpath)s %(matlab)s 
     -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun)s(%(t1)g,%(t2)g,%(dt)g,%(nt)g,%(f0)g,%(Q1)g,%(Q2)g,%(nq)d,'${TARGETS[0]}');quit"
     '''%vars(),stdin=0,stdout=-1)
Flow('data','data-t','put d1=%g d2=1 o1=0 o2=0 n2=%d'%(dt,nq))
#Graph('linear','max2=1.25 min2=-0.75')

Flow('linear-st','data','window n2=1 f2=19|noise seed=201516 var=0.02 | st fhi=%g flo=%g | cabs | transp'%(fhi,flo))
Grey('linear-st','label1=Frequency unit1=Hz label2=Time unit2=s color=j clip=0.4')

Flow('linear-f-t1','linear-st','window n2=1 f2=%g '%(k1))
Flow('linear-f-t2','linear-st','window n2=1 f2=%g '%(k2))
Flow('linear-f-t','linear-f-t1 linear-f-t2','cat axis=2 ${SOURCES[1]}')
Graph('linear-f-t','label1=Frequency unit1=Hz label2=Amplitude ')

Flow('linear-ratio','linear-f-t2 linear-f-t1','divn rect1=3 den=${SOURCES[1]} eps=0.01 | math output="log(input)" |  window min1=20 max1=80')
Graph('linear-ratio','label1=Frequency unit1=Hz label2=Regularized division ratio')

Flow('freq',None,'spike d1=0.992063 n1=127 o1=0 n2=1 d2=1 o2=0 | math output="x1" | window min1=20 max1=80')
Flow('lsfit coef','linear-ratio freq','lsfit coef=${TARGETS[1]} fit=${SOURCES[1]}')
Result('linear-lsfit','linear-ratio lsfit',
       '''
       cat axis=2 ${SOURCES[1]} | 
       graph dash=0,1 title="Least-squares fitted line" 
       label1="Frequency" unit1=Hz label2="Division ratio" unit2=
       ''')
Flow('linear-Q-esti','coef','window n1=1 | math output="-%g*%g/input"'%(pi,t2-t1))



i=0
sts=[]
linears=[]
N=40
for var in range(N):
        s=0.01
        i=i+1
        Flow('linear-%d'%i,'data','window n2=1 f2=%d | noise seed=%d var=%g'%(i-1,i+201516,s))
        Flow('linear-st-%d'%i,'linear-%d'%i,'st fhi=%g flo=%g | cabs | transp'%(fhi,flo))

        linears.append('linear-%d'%i)
        sts.append('linear-st-%d'%i)
        Grey('linear-st-%d'%i,'label1=Frequency unit1=Hz label2=Time unit2=s color=j clip=0.4')
Flow('linears',linears,'cat axis=2 ${SOURCES[1:%d]} | put o2=1 d2=1'%len(linears))

Wig('linears','screenratio=1.8 wherexlabel=b title="Synthetic Data" label2="Trace"')
Graph('linear-15','title="Single linear" transp=y max1=1 min1=0 max2=1.1 min2=-1.1 screenratio=1.8 yreverse=y ')

#single channel
Q1s=[]
i=0
for var in range(N):
        s=0.01
        i=i+1
#       Flow('linear-%d'%i,'data','window n2=1 f2=%d | noise seed=%d var=%g'%(i-1,i+201516,s))
#       Flow('linear-st-%d'%i,'linear-%d'%i,'st fhi=%g flo=%g | cabs | transp'%(fhi,flo))

#       linears.append('linear-%d'%i)
#       sts.append('linear-st-%d'%i)

        Flow('linear-f-t1-%d'%i,'linear-st-%d'%i,'window n2=1 f2=%g '%(k1))
        Flow('linear-f-t2-%d'%i,'linear-st-%d'%i,'window n2=1 f2=%g '%(k2))
        Flow('linear-f-t-%d'%i,['linear-f-t1-%d'%i,'linear-f-t2-%d'%i],'cat axis=2 ${SOURCES[1]}')
#       Graph('linear-f-t','label1=Frequency unit1=Hz label2=Amplitude ')

        Flow('ratio-%d'%i,['linear-f-t2-%d'%i,'linear-f-t1-%d'%i],'divn rect1=3 den=${SOURCES[1]} eps=0.01 | math output="log(input)" |  window min1=20 max1=80')
#Graph('linear-ratio','label1=Frequency unit1=Hz label2=Regularized division ratio')

        Flow(['lsfit-%d'%i,'coef-%d'%i],['ratio-%d'%i,'freq'],'lsfit coef=${TARGETS[1]} fit=${SOURCES[1]}')
        Flow('linear-Q-esti-%d'%i,'coef-%d'%i,'window n1=1 | math output="-%g*%g/input"'%(pi,t2-t1))
        Q1s.append('linear-Q-esti-%d'%i)
Flow('linear-Q1s-esti',Q1s,'cat axis=2 ${SOURCES[1:%d]} |window'%len(Q1s))
Graph('linear-Q1s-esti','min2=20 max2=230 label2="Quality factor" label1=Trace unit1= plotfat=10')


#multi channel
Flow('linear-st-N',sts,'cat axis=3 ${SOURCES[1:%d]}'%len(sts))

Flow('linear-f-t1-N','linear-st-N','window n2=1 f2=%g '%(k1))
Flow('linear-f-t2-N','linear-st-N','window n2=1 f2=%g '%(k2))
Flow('linear-ratio-N','linear-f-t2-N linear-f-t1-N','divn rect1=3 rect2=20 den=${SOURCES[1]} eps=0.1 | math output="log(input)" |  window min1=20 max1=80')

Flow('linear-ratio-1','linear-ratio-N','window n2=1')
Graph('linear-ratio-1','label1=Frequency unit1=Hz label2=Regularized division ratio')

Flow('freq-N',None,'spike d1=0.992063 n1=127 o1=0 n2=1 d2=1 o2=0 | math output="x1" | window min1=20 max1=80|spray axis=2 n=%d o=1 d=1'%N)

Grey3('linear-st-N','title="Time-frequency Cube"')

Grey('linear-f-t2-N','label1=Frequency label2=Trace unit1=Hz title="Numerator" color=j allpos=y')
Grey('linear-f-t1-N','label1=Frequency label2=Trace  unit1=Hz title="Denominator" color=j allpos=y')

lsfits=[]
coefs=[]

i=0
for var in range(N):
        i=i+1
        Flow('linear-ratio-%d'%i,'linear-ratio-N','window n2=1 f2=%d'%(i-1))
        Flow('freq-%d'%i,'freq-N','window n2=1 f2=%d'%(i-1))
        Flow(['lsfit-N-%d'%i,'coef-N-%d'%i],['linear-ratio-%d'%i,'freq-%d'%i],'lsfit coef=${TARGETS[1]} fit=${SOURCES[1]}')
        lsfits.append('lsfit-N-%d'%i)
        coefs.append('coef-N-%d'%i)

Flow('linear-lsfits',lsfits,'cat axis=2 ${SOURCES[1:%d]}'%len(lsfits))

Flow('linear-coef-N',coefs,'cat axis=2 ${SOURCES[1:%d]} |window'%len(coefs))

#Flow('lsfit-N coef-N','linear-ratio-N freq-N','lsfit coef=${TARGETS[1]} fit=${SOURCES[1]}')
#Flow('linear-Q-esti-N','coef-N','window n1=1 | math output="-%g*%g/input"'%(pi,t2-t1))
Flow('linear-Q-esti-N','linear-coef-N','math output="-%g*%g/input"'%(pi,t2-t1))

Grey('linear-ratio-N','label1=Frequency unit1=Hz label2=Trace color=j scalebar=y bias=-0.6 clip=0.6')
Grey('linear-lsfits','label1=Frequency unit1=Hz label2=Trace color=j scalebar=y bias=-0.6 clip=0.6')


Graph('linear-coef-N','min2=-0.02 max2=-0.01 label2="Value" label1=Trace unit1= plotfat=10')
Graph('linear-Q-esti-N','min2=20 max2=230 label2="Quality factor" label1=Trace unit1= plotfat=10')

#sfdisfil <linear-Q-esti-N.rsf
#sfdisfil <linear-Q-esti.rsf


Flow('linear-true',None,'spike n1=40 d1=1 o1=0 |math output="41+x1"')
Graph('linear-true','min2=20 max2=230 label2="Quality factor" label1=Trace unit1= plotfat=10')

Flow('linear-comp','linear-Q-esti-N linear-true linear-Q1s-esti','cat axis=2 ${SOURCES[1:3]}')
Graph('linear-comp','min2=20 max2=230 label2="Quality factor" label1=Trace unit1= plotfat=10')





End()