from rsf.proj import *

def Grey3(data,other):
        Result(data,
       '''
       byte clip=0.0024 |
       transp plane=23 |
       grey3 flat=n frame1=750 frame2=100 frame3=10 
       title=Data point1=0.8 point2=0.8  %s
       '''%other)

def Vel(data,other):
        Result(data,
     '''
     grey color=j allpos=y bias=1.5 clip=0.7
     scalebar=y barreverse=y barunit=km/s
     label2=Midpoint unit2=km label1=Time unit1=s
     title="NMO Velocity"  %s
     '''%other )

def Pick3(data,other):
        Result(data,
       '''
       byte allpos=y gainpanel=all |
       transp plane=23 |
       grey3 flat=n frame1=750 frame2=100 frame3=25 
       label1=Time unit1=s color=j framelabelcol=VP_BLACK
       label3=Velocity unit3=km/s 
       label2=Midpoint unit2=km
       title="Velocity Scan" point1=0.8 point2=0.8  %s
       '''%other)

def Greymig(data,other):
        Result(data,
       '''
                window max1=5 min2=2.5 max2=7.5 |
       grey title="Prestack Kirchhoff Time Migration" 
           label1=Depth unit1=km  %s
       '''%other)

labels=[]
# 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.001 n1=10001'

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 
'''
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.36*x1"
     d2=0.01 n2=1001 d1=0.01 n1=501
     label1=Depth unit1=km
     label2=Distance unit2=km
     label=Velocity unit=km/s
     ''')
Plot('vofz',
     '''
     window min2=2.5 max2=7.5 |
     grey color=j allpos=y bias=1.5
     title=Model 
     ''')

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

# Model data

Flow('dips','lays','deriv scale=y')
Flow('modl','lays dips',
     '''
     kirmod cmp=y dip=${SOURCES[1]} 
     nh=51  dh=0.1  h0=0
     ns=201 ds=0.05 s0=0
     freq=10 dt=0.004 nt=1501
     vel=1.5 gradz=0.36 verb=y |
     pow pow1=1 |
     put d2=0.05 
     label1=Time        unit1=s
     label2=Half-Offset unit2=km 
     label3=Midpoint    unit3=km 
     ''',split=[1,10001], reduce='add')

# Add random noise
Flow('data','modl','noise var=1e-20 seed=101013')
#Flow('data','modl','noise var=1e-20 seed=101013')
Grey3('data','title="Data (Orignal)"')

# Velocity estimation
#####################
Flow('voft','vofz',
     'depth2time velocity=$SOURCE dt=0.004 nt=1501')
Flow('vrms','voft',
     '''
     add mode=p $SOURCE | causint | 
     math output="sqrt(input*0.004/(x1+0.004))" 
     ''')

# Velocity scan
Flow('svscan','data',
     'noise var=1e-6 seed=101013 | vscan v0=1.5 dv=0.02 nv=51 semblance=y',
     split=[3,201], reduce='cat')
Pick3('svscan','title="Velocity Scan (Original)"')

# Using vnmo instead of vrms
########################
# Velocity picking
Flow('vnmo','svscan','pick rect1=100 rect2=10')

for vel in ('vrms','vnmo'):
    Plot(vel,
     '''
     grey color=j allpos=y bias=1.5 clip=0.7
     scalebar=y barreverse=y barunit=km/s
     label2=Midpoint unit2=km label1=Time unit1=s
     title="%s Velocity" 
     ''' % vel[1:].upper())
Result('vnmo','vrms vnmo','SideBySideIso')

# Stacking
##########

Flow('nmo','data vnmo','nmo velocity=${SOURCES[1]}')
Flow('stack','nmo','stack')

# Using vnmo instead of vrms
########################
Flow('nmo0','data vnmo','nmo velocity=${SOURCES[1]}')
Flow('dstack','nmo0',
     '''
     window f1=250 | 
     logstretch | fft1 | 
     transp plane=13 memsize=1000 |
     finstack | 
     transp memsize=1000 |
     fft1 inv=y | logstretch inv=y | 
     pad beg1=250 | put unit1=s
     ''')

Flow('zoff','data','window n2=1')

stacks = {
    'stack': 'Stack with NMO Velocity',
    'dstack': 'DMO Stack',
    'zoff': 'Zero Offset'
    }

for stack in stacks.keys():
    Result(stack,
           '''
           window min1=1.5 max1=5 min2=1 max2=9 | 
           grey title="%s" 
           ''' % stacks[stack])

# Common-offset domain prestack kirchhoff time migration (PSTM)
Flow('spstm','data vnmo',
     '''
         transp memsize=1000 plane=23 |
         prekirch vel=${SOURCES[1]} | stack axis=3 | halfint inv=y adj=y
     ''',split=[3,51,[0]])

## Dix inversion
Flow('semblance','svscan vnmo','slice pick=${SOURCES[1]}')
Flow('vesti','vnmo semblance','dix weight=${SOURCES[1]} rect1=100 rect2=10')

# Time-to-Depth conversion
Flow('pstmz','spstm vesti','time2depth intime=y velocity=${SOURCES[1]} dz=0.01 nz=551')
Greymig('pstmz','title="PSKTM (Original)"')

