from rsf.proj import *

def graph(transp,o2,d2,n2,col,fat,extra=''):
    return '''
    graph transp=%d yreverse=y pad=n min2=%g max2=%g
    plotcol=%d plotfat=%d wantaxis=n wanttitle=n %s
    ''' % (transp,o2,o2+(n2-1)*d2,col,fat,extra)

## Kirchhoff modeling
Flow('modl',None,
     '''
     spike n1=200 o1=-8 d1=0.08 n2=4
     nsp=4 k2=1,2,3,4 mag=1,2.2,3.5,5
     ''')
Flow('refl',None,
     '''
     spike n1=200 n2=4 nsp=4 k2=1,2,3,4
     mag=0.1111110,0.1666667,0.1764706,0.23076920
     ''')

Flow('data','modl refl',
     '''
     kirmod nt=501 dt=0.004 freq=25 refl=${SOURCES[1]}
     ns=12 s0=1.4 ds=-0.1
     nh=64 h0=0.1 dh=0.1
     vel=5 gradz=2 type=v
     ''')

v0=4
dv=0.075
nv=120

###############################
# Create offset and cmp headers
###############################
Flow('off','data','window n1=1 | math output="0.5*x1" ')
Flow('cmp','off','math output="input+x2" ')

###############################
#
# Processing individual CMP
#
###############################
# Extract one CMP 
x = 1500
y = 0.001*x

Flow('msk','cmp','mask min=%g max=%g' % (y-0.001,y+0.001))
Flow('cdpoff','off msk',
     '''
     spray axis=1 n=1 | 
     headerwindow mask=${SOURCES[1]}
     ''')
Flow('cdp','data msk','headerwindow mask=${SOURCES[1]}')

# AVO anomaly
Flow('cdpavo','cdp',
     'math output="input*(1-x2*%g)" ' %(2.0/1.2))

Plot('cdpavo',
     '''put o2=1 d2=1 label2=Trace unit2="" | 
     wiggle poly=y yreverse=y transp=y title="(a)"
     labelsz=14
     titlesz=18 plotcol=7 gridcol=7
     ''')

# Add random noise       
Flow('cdpavonoise','cdpavo',
    '''
    noise rep=y seed=2004 range=0.002 |
    ricker1 frequency=25 |
    add $SOURCE 
    ''')

#Flow('cdpavonoise','cdpavo','noise seed=2014 type=y mean=0 var=0.00000003')

Plot('cdpavonoise',
     '''put o2=1 d2=1 label2=Trace unit2="" | 
     wiggle poly=y yreverse=y transp=y title="(b)"
     labelsz=14
     titlesz=18 plotcol=7 gridcol=7
     ''')

# SNR
#Flow('noise','cdpavo cdpavonoise','math output="${SOURCES[1]}-input"')
#Flow('snr','cdpavo noise','math ')

#Flow('ref','abnmo','stack')
Flow('ref','cdp','window n2=1')

##############################
# Loop for semblances
##############################

dat = 'cdpavonoise'
for typ in ('semblance','ab','weighted','simi'):
    ##################
    # Velocity Scan
    ##################
    scn = typ + 'scn'
    if typ == 'simi':
            Flow(scn,[dat,'ref'],'simivscan thr=0.5 v0=4 dv=0.075 nv=120 type=avo weight=y nb=5 ref=${SOURCES[1]} | clip2 lower=0')
            Plot(scn+'0',scn,
                 '''
                         grey color=j allpos=y pclip=100
                         labelsz=14  titlesz=14
                         label1="Time" unit1="s"
                         label2="Velocity" unit2="kft/s" title=""
                         wherexlabel=bottom wheretitle=top
                         ''')
    elif typ == 'semblance':
            Flow(scn,[dat,'cdpoff'],'vscan offset=${SOURCES[1]} v0=4 dv=0.075 nv=120 type=semb weight=y nb=3 | clip2 lower=0')
            Plot(scn+'0',scn,
                 '''
                         grey color=j allpos=y pclip=100
                         labelsz=14  titlesz=18
                         label1="Time" unit1="s"
                         label2="Velocity" unit2="kft/s" title="(c)"
                         wherexlabel=bottom wheretitle=top
                         ''')
    elif typ == 'weighted':
            Flow(scn,[dat,'cdpoff'],'vscan offset=${SOURCES[1]} v0=4 dv=0.075 nv=120 type=weighted weight=y nb=3 | clip2 lower=0')
            Plot(scn+'0',scn,
                 '''
                         grey color=j allpos=y pclip=100
                         labelsz=14  titlesz=18
                         label1="Time" unit1="s"
                         label2="Velocity" unit2="kft/s" title=""
                         wherexlabel=bottom wheretitle=top
                         ''')
    else:
            Flow(scn,[dat,'cdpoff'],'vscan offset=${SOURCES[1]} v0=4 dv=0.075 nv=120 avosemblance=y nb=3 | clip2 lower=0')
            Plot(scn+'0',scn,
                 '''
                         grey color=j allpos=y pclip=100
                         labelsz=14  titlesz=18
                         label1="Time" unit1="s"
                         label2="Velocity" unit2="kft/s" title="(d)"
                         wherexlabel=bottom wheretitle=top
                         ''')

