from rsf.proj import *
import os

# Size of the model
n1=200
n2=256

# Plot trace
def ref(trace, extra=''):
    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=256 min2=0 max2=1 wanttitle=n wantaxis=n scalebar=y %s
         '''% extra)
    return out

# Generate and display the model
Flow('sigmoid',None,
        '''
        sigmoid n1=%d n2=%d d2=.008 |
        smooth rect1=3 diff1=1 adj=1| smooth rect1=3 |
        put label2=Distance
        ''' % (n1,n2) )
Result('sigmoid', 'grey title=')


# Prepare for predictive painting (estimate dips)
Flow('pad','sigmoid','math output=1 | pad beg1=50 end1=50')
Flow('sigmoid-pad','sigmoid','pad beg1=50 end1=50 | bandpass fhi=60')

Flow('dip-pad','sigmoid-pad pad',
     '''
     dip order=2 p0=0 verb=y niter=10 rect1=3 rect2=3
     mask=${SOURCES[1]}
     ''')
Flow('seed','dip-pad','window n2=1 | math output=x1')
Result('dip-pad',
       '''
       window n1=200 min1=0 |
       grey color=j title="" barlabel="Slope" scalebar=y
       ''')

# RT with one reference trace
Flow('pick-1','dip-pad seed',
     'pwpaint order=2 seed=${SOURCES[1]} i0=128 eps=0.1')
Plot('pick-1',
     '''
     window n1=200 min1=0 |
     grey color=j allpos=y title="Single reference" scalebar=y barlabel="RT"
     ''')
Result('pick-1',['pick-1',ref(128,'scalebar=y')],'Overlay')

# RT with multiple reference traces
refs =[]
picks=[]
for i0 in (1,64,128,192,255):
    pick = 'pick%d' % i0
    picks.append(pick)
    refs.append(ref(i0,'scalebar=y'))

    # RT with single reference trace
    Flow(pick,'dip-pad seed',
         'pwpaint order=2 seed=${SOURCES[1]} i0=%d eps=1' % i0)

np = len(picks)
Flow('pick-5',picks,
     'add ${SOURCES[1:%d]} | scale dscale=%g' % (np,1.0/np))
Plot('pick-5',
     '''
     window n1=200 min1=0 |
     grey color=j allpos=y title="Multiple reference" scalebar=y barlabel="RT"
     ''')
Result('pick-5',['pick-5']+refs,'Overlay')

# Flatten
Flow('flat-1','sigmoid-pad pick-1','iwarp warp=${SOURCES[1]} eps=1 n1=200 o1=0')
Flow('flat-5','sigmoid-pad pick-5','iwarp warp=${SOURCES[1]} eps=1 n1=200 o1=0')
Result('flat-1','grey wanttitle=n')
Result('flat-5','grey wanttitle=n')


# Change the reference trace
Flow('pick-150','dip-pad seed',
     'pwpaint order=2 seed=${SOURCES[1]} i0=150 eps=0.1')
Plot('pick-150',
     '''
     window n1=200 min1=0 |
     grey color=j allpos=y title="Relative Time" scalebar=y barlabel="RT"
     ''')
Result('pick-150',['pick-150',ref(150)],'Overlay')

Flow('time','sigmoid-pad','math output=x1')
Result('time','grey color=j allpos=y')

# Time warping
Flow('invint','time pick-150','iwarp warp=${SOURCES[1]}')
Result('invint','grey color=j allpos=y title="Time Warping" scalebar=y barlabel=RT')

Flow('invint-50','invint','window n2=1 f2=50 | spray axis=2 n=256')
Result('invint-50','window n1=200 min1=0 | grey color=j allpos=y scalebar=y barlabel="RT"')

Flow('pick-50','invint-50 pick-150','iwarp warp=${SOURCES[1]} inv=n')
Plot('pick-50',
     '''
     window n1=200 min1=0 |put o2=0 |
     grey color=j allpos=y title="Relative Time" scalebar=y barlabel="RT"
     ''')
Result('pick-50',['pick-50',ref(50)],'Overlay')

Flow('pick-50-true','dip-pad seed',
     'pwpaint order=2 seed=${SOURCES[1]} i0=50 eps=0.1')
Plot('pick-50-true',
     '''
     window n1=200 min1=0 |
     grey color=j allpos=y title="Relative Time" scalebar=y barlabel="RT" 
     ''')
Result('pick-50-true',['pick-50-true',ref(50)],'Overlay')

Result('diff','pick-50-true pick-50',
    '''
    add ${SOURCES[1]} scale=1,-1 | 
    window n1=200 min1=0 |
    grey color=j title="Diff" scalebar=y
    ''')
Flow('flat1','sigmoid-pad pick-50','iwarp warp=${SOURCES[1]} eps=1 n1=200 o1=0')
Flow('flat2','sigmoid-pad pick-50-true','iwarp warp=${SOURCES[1]} eps=1 n1=200 o1=0')
Flow('flat3','sigmoid-pad pick-150','iwarp warp=${SOURCES[1]} eps=1 n1=200 o1=0')
Result('flat-diff','flat1 flat2',
    '''
    add ${SOURCES[1]} scale=1,-1 | 
    grey title="Difference" maxval=0.00338088
    ''')
Result('flat1','grey title=""')
Result('flat2','grey title=""')
Result('flat3','grey title=""')

End()