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.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 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.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 labelsz=14 labelfat=2 font=4 titlesz=20 titlefat=2
     title=Model screenratio=1
     ''')

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=5 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('data1','modl','noise var=1e-20 seed=101013')              #clean data
Flow('data2','modl','noise var=1e-6 seed=101013')               #noisy data

Flow('vrms',None,
     'math d1=0.004 n1=1001 o1=0 output="x1*125+2000+50*sin(9*x1)" ')

Flow('synt',None,
     '''
     spike d1=0.004 n1=1001 |
     noise rep=y seed=2006 |
     cut n1=100 | 
     bandpass flo=4 fhi=20 |
     spray axis=2 n=141 d=25 o=0 label=Offset unit=m 
     ''')

Flow('cmpa','data1',
     'window n3=1')
Flow('cmpb','data2',
     'window n3=1')

Result('cmpa','grey title="%s" clip=0.008 labelsz=14 labelfat=2 font=4 titlesz=20 titlefat=2  label2="Offset"' % 'Clean CMP')
Result('cmpb','grey title="%s" clip=0.008 labelsz=14 labelfat=2 font=4 titlesz=20 titlefat=2  label2="Offset"' % 'Noisy CMP')

#Flow('refa','cmpa','window n2=1')
Flow('refa','abnmoa','stack')
#Flow('refa','modl','window n3=1 | window n2=1')

dat='cmpa'
for typ in ('semblance','ab','similarity-weighted','similarity-weighted'):
        scn=typ+'scna'
        if typ == 'similarity-weighted':
                Flow(scn,[dat,'refa'],'simivscan thr=0.6 v0=1.5 dv=0.02 nv=51 type=weighted weight=y nb=5 ref=${SOURCES[1]} | clip2 lower=0')
        else:
                Flow(scn,dat,'vscan v0=1.5 dv=0.02 nv=51 weight=y type=%s nb=5 | clip2 lower=0'%typ)
        Plot(scn+'0',scn,
                '''
                grey color=j allpos=y title="%s" pclip=100 labelsz=14 labelfat=2 font=4 titlesz=20 titlefat=2
                '''%typ)
        pik=typ+'pika'
        Flow(pik,scn,'scale axis=2 | pick rect1=10')
        Plot(pik,
             '''
             graph yreverse=y transp=y min2=1.5 max2=2.5 pad=n
             plotcol=7 plotfat=15 wanttitle=n wantaxis=n
             ''')
        Result(scn,[scn+'0',pik],'Overlay')
        nmo = typ + 'nmoa'
        Flow(nmo,[dat,pik],'nmo half=y velocity=${SOURCES[1]} str=0.1 ')
        Result(nmo,
                        '''
                        grey title="%s"
                        grid2=y gridcol=6 gridfat=7 g2num=0.5 clip=0.008 labelsz=14 labelfat=2 font=4 titlesz=20 titlefat=2
                        ''' % typ)

        inmo = typ + 'inmoa'
        Flow(inmo,[nmo,pik],'inmo half=y velocity=${SOURCES[1]} str=0.4')
        Result(inmo,
                        '''
                        grey title="%s" clip=0.008 labelsz=14 labelfat=2 font=4 titlesz=20 titlefat=2
                        ''' % typ)

#Flow('refb','cmpb','window n2=1')
Flow('refb','abnmoa','stack')
#Flow('refb','modl','window n3=1 | window n2=1')
dat='cmpb'
for typ in ('semblance','ab','weighted','similarity-weighted'):
        scn=typ+'scnb'
        if typ == 'similarity-weighted':
                Flow(scn,[dat,'refb'],'simivscan thr=0.6 v0=1.5 dv=0.02 nv=51 type=weighted weight=y nb=5 ref=${SOURCES[1]} | clip2 lower=0')
        else:
                Flow(scn,dat,'vscan v0=1.5 dv=0.02 nv=51 weight=y type=%s nb=5 | clip2 lower=0'%typ)
        Plot(scn+'0',scn,
                '''
                grey color=j allpos=y title="%s" pclip=100 labelsz=14 labelfat=2 font=4 titlesz=20 titlefat=2
                '''%typ)
        pik=typ+'pikb'
        Flow(pik,scn,'scale axis=2 | pick rect1=10')
        Plot(pik,
             '''
             graph yreverse=y transp=y min2=1.5 max2=2.5 pad=n
             plotcol=7 plotfat=15 wanttitle=n wantaxis=n
             ''')
        Result(scn,[scn+'0',pik],'Overlay')
        nmo = typ + 'nmob'
        Flow(nmo,[dat,pik],'nmo half=y velocity=${SOURCES[1]} str=0.1 ')
        Result(nmo,
                        '''
                        grey title="%s"
                        grid2=y gridcol=6 gridfat=7 g2num=0.5 clip=0.008 labelsz=14 labelfat=2 font=4 titlesz=20 titlefat=2
                        ''' % typ)

        inmo = typ + 'inmob'
        Flow(inmo,[nmo,pik],'inmo half=y velocity=${SOURCES[1]} str=0.4')
        Result(inmo,
                        '''
                        grey title="%s" clip=0.008 labelsz=14 labelfat=2 font=2 titlesz=1t0 titlefat=2
                        ''' % typ)

Result('hw-comp1','./Fig/cmpa.vpl ./Fig/semblancescna.vpl ./Fig/abscna.vpl ./Fig/similarity-weightedscna.vpl','SideBySideAniso')
Result('hw-comp2','./Fig/cmpb.vpl ./Fig/semblancescnb.vpl ./Fig/abscnb.vpl ./Fig/similarity-weightedscnb.vpl','SideBySideAniso')






End()