from rsf.proj import *
from rsf.recipes.beg import server
from rsf.recipes.tpx import FPX
import math
import tail_el
import velcon2d

### Model building

Flow('surf',None,'spike n1=1001 d1=0.001 k1=500 mag=0.375')#0.375; 1 for 4km/s

Flow('refl',None,'spike n1=1001 d1=0.001 k1=500 mag=0.7')

Flow('diffr_data','surf refl',
     '''
     kirmod nt=250 dt=0.004 freq=20 
     nh=1 dh=0.5 h0=0.  
     ns=1001 ds=0.001 s0=0. 
     cmp=y refl=${SOURCES[1]}
     vel=1.5 type=c |
     put label2=Offset unit2=km label3=Midpoint unit3=km |
     window | 
     costaper nw1=100 nw2=100
     ''',split=[1,1001],reduce='add')

#Result('diffr_data','grey pclip=100 gainpanel=each title="Diffr Data"')

### 2D Post-stack velocity continuation

nv_vc=100
dv_vc=.01
v0_vc=1.0

pad_time=200
pad_af1=1000
pad_af2=1000

taper = '''
    costaper nw1=50 nw2=50
    '''
pad = '''
    pad end1=%d beg2=%d end2=%d
    ''' % (pad_time,pad_af1,pad_af2)

wind = '''
    window n1=%d n3=%d f3=%d | put o3=%g
    ''' %(250,1001,pad_af1,0)

Flow('vczo_input','diffr_data',taper + ' | ' + pad + ' | cosft sign2=1')
Flow('vczo_1st_step','vczo_input','vczo nv=%g dv=%g v0=%g | window' %(1,v0_vc,0.0))
Flow('vczo_confft','vczo_1st_step',
                   '''
                   vczo nv=%g dv=%g v0=%g
                   '''% (nv_vc,dv_vc,v0_vc),split=[2,'omp'],reduce='cat axis=3' )
Flow('vczo_con','vczo_confft','cosft sign3=-1 | ' + wind + ' | transp plane=23')

Flow('vczo_125','vczo_con','window n3=1 min3=1.25')
Flow('vczo_15','vczo_con','window n3=1 min3=1.5')
Flow('vczo_175','vczo_con','window n3=1 min3=1.75')

Flow('vczo_con_nw','vczo_confft','cosft sign3=-1 | transp plane=23')

### Stacking path summation

Flow('path_integral','vczo_con','stack axis=3')

### Tail models

Flow('path1','vczo_con','window n3=1')
Flow('path2','vczo_con','window n3=1 f3=-1')

### Adaptive tail elimination

shift_num=3
tail_el.tail_el('path_integral','path2','path1',
              nt = 250,
              nsh=shift_num,
              lrect1=20,
              lrect2=30,
              lniter=400)

### Fourier Transforms

sym='y'

invft = '''
        fft3 axis=2 inv=y |
        fft1 inv=y |
        t2warp inv=y
        '''

fwdft = '''
        t2warp pad=1000 |
        fft1 |
        fft3 axis=2
        '''

opt='n'

invftexp = '''
           fft3 axis=2 inv=y sym=%s opt=%s |
           transp |
           fft3 axis=2 inv=y sym=%s opt=%s |
           transp |
           real |
           t2warp inv=y
           '''%(sym,opt,sym,opt)
fwdftexp = '''
           t2warp pad=1000 |
           transp plane=12 | rtoc |
           fft3 axis=2 sym=%s opt=%s |
           transp plane=12 |
           fft3 axis=2 sym=%s opt=%s
           '''%(sym,opt,sym,opt)

Flow('fft','diffr_data','costaper nw1=10 nw2=100 | ' + fwdft)

Flow('fft-exp','diffr_data','costaper nw1=10 nw2=100 | ' + fwdftexp)

Flow('erfi_input','fft','math output="1.0 + I*0.0"')

Flow('erfi_input-exp','fft-exp','math output="1.0 + I*0.0"')

### GPI

Flow('pi_gaussian_erfi','erfi_input','gpi3dzo v_0=1.5 v_a=1.0 v_b=2.0 beta=10')#beta=10

Flow('pi_gaussian_fft','fft','gpi3dzo v_0=1.5 v_a=1.0 v_b=2.0 beta=10')
Flow('pi_gaussian','pi_gaussian_fft',
                  '''
                  fft3 axis=2 inv=y |
                  fft1 inv=y |
                  t2warp inv=y
                  ''')

