# 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','synth','revision1','new_data')
filesloc1 = os.path.join('users','dwang','Geophysics','Matched','synth','revision1','new_data','xdr')

flist = ''' 
        model1_shot181_p0.rsf       model1_shot181_srme_iter1_local.rsf
        model1_multwav_iter1.rsf    model1_shot181_data.rsf
  '''
#fetch data from FTP server
for fil in str.split(flist):
    Fetch(fil,filesloc,FTPserver)
Alias('fetch',['model1_shot181_p0.rsf','model1_shot181_srme_iter1_local.rsf', \
               'model1_multwav_iter1.rsf','model1_shot181_data.rsf'])


Flow('model1_shot181_data_n.rsf','model1_shot181_data.rsf','sfdd form=native')
Flow('model1_shot181_p0_n.rsf','model1_shot181_p0.rsf','sfdd form=native')
Flow('model1_multwav_iter1_n.rsf','model1_multwav_iter1.rsf','sfdd form=native')
Flow('model1_shot181_srme_iter1_local_n.rsf','model1_shot181_srme_iter1_local.rsf','sfdd form=native')
###################################################################
#                                                                 #
#               Diagonal estimate for different Lambda             #
#                                                                 #
###################################################################
cluster = ARGUMENTS.get('cluster',0)
if int(cluster) is 1:
        filein = 'model1_shot181_data_n.rsf'
        lsqrmult = 'model1_multwav_iter1_n.rsf'
        lambdas = [0.0,0.1,0.5,1.0,2.0,5.0]
        maxiter = 100
        for lamb in lambdas:
                fileout = 'scaler_iter1_%3.1f.rsf' % lamb
                scalemult = '%3.1fmodel1_iter1_scalmult.rsf' % lamb
                command = python+' ../code/Mdiag_bfgs.py lsqrmult=%s scalemult=%s  \
                      lamb=%3.1f maxiter=%d' %(lsqrmult,scalemult,lamb,maxiter)
                Flow([fileout,scalemult],[filein,lsqrmult],command)
