from rsf.proj import *

def rand(seed,n=128):
    name = 'rand%d' % seed
    Flow(name,None,
         '''
         spike n1=%d n2=%d d1=1 d2=1 |
         noise rep=y seed=%d 
         ''' % (n,n,seed))
    return name

def grey(title):
    return '''
    grey title="%s" crowd=0.86 wantaxis=n
    ''' % title

Flow('rand2d',rand(2010),'smooth rect1=3 rect2=3')
Flow('randph',rand(2011),'rtoc | math output="exp(I*input)" ')

for txtr in ('ridges','wood','herr','skull'):
    Fetch(txtr + '.h','textures')
    Flow(txtr,txtr + '.h','dd form=native | put d1=1 o1=0 d2=1 o2=0')

Fetch('stanford.tree.H','textures')
Flow('tree-hole','stanford.tree.H herr',
     '''
     dd form=native | put d1=1 o1=0 d2=1 o2=0 |
     mul ${SOURCES[1]}
     ''')
Flow('skull-hole','skull herr','mul ${SOURCES[1]}')

Fetch('WGstack.H','book/iee')
Flow('WGstack','WGstack.H','window f2=500 j2=2 n2=125 n1=250 j1=2')

for case in ('rand2d','ridges','wood','herr',
             'tree-hole','skull-hole','WGstack'):
    Plot(case,grey('Training Image'))

    if case != 'WGstack':
        # FFT method

        fft = 'fft-'+case

        Flow(fft,case,'rtoc | fft3 axis=1 pad=1 | fft3 axis=2 pad=1')

        phase = 'phase-'+case

        Flow(phase,fft,'math output="log(input)*abs(input)" | imag')
        Plot(phase,grey('Phase Spec.(TI)'))

        campl = 'campl-'+case

        Flow(campl,fft,'math output="abs(input)" ')
        Plot(campl,'real | ' + grey('Ampl. Spec.(TI)') + ' allpos=y')

        synth = 'synth-'+case
        Flow(synth,[campl,'randph'],
             '''
             mul ${SOURCES[1]} |
             fft3 inv=y axis=2 |
             fft3 inv=y axis=1 |
             real
             ''')
        Plot(synth,grey('FT synthesis'))

        Result(case+'-ftsyn',[case,synth,phase,campl],'TwoRows')

    #------------------------------------------------------------ 
    # Texture Synthesis via PEF method.
    #------------------------------------------------------------ 
    # Minimize ||Data*PEF||^2 for unknown PEF.
    #------------------------------------------ 
    pef = 'pef-'+case
    lag = 'lag-'+case
    mis = 'mis-'+case
    Flow([pef,lag,mis+'i'],case,
         '''
         pef niter=40 a=10,10 center=4,0
         lag=${TARGETS[1]} maskin=$SOURCE maskout=${TARGETS[2]}
         ''')
    Flow(mis,mis+'i','dd type=float')

    # Impulse response of inverse PEF - decon of a spike.
    #----------------------------------------------------
    Flow('inv'+pef,pef,
         '''
         spike k1=64 k2=32 n1=128 n2=128 |
         helicon filt=$SOURCE div=1
         ''',stdin=0)
    Plot(pef,'inv'+pef,grey('spike/PEF'))

    # Compute residual for known PEF: res = Data*PEF
    #-----------------------------------------------
    wht = 'wht-'+case
    Flow(wht,[case,pef,lag,mis],
         'helicon filt=${SOURCES[1]} | mul ${SOURCES[3]}')
    Plot(wht,grey('Residual = TI*PEF'))

    # Synthesized texture = Random #'s/PEF  ("/" = poly. div.)
    #---------------------------------------------------------
    syn = 'syn-'+case
    Flow(syn,[rand(2012),pef,lag],'helicon filt=${SOURCES[1]} div=1')
    Plot(syn,grey('Synthesized Image'))

    Plot(case+'-pefsyn',[case,wht,syn],'SideBySideAniso',vppen='txscale=3')
    Result(case+'-pefsyn',[case,wht,syn],'SideBySideIso')

Result('holes','herr-pefsyn tree-hole-pefsyn skull-hole-pefsyn',
       'OverUnderAniso')

#---------------------------
# Fill missing data with PEF
#---------------------------
Flow('fillpef','tree-hole pef-tree-hole',
     'miss filt=${SOURCES[1]} padin=20')
Plot('fillpef',grey('PEF Infilled data'))

#-----------------------------
# Fill missing data with Lapl. 
#-----------------------------
Flow('filllap','tree-hole','lapfill')
Plot('filllap',grey('Laplacian Infilled data'))

Result('tree-hole-filled','tree-hole pef-tree-hole fillpef filllap',
       'TwoRows')

spectra = 'rtoc | fft3 axis=1 pad=1 | math output="abs(input)" | real'

Flow('rand1d',None,
     '''
     spike n1=64 | noise rep=y seed=287293 |
     smooth rect1=3
     ''')
Flow('spc','rand1d',spectra)

spcs = []
label = ''
for a in (4,10,20):
    pef = 'pef%d' % a
    lag = 'lag%d' % a
    Flow([pef,lag],'rand1d','pef lag=${TARGETS[1]} a=%d' % a)
    imp = 'imp%d' % a
    Flow(imp,pef,
         '''
         spike n1=64 k1=1 |
         helicon filt=$SOURCE div=1
         ''',stdin=0)
    spc = 'spc%d' % a
    Flow(spc,imp,spectra)
    spcs.append(spc)

    label += ':Inv. PEF Spectrum, Na=%d' % a

Result('rand1d-spec',['rand1d','spc']+spcs,
       '''
       cat axis=2 ${SOURCES[1:%d]} |
       dots dots=0 connect=1 gaineach=1 constsep=1 strings=0 label1=" " 
       yreverse=y labelsz=7
       labels="Input:Input Spectrum%s"
       ''' % (len(spcs)+2,label))



End()