### checking Bona's idea for tapering pi

Flow('left_taper_fft','fft','gpi3dzo v_0=1.0 v_a=0.5 v_b=1.0 beta=10')
Flow('right_taper_fft','fft','gpi3dzo v_0=2.0 v_a=2.0 v_b=2.5 beta=10')
Flow('flat_fft','fft','gpi3dzo v_0=1.5 v_a=1.0 v_b=2.0 beta=0.0')

Flow('left','left_taper_fft',
                  '''
                  fft3 axis=2 inv=y |
                  fft1 inv=y |
                  t2warp inv=y
                  ''')
Flow('right','right_taper_fft',
                  '''
                  fft3 axis=2 inv=y |
                  fft1 inv=y |
                  t2warp inv=y
                  ''')
Flow('flat','flat_fft',
                  '''
                  fft3 axis=2 inv=y |
                  fft1 inv=y |
                  t2warp inv=y
                  ''')

### combine all the parts

Flow('gpi_taper','left flat right','add scale=1,1 ${SOURCES[1]} | add scale=1,1 ${SOURCES[2]}')
#Result('gpi_taper','grey pclip=100 title="gpi taper = left + flat + right"')

### regular pi

eps=0.0001
v_a=1.0#v0_vc
v_b=2.0#v0_vc + (nv_vc-1)*dv_vc

exp0=0
exp1=1

velcon2d.pi2d(eps,'erfi_input','fft','path-sum',v_a,v_b,exp0,opt,sym)
velcon2d.pi2d(eps,'erfi_input-exp','fft-exp','path-sum-exp',v_a,v_b,exp1,opt,sym)

#?# Not sure why but erfi is not centered at 0 wavenumber

### velocity continuation through 2d fft

cube2append = []
cube2appendfft = []

velcon2d.vc2d('fft','all','vc-2d-fft','vc-2d',cube2append,cube2appendfft,
          v0_vc,
          nv_vc,
          dv_vc,
          eps)

### Figures for article
def grey(title):
        return '''
        grey title="%s" gainpanel=each pclip=100
        label1=Time unit1="s" label2=Distance unit2="km"
        titlefat=3 labelfat=3 font=2
        ''' % (title)
def grey_ng(title):
        return '''
        math output="(input)" |
        grey title="%s" clip=2.0e-04  screenratio=1
        label1=Time unit1="s" label2=Distance unit2="km"
        titlefat=3 labelfat=3 font=2
        ''' % (title)


Result('a-data','diffr_data',grey_ng(''' Diffraction '''))
Result('a-path-integral','pi_4_path-sum',grey_ng(''' Path-Summation Image '''))#4
#Result('a-path-int-te','path_int_te_nsh%d'%(shift_num),grey_ng(''' PI Image: 'Tails' Subtracted '''))#5
#Result('a-tails','tails_nsh%d'%(shift_num),grey_ng(''' Predicted 'Tails' '''))#6
Result('a-pi-gaussian','pi_gaussian',
                                     '''
                                     math output="(input)" |
                                     grey title="%s" clip=2.0e-04 screenratio=1  
                                     label1=Time unit1="s" label2=Distance unit2="km"
                                     titlefat=3 labelfat=3 font=2
                                     '''%('PS Image: Gaussian Weighting'))

### Plotting magnitude and phase responses

def fftgrey(title,extra=''):
        return '''
        math output="abs(input)" | real | scale axis=2 |
        grey title="%s" min2=-100 max2=100 max1=80
        label1=Frequency unit1="Hz" label2=Wavenumber unit2="1/km" color=i bias=0.2 clip=0.3
        titlefat=3 labelfat=3 font=2 scalebar=y allpos=n minval=0.0 maxval=1.0 barlabel="magnitude" %s
        ''' % (title,extra)

def fftgreyph(title,extra=''):
        return '''
        math output="arg(input)" | real |
        grey title="%s" min2=-100 max2=100 max1=80
        label1=Frequency unit1="Hz" label2=Wavenumber unit2="1/km" color=x
        titlefat=3 labelfat=3 font=2 scalebar=y allpos=n bias=0.0 barlabel="phase" barunit="radians" %s
        ''' % (title,extra)

def fftgraph(title,freq=20,extra=''):
        return '''
        window n1=1 min1=%g min2=-50 max2=50 | math output="abs(input)" | real |
        graph title="%s" %s
        ''' % (freq,title,extra)

