from rsf.proj import *
from math import pi

# Synthetic signal
Flow('zero',None,'spike n1=300 nsp=2 k1=100,200 mag=1,2 | ricker1 frequency=10')
Flow('hilb','zero','envelope hilb=y')
Flow('envp','zero','envelope')

Plot('zero','zero envp','cat axis=2 ${SOURCES[1]} | graph title="Zero Phase Signal" dash=0,1')
Plot('hilb','hilb envp','cat axis=2 ${SOURCES[1]} | graph title="Hilbert Transform" dash=0,1')

Result('hilb','zero hilb','OverUnderAniso')

# Impulse response
shift = 2*pi*256
Flow('nfreq',None,
     '''
     math type=complex
     n1=512 d1=0.00195312 o1=-0.5 label1=Frequency unit1=cycle
     output="-I*x1/abs(x1)"
     ''')
Flow('ntime','nfreq',
     '''
     math output="input*exp(I*x1*%g)" |
     put fft3_n1=512 fft3_o1=-256 |
     fft3 axis=1 inv=y |
     real
     ''' % shift)

Flow('spike',None,'spike n1=512 k1=257 d1=1 o1=-256')

fft = '''
rtoc |
fft3 axis=1 pad=1 opt=n |
math output="input*exp(-I*x1*%g)"
''' % shift

for order in (10,100):
    time = 'time%d' % order
    freq = 'freq%d' % order
    Flow(time,'spike',
         'envelope hilb=y order=%d' % order)
    Flow(freq,time,fft)

Result('freq','nfreq freq10 freq100',
       '''
       cat axis=2 ${SOURCES[1:3]} | imag |
       reverse which=1 |
       dots labels=ideal:order=10:order=100
       ''')
Result('time','ntime time10 time100',
       '''
       cat axis=2 ${SOURCES[1:3]} |
       window min1=-33 max1=33 |
       dots unit1=sample labels=ideal:order=10:order=100
       ''')

ref = [0.6,0.8,1.0]
labels = ':'.join(map(lambda x: 'ref=%g' % x,ref))

for r in range(3):
    time = 'time%d' % r
    freq = 'freq%d' % r
    Flow(time,'spike','envelope hilb=y ref=%g' % ref[r])
    Flow(freq,time,fft)

Result('rfreq','freq0 freq1 freq2',
       '''
       cat axis=2 ${SOURCES[1:3]} | imag |
       reverse which=1 |
       dots labels=%s
       ''' % labels)
Result('rtime','time0 time1 time2',
       '''
       cat axis=2 ${SOURCES[1:3]} |
       window min1=-33 max1=33 |
       dots unit1=sample labels=%s
       ''' % labels)

SU = os.environ.get('CWPROOT')

if SU:
    suhilb = os.path.join(SU,'bin','suhilb')

    Flow('tspike','spike','segyheader')
    Flow('spike.su','spike tspike',
         'suwrite suxdr=y tfile=${SOURCES[1]}')
    Flow('hilb.su','spike.su',suhilb)
    Flow('sutime','hilb.su',
         'suread suxdr=y read=d | put d1=1 o1=-256')
    Flow('sufreq','sutime',fft)

    Result('sufreq','sufreq freq100',
       '''
       cat axis=2 ${SOURCES[1]} | imag |
       reverse which=1 |
       dots labels=SU:Madagascar
       ''')
    Result('sutime','sutime time100',
       '''
       cat axis=2 ${SOURCES[1]} |
       window min1=-33 max1=33 |
       dots unit1=sample labels=SU:Madagascar
       ''')

End()