from rsfproj import *

# Download data
Fetch('seabin.hh','seab')

# Create mesh
Flow('mesh','seabin.hh','dd form=native | pad n1=200 beg1=20')

# Mask for known values
Flow('mask','mesh',
     'add mode=p $SOURCE | mask min=1e-6 | dd type=float')

def plotdata(title):
    return '''
    grey title="%s" transp=n yreverse=n clip=0.65
    label1=Longitude label2=Latitude 
    ''' % title

# Plot input data
Plot('mesh',plotdata('(a) Input'))

# Laplacian filter
Flow('lag.asc',None,
     '''
     echo 1 99 100 101 n1=4 n=100,100 
     data_format=ascii_int in=$TARGET
     ''')
Flow('lag','lag.asc','dd form=native')

Flow('flt.asc','lag',
     '''
     echo -4 -1 -4 -1 a0=10 n1=4 lag=$SOURCE 
     data_format=ascii_float in=$TARGET
     ''',stdin=0)
Flow('flt','flt.asc','dd form=native')

# Spectral factorization on a helix
Flow('lapflt laplag','flt',
     'wilson eps=1e-3 lagout=${TARGETS[1]}')

def plotfilt(title):
    return '''
    grey wantaxis=n title="%s" pclip=100 
    crowd=0.85 screenratio=1
    ''' % title

# Filter impulse response
Flow('spike',None,'spike n1=15 n2=15 k1=8 k2=8')
Flow('imp0','spike flt','helicon filt=${SOURCES[1]} adj=0')
Flow('imp1','spike flt','helicon filt=${SOURCES[1]} adj=1')
Flow('imp','imp0 imp1','add ${SOURCES[1]}')
Plot('imp',plotfilt('(a) Laplacian'))

# Test factorization
Flow('fac0','imp lapflt',
     'helicon filt=${SOURCES[1]} adj=0 div=1')
Flow('fac1','imp lapflt',
     'helicon filt=${SOURCES[1]} adj=1 div=1')
Plot('fac0',plotfilt('(b) Laplacian/Factor'))
Plot('fac1',plotfilt('(c) Laplacian/Factor\''))
Flow('fac','fac0 lapflt',
     'helicon filt=${SOURCES[1]} adj=1 div=1')
Plot('fac',plotfilt('(d) Laplacian/Factor/Factor\''))

Result('laplace','imp fac0 fac1 fac','TwoRows',
       vppen='gridsize=5,5 xsize=11 ysize=11')

seed  = 1203 # MODIFY ME !!!

# Random initial model
Flow('rand','mesh lapflt',
     '''
     noise rep=y seed=%d var=0.02 | 
     helicon filt=${SOURCES[1]} div=y      
     ''' % seed)
Plot('rand',plotdata('(b) Random'))
Result('mesh','mesh rand','SideBySideAniso')

# Missing data interpolation

for prec in (0,1):
    errors = []
    # Loop over iterations
    for niter in range(101):
        interp = 'interp%d-%d' % (niter,prec)
        # K[x] =~ K[d-K[m0]]; D[x] =~ 0; m = m+m0
        Flow(interp,'rand mask mesh lapflt',
             '''
             add mode=p ${SOURCES[1]} |
             add scale=-1,1 ${SOURCES[2]} |
             miss filt=${SOURCES[3]} mask=${SOURCES[1]} 
             niter=%d prec=%d exact=n |
             add ${SOURCES[0]}
             ''' % (niter,prec))
        title = ('(a) Multiplication','(b) Division')[prec]
        Plot(interp,plotdata(title))

        error = 'err-'+interp
        Flow(error,[interp,'interp100-%d' % prec],
             '''
             add scale=1,-1 ${SOURCES[1]} | 
             math output="input*input" |
             stack axis=2 norm=n |
             stack axis=1 norm=n |
             math output="sqrt(input)"
             ''')
        errors.append(error)
    error = 'error%d' % prec
    Flow(error,errors,'cat axis=1 ${SOURCES[1:100]}')

Result('interp','interp50-0 interp50-1','SideBySideAniso')
Result('error','error0 error1',
       '''
       cat axis=2 ${SOURCES[1]} | 
       put o1=0 d1=1 label1=Iterations |
       graph dash=0,1 title="Model Convergence"
       label2="|Model(n)-Model(N)|" 
       ''')

End()