from rsf.proj import *

Flow('sigmoid',None,
     '''
     sigmoid n1=200 n2=210 d2=0.008 |
     put label2=Distance |
     smooth rect1=3 diff1=1 adj=1 | smooth rect1=3
     ''')
Result('sigmoid','grey title="Input" ')
Plot('sig','sigmoid','grey wanttitle=n clip=0.003')

Flow('sdip','sigmoid',
     'dip order=2 p0=0 verb=y niter=10 rect1=3 rect2=3')
Result('sdip','grey color=j title="Dip" ')

Flow('spwd','sigmoid sdip',
     'pwd dip=${SOURCES[1]} order=2 | envelope | sharpen perc=99')
Result('sigpwd','spwd','grey title=Residual allpos=y')

Flow('sdip0','sdip','math output=0')
Flow('sdip1','sdip','math output=-0.285')

for t in range(1,210):
    left = 'left%d' % t
    Flow(left,'sigmoid','cut n2=%d' % t)

for x in ('','0','1'):
    pwd = 'pwd%s' % x
    Flow(pwd, ['sigmoid','sdip%s' % x],
         'pwd dip=${SOURCES[1]} order=2')
    blds = ['sig']
    for t in range(1,210):
        bld = 'bld%s-%d' % (x,t)
        Flow(bld,[pwd,'left%d' % t],
             'cut f2=%d | add ${SOURCES[1]}' % t)
        Plot(bld,'grey wanttitle=n clip=0.003')
        blds.append(bld)
    blend = 'blend%s' % x
    Plot(blend,blds,'Movie')

Flow('pad','sigmoid','math output=1 | pad beg1=50 end1=50')
Flow('sigmoid-pad','sigmoid','pad beg1=50 end1=50 | bandpass fhi=60')

Flow('sdip-pad','sigmoid-pad pad',
     '''
     dip order=2 p0=0 verb=y niter=10 rect1=3 rect2=3 mask=${SOURCES[1]}
     ''')

Flow('trace','sigmoid-pad','window n2=1 f2=105')
Flow('pred','sdip-pad trace',
     'pwpaint order=2 seed=${SOURCES[1]} i0=105 eps=0.1')
Plot('pred','window n1=200 min1=0 | grey wanttitle=n clip=0.003')

preds=[]
for n2 in range(104):
    pred = 'pred%d' % n2
    Flow(pred,'pred','cut n2=%d | cut f2=%d' % (105-n2,105+n2+1))
    Plot(pred,'window n1=200 min1=0 | grey wanttitle=n clip=0.003')
    preds.append(pred)
preds.append('pred')
Plot('preds',preds,'Movie')

Flow('seed','sdip-pad','window n2=1 | math output=x1')
Flow('spick','sdip-pad seed',
     'pwpaint order=2 seed=${SOURCES[1]} i0=105 eps=0.1')
Plot('spick0','spick',
     '''
     window n1=200 min1=0 |
     grey color=j allpos=y
     Xscalebar=y Xbarreverse=y
     title="Relative Age" clip=0.8
     minval=0 maxval=0.8
     ''')
Plot('spick1','spick',
     '''
     window n1=200 min1=0 |
     clip clip=0.8 |
     contour c0=0 dc=0.04 nc=20
     transp=y yreverse=y plotcol=7 plotfat=5
     Xscalebar=y Xbarreverse=y barlabel=" " 
     wanttitle=n wantaxis=n
     minval=0 maxval=0.8
     ''')
Result('spick','spick0 spick1','Overlay')

Flow('sflat','sigmoid-pad spick','iwarp warp=${SOURCES[1]} eps=1 n1=200 o1=0')

def ref(trace):
    out = 'ref%d' % trace
    Flow(out+'.asc',None,
         '''
         echo %d 0 %d 1 n1=4 in=$TARGET data_format=ascii_float
         ''' % (trace,trace))
    Plot(out,out+'.asc',
         '''
         dd form=native type=complex | 
         graph min1=0 max1=210 min2=0 max2=1 wanttitle=n wantaxis=n
         ''')
    return out