#    Plot(scn+'0',scn,
#         '''
#         grey color=j allpos=y pclip=100
#         labelsz=14  titlesz=14
#         label2=Velocity unit2=kft/s title=""
#         wherexlabel=bottom wheretitle=top
#         ''')
    ##################
    # Velocity Picking
    ##################
    pik = typ + 'pik'
    if typ == 'semblance':
            Flow(pik,scn,
                 '''
                 mutter x0=4 inner=y v0=3 half=y |
                 mutter x0=4 inner=y v0=4 half=y |
                 pick rect1=50 rect2=100 vel0=5 gate=10
                 ''')
    else:
            Flow(pik,scn,
                 '''
                 mutter x0=4 inner=y v0=3 half=n |
                 mutter x0=4 inner=y v0=4 half=n |
                 pick rect1=50 rect2=100 vel0=5 gate=10
                 ''')
    Plot(pik+'a',pik,graph(1,v0,dv,nv,0,20))
    Plot(pik+'b',pik,graph(1,v0,dv,nv,7,5))
    Plot(scn,[scn+'0',pik+'a', pik+'b'],'Overlay')
    ##################
    # NMO
    ##################
    nmo = typ + 'nmo'
    Flow(nmo,[dat,'cdpoff',pik],
         'nmo offset=${SOURCES[1]} velocity=${SOURCES[2]} str=0.2')
    if typ == 'semblance':
        Plot(nmo,
                '''
                 put o2=1 d2=1 label2=Trace unit2="" | 
                 wiggle poly=y yreverse=y transp=y title="(a)"
                 labelsz=12
                 titlesz=12 plotcol=7 gridcol=7
                ''')
    elif typ == 'ab':
        Plot(nmo,
                '''
                 put o2=1 d2=1 label2=Trace unit2="" | 
                 wiggle poly=y yreverse=y transp=y title="(b)"
                 labelsz=12 
                 titlesz=12 plotcol=7 gridcol=7
                ''')
#    Plot(nmo,
#         '''
#         put o2=1 d2=1 label2=Trace unit2="" | 
#         wiggle poly=y yreverse=y transp=y title=""
#         labelsz=10 
#         titlesz=14 plotcol=7
#         ''')
    #comp1 = typ + 'comp1'
    #Result(comp1,['cdpavonoise',scn,nmo],'SideBySideAniso')

###################
# Semblance + Stack
###################
Flow('stack','semblancenmo','stack')
Plot('stack',
     '''
     put o2=0 label2=Amplitude unit2="" | 
     wiggle poly=y yreverse=y transp=y title="(a)" labelsz=14 titlesz=14 plotcol=7 gridcol=7
     ''')

###################
# Semblance + SNR stack
###################
Flow('snrstack','semblancenmo','snrstack')
Plot('snrstack',
     '''
     put o2=0 label2=Amplitude unit2="" | 
     wiggle poly=y yreverse=y transp=y title="(b)" labelsz=14 titlesz=14 plotcol=7 gridcol=7
     ''')

######################
# AB semblance + stack
###################### 
Flow('abstack','abnmo','stack')
Plot('abstack',
     '''
     put o2=0 label2=Amplitude unit2="" | 
     wiggle poly=y yreverse=y transp=y title="ab-stack" labelsz=10 titlesz=14 plotcol=7 gridcol=7
     ''')

######################
# AB semblance + SNR stack
###################### 
Flow('absnrstack','abnmo','snrstack')
Plot('absnrstack',
     '''
     put o2=0 label2=Amplitude unit2="" | 
     wiggle poly=y yreverse=y transp=y title="ab-stack" labelsz=10 titlesz=14 plotcol=7 gridcol=7
     ''')

