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

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

par = dict(
    nx=101, ox=-1000.0, dx=20.0,  lx='x', ux='m',
    ny=101, oy=-1000.0, dy=20.0,  ly='y', uy='m', 
    nz=41,  oz=+0000.0, dz=20.0,  lz='z', uz='m',
    nit=10
)   
wplot.param(par)

# ------------------------------------------------------------
# receiver subsampling
par['jx']=2;
par['njx']=(par['nx']-1)/par['jx']+1
par['djx']=par['dx']*par['jx']

par['jy']=2;
par['njy']=(par['ny']-1)/par['jy']+1
par['djy']=par['dy']*par['jy']

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

def gplot2d(custom,par):
    return '''
    scale rscale=1e5 |
    graph title="" symbol=. plotfat=20
    min1=%(xmin)g max1=%(xmax)g label1=%(lx)s unit1=%(ux)s
    label2="\F10 D\F3 g\_z\^" unit2="mgal"
    %(labelattr)s
    '''%par
    + ' ' + custom

def shape3d(par):
    return '''
    put 
    n1=%(njx)d o1=%(ox)g d1=%(djx)g 
    n2=%(njy)d o2=%(oy)g d2=%(djy)g
    '''%par

def gplot3d(custom,par):
    return '''
    transp |
    grey title=""
    labelrot=n wantaxis=y yreverse=n 
    min2=%g max2=%g label2=%s unit2=%s
    min1=%g max1=%g label1=%s unit1=%s
    screenratio=%g screenht=%g wantscalebar=%s
    %s
    ''' % (par['xmin'],par['xmax'],par['lx'],par['ux'],
           par['ymin'],par['ymax'],par['ly'],par['uy'],
           par['mapratio'],par['mapheight'],par['scalebar'],
           par['labelattr']+' '+custom)

# ------------------------------------------------------------
# receivers
geom.horizontal2d('rr2d',par['oz'],'',par,jx=par['jx'])
Plot('rr2d',wplot.rrplot2d('',par))

geom.horizontal3d('rr3d',par['oz'],'',par,jx=par['jx'],jy=par['jy'])

Plot('rr3d',
    'window n1=2 | dd type=complex | window |'
    + wplot.mapXY('symbol=. plotfat=10 plotcol=1 wantaxis=n',par))

# ------------------------------------------------------------
# true model (delta density; units kg / m^3)
Flow('spk3d',None,
     '''
     spike nsp=2 mag=1.25e+3,-1.0e3
     n1=%(nz)d o1=%(oz)g d1=%(dz)g k1=10,20 l1=21,31
     n2=%(nx)d o2=%(ox)g d2=%(dx)g k2=20,70 l2=31,81
     n3=%(ny)d o3=%(oy)g d3=%(dy)g k3=10,50 l3=51,91 |
     put 
     label1=%(lz)s label2=%(lx)s label3=%(ly)s
     unit1=%(uz)s   unit2=%(ux)s  unit3=%(uy)s
     ''' % par)
Result('spk3d','byte gainpanel=a pclip=100 |'
           + wplot.igrey3d('',par))
Plot('cut3d','spk3d','window n1=1 min1=400 |'
           + gplot3d('grid=y',par))
Result('cut3d',['cut3d','rr3d'],'Overlay')


Flow('spk2d','spk3d','window n3=1 min3=0')
Plot(  'spk2d',wplot.igrey2d('mean=y',par))
Result('spk2d','spk2d rr2d','Overlay')

# ------------------------------------------------------------
# initialize operators
# ------------------------------------------------------------
G2D = gravop2d(par,'rr2d','verb=y')
if(do3D=='y'):
    G3D = gravop3d(par,'rr3d','verb=y')

# ------------------------------------------------------------
# FORWARD/ADJOINT operators
# ------------------------------------------------------------
G2D.FORW('spk2d','dat2d')
G2D.ADJT('mod2d','dat2d')
Result('dat2d',gplot2d('',par))
Result('mod2d',wplot.igrey2d('color=e',par))

if(do3D=='y'):
    G3D.FORW('spk3d','dat3d')
    G3D.ADJT('mod3d','dat3d')
    
    Result('dat3d',shape3d(par) + '|' + gplot3d('color=e',par))
    Result('mod3d','byte gainpanel=a pclip=100 |'
            + wplot.igrey3d('color=e',par))

# ------------------------------------------------------------
# DOT PRODUCT TEST
# ------------------------------------------------------------
Flow('m2d','mod2d','math output=0') # init m
Flow('d2d','dat2d','math output=0') # init d
DP2 = DPTEST(G2D,['m2d'],['d2d'])   # init DP test
DP2.RUN()                             

if(do3D=='y'):
    Flow('m3d','mod3d','math output=0') # init m
    Flow('d3d','dat3d','math output=0') # init d
    DP3 = DPTEST(G3D,['m3d'],['d3d'])   # init DP test
    DP3.RUN()

End()