Plot('sflat','grey title="Single reference" ')
Result('sflat',['sflat',ref(105)],'Overlay')

# Add painting

Flow('cont','trace',
     'envelope | max1 | window n1=5 | real')
Flow('k1.p','cont',
     '''
     math output="1.5+(input+0.2)/0.004" | 
     dd type=int form=ascii format="%d," line=25 --out=$TARGET
     ''',stdout=0)
Command('k1.par','k1.p',
         'printf "k1=" > $TARGET && cat $SOURCE >> $TARGET')

Flow('spikes','k1.par',
     '''
     spike n1=300 nsp=5 par=$SOURCE |
     smooth rect1=5  
     ''',stdin=0)
Flow('spaint','sdip-pad spikes',
     'pwpaint order=2 seed=${SOURCES[1]} i0=105 eps=0.1')

Result('spaint','spaint sigmoid',
     '''
     window n1=200 min1=0 |
     add ${SOURCES[1]} scale=0.15,1 | 
     grey color=G title="Painted" allpos=y
     ''')

paints=[]
for n2 in range(104):
    paint = 'paint%d' % n2
    Flow(paint,'spaint sigmoid',
         '''
         cut n2=%d | cut f2=%d | window n1=200 min1=0 | 
         add scale=0.15,1 ${SOURCES[1]}
         ''' % (105-n2,105+n2+1))
    Plot(paint,'grey allpos=y color=G wanttitle=n clip=0.03')
    paints.append(paint)
paints.append('spaint')
Plot('paints',paints,'Movie')

# Add coherence

Flow('scor a2 b2','sigmoid sdip',
     'pwcoh dip=${SOURCES[1]} a2=${TARGETS[1]} b2=${TARGETS[2]} rect=3')

Flow('ab0','a2','math output=1')
Flow('ba0','a2','math output=1')

eps=2.0e-8
perc=97

for i in range(20):
    Flow('ab%d' % (i+1),'ab%d scor a2' % i,
         '''
         math c=${SOURCES[1]} a2=${SOURCES[2]} 
         output="(a2+%g)*input/(c*input+%g)" |
         sharpen perc=%g
         ''' % (eps,eps,perc))
    Flow('ba%d' % (i+1),'ba%d scor b2' % i,
         '''
         math c=${SOURCES[1]} a2=${SOURCES[2]}
         output="(a2+%g)*input/(c*input+%g)" |
         sharpen perc=%g
         ''' % (eps,eps,perc))

    Flow('scoh%d' % (i+1),'ab%d ba%d' % (i+1,i+1),'add mode=p ${SOURCES[1]}')

Plot('scoh','scoh20',
       'grey allpos=y title="(f) Coherence" ')

# Multiple references

picks=[]
refs=[]
for i0 in (0,49,99,149,199):
    pick = 'pick%d' % i0
    picks.append(pick)
    refs.append(ref(i0))

    Flow(pick,'sdip-pad seed',
         'pwpaint order=2 seed=${SOURCES[1]} i0=%d eps=0.1' % i0)

np = len(picks)
Flow('spicks',picks,
     'add ${SOURCES[1:%d]} | scale dscale=%g' % (np,1.0/np))

Flow('flat','sigmoid-pad spicks','iwarp warp=${SOURCES[1]} eps=1 n1=200 o1=0')
Plot('flat','grey title="Multiple references" ')
Result('flat',['flat']+refs,'Overlay')

#Result('spaint','sigmoid sdip spaint spick','TwoRows')
#Result('sflat','sflat flat','SideBySideIso')

flats=[]
for n2 in range(104):
    flat = 'flat%d' % n2
    Flow(flat,'sflat','cut n2=%d | cut f2=%d' % (105-n2,105+n2+1))
    orig = 'orig%d' % n2
    Flow(orig,'sigmoid','cut f2=%d n2=%d' % (105-n2,2*n2+1))
    Plot(flat,[flat,orig],
         'add ${SOURCES[1]} | grey wanttitle=n clip=0.003')
    flats.append(flat)
flats.append('sflat')
Plot('flats',flats,'Movie')


###########################################################################
End()