Result('spstm',['Fig/pstmz.vpl']+labels, 'Overlay')

##########################################
## The following is for blending test
##########################################
Flow('dithert',None,'math output=1 n1=51 n2=201 | noise rep=y seed=2013 var=0.5 ')
Flow('data1-dither','data dithert','datstretch inv=y datum=${SOURCES[1]} | math output="input*2"')

## Blending
Flow('data-b','data data1-dither','add scale=1,2 ${SOURCES[1]}')
Grey3('data-b','title="Data (Blended)"')

# Velocity scan
Flow('svscan-b','data-b',
     'noise var=1e-7 seed=101013 | vscan v0=1.5 dv=0.02 nv=51 semblance=y',
     split=[3,201], reduce='cat')
Pick3('svscan-b','title="Velocity Scan (Blended)"')

# Velocity picking for blended data
Flow('vnmo-b','svscan-b','pick rect1=100 rect2=10')
Vel('vnmo-b','')

# Stacking
##########
Flow('nmo-b','data-b vnmo-b','nmo velocity=${SOURCES[1]}')
Flow('stack-b','nmo-b','stack')

# DMO
########################
Flow('nmo0-b','data-b vnmo-b','nmo velocity=${SOURCES[1]}')
Flow('dstack-b','nmo0-b',
     '''
     window f1=250 | 
     logstretch | fft1 | 
     transp plane=13 memsize=1000 |
     finstack | 
     transp memsize=1000 |
     fft1 inv=y | logstretch inv=y | 
     pad beg1=250 | put unit1=s
     ''')

Grey3('nmo-b','')
Plot('stack-b','grey title="Stack (Blended)"')
Plot('dstack-b','grey title="DMO Stack (Blended)"')
Plot('spstm-datab','spstm-b','grey title="PSTM (Blended)"')
Result('comp-b','stack-b dstack-b spstm-datab','SideBySideAniso')


# Common-offset domain prestack kirchhoff time migration (PSTM)
Flow('spstm-b','data-b vnmo-b',
     '''
         transp memsize=1000 plane=23 |
         prekirch vel=${SOURCES[1]} | stack axis=3 | halfint inv=y adj=y
     ''',split=[3,51,[0]])

## Dix inversion
Flow('semblance-b','svscan-b vnmo-b','slice pick=${SOURCES[1]}')
Flow('vesti-b','vnmo-b semblance-b','dix weight=${SOURCES[1]} rect1=100 rect2=10')

# Time-to-Depth conversion
Flow('pstmz-b','spstm-b vesti-b','time2depth intime=y velocity=${SOURCES[1]} dz=0.01 nz=551')
Greymig('pstmz-b','title="PSKTM (Blended)"')

Result('spstm-b',['Fig/pstmz-b.vpl']+labels, 'Overlay')

##########################################
## The following is for deblending test
##########################################

# Stacking
##########
Flow('nmo-db','nmo-b','transp | mf nfw=9 | transp')
Flow('stack-db','nmo-db','stack')

# DMO
########################

Flow('dstack-db','nmo-db',
     '''
     window f1=250 | 
     logstretch | fft1 | 
     transp plane=13 memsize=1000 |
     finstack | 
     transp memsize=1000 |
     fft1 inv=y | logstretch inv=y | 
     pad beg1=250 | put unit1=s
     ''')

Grey3('nmo-db','')
Plot('stack-db','grey title="Stack (Blended)"')
Plot('dstack-db','grey title="DMO Stack (Blended)"')
Plot('spstm-datadb','spstm-db','grey title="PSTM (Deblended)" clip=500000')
Result('comp-db','stack-db dstack-db spstm-datadb','SideBySideAniso')


# Common-offset domain prestack kirchhoff time migration (PSTM)
Flow('data-db','nmo-db vnmo-b','inmo velocity=${SOURCES[1]}')
Grey3('data-db','title="Data (Deblended)"')

# Velocity scan
Flow('svscan-db','data-db',
     'noise var=1e-6 seed=101013 | vscan v0=1.5 dv=0.02 nv=51 semblance=y',
     split=[3,201], reduce='cat')
Pick3('svscan-db','title="Velocity Scan (Deblended)"')


# Common-offset domain prestack kirchhoff time migration (PSTM)
Flow('spstm-db','data-db vnmo-b',
     '''
         transp memsize=1000 plane=23 |
         prekirch vel=${SOURCES[1]} | stack axis=3 | halfint inv=y adj=y
     ''',split=[3,51,[0]])

# Time-to-Depth conversion
Flow('pstmz-db','spstm-db vesti-b','time2depth intime=y velocity=${SOURCES[1]} dz=0.01 nz=551')
Greymig('pstmz-db','title="PSKTM (Deblended)"')

Result('spstm-db',['Fig/pstmz-db.vpl']+labels, 'Overlay')


Result('data-comp','Fig/data.vpl Fig/data-b.vpl Fig/data-db.vpl','SideBySideIso')
Result('svscan-comp','Fig/svscan.vpl Fig/svscan-b.vpl Fig/svscan-db.vpl','SideBySideIso')
Result('spstm-comp','Fig/spstm.vpl Fig/spstm-b.vpl Fig/spstm-db.vpl','SideBySideIso')

End()