#49
def fftgraphm(title,extra=''):
        return '''
        window n1=1 min1=30 min2=-50 max2=50 squeeze=y | math output="arg(input)" | real |
        graph3 title="%s" %s frame2=45 frame3=9 movie=2
        ''' % (title,extra)

def fftgraphmabs(title,extra=''):
        return '''
        window n1=1 min1=20 min2=-50 max2=50 squeeze=y | math output="abs(input)" | real |
        graph3 title="%s" %s frame2=45 frame3=9 movie=3
        ''' % (title,extra)


#def fftgraphph(title,extra=''):
#        return '''
#        math output="arg(input)" | real |
#        graph title="%s" min2=-200 max2=200 max1=80
#        label1=Frequency unit1="Hz" label2=Wavenumber unit2="1/km" color=j
#        titlefat=3 labelfat=3 font=2 scalebar=y allpos=n bias=0.0 barlabel="phase" %s
#        ''' % (title,extra)

Result('a-erfi-fft-mag','erfi_ba_4_path-sum',fftgrey('PS Magnitude'))

Result('a-erfi-fft-phase','erfi_ba_4_path-sum',fftgreyph('PS Phase'))

Result('a-gaussian-erfi-mag','pi_gaussian_erfi',fftgrey('GPS Magnitude'))

Result('a-gaussian-erfi-phase','pi_gaussian_erfi',fftgreyph('GPS Phase'))

#Result('erfi-fft-exp-mag','erfi_ba_4_path-sum-exp',fftgrey('PI Magnitude'))

#Result('erfi-fft-exp-phase','erfi_ba_4_path-sum-exp',fftgreyph('PI Phase'))

#,'bias=0.3 clip=0.2 allpos=n')
Result('a-fft-mag','fft',fftgrey('Data FFT Magnitude','screenratio=1'))

#Result('fft-phase','fft',fftgreyph('Data FFT Phase'))

#'bias=0.3 clip=0.2 allpos=n'
Result('a-pi-fft-mag','pi_erfi_fft_4_path-sum',fftgrey('PS Magnitude'))

#Result('pi-fft-phase','pi_erfi_fft_4_path-sum',fftgreyph('PI Phase'))

# ,'bias=0.3 clip=0.2 allpos=n'
Result('a-gpi-fft-mag','pi_gaussian_fft',fftgrey('GPS Magnitude'))

#Result('gpi-fft-phase','pi_gaussian_fft',fftgreyph('GPI Phase'))

prog = Program('filter.c')

v1=2.0#1.0#1.7
v2=2.5#3.5#2.0

Flow('pi_erfi_fft-cut','pi_erfi_fft_4_path-sum %s' % prog[0],'./${SOURCES[1]} logis=30 v1=%g v2=%g'% (v1,v2))

#Result('pi-fft-mag-cut','pi_erfi_fft-cut',fftgrey('PI Magnitude Cut','bias=0.3 clip=0.2 allpos=n'))

#Result('pi-fft-phase-cut','pi_erfi_fft-cut',fftgreyph('PI Phase Cut'))

Flow('pi-cut','pi_erfi_fft-cut',invft)

#Result('a-path-integral-cut','pi-cut',grey_ng(''' Path Integral Image Cut '''))#4

### Trying to smooth erfi spectrum - equivalent to GPI?

sm2 = 5
rep = 10

Flow('re-erfi-ps-sm','erfi_ba_4_path-sum','real | smooth rect2=%g repeat=%d | rtoc'%(sm2,rep))

Flow('erfi-ps-sm','erfi_ba_4_path-sum re-erfi-ps-sm',
        '''
        imag |
        smooth rect2=%g repeat=%d |
        rtoc | math output="I*real(input)" |
        add scale=1,1 ${SOURCES[1]}
        '''%(sm2,rep))

#Result('erfi-ps-sm-mag','erfi-ps-sm',fftgrey('PI erfi smooth'))

#Result('erfi-ps-sm-phase','erfi-ps-sm',fftgreyph('PI erfi smooth'))

Flow('pi-erfi-smooth','fft erfi-ps-sm',
        '''
        math  K=${SOURCES[1]} output="input*K" |
        ''' + invft)

#Result('a-path-integral-sm','pi-erfi-smooth',grey_ng(''' Path Integral Image Smooth '''))#4

### Make graphs slicing error functions at 20 Hz

Plot('graph-erfi-fft-mag','erfi_ba_4_path-sum',fftgraph('PI erfi graph',20))

