from rsf.proj import *
import string

# Vs and density
Flow('vs','time1','math output=0')
Flow('rho','time1','math output=1')

g0=0.98  # starting change 
g1=2-g0  # last change
ng=101   # number of changes to scan
dg = (g1-g0)/(ng-1)
niter = 100 # number of iterations
rect1 = 25 # vertical smoothing

# time sampling
dt = 0.002
nt = int(1.75/dt+1.5)

# anisotropy for picking
an = 1000

under = string.maketrans('_','-')
for c in range(4):
    case = ('time1','time2','time2_over','time2_over_compaction')[c]
    src = 'vel_%s.rsf' % case
    case = case.translate(under)

    # get model from data server
    Fetch(src,'long')
    Flow(case,src,
         '''
         dd form=native |
         window f2=1 |
         put label=Velocity unit=km/s label1=Depth unit1=km
         ''')

    # display
    Result(case,
           '''
           grey color=j allpos=y bias=1.5
           scalebar=y label2=Trace wanttitle=n
           barreverse=y parallel2=n font=2 format2="%3.1f"
           ''')

    # model seismic (+ Gaussian random noise)
    seis = 'seis-'+case
    Flow(seis,[case,'vs','rho'],
         '''
         cat axis=3 ${SOURCES[1:3]} |
         transp plane=23 |
         modrefl2 dt=%g nt=%d |
         window n2=1 |
         ricker1 frequency=30 |
         noise seed=%d var=5e-8
         ''' % (dt,nt,10*c))

    # display
    Result(seis,'grey label2=Trace label1=Time unit1=s wanttitle=n parallel2=n font=2 format2="%3.1f"')

    # Select one trace from the middle
    trace = 'trace-'+case
    Flow(trace,seis,'window n2=1 f2=50')

    if c > 0:
        # Model difference
        mdif = 'mdif-'+case
        Flow(mdif,[case,'time1'],'add scale=1,-1 ${SOURCES[1]}')
        Result(mdif,
               '''
               grey color=j 
               scalebar=y label2=Trace wanttitle=n
               barreverse=y barlabel="Velocity Change" minval=0
               parallel2=n font=2 format2="%3.1f" formatbar="%3.2f"
               ''')

        # Analyze one trace
        ###################

        # Take time-lapse difference
        dif = 'tdif-'+case
        Flow(dif,[trace,'trace-time1'],'add scale=1,-1 ${SOURCES[1]}')

        Result(dif,['trace-time1',trace,dif],
               '''
               cat axis=2 ${SOURCES[1:3]} |
               dots labels=Before:After:Difference label1=Time unit1=s
               yreverse=y gaineach=n 
               ''')

        # Separate bottom part
        n1 = (400,250,250)[c-1]

        bot = 'tbot-'+case
        Flow(bot,trace,'window f1=%d' % n1)

        bot0 = 'tbot0-'+case
        Flow(bot0,'trace-time1','window f1=%d' % n1)

        top0 = 'ttop0-'+case
        Flow(top0,trace,'window n1=%d' % n1)

        # Scan shifts computing local similarity
        scan = 'tscan-'+case
        Flow(scan,[bot,bot0],
             '''
             warpscan other=${SOURCES[1]} niter=%d
             ng=%d g0=%g dg=%g rect1=%d | 
             math output='(1+input)^4'
             ''' % (niter,ng,g0,dg,rect1))

        Plot(scan,
             '''
             grey allpos=y
             color=j wanttitle=n
             label1=Time unit1=s
             label2="Relative stretch"
             ''')

        # Pick the stretch
        pick = 'tpick-'+case
        Flow(pick,scan,'pick rect1=%d vel0=1 | window' % rect1)
        Plot(pick,
             '''
             graph transp=y min2=%g max2=%g 
             yreverse=y plotcol=7 plotfat=5 
             wantaxis=n wanttitle=n pad=n
             ''' % (g0,g1))
        Result(scan,[scan,pick],'Overlay')

        # Convert stretch to shift
        shift = 'tshift-'+case
        Flow(shift,pick,'math output="(input-1)*x1" ')


        # Smooth pick
        pick = 'tpicks-'+case
        Flow(pick,scan,
             '''
             pick rect1=%d vel0=1 an=%d |
             window''' % (10*rect1,an))

        # Interval velocity ratio
        rat = 'trat-'+case
        Flow(rat,[top0,pick],
             '''
             math output=1 |
             cat axis=1 ${SOURCES[1]} |
             math output="input*x1" |
             smoothder eps=10
             ''')

        # Apply the stretch
        warp = 'twarp-'+case
        Flow(warp,[bot,bot0,shift],
             '''
             warp1 other=${SOURCES[1]} warpin=${SOURCES[2]}
             verb=1 nliter=0 
             ''')

        warp2 = 'twarp2-'+case
        Flow(warp2,[top0,warp],'cat axis=1 ${SOURCES[1]}')

        # Take time-lapse difference
        dif2 = 'tdif1-'+case
        Flow(dif2,[warp2,'trace-time1'],'add scale=1,-1 ${SOURCES[1]}')

        Result(dif2,['trace-time1',warp2,dif2],
               '''
               cat axis=2 ${SOURCES[1:3]} |
               dots labels=Before:After:Difference label1=Time unit1=s
               yreverse=y gaineach=n 
               ''')

        ######################################
        # 2-D work
        ######################################

        # Take time-lapse difference
        dif = 'dif-'+case
        Flow(dif,[seis,'seis-time1'],'add scale=1,-1 ${SOURCES[1]}')

        Result(dif,
               'grey label2=Trace label1=Time unit1=s wanttitle=n parallel2=n font=2 format2="%3.1f"')

        bot = 'bot-'+case
        Flow(bot,seis,'window f1=%d' % n1)

        bot0 = 'bot0-'+case
        Flow(bot0,'seis-time1','window f1=%d' % n1)

        top0 = 'top0-'+case
        Flow(top0,seis,'window n1=%d' % n1)

        # Scan shifts computing local similarity
        scan = 'scan-'+case
        Flow(scan,[bot,bot0],
             '''
             warpscan other=${SOURCES[1]} niter=%d
             ng=%d g0=%g dg=%g rect1=%d rect3=2 | 
             math output='(1+input)^4'
             ''' % (niter,ng,g0,dg,rect1))

        Result(scan,
               '''
               byte allpos=y |
               grey3 flat=n frame1=%d frame2=50 frame3=49
               color=j wanttitle=n
               label1=Time unit1=s
               label2="Relative stretch"
               label3=Trace
               parallel2=n font=2 format2="%%3.1f"
               ''' % (600-n1))

        # Pick the stretch
        pick = 'pick-'+case
        Flow(pick,scan,
             'pick rect1=%d rect2=3 gate=100 vel0=1 | window' % (rect1/2))

        Result(pick,[top0,pick],
               '''
               math output=1 | cat axis=1 ${SOURCES[1]} |
               grey color=j scalebar=y wanttitle=n
               label1=Time unit1=s label2=Trace bias=1
               barlabel="Relative stretch" barunit=
               parallel2=n font=2 format2="%3.1f"
               ''')

        # Convert stretch to shift
        shift = 'shift-'+case
        Flow(shift,pick,'math output="(input-1)*x1" ')


        # Apply the stretch
        warp = 'warp-'+case
        Flow(warp,[bot,bot0,shift],
             '''
             warp1 other=${SOURCES[1]} warpin=${SOURCES[2]}
             verb=1 nliter=0 
             ''')

        warp2 = 'warp2-'+case
        Flow(warp2,[top0,warp],'cat axis=1 ${SOURCES[1]}')

        # display
        Result(warp2,'grey label2=Trace label1=Time unit1=s wanttitle=n parallel2=n font=2 format2="%3.1f" ')

        # Take time-lapse difference
        dif2 = 'dif2-'+case
        Flow(dif2,[warp2,'seis-time1'],'add scale=1,-1 ${SOURCES[1]}')

        Result(dif2,
               'grey label2=Trace label1=Time unit1=s wanttitle=n parallel2=n font=2 format2="%3.1f"')

End()