from rsf.proj import *
import os

# Straight line
Flow('clean',None,'math n1=50 o1=1 d1=1 n2=1 output="3+0.1*x1"')
Flow('noisy','clean','noise var=0.01 seed=2007')
Plot('noisy',
     'graph wanttitle=n symbol=* Xplotcol=5 label1="\F2 X" label2="\F2 Y" format2="%3.1f" min2=2 max2=9 plotfat=3 parallel2=n')

Result('noisy','Overlay')

Flow('t0','noisy','math output=1')
Flow('t1','noisy','math output="%g*x1"' % (1/29.3002))
Flow('t2','noisy','math output="%g*x1*x1"' % (1/1146.01))
Flow('t','t0 t1','cat axis=2 ${SOURCES[1]}')
Flow('tt','t t2','cat axis=2 ${SOURCES[1]}')

Flow('flt pred','t noisy',
     'lpf match=${SOURCES[1]} pred=${TARGETS[1]} rect1=1000')

Plot('pred','graph title="\F2 Y(X) = a+b*X" wantaxis=n min2=2 max2=9 Xplotcol=3 plotfat=3')
Result('pred','noisy pred','Overlay')

# Curve
Flow('clean2',None,'math n1=50 o1=1 d1=1 output="sqrt(9+0.025*x1*x1)"')
Flow('noisy2','clean2','noise var=0.01 seed=2007')
Plot('noisy2',
     'graph wanttitle=n symbol=* Xplotcol=5 label1="\F2 X" label2="\F2 Y" format2="%3.1f" min2=2 max2=9 plotfat=3 parallel2=n')

Result('noisy2','Overlay')

Flow('flt2 pred2','t noisy2',
     'lpf match=${SOURCES[1]} pred=${TARGETS[1]} rect1=1000')

Plot('pred2','graph title="\F2 Y(X) = a+b*X" wantaxis=n min2=2 max2=9 Xplotcol=3 plotfat=3')
Result('pred2','noisy2 pred2','Overlay')

Flow('flt3 pred3','tt noisy2',
     'lpf match=${SOURCES[1]} pred=${TARGETS[1]} rect1=1000')

Plot('pred3',
     'graph title="\F2 Y(X) = a+b*X+c*X\^2" wantaxis=n plotfat=3 min2=2 max2=9 Xplotcol=3')
Result('pred3','noisy2 pred3','Overlay')

for case in 'nf':
    noisy = 'noisy2'+case
    window = 'window %c1=25' % case
    Flow(noisy,'noisy2',window)

    pred = 'pred'+case
    Flow(['flt2'+case,pred],['t',noisy],
         window + '''|
         lpf match=${SOURCES[1]} pred=${TARGETS[1]} rect1=1000
         ''')
    Flow(pred+'2',[pred,'noisy2'],'remap1 order=2 pattern=${SOURCES[1]}')

    Plot(pred+'2',
         '''
         graph title="\F2 Y\_k\^(X) = a\_k\^+b\_k\^*X"
         wantaxis=n plotfat=3 min2=2 max2=9 Xplotcol=3
         ''')

Result('pred4','noisy2 predn2 predf2','Overlay')

pres = []
errs = []
for rect in range(2,101):
    flt = 'prf%d' % rect
    pre = 'pre%d' % rect
    err = 'err%d' % rect

    Flow([flt,pre],'t noisy2',
         'lpf match=${SOURCES[1]} pred=${TARGETS[1]} rect1=%d' % rect)
    Plot(pre,
         '''
         graph title="\F2 Y(X) = a(X)+b(X)*X"
         wantaxis=n plotfat=3 min2=2 max2=9 Xplotcol=3
         ''')
    pres.append(pre)

    Flow(err,[pre,'noisy2'],
         '''
         math d=${SOURCES[1]} output="(input-d)^2" |
         stack axis=1 norm=n
         ''')
    errs.append(err)


Flow('errs',errs,'cat axis=1 ${SOURCES[1:%d]} | put o1=2 d1=1' % len(errs))
Plot('pres',pres,'Movie',view=1)

