from rsf.proj import *
import math

# Download data

horizon = 'Penobscot_HorB.txt'

Fetch(horizon,'data',
      server='https://raw.githubusercontent.com',
      top='agile-geoscience/notebooks/master')

# Convert from xyz triples to horizon

Flow('xyz',horizon,
     '''
     echo n1=3 n2=254674 data_format=ascii_float in=$SOURCE |
     dd form=native
     ''')

Flow('xy','xyz','window n1=2 | dd type=int')
Flow('horizon','xyz xy',
     '''
     window f1=2 squeeze=n | 
     intbin head=${SOURCES[1]} xkey=0 ykey=1 |
     window | lapfill grad=y verb=y niter=500 |
     put label=Time unit=ms label1=Inline label2=Crossline
     ''')

Result('horizon',
     '''
     grey color=x bias=950 clip=133 scalebar=y title=Horizon 
     yreverse=n transp=n screenwd=11.9 screenht=9.26
     ''')

Flow('sobel','horizon','grad2')

Result('sobel',
       '''
       grey allpos=y title="Sobel filter"
       yreverse=n transp=n screenwd=11.9 screenht=9.26
       ''')

# Canny edge detection

Flow('square',None,
     'spike n1=128 n2=128 k1=32 l1=96 k2=32 l2=96 d1=1 d2=1')

a=math.pi/12 # 15 degrees
c=math.cos(a)
s=math.sin(a)

Flow('x','square','math output="64+(%g)*(x1-64)+(%g)*(x2-64)" ' % (c,-s))
Flow('y','square','math output="64+(%g)*(x1-64)+(%g)*(x2-64)" ' % (s, c))
Flow('xyr','x y','cat axis=3 ${SOURCES[1]}')

Flow('rotate','square xyr',
     '''
     iwarp2 warp=${SOURCES[1]} |
     smooth rect1=3 rect2=3 repeat=3 |
     noise type=n range=0.2
     ''')

Plot('rotate','grey screenratio=1 wantaxis=n title="noisy image" wheretitle=t')

Flow('canny','rotate','canny')

Plot('canny',
     '''
     dd type=float | math output=1-input | 
     grey screenratio=1 wantaxis=n title="Canny filter 1" wheretitle=t
     ''')

Flow('clean','rotate','impl2 rect1=10 rect2=10 tau=1 pclip=90')

Flow('canny2','clean','canny')

Plot('canny2',
     '''
     dd type=float | math output=1-input | 
     grey screenratio=1 wantaxis=n title="Canny filter 2" wheretitle=t
     ''')

Result('canny','rotate canny canny2','SideBySideIso',vppen='xsize=11')

# Apply to horizon

Flow('hcanny','horizon','canny')

Result('hcanny',
     '''
     dd type=float | 
     grey title="Canny filter 1" 
     yreverse=n transp=n screenwd=11.9 screenht=9.26
     ''')

# expose a parameter

rect = float(ARGUMENTS.get('rect',10))

Flow('hclean','horizon','impl2 rect1=%g rect2=%g tau=1' % (rect,rect))

Flow('hcanny2','hclean','canny | dd type=float')

Result('hcanny2',
     '''
     grey title="Canny filter 2" 
     yreverse=n transp=n screenwd=11.9 screenht=9.26
     ''')

# Overlay plot

Result('overlay','horizon hcanny2',
       '''
       math edge=${SOURCES[1]} output="input+500*(2-edge)" |
       grey color=xC bias=950 clip=170 scalebar=y title=Horizon 
       yreverse=n transp=n screenwd=11.9 screenht=9.26 cliprgb=0
       maxval=1120
       ''')

# Playing with parameters

cannys = []
for rect in range(1,21):
    canny = 'canny-%d' % rect
    Flow(canny,'horizon',
         '''
         impl2 rect1=%d rect2=%d tau=1 | canny 
         ''' % (rect,rect))
    cannys.append(canny)
Flow('cannys',cannys,
     '''
     cat axis=3 ${SOURCES[1:%d]} | 
     dd type=float | put o3=1 d3=1
     ''' % len(cannys))
Plot('cannys',
     '''
     grey title="Canny filter" 
     yreverse=n transp=n screenwd=11.9 screenht=9.26
     ''',view=1)

# Stack them

Result('cannys',
       '''
       stack axis=3 norm=n |
       grey allpos=y title="stacked Canny filter"
       yreverse=n transp=n screenwd=11.9 screenht=9.26
       ''')

End()