# Author: Deli Wang
#         College of Geoexploration Science and Technology,
#         Jilin University,Changchun,China
#         Seismic Laboratory for Imaging and Modeling
#         Department of Earch & Ocean Sciences
#         The University of British Columbia
#         
# Date  : November, 2007

import os
import string as str
from rsfproj import *

python = WhereIs('python')

###################################################################
#                                                                 #
#                         Data preparation                        #
#                                                                 #
###################################################################
# define SLIM FTP server information
FTPserver = {
    'server': os.environ.get('SLIM_SERVER'),
    'login': os.environ.get('SLIM_LOGIN'),
    'password':os.environ.get('SLIM_PASSWORD')}

filesloc = os.path.join('users','dwang','Geophysics','Matched','real','revision1')
filesloc1 = os.path.join('users','dwang','Geophysics','Matched','real','revision1','xdr')


flist = ''' 
        saga_cop200_data.rsf   saga_cop200_lsqrmult.rsf
        saga_cop200_srme.rsf
  '''

#fetch data from FTP server
for fil in str.split(flist):
    Fetch(fil,filesloc,FTPserver)

Alias('fetch',['saga_cop200_data.rsf','saga_cop200_srme.rsf', \
               'saga_cop200_lsqrmult.rsf'])

Flow('saga_cop200_data_n.rsf','saga_cop200_data.rsf','sfdd form=native')
Flow('saga_cop200_srme_n.rsf','saga_cop200_srme.rsf','sfdd form=native')
Flow('saga_cop200_lsqrmult_n.rsf','saga_cop200_lsqrmult.rsf','sfdd form=native')

###################################################################
#                                                                 #
#               Diagonal estimate for different Lambda            #
#                                                                 #
###################################################################
cluster = ARGUMENTS.get('cluster',0)
if int(cluster) is 1:
        command = python+' ../code/Mdiag_bfgs.py '

        Flow('2.0saga_cop200_scaler_init.rsf 2.0saga_cop200_scalmult_init.rsf', \
        'saga_cop200_data_n.rsf  saga_cop200_lsqrmult_n.rsf',command +  \
        'lsqrmult=${SOURCES[1]} scalemult=${TARGETS[1]} lamb=2.0 maxiter=75')

        Flow('1.0saga_cop200_scaler_init.rsf 1.0saga_cop200_scalmult_init.rsf', \
        'saga_cop200_data_n.rsf 2.0saga_cop200_scalmult_init.rsf',command +  \
        'lsqrmult=${SOURCES[1]} scalemult=${TARGETS[1]} lamb=1.0 maxiter=75')

        Flow('0.7saga_cop200_scaler_init.rsf 0.7saga_cop200_scalmult_init.rsf', \
        'saga_cop200_data_n.rsf 1.0saga_cop200_scalmult_init.rsf',command +  \
        'lsqrmult=${SOURCES[1]} scalemult=${TARGETS[1]} lamb=0.7 maxiter=75')

        Flow('0.5saga_cop200_scaler_init.rsf 0.5saga_cop200_scalmult_init.rsf', \
        'saga_cop200_data_n.rsf 0.7saga_cop200_scalmult_init.rsf',command +  \
        'lsqrmult=${SOURCES[1]} scalemult=${TARGETS[1]} lamb=0.5 maxiter=75')

        Flow('0.3saga_cop200_scaler_init.rsf 0.3saga_cop200_scalmult_init.rsf', \
        'saga_cop200_data_n.rsf 0.5saga_cop200_scalmult_init.rsf',command +  \
        'lsqrmult=${SOURCES[1]} scalemult=${TARGETS[1]} lamb=0.3 maxiter=75')

        Flow('0.0saga_cop200_scaler_init.rsf 0.0saga_cop200_scalmult_init.rsf', \
        'saga_cop200_data_n.rsf 0.3saga_cop200_scalmult_init.rsf',command +  \
        'lsqrmult=${SOURCES[1]} scalemult=${TARGETS[1]} lamb=0.0 maxiter=75')
else:
        Fetch('0.3saga_cop200_scalmult_init_xdr.rsf',filesloc1,FTPserver)
        Fetch('0.0saga_cop200_scalmult_init_xdr.rsf',filesloc1,FTPserver)
        Flow('0.3saga_cop200_scalmult_init.rsf','0.3saga_cop200_scalmult_init_xdr.rsf',\
            'sfdd form=native')
        Flow('0.0saga_cop200_scalmult_init.rsf','0.0saga_cop200_scalmult_init_xdr.rsf',\
            'sfdd form=native')
###################################################################
#                                                                 #
#               Bayesian primaries-multiples separation           #
#                                                                 #
###################################################################
#pms_bayes = WhereIs('pms_bayes.py')
pms_bayes = ' ../code/pms_bayes.py'
# Number of interations
interBayes = 5
padlen = 10