Plot('errs1','errs',
     'graph wanttitle=n label2="Data Misfit" label1="Smoothing Radius" ')
Plot('errs2','errs',
     '''
     window n1=12 |
     graph wanttitle=n label2="Data Misfit" label1="Smoothing Radius"
     ''')
Result('errs','errs1 errs2','SideBySideIso')

Result('pred5','noisy2 pre7','Overlay')

rows = {'t': ([],[],[]),   's': ([],[],[])}
pars = {'t': (0.1,1,10), 's': (3,15,75)}

for i in range(100):
    spike = 'spike%d' % i
    Flow(spike,None,'spike n1=100 o1=1 d1=1 k1=%d' % (i+1))

for j in range(3):
    for i in range(100):
        spike = 'spike%d' % i

        for case in ('nf'):
            x = 'x%c%d' % (case,i)
            Flow(x,spike,'window %c1=50' % case)

        tn = 'tn%d-%d' % (i,j)
        tf = 'tf%d-%d' % (i,j)
        trow = 'trow%d-%d' % (i,j)

        eps = pars['t'][j]
        eps = eps*eps

        Flow(tn,'xn%d xf%d t1' % (i,i),
             '''
             igrad | igrad adj=y |
             math top=${SOURCES[0]} bot=${SOURCES[1]} t1=${SOURCES[2]}
             output="%g*input+top+t1*bot"
             ''' % eps)
        Flow(tf,'xf%d xn%d t1' % (i,i),
             '''
             igrad | igrad adj=y |
             math top=${SOURCES[1]} bot=${SOURCES[0]} t1=${SOURCES[2]}
             output="%g*input+t1*top+t1*t1*bot"
             ''' % eps)

        Flow(trow,[tn,tf],'cat ${SOURCES[1]} axis=1')
        rows['t'][j].append(trow)

        sn = 'sn%d-%d' % (i,j)
        sf = 'sf%d-%d' % (i,j)
        srow = 'srow%d-%d' % (i,j)

        rect = pars['s'][j]

        Flow(sn,'xf%d xn%d t1' % (i,i),
             '''
             smooth rect1=%d |
             math  t1=${SOURCES[2]} output="input*t1" |
             smooth rect1=%d |
             add ${SOURCES[1]}
             ''' % (rect,rect))
        Flow(sf+'t','xn%d t1' % i,
             '''
             smooth rect1=%d |
             math  t1=${SOURCES[1]} output="input*t1" |
             smooth rect1=%d
             ''' % (rect,rect))
        Flow(sf,['xf%d' % i, 't1',sf + 't'],
             '''
             smooth rect1=%d |
             math  t1=${SOURCES[1]} output="input*(t1*t1-1)" |
             smooth rect1=%d |
             add ${SOURCES[0]} ${SOURCES[2]}
             ''' % (rect,rect))

        Flow(srow,[sn,sf],'cat ${SOURCES[1]} axis=1')
        rows['s'][j].append(srow)

    for case in 'ts':
        mat = case+'mat%d' % j

        title = '%s Regularization (%s)' % \
                (('Tikhonov','Shaping')[case=='s'],pars[case][j])

        Flow(mat,rows[case][j],
             '''
             cat ${SOURCES[1:100]} axis=2 |
             put o1=1 o2=1 d2=1 label1= unit1=
             ''')
        Result(mat,
               '''
               grey title="\F2 %s" wanttitle=n screenratio=1 scalebar=y
               label1="\F2 Row" label2="\F2 Column" parallel2=n
               bartype=h 
               ''' % title)

        eig = case+'eig%d' % j
        Flow(eig,mat,'jacobi niter=20')
        Result(eig,
               '''
               put o1=1 d1=1 | sort | 
               graph symbol=* title="\F2 %s" screenratio=1 plotfat=3
               wheretitle=b wherexlabel=t wanttitle=n  parallel2=n 
               label1="\F2 Eigenvalue number" label2="\F2 Magnitude"
              ''' % title)

End()