from rsf.proj import *

####################### EDIT TO CHOOSE BETWEEN ########################
################ MADAGASCAR AND STANDALONE INSTALLS ###################
asg = 'asg'
stdmdl='standardmodel'
################################# OR ##################################
RSFSRC          = os.getenv('RSFSRC')
#asg             = os.path.join(RSFSRC,'trip/iwave/asg/main/asg.x')
#stdmdl          = os.path.join(RSFSRC,'trip/iwave/demo/main/standardmodel.x')
#######################################################################

# dummy target
Flow('cout0.txt',None,'touch $TARGET')

choose = {
'bm': '3',
'by': '4'
}

label = {
'by': 'Buoyancy',
'bm': 'Bulk Modulus'
}

unit = {
'by': 'cm\^3\_/g',
'bm': 'GPa'
}

for m in ('bm','by'):
    # Get model
    mod = m + '0.rsf'
    # Fetch(mod,'iwave')

    Flow([mod, mod +'@'], None,
         stdmdl + ' model=8 hfile=' + mod + ' choose='
         + choose[m] +
         '''
         d1=2.5 d2=2.5 d3=1.0 
         n1=721 n2=3121 n3=1 
         f3=3300
         label1=Depth unit1=m
         label2=Distance unit2=m
         label=
         '''
         + label[m] +
         '''
         unit=
         '''
         + unit[m] +
         '''
         datapath=.
         ''',stdin=0,stdout=0)

    j = 2 # subsampling
    for r in range(1,4):
        modrsf = '%s%d' % (m,r)
        Flow(modrsf,mod,
             '''
             dd form=native | window j1=%d j2=%d |
             put label1=Depth unit1=m label2=Distance unit2=m
             label=%s unit="%s"
             ''' % (j,j,
                    ('Bulk Modulus','Buoyancy')[m=='by'],
                    ('GPa','cm\^3\_/g')[m=='by']))

        Result(modrsf,'grey color=c mean=y title="%g-m model" scalebar=y barreverse=y' % (2.5*j))

        j *= 2

Flow('data',None,'spike n1=1501 n2=301 d1=0.002 d2=1 o2=0')
Flow('slice','data','window n1=1')

hkeys = dict(sx=3300,gx='100+20*x1',delrt=0,selev=-40,gelev=-20)
hks = hkeys.keys()

i=0
hstr = ''
for hk in hks:
    Flow(hk,'slice','math output=%s | dd type=int' % str(hkeys[hk]))
    i += 1
    hstr += '%s=${SOURCES[%d]} ' % (hk,i)

Flow('tdata',['data']+hks,'segyheader ' + hstr)
Flow('hdr.su','data tdata',
     'suwrite tfile=${SOURCES[1]} endian=0')

### create wavelet, src file

import math

fpeak=5
t0 = 1.0/fpeak
s  = 1.0 /(math.sqrt(2)*math.pi*fpeak)

# Gaussian derivative wavelet (type=gaussd in suwaveform)
Flow('wavelet',None,
     '''
     math n1=101 d1=0.004 
     output="%g*(x1-%g)*exp(-%g*(x1-%g)^2)" | put o1=-0.1
     ''' % (-1.0/(s*s*s*math.sqrt(2*math.pi)),t0,1/(2*s*s),t0))
Flow('twavelet','wavelet','segyheader')
Flow('wavelet_base.su','wavelet twavelet','suwrite endian=0 tfile=${SOURCES[1]}')

Flow('src.su', ['wavelet_base.su', 'hdr.su'],
      '''   
      towed_array
      data=${SOURCES[1]} 
      src=${SOURCES[0]}
      towed=${TARGETS[0]}
      ''',
      stdin=0, stdout=0)

# movie for PML demo
Flow('bmh', None,
     '''
     makevel n1=91 n2=391 
     d1=20.0 d2=20.0 v000=2.25 | 
     sfput dim=2 gdim=2 id1=0 id2=1
     ''',
     stdin=0)

Flow('byh', None,
     '''
     makevel n1=91 n2=391 
     d1=20.0 d2=20.0 v000=1.0 | 
     sfput dim=2 gdim=2 id1=0 id2=1
     ''',
     stdin=0)

Flow('movie1','bmh byh src.su',
     '''
     makevel n1=91 n2=391 n3=41 
     d1=20.0 d2=20.0 d3=100.0 v000=0.0 | 
     sfput dim=2 gdim=3 id1=0 id2=1 id3=2 > ${TARGETS[0]} &&
     ''' +
     asg +
     '''
     bulkmod=${SOURCES[0]} buoyancy=${SOURCES[1]} 
     source_p=${SOURCES[2]} movie_p=${TARGETS[0]} 
     deriv=0 adjoint=0 order=1 cfl=0.5 cmin=1.0 cmax=6.0 
     dmin=0.8 dmax=3.0
     ''', stdin=0, stdout=0)

