from rsf.proj import *

from math import pi
import os.path

env = Environment()

gray3 = 'byte gainpanel=all | grey3 flat=y frame1=27 frame2=125 frame3=100 point1=.2 label1=Time label2=IL label3=XL unit1=s unit2= unit3= title=%s'
color3 = 'byte bar=bar.rsf gainpanel=all | grey3 color=j flat=y scalebar=y frame1=27 frame2=125 frame3=100 point1=.2 label1=Time label2=IL label3=XL unit1=s unit2= unit3= title=%s barlabel=%s'
white3 = 'byte gainpanel=all allpos=y polarity=y | grey3 flat=y frame1=27 frame2=125 frame3=100 point1=.2 label1=Time label2=IL label3=XL unit1=s unit2= unit3= title=%s'
black3 = 'byte gainpanel=all allpos=y | grey3 flat=y frame1=27 frame2=125 frame3=100 point1=.2 label1=Time label2=IL label3=XL unit1=s unit2= unit3= title=%s'
#gray = 'window f1=27 n1=1 | grey label1=IL label2=XL unit1= unit2= title=%s'
#white = 'window f1=27 n1=1 | grey allpos=y polarity=y label1=IL label2=XL unit1= unit2= title=%s'
#black = 'window f1=27 n1=1 | grey allpos=y label1=IL label2=XL unit1= unit2= title=%s'
#color = 'window f1=27 n1=1 | grey color=j scalebar=y label1=IL label2=XL unit1= unit2= barunit= title=%s barlabel=%s'

Fetch('penobscot_subset.bin','data',
      server='https://github.com',
      top='seg/tutorials-2015/raw/master/1512_Semblance_coherence_and_discontinuity/')

Flow('penobs','penobscot_subset.bin',
     '''
     echo in=$SOURCE n1=55 n2=200 n3=250
     o1=0.116 d1=0.004 unit1=s label1=Time
     o2=1253 d2=1 label2=Crossline
     o3=1207 d3=1 label3=Inline
     data_format=native_short |
     dd type=float | transp plane=23
     ''',stdin=0)
Result('penobs',gray3 % '"Seismic data"')

Flow('dip','penobs','fdip n4=2 rect1=3 rect2=20 rect3=10')
Flow('idip','dip','window n4=1')
Flow('xdip','dip','window f4=1')
Flow('xdip23','xdip','transp plane=23')
Result('idip',color3 % ('"IL Dip"','"Slope"') )
Result('xdip',color3 % ('"XL Dip"','"Slope"') )

Flow('pen','penobs dip','pwspray2 dip=${SOURCES[1]} ns2=1 ns3=4 | transp | median')
Result('pen',gray3 % '"Seismic data"')

#############
### SOBEL ###
#############
Flow('is','pen','grad3 dim=2 | smooth rect3=3')
Flow('s','pen is','grad3 dim=3 | smooth rect2=3 | math x=${SOURCES[1]} output="input*input+x*x"')
Result('s',white3 % '"Flat Sobel"')

###############
### PWSOBEL ###
###############
Flow('ipwd','pen idip','pwd dip=${SOURCES[1]} n4=0')
Flow('xpwd','pen xdip','pwd dip=${SOURCES[1]} n4=1')
Result('ipwd',gray3 % '"IL PWD"' )
Result('xpwd',gray3 % '"XL PWD"' )

Flow('isobel','ipwd xdip23','transp plane=23 | pwsmooth dip=${SOURCES[1]} ns=3 | transp plane=23')
Flow('xsobel','xpwd idip','pwsmooth dip=${SOURCES[1]} ns=3')
Result('isobel',white3 % '"IL Sobel"' )
Result('xsobel',white3 % '"XL Sobel"' )

Flow('sobel','isobel xsobel','math x=${SOURCES[1]} output="input*input+x*x" | smooth rect1=3')
Result('sobel',white3 % '"Plane-wave Sobel"' )

Flow('ipwd2','sobel idip','pwd dip=${SOURCES[1]} n4=0')
Flow('xpwd2','sobel xdip','pwd dip=${SOURCES[1]} n4=1')
Flow('isobel2','ipwd2 xdip23','transp plane=23 | pwsmooth dip=${SOURCES[1]} ns=3 | transp plane=23 | clip2 lower=0')
Flow('xsobel2','xpwd2 idip','pwsmooth dip=${SOURCES[1]} ns=3 | clip2 lower=0')
Flow('sobel2','isobel2 xsobel2','math x=${SOURCES[1]} output="input*input+x*x" | smooth rect1=3')
Result('sobel2',white3 % '"Cascaded Sobel"' )