#-------------Over matched multiples Bayessian separation---------                     
Flow('0.0saga_cop200_prim_init.rsf',['saga_cop200_data_n.rsf', \
     '0.0saga_cop200_scalmult_init.rsf'],\
    'sfmath a=${SOURCES[1]} output="input-a"')

Flow('0.0saga_cop200_bayes2D_init_prim.rsf 0.0saga_cop200_bayes2D_init_mult.rsf',['0.0saga_cop200_scalmult_init.rsf','0.0saga_cop200_prim_init.rsf'],
                python+' '+pms_bayes+"""   
                        predmult=${SOURCES[0]}
                        predprim=${SOURCES[1]}
                        niter=%(interBayes)s
                        padlen=%(padlen)s
                        transparams=6,64,1
                        mu=1.0
                        lamb1=2.0
                        lamb2=1.2
                        poutput=${TARGET}
                        moutput=${TARGETS[1]}
                        --tmpdatapath=./
                        --localtmpdir=./
                """ %vars(),stdin=0,stdout=-1)
#-------------Correctly matched multiples Bayessian separation---------                  
Flow('0.3saga_cop200_prim_init.rsf',['saga_cop200_data_n.rsf', \
     '0.3saga_cop200_scalmult_init.rsf'],\
    'sfmath a=${SOURCES[1]} output="input-a"')

Flow('0.3saga_cop200_bayes2D_init_prim.rsf 0.3saga_cop200_bayes2D_init_mult.rsf',['0.3saga_cop200_scalmult_init.rsf','0.3saga_cop200_prim_init.rsf'],
                python+' '+pms_bayes+"""   
                        predmult=${SOURCES[0]}
                        predprim=${SOURCES[1]}
                        niter=%(interBayes)s
                        padlen=%(padlen)s
                        transparams=6,64,1
                        mu=1.0
                        lamb1=2.0
                        lamb2=1.2
                        poutput=${TARGET}
                        moutput=${TARGETS[1]}
                        --tmpdatapath=./
                        --localtmpdir=./
                """ %vars(),stdin=0,stdout=-1)
#-------------No matching multiples(SRME multiples) Bayessian separation--------- 
Flow('saga_cop200_lsqrprim.rsf',['saga_cop200_data_n.rsf', \
     'saga_cop200_lsqrmult_n.rsf'],'sfmath a=${SOURCES[1]} output="input-a"')

Flow('saga_cop200_bayes2D_lsqrmult_prim.rsf saga_cop200_bayes2D_lsqrmult_mult.rsf',['saga_cop200_lsqrmult_n.rsf','saga_cop200_lsqrprim.rsf'],
                python+' '+pms_bayes+"""   
                        predmult=${SOURCES[0]}
                        predprim=${SOURCES[1]}
                        niter=%(interBayes)s
                        padlen=%(padlen)s
                        transparams=6,64,1
                        mu=1.0
                        lamb1=2.0
                        lamb2=1.2
                        poutput=${TARGET}
                        moutput=${TARGETS[1]}
                        --tmpdatapath=./
                        --localtmpdir=./
                """ %vars(),stdin=0,stdout=-1)

###################################################################
#                                                                 #
#               Make figures needed in the paper                  #
#                                                                 #
###################################################################
# Define plot functions
def greym(title):
     return '''
     sfdd form=native |sfcut n1=60 |sfwindow n1=938|sfput d1=0.004 d2=25 o2=0| sfagc| 
     sfgrey xinch=6 yinch=8  clip=2.5 grid2=y g2num=1.0
     label1="Time" unit1="s"  label2="Shot location" unit2="m"
     labelsz=6.0  titlesz=6.0 title="%s"
      ''' % (title)

#-------Basic plot-------------------
Plot('F2a','saga_cop200_data.rsf',greym(''))
Plot('F2b','saga_cop200_srme.rsf',greym(''))
Plot('F2c','saga_cop200_bayes2D_lsqrmult_prim.rsf',greym(''))
Plot('F2d','0.3saga_cop200_bayes2D_init_prim.rsf',greym(''))

Plot('overm','0.0saga_cop200_bayes2D_init_prim.rsf',greym(''))
#-------annotations-----------------    
labels = (
    (6.17,7.4,''),
    (6.17,6.52,''),
    (6.17,2.5,''),
    (8.8,3.4,'')
    )
nlabels = len(labels)
annot = []
for i in range(nlabels):
    annot.append('annot%d' % i)
    Plot(annot[i],None,
         '''
         box x0=%g y0=%g label="%s"
         xt=0.5 yt=0.4 lab_fat=0 
         boxit=1 size=0.2 pointer=1 x_oval=0.0 y_oval=0.0
         ''' % labels[i])

#------Final result----------------
Result('Figure2-a',['F2a']+annot,'Overlay')
Result('Figure2-b',['F2b']+annot,'Overlay')
Result('Figure2-c',['F2c']+annot,'Overlay')
Result('Figure2-d',['F2d']+annot,'Overlay')

Result('overmatch',['overm']+annot,'Overlay')
End()