from rsf.proj import *

# cgiter dircvv dirrst

# Mobil AVO CMP gather 807 at well4 location
Fetch('cmp807_raw.HH','rad')

clip=187.628

# Preprocessing
Flow('cmp','cmp807_raw.HH',
     '''
     dd form=native | tpow tpow=2 | mutter half=n v0=1.3 tp=0.2 |
     put label1=Time unit1=s label2=Offset unit2=km
     ''')
Plot('cmp','grey title="Input CMP Gather" clip=%g' % clip)

# Velocity Transform
Flow('veltran','cmp','veltran anti=0.1 v0=1.250 dv=0.125 nv=60 adj=y')
Plot('veltran','grey title="Velocity Scan" ')

# Display Side by Side
Result('veltran','cmp veltran','SideBySideAniso')

# Least-Squares Optimization
Flow('velinv','cmp veltran',
     '''
     conjgrad veltran anti=0.1 v0=1.250 dv=0.125 nv=60 niter=20
     mod=${SOURCES[1]}
     ''')
Plot('velinv','grey title="Least-Squares Velocity Scan" ')

# Data Prediction Error

for case in ('tran','inv'):
    err = 'err'+case
    Flow(err,['vel'+case,'cmp'],'veltran anti=0.1 adj=n | add scale=1,-1 ${SOURCES[1]}')
    Plot(err,'grey title="Data Residual" clip=%g' % clip)
    Result(err,[err,'vel'+case],'SideBySideAniso')

def iscan(miter,psun1,psun2,niter=10):
    return '''
    cgscan miter=%d psun1=%d psun2=%d error=${TARGETS[1]}
    niter=%d v0=1.250 dv=0.125 nv=60 anti=0.1
    ''' % (miter,psun1,psun2,niter)

for psun in (0,1):
    scan = 'scan%d' % psun
    err  = 'err%d' % psun
    Flow([scan,err],'cmp',iscan(2,psun,psun))

Plot('vscan','scan1','grey title="Least-squares Velocity Scan" ')

Flow('scan err','cmp',iscan(0,1,1,0))
Flow('scan2 err2','scan',iscan(0,1,1,0) + ' adj=1 | ' + iscan(0,1,1,0))

nx = 1000
ny = 60

shifts = []
shifts2 = []
for x  in range(-3,4):
    for y in range(-3,4):
        shift = 'shift%d%d' % (x,y)
        if x > 0:
            if y > 0:
                tra = 'window f1=%d f2=%d | pad end1=%d end2=%d' % (x,y,x,y)
            else:
                tra = 'window f1=%d n2=%d | pad end1=%d beg2=%d' % (x,ny+y,x,-y)
        else:
            if y > 0:
                tra = 'window n1=%d f2=%d | pad beg1=%d end2=%d' % (nx+x,y,-x,y)
            else:
                tra = 'window n1=%d n2=%d | pad beg1=%d beg2=%d' % (nx+x,ny+y,-x,-y)
        Flow(shift,'scan',tra)
        Flow(shift+'2','scan2',tra)
        shifts.append(shift)
        shifts2.append(shift+'2')

Flow('shifts',shifts,'cat ${SOURCES[1:%d]}' % len(shifts))
Flow('shifts2',shifts2,'cat ${SOURCES[1:%d]}' % len(shifts))

Flow('flt pre','shifts2 scan',
     'lpf match=${SOURCES[1]} pred=${TARGETS[1]} rect1=20 rect2=10')

Flow('inv','shifts flt','add mode=p ${SOURCES[1]} | stack axis=3 norm=n')
Flow('cmp2','scan1',iscan(0,1,1,0) + ' adj=1')

Result('cgiter','err0 err1',
       '''
       cat axis=2 ${SOURCES[1]} | scale axis=2 |
       graph min2=0 max2=1.1 labelsz=10 titlesz=12 dash=1,0
       title="Iterative Convergence" unit1= unit2=
       ''')

Result('dircvv','cmp vscan','SideBySideAniso')

Flow('cmp2','scan1',iscan(0,1,1,0) + ' adj=1')

Plot('cmp2','grey title="Predicted Data" clip=%g' % clip)

Plot('res','cmp2 cmp',
     'add scale=1,-1 ${SOURCES[1]} | grey title="Residual Error" clip=%g' % clip)

Result('dirrst','cmp2 res','SideBySideAniso')

End()