###########################
# Semblance + Simi stack
########################### 
# Reference trace
Flow('ref3','semblancenmo','window n2=1')

Flow('refs3','ref3','spray axis=2 n=12')

Flow('weight3','refs3 abnmo',
     '''
     similarity other=${SOURCES[1]} niter=30 rect1=5 rect2=7
     ''')

Flow('tweight3','weight3','threshold pclip=20')

Flow('simistack', 'tweight1 refs1 semblancenmo',
     '''
     sfmath y=${SOURCES[1]} output=input*y | stack

     ''')

Plot('simistack',
     '''
     put o2=0 label2=Amplitude unit2="" | 
     wiggle poly=y yreverse=y transp=y title="simistack" labelsz=10 titlesz=14
     ''')

######################
# AB semblance + stack
###################### 
Flow('abstack','abnmo','stack')
Plot('abstack',
     '''
     put o2=0 label2=Amplitude unit2="" | 
     wiggle poly=y yreverse=y transp=y title="ab-stack" labelsz=10 titlesz=14
     ''')

######################
# AB semblance + SNR stack
###################### 
Flow('absnrstack','abnmo','snrstack')
Plot('absnrstack',
     '''
     put o2=0 label2=Amplitude unit2="" | 
     wiggle poly=y yreverse=y transp=y title="ab-snrstack" labelsz=10 titlesz=14
     ''')

###########################
# AB semblance + Simi stack
########################### 
# Reference trace
Flow('ref1','cdpavonoise','window f2=0 n2=1')

Flow('refs1','ref1','spray axis=2 n=12')

Flow('weight1','refs1 abnmo',
     '''
     similarity other=${SOURCES[1]} niter=30 rect1=5 rect2=7
     ''')

Flow('tweight1','weight1','threshold pclip=20')

Flow('absimistack', 'tweight1 refs1 abnmo',
     '''
     sfmath y=${SOURCES[1]} output=input*y | stack

     ''')

Plot('absimistack',
     '''
     put o2=0 label2=Amplitude unit2="" | 
     wiggle poly=y yreverse=y transp=y title="(c)" labelsz=14 titlesz=14 plotcol=7 gridcol=7
     ''')


#############################
# Simi semblance + Simi stack
############################# 
# Reference trace
Flow('ref2','siminmo','window n2=1')

Flow('refs2','ref2','spray axis=2 n=12')

Flow('weight2','refs1 siminmo',
     '''
     similarity other=${SOURCES[1]} niter=30 rect1=5 rect2=7
     ''')

Flow('tweight2','weight2','threshold pclip=20')

Flow('simisimistack', 'tweight2 refs2 siminmo',
     '''
     sfmath y=${SOURCES[1]} output=input*y | stack

     ''')
Plot('simisimistack',
     '''
     put o2=0 label2=Amplitude unit2="" | 
     wiggle poly=y yreverse=y transp=y title="simi-simistack" labelsz=10 titlesz=14
     ''')

#####################
# Putting it together
#####################
#Result('comb-1','cdpavonoise semblancescn semblancenmo stack','SideBySideAniso')
#Result('comb-2','cdpavonoise abscn abnmo abstack','SideBySideAniso')
#Result('comb-3','cdpavonoise abscn abnmo absimistack','SideBySideAniso')
#Result('comb-4','cdpavonoise weightedscn weightednmo','SideBySideAniso')
#Result('comb-5','cdpavonoise simiscn siminmo simisimistack','SideBySideAniso')

Result('lcomb-1','cdpavo cdpavonoise semblancescn abscn','SideBySideAniso')
# increse title and label by ,vppen='txscale=1.8'
Result('lcomb-2','semblancenmo abnmo ','SideBySideAniso')
Result('lcomb-3','stack snrstack absimistack','SideBySideAniso')


# Data format conversion
# Make additional header keys 
Flow('offset','cdpavonoise',
     'window n1=1 | math output="1000*x1" | dd type=int')

# Generate SEGY trace headers
Flow('tdata','cdpavonoise offset','segyheader offset=${SOURCES[1]}')

# Convert RSF to SU
Flow('cdpavonoise.su','cdpavonoise tdata','suwrite tfile=${SOURCES[1]} suxdr=y su=y')
# Convert RSF to SEGY
Flow('cdpavonoise.sgy','cdpavonoise tdata','segywrite tfile=${SOURCES[1]}')

Flow('abnmo.su','abnmo tdata','suwrite tfile=${SOURCES[1]} suxdr=y su=y')

End()