# 
# convolution/correlation demo
#
from rsf.proj import *
import sys,wplot,pplot,awe
sys.path.append('../../PYUTIL')
from corLOP import corop
from conLOP import conop

from DPTEST import *
from LICGsolver import *
from LICGutils  import *

# ------------------------------------------------------------
do3D = 'y'
doINV = 'y'
# ------------------------------------------------------------

parT = dict( nt=61, ot=+0, dt=1, lt='t', ut='s', nc=30)
parU = dict( nt=31, ot=-5, dt=1, lt='t', ut='s', nc=30)
wplot.param(parT)
wplot.param(parU)

par = dict( nt=1001, ot=-0.1, dt=0.001, lt='t', ut='s', kt=200)
wplot.param(par)

# ------------------------------------------------------------
Flow('f',None,
    '''
    spike nsp=3 mag=+2,-1,+1 k1=15,24,36 l1=17,28,36
    n1=%(nt)d o1=%(ot)g d1=%(dt)g
    '''%parT)

Flow('g',None,
    '''
    spike nsp=3 mag=3,2,1 k1=10,11,12
    n1=%(nt)d o1=%(ot)g d1=%(dt)g
    '''%parU)

# ------------------------------------------------------------

CON = conop('g',parT,'')
COR = corop('g',parT,'')

CON.FORW('f',   'fcon')
CON.ADJT('acon','fcon')

COR.FORW('f',   'fcor')
COR.ADJT('acor','fcor')

# ------------------------------------------------------------
def plt(col,custom,par):
    return ''' 
    graph title=""
    symbol=. plotfat=15 
    screenratio=0.15 screenht=4 
    grid1=y g1num=5
    min1=-50 max1=+50
    label1=%(lt)s unit1=%(ut)s
    '''%par + 'plotcol=%d'%col + par['labelattr'] + 'labelfat=3' + ' ' +custom

con=['f','g','fcon','acon']
for f in con:
    c = con.index(f)+1
    Plot(f,plt(c,'',parT))
pplot.multip('con',con,len(con),1, 0.5,0.5, -4,0);

cor=['f','g','fcor','acor']
for f in cor:
    c = cor.index(f)+1
    Plot(f,plt(c,'',parT))
pplot.multip('cor',cor,len(cor),1, 0.5,0.5, -4,0);

#Flow('tcor',['f','g'],'./npyCorr.py flt=${SOURCES[1]} nc=30')
#Result('tcor',plt(3,'xll=+1 screenht=2',parT))
#
#Flow('tcon',['f','g'],'./npyConv.py flt=${SOURCES[1]}')
#Result('tcon',plt(3,'xll=+1 screenht=2',parT))

# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
Flow('dcon','fcon','window squeeze=n') 
Flow('mcon','acon','window squueze=n') 
DCON = DPTEST(CON,['mcon'],['dcon'])           
DCON.RUN()

Flow('dcor','fcor','window squeeze=n') 
Flow('mcor','acor','window squueze=n') 
DCOR = DPTEST(COR,['mcor'],['dcor'])           
DCOR.RUN()

# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
if(doINV == 'y'):
    NITER = 15
    
    Flow('zcon','acon','math output=0')
    LCON = LICG('icon',par, Lop=CON, d=['fcon'], xo=['zcon'], xf=['icon'], niter=NITER)
    LCON.iterate()
    plotOFUN('ICON',NITER,par,'min2=-80')
    Plot('jcon','icon',plt(5,'',parT))
    pplot.multip('icon',con+['jcon'],len(con)+1,1, 0.5,0.5, -4,0);
    
    Flow('zcor','acor','math output=0')
    LCOR = LICG('icor',par, Lop=COR, d=['fcor'], xo=['zcor'], xf=['icor'], niter=NITER)
    LCOR.iterate()
    plotOFUN('ICOR',NITER,par,'min2=-80')
    Plot('jcor','icor',plt(5,'',parT))
    pplot.multip('icor',cor+['jcor'],len(con)+1,1, 0.5,0.5, -4,0);
    
# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------

Flow('a',None,
    '''
    spike nsp=1 mag=0 
    n1=%(nt)d o1=%(ot)g d1=%(dt)g |
    noise rep=y |
    ricker1 frequency=20 |
    scale axis=123
    '''%par)
Flow('b','a','put o1=0')

WCON = conop('a',par,'norm=n')
WCOR = corop('a',par,'norm=n')

WCON.FORW('b','wcon')
WCOR.FORW('b','wcor')

for i in (['a','b','wcon','wcor']):
    Result(i,'window | scale axis=123|'
        + wplot.waveplot('grid=y min1=-1.5 max1=+1.5',par))
    
End()
