from rsf.proj import *

# Generate a reflector model

layers = (
     ((0,2),(3.5,2),(4.5,2.5),(5,2.25),
       (5.5,2),(6.5,2.5),(10,2.5)),
     ((0,2.5),(10,3.5)),
     ((0,3.2),(3.5,3.2),(5,3.7),
       (6.5,4.2),(10,4.2)),
     ((0,4.5),(10,4.5))
)

nlays = len(layers)
for i in range(nlays):
     inp = 'inp%d' % i
     Flow(inp+'.asc',None,
          '''
          echo %s in=$TARGET
          data_format=ascii_float n1=2 n2=%d
          ''' % \
          (' '.join(map(lambda x: ' '.join(map(str,x)),
                           layers[i])),len(layers[i])))

dim1 = 'o1=0 d1=0.05 n1=201'
Flow('lay1','inp0.asc',
     'dd form=native | spline %s fp=0,0' % dim1)
Flow('lay2',None  ,
     'math %s output="2.5+x1*0.1" '      % dim1)
Flow('lay3','inp2.asc',
     'dd form=native | spline %s fp=0,0' % dim1)
Flow('lay4',None  ,'math %s output=4.5'  % dim1)

Flow('lays','lay1 lay2 lay3 lay4',
     'cat axis=2 ${SOURCES[1:4]}')

graph = '''
graph min1=2.5 max1=7.5 min2=0 max2=5
yreverse=y wantaxis=n wanttitle=n screenratio=1
'''
Plot('lays0','lays',graph + ' plotfat=10 plotcol=0')
Plot('lays1','lays',graph + ' plotfat=2 plotcol=7')
Plot('lays2','lays',graph + ' plotfat=2')

# Velocity

Flow('vofz',None,
     '''
     math output="1.5+0.25*x1"
     d2=0.05 n2=201 d1=0.01 n1=501
     label1=Depth unit1=km
     label2=Distance unit2=km
     label=Velocity unit=km/s
     ''')
Plot('vofz',
     '''
     window min2=2.75 max2=7.25 |
     grey color=j allpos=y bias=1.5
     title=Model screenratio=1
     ''')

Result('lmodel','vofz lays0 lays1','Overlay')

# Exploding reflector

Flow('expl','lays vofz',
     '''
     unif2 d1=0.01 n1=501 v00=1,2,3,4,5 |
     depth2time velocity=${SOURCES[1]} |
     ai2refl | ricker1 frequency=10 |
     time2depth velocity=${SOURCES[1]} |
     put label1=Depth unit1=km 
     label2=Distance unit2=km
     ''')

Result('expl',
       '''
       window max1=5 min2=2.5 max2=7.5 | 
       grey title="Exploding Reflector" screenratio=1
       ''')

# Modeling
##########

proj = Project()
prog = proj.Program('phaseshift.c')

# Cosine Fourier transform
Flow('kexpl','expl','cosft sign2=1 | rtoc')

Flow('kmodl','kexpl %s vofz' % prog[0],
     '''
     ./${SOURCES[1]} vel=${SOURCES[2]} 
     nw=247 dw=0.16276
     ''')

Flow('modl','kmodl',
     'pad n1=769 | fft1 inv=y | cosft sign2=-1')

Result('modl',
       '''
       grey title="Exploding Reflector Modeling" 
       label1=Time unit1=s
       ''')

# Migration
###########

Flow('kdata','modl',
     'cosft sign2=1 | fft1 | window n1=247')

Flow('kmigr','kdata %s vofz' % prog[0],
     './${SOURCES[1]} vel=${SOURCES[2]} adj=1')

Flow('migr','kmigr','real | cosft sign2=-1')

# !!! UNCOMMENT BELOW !!!

#Result('migr',
#       '''
#       window max1=5 min2=2.5 max2=7.5 | 
#       grey title="Exploding Reflector Migration" 
#       screenratio=1 label1=Depth unit1=km
#       ''')

# Least-Squares Migration
#########################

Flow('kinvs','kdata %s vofz kmigr' % prog[0],
     '''
     cconjgrad ./${SOURCES[1]} split=2 
     nw=247 dw=0.16276
     vel=${SOURCES[2]} mod=${SOURCES[3]} niter=10
     ''')

Flow('invs','kinvs','real | cosft sign2=-1')

# !!! UNCOMMENT BELOW !!!

#Result('invs',
#       '''
#       window max1=5 min2=2.5 max2=7.5 | 
#       grey title="Exploding Reflector LS Migration" 
#       screenratio=1 label1=Depth unit1=km
#       ''')


End()