Plot('graph-g-erfi-fft-mag','pi_gaussian_erfi',fftgraph('GPI erfi graph',20))

Plot('graph-pi-fft-mag','pi_erfi_fft_4_path-sum',fftgraph('PI graph',20))

Plot('graph-gpi-fft-mag','pi_gaussian_fft',fftgraph('GPI graph',20))

Plot('graph-picut-fft-mag','pi_erfi_fft-cut',fftgraph('PI cut graph',20))

### Lets consider response in time-space domain

Flow('filter-time','erfi_ba_4_path-sum',invft)

Flow('filter-time-exp','erfi_ba_4_path-sum-exp',invftexp)

#Result('filter-time','grey title="Filter in time-space domain" scalebar=y pclip=95')

### Transforming back the filter - to see if we loose information

Flow('filter-fft','filter-time',fwdft)

Flow('filter-fft-exp','filter-time-exp',fwdftexp)

#Result('filter-fft-mag','filter-fft',fftgrey('Magnitude'))
#Result('filter-fft-ph','filter-fft',fftgreyph('Phase'))

#Result('filter-fft-exp-mag','filter-fft-exp',fftgrey('Magnitude'))
#Result('filter-fft-exp-ph','filter-fft-exp',fftgreyph('Phase'))

### Analyzing VC responses

#Result('all-fft-mag','all-fft','stack axis=3 | ' + fftgrey('Magnitude','gainpanel=all'))

#Result('all-fft-ph','all-fft','stack axis=3 | ' + fftgreyph('Phase','gainpanel=all'))

#Result('graph-all-fft-ph','all-fft',fftgraphm('all fft graph',20))

#Result('graph-all-fft-mag','all-fft',fftgraphmabs('all fft graph',20))

### Cutting different parts of the filter - looking at different parts of the spectrum

#Flow('filter-time-cut','filter-time','cut max1=0.6')

#Result('filter-time-cut','grey title="Cut" scalebar=y pclip=99.8')

### Going back - t2 FFT domain and multiply by the data FT

#Flow('pi-cut-fft','filter-time-cut fft',
#               '''
#               t2warp pad=1000 |
#               fft1 |
#               fft3 axis=2 |
#               math  K=${SOURCES[1]} output="input*K"
#               ''')

#Flow('pi-cut-fft','filter-time-cut fft-exp',fwdftexp + ' | math  K=${SOURCES[1]} output="input*K"')

#Flow('pi-fft','filter-time fft',
#               '''
#               t2warp pad=1000 |
#               fft1 |
#               fft3 axis=2 |
#               math  K=${SOURCES[1]} output="input*K"
#               ''')

#Flow('pi-fft-exp','filter-time-exp fft-exp',fwdftexp + ' | math  K=${SOURCES[1]} output="input*K"')

### Explore amplitude

#Result('path-sum-mag','pi_erfi_fft_4_path-sum',fftgrey('Magnitude','mean=y clip=0.5'))
#Result('pi-cut-mag','pi-cut-fft',fftgrey('Magnitude','mean=y clip=0.5'))
#Result('pi-cut-check-mag','pi-cut-check-fft',fftgrey('Magnitude','mean=y clip=0.5'))

### Explore phase

#Result('path-sum-mag','pi_erfi_fft_4_path-sum',fftgrey('Magnitude','mean=y clip=0.5'))
#Result('pi-cut-mag','pi-cut-fft',fftgrey('Magnitude','mean=y clip=0.5'))
#Result('pi-cut-check-mag','pi-cut-check-fft',fftgrey('Magnitude','mean=y clip=0.5'))

#Flow('pi-cut','pi-cut-fft',
#               '''
#               fft3 axis=2 inv=y |
#               fft1 axis=1 inv=y |
#               t2warp inv=y
#               ''')

#Flow('pi-cut','pi-cut-fft',invftexp)

#Flow('pi-cut-check','pi-cut-check-fft',
#               '''
#               fft3 axis=2 inv=y |
#               fft1 axis=1 inv=y |
#               t2warp inv=y
#               ''')

#Flow('pi-cut-check','pi-cut-check-fft',invftexp)

#Result('pi-cut',grey_ng(''' Path Integral Image Cut '''))
#Result('pi-cut-check',grey_ng(''' Path Integral Image Check '''))

# dont know how to convolve it
#Flow('pi-conv',['diffr_data','filter-time'],'convolve2 flt=${SOURCES[1]}')

End()