from rsf.recipes.xproj import *
from rsf.recipes.beg import server
import math

approx=(1,5)

Fetch('filt_mig.sgy','teapot',server)

Flow('mig tmig mig.asc mig.bin','filt_mig.sgy',
     'segyread tfile=${TARGETS[1]} hfile=${TARGETS[2]} bfile=${TARGETS[3]}')

Flow('cube mask','mig','intbin xk=ep yk=tracf mask=${TARGETS[1]}')

Flow('cubew','cube','window n1=751')
Flow('slice','cubew','window n1=1 f1=420')

az = 70*math.pi/180

Flow('x','slice','math output="%g+(%g)*x1+(%g)*x2" ' % (345.0/2,math.cos(az),math.sin(az)))
Flow('y','slice','math output="%g+(%g)*x1+(%g)*x2" ' % (188.0/2,-math.sin(az),math.cos(az)))

Flow('xy','x y','cat axis=3 ${SOURCES[1]} | put n1=64860 n2=1 | window | transp')

Flow('cuber','cubew xy',
     '''
     window min1=0.45 n1=512 | put n2=64860 n3=1 | transp memsize=1000 |
     bin xkey=0 ykey=1 head=${SOURCES[1]} 
     ny=256 nx=64 x0=310 dx=1 y0=-150 dy=1 interp=2 |
     transp plane=13 memsize=1000 | put o2=1 o3=1 unit2= unit3=
     ''')

def cubeplot(frame=[420,199,99],title='',clip='',extra=''):
    return '''
    byte gainpanel=all %s |
    grey3 frame1=%d frame2=%d frame3=%d flat=y 
    label1=Time unit1=s label2=Inline label3=Crossline
    title="%s" %s point1=0.7 point2=0.7 flat=n
    ''' % (clip,frame[0],frame[1],frame[2],title,extra)

Result('cuber',cubeplot([195,125,25]))


def Xseislet(target, data, dip, order=1, niter=5):
        Tflow(target+'dip', data, dip+' order=%d niter=%d '%(order, niter))
        Flow(target+'dip1', target+'dip',
                'window n4=1')
        Flow(target+'dip2',target+'dip',
                'window f4=1')
        Flow(target+'dip2t',target+'dip2',
                ' transp plane=23 ')
        Tflow(target+'seis1', [data, target+'dip1'],
                'seislet dip=${SOURCES[1]} type=b eps=0.01 adj=y inv=y unit=y')
        Flow(target+'seis2', [target+'seis1', target+'dip2t'],
                '''
                transp plane=23 
                | seislet dip=${SOURCES[1]} type=b eps=0.01 adj=y inv=y unit=y 
                | transp plane=23
                ''')
        for i1 in approx:
                Flow(target+'approx%d'%i1,
                        [target+'seis2', target+'dip2t', target+'dip1'],
                        '''
                        threshold pclip=%d 
                        | transp plane=23 
                        | seislet dip=${SOURCES[1]} type=b eps=0.01 adj=n inv=y unit=y 
                        | transp plane=23 
                        | seislet dip=${SOURCES[2]} type=b eps=0.01 adj=n inv=y unit=y
                '''%i1)
        Flow(target+'pwd', [data, target+'dip' ],
                '''
                pwd dip=${SOURCES[1]} order=%d
                | math output="input*input"
                | stack norm=n
                | stack norm=n
                | stack norm=n axis=1
                | math output="sqrt(input/%g)"
                '''%(order, 512*256*64 ) )
        Flow(target+'ncc', [target+'approx%d'%(approx[0]), target+'approx%d'%(approx[1]),  data],
                '''
                cat axis=4 ${SOURCES[1]}
                | comp mode=0 ref=${SOURCES[2]}
                ''' )
        Flow(target, [target+'dip_runtime', target+'pwd', target+'ncc'],
                'cat axis=1 ${SOURCES[1:3]} ')


Xseislet('xfdip', 'cuber', 'fdip rect1=10 rect2=10 rect3=5 ', order=1, niter=0)
Xseislet('xdip1', 'cuber', 'dip rect1=10 rect2=10 rect3=5 ', order=1, niter=6)
Xseislet('xdip2', 'cuber', 'dip rect1=10 rect2=10 rect3=5 ', order=2, niter=5)

Flow('xseis','xfdip xdip1 xdip2',
        'cat axis=2 ${SOURCES[1:3]} o=0 d=1 ')

# 2D Seislet transform
str = 'xdip1'
Result('dip2d',str+'dip1',cubeplot([195,5,25],
        clip='bar=bar1.rsf', extra='color=j'))
Result('seis2d',str+'seis1',cubeplot([195,5,25]))
Flow(str+'approx', [str+'seis1', str+'dip1'],
        '''
        threshold pclip=10
        | seislet dip=${SOURCES[1]} type=b eps=0.01 adj=n inv=y unit=y 
        ''')
Result('approx2d', str+'approx', cubeplot([195,5,25]))



# 3D seislet transform

Result('fdip1', 'xfdipdip1',cubeplot([195,125,25],
        clip='bar=bar1.rsf clip=2.6 minval=-2.6 maxval=2.6',
        extra='''
                color=j scalebar=y 
                barlabel="Inline Slpoe" barunit="dt/dx" 
                '''))
Result('fdip2', 'xfdipdip2',cubeplot([195,125,25],
        clip='bar=bar2.rsf clip=2.6 minval=-2.6 maxval=2.6',
        extra='''
                color=j scalebar=y 
                barlabel="Crossline Slpoe" barunit="dt/dy" 
                '''))


Result('xdip2seis1',cubeplot([195,1,1],extra='label2="Inline Scale"'))
Result('xdip2seis2',cubeplot([195,1,1],
        extra='label2="Inline Scale" label3="Crossline Scale"'))


Result('fseis1','xfdipseis1',
        cubeplot([195,2,1],extra='label2="Inline Scale"'))
Result('fseis2','xfdipseis2',
        cubeplot([195,2,1], extra='label2="Inline Scale" label3="Crossline Scale"'))

for i1 in range(2):
        Result('fapprox%d'%(i1+1), 'xfdipapprox%d'%(approx[i1]) ,
                cubeplot([195,125,25]))

End()