else:
        Fetch('0.0model1_iter1_scalmult_xdr.rsf',filesloc1,FTPserver)
        Fetch('0.5model1_iter1_scalmult_xdr.rsf',filesloc1,FTPserver)
        Flow('0.0model1_iter1_scalmult.rsf','0.0model1_iter1_scalmult_xdr.rsf', \
             'sfdd form=native')
        Flow('0.5model1_iter1_scalmult.rsf','0.5model1_iter1_scalmult_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 = 7
padlen = 10

#-------------Over matched multiples Bayessian separation---------                     
Flow('0.0model1_iter1_prim.rsf',['model1_shot181_data_n.rsf','0.0model1_iter1_scalmult.rsf'],\
         'sfmath a=${SOURCES[1]} output="input-a"')
Flow('0.0Bayes_iter1_prim.rsf 0.0Bayes_iter1_mult.rsf',['0.0model1_iter1_scalmult.rsf','0.0model1_iter1_prim.rsf'],
                python+' '+pms_bayes+"""   
                        predmult=${SOURCES[0]}
                        predprim=${SOURCES[1]}
                        niter=%(interBayes)s
                        padlen=%(padlen)s
                        transparams=6,8,1
                        mu=1.0
                        lamb1=0.5
                        lamb2=0.2
                        poutput=${TARGET}
                        moutput=${TARGETS[1]}
                        --tmpdatapath=./
                        --localtmpdir=./
                """ %vars(),stdin=0,stdout=-1)
#-------------Correctly matched multiples Bayessian separation---------                     
Flow('0.5model1_iter1_prim.rsf',['model1_shot181_data_n.rsf','0.5model1_iter1_scalmult.rsf'],\
         'sfmath a=${SOURCES[1]} output="input-a"')
Flow('0.5Bayes_iter1_prim.rsf 0.5Bayes_iter1_mult.rsf',['0.5model1_iter1_scalmult.rsf','0.5model1_iter1_prim.rsf'],
                python+' '+pms_bayes+"""   
                        predmult=${SOURCES[0]}
                        predprim=${SOURCES[1]}
                        niter=%(interBayes)s
                        padlen=%(padlen)s
                        transparams=6,8,1
                        mu=1.0
                        lamb1=0.5
                        lamb2=0.2
                        poutput=${TARGET}
                        moutput=${TARGETS[1]}
                        --tmpdatapath=./
                        --localtmpdir=./
                """ %vars(),stdin=0,stdout=-1)

#-------------No matching multiples(SRME multiples) Bayessian separation---------                     
Flow('model1_multwav_iter1_prim.rsf',['model1_shot181_data_n.rsf','model1_multwav_iter1_n.rsf'],\
         'sfmath a=${SOURCES[1]} output="input-a"')
Flow('Bayes_multwav_iter1_prim.rsf Bayes_multwav_iter1_mult.rsf',['model1_multwav_iter1_n.rsf','model1_multwav_iter1_prim.rsf'],
                python+' '+pms_bayes+"""   
                        predmult=${SOURCES[0]}
                        predprim=${SOURCES[1]}
                        niter=%(interBayes)s
                        padlen=%(padlen)s
                        transparams=6,8,1
                        mu=1.0
                        lamb1=0.5
                        lamb2=0.2
                        poutput=${TARGET}
                        moutput=${TARGETS[1]}
                        --tmpdatapath=./
                        --localtmpdir=./
                """ %vars(),stdin=0,stdout=-1)

###################################################################
#                                                                 #
#               Make figures needed in the paper                  #
#                                                                 #
###################################################################
def greym(title):
     return '''
     sfdd form=native |sfwindow n1=480|sfput d1=0.004 d2=15 o2=-2715|  
     sfgrey xinch=6 yinch=8 
     clip=2.5
     label1="Time" unit1="s" 
     label2="Offset" unit2="m"
     labelsz=6.0  titlesz=6.5 title="%s"
      ''' % (title)

def scalplot(title):
     return '''
     sfdd form=native |sfput d1=1|  
     sfgraph labelsz=7.0 titlesz=7.0 
     xinch=9 yinch=6 
     label2="Scaling amplitude" 
     label1="Number of coefficient"
     labelsz=6.0  titlesz=6.5 title="%s"
      ''' % (title)
#---------------------basic plot----------------------
Plot('F1a','model1_shot181_data.rsf',greym(''))
Plot('F1b','model1_multwav_iter1.rsf',greym(''))
Plot('F1c','model1_shot181_p0.rsf',greym(''))
Plot('F1d','model1_shot181_srme_iter1_local.rsf',greym(''))
Plot('F1e','Bayes_multwav_iter1_prim.rsf',greym(''))
Plot('F1f','0.5Bayes_iter1_prim.rsf',greym(''))
Plot('F1i','0.0Bayes_iter1_prim.rsf',greym(''))


#---------------------annotations---------------------  
Plot('ant1',None,'box x0=6.85 y0=4.7 label="" xt=-0.5 yt=0.4 lab_fat=0 boxit=1 \
       size=0.2 pointer=1 x_oval=0.0 y_oval=0.0')
Plot('ant2',None,'box x0=8.8 y0=5.7 label="" xt=0.5 yt=0.4 lab_fat=0 boxit=1 \
       size=0.2 pointer=1 x_oval=0.0 y_oval=0.0')
#Plot('ant3',None,'box x0=6.0 y0=2.35 label="" xt=0.0 yt=0.0 lab_fat=0 boxit=1 \
#       size=0.0 pointer=1 x_oval=0.0 y_oval=0.0')
Plot('ant3',None,'box x0=7.3 y0=6.5 label="" xt=-0.5 yt=0.4 lab_fat=0 boxit=1 \
       size=0.2 pointer=1 x_oval=0.0 y_oval=0.0')

Plot('ant4',None,'box x0=7.2 y0=4.7 label="" xt=-0.5 yt=0.4 lab_fat=0 boxit=1 \
       size=0.2 pointer=1 x_oval=0.0 y_oval=0.0')
Plot('ant5',None,'box x0=9.05 y0=5.7 label="" xt=0.5 yt=0.4 lab_fat=0 boxit=1 \
       size=0.2 pointer=1 x_oval=0.0 y_oval=0.0')
Plot('ant',None,'box x0=6.85 y0=5.7 label="" xt=0.0 yt=0.0 lab_fat=0 boxit=1 \
       size=0.0 pointer=1 x_oval=0.0 y_oval=0.0')

#---------------------Final results in paper----------
Result('Figure1-a',['F1a']+['ant'],'Overlay')
Result('Figure1-b',['F1b']+['ant'],'Overlay')
Result('Figure1-c',['F1c']+['ant'],'Overlay')
Result('Figure1-d',['F1d']+['ant1']+['ant2'],'Overlay')
Result('Figure1-e',['F1e']+['ant4']+['ant5'],'Overlay')
Result('Figure1-f',['F1f']+['ant'],'Overlay')


End()