#
# Green operator demo
#
try:    from rsf.cluster import *
except: from rsf.proj    import *
import wplot,awe,geom
sys.path.append('../../PYUTIL')
from greenLOP import *

from DPTEST import *
from LICGsolver import *
from LICGutils  import *

# ------------------------------------------------------------
do3D = 'y'
doINV = 'n'
# ------------------------------------------------------------

par = dict(nx=200, dx=0.05,  ox=-5.00, lx='x', ux='km',
           ny=200, dy=0.05,  oy=-5.00, ly='y', uy='km',
           nz=200, dz=0.05,  oz=+0.00, lz='z', uz='km',
           nt=201, dt=0.001, ot=0,     lt='t', ut='s',
           kt=101,
           vel=2.0,
           apt=15,
           verb='y'
     )
wplot.param(par)

# ------------------------------------------------------------

# orbit parameters
xorb=0.0
yorb=0.0
zorb=5.0
rorb=4.5

# topography parameters
xtop=1.0
ytop=0.0
ztop=4.0

xtop=0.0
ytop=0.0
ztop=5.0
rtop=1.0

nlon=91  # E/W
nlat=91  # N/S

# ------------------------------------------------------------
# mod time
par['nmt']=201
par['omt']=0.0
par['dmt']=par['dt']

# dat time
par['ndt']=2001
par['odt']=1.5
par['ddt']=par['dt']
par['maxdt']=par['odt']+(par['ndt']-1)*par['ddt']

# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
# topo 2D
geom.circle('ss2D',xtop,ztop,rtop,nlat,'',par)

Flow('snx2D','ss2D','window n1=1 f1=0 | add add=%g | scale axis=123'%(-xtop))
Flow('snz2D','ss2D','window n1=1 f1=1 | add add=%g | scale axis=123'%(-ztop))
Flow('sn2D',['snx2D','snz2D'],
         '''
         cat axis=2 space=n ${SOURCES[1]} | transp |
         put o2=0 d2=1
         ''')

# topo 3D
if do3D=='y':
    geom.sphere('ss3D',xtop,ytop,ztop,rtop,nlat,nlon,'',par)

    Flow('snx3D','ss3D','window n1=1 f1=0 | add add=%g | scale axis=123'%(-xtop))
    Flow('sny3D','ss3D','window n1=1 f1=1 | add add=%g | scale axis=123'%(-ytop))
    Flow('snz3D','ss3D','window n1=1 f1=2 | add add=%g | scale axis=123'%(-ztop))
    Flow('sn3D',['snx3D','sny3D','snz3D'],
         '''
         cat axis=2 space=n ${SOURCES[1:3]} | transp |
         put o2=0 d2=1
         ''')

# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
# orbit 2D
geom.circle('rr2D',xorb,zorb,rorb,nlat,'',par)

Flow('rnx2D','rr2D','window n1=1 f1=0 | add add=%g | scale axis=123'%(-xorb))
Flow('rnz2D','rr2D','window n1=1 f1=1 | add add=%g | scale axis=123'%(-zorb))
Flow('rn2D',['rnx2D','rnz2D'],
         '''
         cat axis=2 space=n ${SOURCES[1]} | transp |
         put o2=0 d2=1
         ''')

# orbit 3D
geom.sphere('rr3D',xorb,yorb,zorb,rorb,nlat,nlon,'',par)

Flow('rnx3D','rr3D','window n1=1 f1=0 | add add=%g | scale axis=123'%(-xorb))
Flow('rny3D','rr3D','window n1=1 f1=1 | add add=%g | scale axis=123'%(-yorb))
Flow('rnz3D','rr3D','window n1=1 f1=2 | add add=%g | scale axis=123'%(-zorb))
Flow('rn3D',['rnx3D','rny3D','rnz3D'],
         '''
         cat axis=2 space=n ${SOURCES[1:3]} | transp |
         put o2=0 d2=1
         ''')

# ------------------------------------------------------------
Plot('ss2D',wplot.ssplot2d('wantaxis=y',par))
Plot('rr2D',wplot.rrplot2d('plotfat=10',par))
Result('geom2D',['ss2D','rr2D'],'Overlay')

# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
# ------------------------------------------------------------
# source data
awe.wavelet('tmp',20,'',par)
Flow('wav2D','tmp',
     'window | spray axis=2 n=%d o=0 d=1 label="" unit=""'%nlat)
Result('wav2D','grey title="" pclip=100 %(labelattr)s'%par)

#Result('wav2D',wplot.dgrey2d('',par))

if do3D=='y':
    Flow('wav3D','tmp',
        'window | spray axis=2 n=%d o=0 d=1 label="" unit=""'%(nlat*nlon))
    Result('wav3D','grey title="" pclip=100 %(labelattr)s'%par)

# ------------------------------------------------------------
# initialize operators
# ------------------------------------------------------------
G2D=greenop2d(par,
              mco='ss2D',mno='sn2D',
              dco='rr2D',dno='rn2D',
              vel=par['vel'],
              apt=par['apt'])
if do3D=='y':
    G3D=greenop3d(par,
                  mco='ss3D',mno='sn3D',
                  dco='rr3D',dno='rn3D',
                  vel=par['vel'],
                  apt=par['apt'])

# ------------------------------------------------------------
# FORWARD/ADJOINT operators
# ------------------------------------------------------------
G2D.FORW('wav2D','dat2D')
G2D.ADJT('mod2D','dat2D')
Result('dat2D','grey title="" pclip=100 max1=%(maxdt)g %(labelattr)s'%par)
Result('mod2D','grey title="" pclip=100 %(labelattr)s'%par)

if do3D=='y':
    G3D.FORW('wav3D','dat3D')
    G3D.ADJT('mod3D','dat3D')
    Result('dat3D','grey title="" pclip=100 max1=%(maxdt)g %(labelattr)s'%par)
    Result('mod3D','grey title="" pclip=100 %(labelattr)s'%par)

# ------------------------------------------------------------
# DOT PRODUCT TEST
# ------------------------------------------------------------
Flow('m2D','mod2D','window squeeze=n') # init m
Flow('d2D','dat2D','window squeeze=n') # init d
D2D = DPTEST(G2D,['m2D'],['d2D'])      # init DP test
D2D.RUN()                              #  run DP test

if do3D=='y':
    Flow('m3D','mod3D','window squeeze=n') # init m
    Flow('d3D','dat3D','window squeeze=n') # init d
    D3D = DPTEST(G3D,['m3D'],['d3D'])      # init DP test
    D3D.RUN()                              #  run DP test

# ------------------------------------------------------------
# LS inverse
# ------------------------------------------------------------
if(doINV=='y'):
    NITER=10

    Flow('zer2D','mod2D','math output=0')
    S2D = LICG('lsq2D',par, Lop=G2D, d=['dat2D'], xo=['zer2D'], xf=['lsq2D'], niter=NITER)
    S2D.iterate()
    plotOFUN('lsq2D',NITER,par)
    Result('lsq2D','grey title="" pclip=100 %(labelattr)s'%par)

    if do3D=='y':
        Flow('zer3D','mod3D','math output=0')
        S3D = LICG('lsq3D',par, Lop=G3D, d=['dat3D'], xo=['zer3D'], xf=['lsq3D'], niter=NITER)
        S3D.iterate()
        Result('lsq3D','grey title="" pclip=100 %(labelattr)s'%par)

End()
