from rsf.proj import *
from rsf.prog import RSFROOT

# Program compilation
#####################

proj = Project()

# To do the coding assignment in Fortran,
# comment the next line and uncomment the lines below
exe = proj.Program('wave.c')

# UNCOMMENT BELOW IF YOU WANT TO USE FORTRAN
#exe = proj.Program('wave.f90',
#                   F90PATH=os.path.join(RSFROOT,'include'),
#                   LIBS=['rsff90']+proj.get('LIBS'))

# UNCOMMENT BELOW IF YOU WANT TO USE PYTHON
#exe = proj.Command('wave.exe','wave.py','cp $SOURCE $TARGET')
#AddPostAction(exe,Chmod(exe,0o755))


# Constant velocity test
########################

# Source wavelet
Flow('wavelet',None,
     '''
     spike n1=1000 d1=0.001 k1=201 |
     ricker1 frequency=10
     ''')

# Source location
Flow('source',None,
     '''
     spike n1=201 n2=301 d1=0.01 d2=0.01
     label1=x1 unit1=km label2=x2 unit2=km
     k1=101 k2=151     
     ''')

# Velocity model
Flow('v1','source','math output=1')
Flow('v2','source','math output=1.5')

# Modeling
Flow('wave','source %s wavelet v1 v2' % exe[0],
     '''
     ./${SOURCES[1]} wav=${SOURCES[2]}
     v=${SOURCES[3]}  vx=${SOURCES[4]}
     ft=200 jt=5 
     ''')

Plot('wave','grey gainpanel=all title=Wave',view=1)

Result('wave',
       '''
       window n3=1 min3=0.9 |
       grey title=Wave screenht=8 screenwd=12
       ''')

# Download Hess VTI model
#########################
zcat = WhereIs('gzcat') or WhereIs('zcat')
for case in ('vp','epsilon'):
    sgy = 'timodel_%s.segy' % case
    sgyz = sgy + '.gz'
    Fetch(sgyz,dir='hessvti',
          server='https://s3.amazonaws.com',
          top='open.source.geoscience/open_data')

    # Uncompress
    Flow(sgy,sgyz,zcat + ' $SOURCE',stdin=0)
    # Convert to RSF format
    Flow(case,sgy,
         '''
         segyread read=data | 
         window j1=2 j2=2 | put d1=40 d2=40 
         unit1=ft label1=Depth unit2=ft label2=Distance 
         ''')

# Horizontal velocity
Flow('vx','vp epsilon',
     'math e=${SOURCES[1]} output="input*sqrt(1+2*e)"')

for case in ('vp','vx'):
    Result(case,
           '''
           grey color=j pclip=100 allpos=y bias=5000 
           scalebar=y barreverse=y wanttitle=n 
           barlabel=Velocity barunit=ft/s 
           screenht=5 screenwd=12 labelsz=6
           ''')

Flow('hsource','vp','spike k1=300 k2=900')
Flow('hess','hsource %s wavelet vp vx' % exe[0],
     '''
     ./${SOURCES[1]} wav=${SOURCES[2]}
     v=${SOURCES[3]}  vx=${SOURCES[4]}
     ft=200 jt=5 
     ''')

End()