from rsf.proj import *

# Prepare source and velocity
#############################

# Source wavelet
Flow('wavelet',None,
     '''
     spike n1=1500 d1=0.001 k1=201 |
     ricker1 frequency=10
     ''')
# Source spectra
Result('spectra','wavelet',
    '''
        spectra |window n1=50 |
        graph title="Source Spectra"
        ''')

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

# Background velocity
Flow('v0','source','math output=2.')

# Velocity perturbation
Flow('detv',None,
    '''
        spike n1=501 d1=0.01 o1=0 k1=300,301,302 mag=0.02 |
        spray axis=2 n=501 d=0.01 o=0 |
        costaper nw2=25
        ''')

# True velocity
Flow('vel','v0 detv','add ${SOURCES[1]}')
Result('vel',
    '''
        grey scalebar=y barreverse=y 
        color=j bias=2.01 title="Velocity"
        ''')

# Define two string functions
#############################

# program execution script
def modeling(born):
    return'''
    ./${SOURCES[1]} wav=${SOURCES[2]}
    v=${SOURCES[3]} detv=${SOURCES[4]}
        snapshot=${TARGETS[1]} 
        born=%s ft=200 jt=20 
        nr=301 r0=100 rz=200
        ''' %born

# plot curves
def plot(title):
    return'''
        cat ${SOURCES[1]} axis=3 |
        window n2=1 f2=150 min1=1. max1=1.4 |
        graph dash=0,1 title="%s"
        label2=Amplitude unit2= 
        wheretitle=bottom wherexlabel=top 
        screenratio=0.33 screenht=4.5
        labelsz=7 titlesz=8 
        labelfat=3 plotfat=9 titlefat=4 plotcol=4,7
        ''' %title

# C code test: Born modeling
############################

# Program compilation
proj=Project()
exe_wave = proj.Program('born_c.c')

# Forward modeling (born=n, vel=vel) 
Flow('full_c full_snap_c','source %s wavelet vel detv' %exe_wave[0], modeling('n'));

# Diving wave (born=n, vel=v0)
Flow('diving_c diving_snap_c','source %s wavelet v0 detv' %exe_wave[0], modeling('n'));

# U-U0
Flow('diff_c','full_c diving_c','add ${SOURCES[1]} scale=1,-1')

# Born modeling (born=y, vel=v0)
Flow('born_c born_snap_c','source %s wavelet v0 detv' %exe_wave[0], modeling('y'));

# Comparison
Result('comp_c','diff_c born_c', plot('C: U-U\_0\^ (solid) vs. \F10 d\F3 U (dashed)'))

# C++ code test: Born modeling
############################

# Program compilation
exe_wave = proj.Program('born_cc.cc',
                     LIBS=['rsf++']+proj.get('LIBS'))

# Forward modeling (born=n, vel=vel) 
Flow('full_cc full_snap_cc','source %s wavelet vel detv' %exe_wave[0], modeling('n'));

# Diving wave (born=n, vel=v0)
Flow('diving_cc diving_snap_cc','source %s wavelet v0 detv' %exe_wave[0], modeling('n'));

# U-U0
Flow('diff_cc','full_cc diving_cc','add ${SOURCES[1]} scale=1,-1')

# Born modeling (born=y, vel=v0)
Flow('born_cc born_snap_cc','source %s wavelet v0 detv' %exe_wave[0], modeling('y'));

# Comparison
Result('comp_cc','diff_cc born_cc', plot('C++: U-U\_0\^ (solid) vs. \F10 d\F3 U (dashed)'))

# F90 code test: Born modeling
############################

# Program compilation
exe_wave = proj.Program('born_f90.f90',
                     LIBS=['rsff90']+proj.get('LIBS'))

# Forward modeling (born=n, vel=vel) 
Flow('full_f90 full_snap_f90','source %s wavelet vel detv' %exe_wave[0], modeling('n'));

# Diving wave (born=n, vel=v0)
Flow('diving_f90 diving_snap_f90','source %s wavelet v0 detv' %exe_wave[0], modeling('n'));

# U-U0
Flow('diff_f90','full_f90 diving_f90','add ${SOURCES[1]} scale=1,-1')

# Born modeling (born=y, vel=v0)
Flow('born_f90 born_snap_f90','source %s wavelet v0 detv' %exe_wave[0], modeling('y'));

# Comparison
Result('comp_f90','diff_f90 born_f90', plot('F90: U-U\_0\^ (solid) vs. \F10 d\F3 U (dashed)'))

# Concatenate the comparisons
###############################
Result('comps','Fig/comp_c.vpl Fig/comp_cc.vpl Fig/comp_f90.vpl','OverUnderAniso')

End()