n1 = 55
na = 45
picks = []
slices = []

for i in range(n1):
    isobelt = 'isobelt-%d' % i
    xsobelt = 'xsobelt-%d' % i
    Flow(isobelt,'isobel2','window f1=%d n1=1' % i )
    Flow(xsobelt,'xsobel2','window f1=%d n1=1' % i )
    asobels = []
    for j in range(0,360,8):
        a = j*pi/180
        asobel = 'asobel-%d-%d' % (i,j)
        Flow(asobel,[isobelt,xsobelt],'math x=${SOURCES[1]} output="abs(input*cos(%f)+x*sin(%f))"' % (a,a) )
        asobels.append(asobel)

    asobelt = 'asobel-%d' % i
    Flow(asobelt,asobels,'cat ${SOURCES[1:%d]} axis=3 | put d3=8 o3=0 | scale axis=3 | transp plane=23' % na )
#    env.AddPostAction(asobelt+'.rsf','/bin/bash ./ignore.sh %s' % asobelt )

    ################
    pick = 'pick-%d' % i
    Flow(pick,asobelt,'pick vel0=180 rect1=3 rect2=10 an=1')
    picks.append(pick)

    slice = 'slice-%d' % i
    Flow(slice,[asobelt,pick],'slice pick=${SOURCES[1]}')
    slices.append(slice)

Flow('picks',picks,'cat ${SOURCES[1:%d]} axis=3 | transp plane=23 | transp' % n1 )
Result('picks','byte bar=bar.rsf gainpanel=all bias=180 | grey3 color=j scalebar=y frame1=104 frame2=200 frame3=100 flat=y point1=.3 label1=Time label2=IL label3=XL unit1=s unit2= unit3= title=Azimuth barlabel=Azimuth')
#env.AddPostAction('picks.rsf','/bin/bash ./ignore.sh isobelt')
#env.AddPostAction('picks.rsf','/bin/bash ./ignore.sh xsobelt')
#env.AddPostAction('picks.rsf','/bin/bash ./ignore.sh asobel')
#env.AddPostAction('picks.rsf','/bin/bash ./ignore.sh pick')

Flow('slices',slices,'cat ${SOURCES[1:%d]} axis=3 | transp plane=23 | transp | smooth rect1=3' % n1 )
Result('slices',white3 % '"Azimuthal Sobel"')
#env.AddPostAction('slices.rsf','/bin/bash ./ignore.sh slice')

#####################
##### COHERENCE #####
#####################
for case in range(3):
    coh = 'coh%d' % case
    Flow(coh,'pen','coherence mode=c%d twod=n' % (case+1))
    Result(coh,black3 % ('"Cross-correlation"','Semblance','Eigenstructure')[case])

# Gradient structure tensor
Flow('dz','pen','deriv order=1')
Flow('dx','pen','transp | deriv order=1 | transp')
Flow('dy','pen','transp plane=23 | deriv order=1 | transp plane=23')
Flow('grad','dz dx dy',
     '''
     cat axis=4 ${SOURCES[1:3]} | 
     transp plane=34 | transp plane=23 | transp plane=12 | spray axis=2 n=1
     ''')
Flow('gradt','grad','transp')
Flow('gst','grad gradt','matrix B=${SOURCES[1]} | smooth rect3=3 rect4=3 rect5=3')

# Singular values
Flow('eig','gst','svd')
Flow('eig1','eig','window n1=1')
Flow('eig2','eig','window n1=1 f1=1')
Flow('coh','eig1 eig2','math l1=${SOURCES[0]} l2=${SOURCES[1]} output="(l1-l2)/(l1+l2)" ')
Result('coh',black3 % '"Gradient structure tensor"')

# Predictive coherence
Flow('dip2','pen','dip rect1=5 rect2=20 rect3=20')
Flow('pwspray','pen dip2','pwspray2 ns2=2 ns3=2 dip=${SOURCES[1]}')

box=3

Flow('pwdiff','pen pwspray',
     'spray axis=2 n=25 | math s=${SOURCES[1]} output="(s-input)^2" ')

Flow('pcoh','pwdiff',
     '''
     transp plane=12 | spray axis=2 n=1 | put n1=5 n2=5 | 
     boxsmooth rect1=%d | boxsmooth rect2=%d | 
     window f1=%d f2=%d | 
     stack axis=2 min=y | stack axis=1 min=y 
     ''' % (box,box,box-1,box-1))

Result('pcoh',white3 % 'Predictive' )

End()