from rsf.proj import *
import math

def numpy_load(target=None,source=None,env=None):
    'convert from numpy to RSF format'
    import numpy, m8r
    data = numpy.load(str(source[0]))
    rsf = m8r.Output(str(target[0]))
    rsf.put("n1",len(data))
    rsf.put("o1",0)
    rsf.put("d1",0.002)
    rsf.put("label1","Time")
    rsf.put("unit1","s")
    rsf.write(data)
    return 0

for case in ('rpp','seismic'):
    npy = case+'.npy'
    Fetch(npy,'1606_Wavelet_estimation',
          server='https://raw.githubusercontent.com',
          top='seg/tutorials-2016/master')

    Command(case+'.rsf',npy,action=Action(numpy_load))
    Plot(case,
         '''
         window max1=1.5 | 
         graph title="%s"
         ''' % ('Well-Log Reflectivity','Seismic')[case=='seismic'])

Result('data','rpp seismic','OverUnderAniso')

# Compute spectrum

Flow('spec','seismic','spectra')

Plot('spec',
     '''
     math output="20*log(input)/log(10)" | 
     graph label2=power unit2=db title="Amplitude Spectrum" 
     min1=0 max1=250 min2=-110 max2=-25
     ''')

Flow('line.asc',None,
     '''
     echo 5 -80 80 -35 130 -35 160 -80 n1=2 n2=4 data_format=ascii_int in=$TARGET
     ''')
Flow('line','line.asc','dd type=complex form=native | window')

Plot('line',
     '''
     graph plotcol=5 plotfat=5 wantaxis=n wanttitle=n 
     min1=0 max1=250 min2=-110 max2=-25
     ''')

Result('spec','spec line','Overlay')

freqs = [5, 80, 130, 160]

Flow('ormsby',None,
     '''
     spike n1=129 d1=0.002 k1=65 o1=-0.128 | 
     trapez frequency=5,80,130,160 
     ''')
Result('ormsby','graph title="Ormsby wavelet" label2=Amplitude')

Flow('formsby','ormsby',
     'pad n1=1024 beg1=460 | fft1 | real | put fft_o1=0')

# Compare with autocorrelation

Flow('auto','seismic','fft1 | math output="input*conj(input)" | fft1 inv=y')
Flow('auto2','auto','window n1=65 | reverse which=1 | window n1=64 | cat $SOURCE axis=1 | window n1=129')

Result('auto2','ormsby auto2',
       'cat axis=2 ${SOURCES[1]} | scale axis=1 | graph title=Autocorrelation dash=1,0 label2=Amplitude')

# Wavelet estimation by spectral division

Flow('specrpp','rpp','spectra')
Flow('specdiv','spec specrpp',
     'divn rect1=10 den=${SOURCES[1]}')
Flow('wavdiv','specdiv',
     '''
     rtoc | math output="input*exp(I*x1*%g)" | fft1 inv=y
     ''' % (2*math.pi))
Result('wavdiv','window f1=460 n1=129 | graph title="Division Wavelet" label2=Amplitude')

Result('specdiv',
       '''
       math output="20*log(input)/log(10)" | 
       graph label2=power unit2=db title="Division Amplitude Spectrum" 
       min1=0 max1=250 min2=-60 max2=10
       ''')

Result('wavelet','wavdiv ormsby',
       '''
       window f1=460 n1=129 | cat axis=2 ${SOURCES[1]} | graph title=Wavelet label2=Amplitude
       ''')

# Create synthetic

Flow('frpp','rpp','fft1')
Plot('wrpp','rpp','window min1=0.75 max1=1.5 | graph transp=y yreverse=y title=Reflectivity')
Plot('wseis','seismic',
     '''
     window min1=0.75 max1=1.5 | spray axis=2 n=4 | 
     wiggle poly=y transp=y yreverse=y title=Data
     ''')

for case in range(2):
    wavelet = ('formsby','specdiv')[case]

    synth = 'synth%d' % case
    Flow(synth,[wavelet,'frpp'],
         'rtoc | math freq=${SOURCES[1]} output="-input*freq" | fft1 inv=1')
    Result(synth,'window max1=1.5 | graph title="Synthetic" ')

    Plot('w'+synth,synth,
         '''
         window min1=0.75 max1=1.5 | spray axis=2 n=4 | 
         wiggle poly=y transp=y yreverse=y title=Synthetic
         ''')
    Result('w'+synth,['wrpp','w'+synth,'wseis'],'SideBySideAniso',vppen='txscale=1.5')

End()