Flow('movie2','bmh byh src.su',
     '''
     makevel n1=91 n2=391 n3=41 
     d1=20.0 d2=20.0 d3=100.0 v000=0.0 | 
     sfput dim=2 gdim=3 id1=0 id2=1 id3=2 > ${TARGETS[0]} &&
     ''' +
     asg +
     '''
     bulkmod=${SOURCES[0]} buoyancy=${SOURCES[1]} 
     source_p=${SOURCES[2]} movie_p=${TARGETS[0]} 
     deriv=0 adjoint=0 order=1 cfl=0.5 cmin=1.0 cmax=6.0 
     dmin=0.8 dmax=3.0 nr1=250 nl2=250 nr2=250 pmlampl=1.0
     ''', stdin=0, stdout=0)

Flow('movie3','bmh byh src.su',
     '''
     makevel n1=91 n2=391 n3=41 
     d1=20.0 d2=20.0 d3=100.0 v000=0.0 | 
     sfput dim=2 gdim=3 id1=0 id2=1 id3=2 > ${TARGETS[0]} &&
     ''' +
     asg +
     '''
     bulkmod=${SOURCES[0]} buoyancy=${SOURCES[1]} 
     source_p=${SOURCES[2]} movie_p=${TARGETS[0]} 
     deriv=0 adjoint=0 order=1 cfl=0.5 cmin=1.0 cmax=6.0 
     dmin=0.8 dmax=3.0 nr1=100 nl2=100 nr2=100 pmlampl=1.0
     ''', stdin=0, stdout=0)

Flow('frame13','movie1','window f3=12 n3=1')
Flow('frame40-1','movie1','window f3=40 n3=1')
Flow('frame40-2','movie2','window f3=40 n3=1')
Flow('frame40-3','movie3','window f3=40 n3=1')

Result('frame13',  'grey mean=y clip=0.01 scalebar=y barreverse=y wanttitle=n')
Result('frame40-1','grey mean=y clip=0.01 scalebar=y barreverse=y wanttitle=n')
Result('frame40-2','grey mean=y clip=0.01 scalebar=y barreverse=y wanttitle=n')
Result('frame40-3','grey mean=y clip=0.01 scalebar=y barreverse=y wanttitle=n')

# 2-4 results
j = 1
traces = []
for r in range(4):

    by = 'by%d.rsf' % r
    bm = 'bm%d.rsf' % r

    data = 'data%d' % r
    sudata = data + '.su'

    Flow(sudata,'data tdata' + ' ' + bm + ' ' + by + ' src.su',
         'suwrite < ${SOURCES[0]} tfile=${SOURCES[1]} endian=0 > ${TARGETS[0]} && ' +
         asg +
         ''' 
         bulkmod=${SOURCES[2]} buoyancy=${SOURCES[3]} 
         source_p=${SOURCES[4]} data_p=${TARGETS[0]} 
         deriv=0 adjoint=0 order=2 cfl=0.75 cmin=1.0 
         cmax=4.5 dmin=0.8 dmax=3.0
         nr1=500 nl2=500 nr2=500 pmlampl=0.1
         ''',
         stdin=0, stdout=0)

    Flow(data,sudata,
         'suread read=data endian=0')

    Result(data,'grey title="%g-m data, 2-4 scheme" ' % (2.5*j))

    j *= 2

    trace = 'trace%d' % r
    Flow(trace,data,'window n2=1 f2=100')
    traces.append(trace)

Result('trace',traces,
       '''
       cat axis=2 ${SOURCES[1:4]} | 
       graph plotcol=4,2,6,7 wanttitle=n label2=Pressure unit2=MPa
       ''')

Result('wtrace',traces,
       '''
       cat axis=2 ${SOURCES[1:4]} | window min1=1.8 max1=2.5 |
       graph plotcol=4,2,6,7 wanttitle=n label2=Pressure unit2=MPa
       ''')

# 2-8 results
j = 1
traces8k = []
for r in range(4):

    by = 'by%d' % r
    bm = 'bm%d' % r

    data = 'data8k%d' % r
    sudata = data + '.su'

    Flow(sudata,'data tdata' + ' ' + bm + ' ' + by + ' src.su',
         'suwrite < ${SOURCES[0]} tfile=${SOURCES[1]} endian=0 > ${TARGETS[0]} && ' +
         asg +
         ''' 
         bulkmod=${SOURCES[2]} buoyancy=${SOURCES[3]} 
         source_p=${SOURCES[4]} + data_p=${TARGETS[0]} 
         deriv=0 adjoint=0 order=2 cfl=0.75 cmin=1.0 
         cmax=4.5 dmin=0.8 dmax=3.0
         nr1=500 nl2=500 nr2=500 pmlampl=0.1
         ''',
         stdin=0, stdout=0)

    Flow(data,sudata,
         'suread read=data endian=0')

    trace = 'trace' + str(r+4)
    Flow(trace,data,'window n2=1 f2=100')
    traces8k.append(trace)

    Result(data,'grey title="%g-m data, 2-8 scheme" ' % (2.5*j))

    j *= 2

Result('trace8k',traces8k,
   '''
   cat axis=2 ${SOURCES[1:4]} | 
   graph plotcol=4,2,6,7 wanttitle=n label2=Pressure unit2=MPa
   ''')
Result('wtrace8k',traces8k,
   '''
   cat axis=2 ${SOURCES[1:4]} | window min1=1.8 max1=2.5 |
   graph plotcol=4,2,6,7 wanttitle=n label2=Pressure unit2=MPa
   ''')



End()