# Author: G. Hennenfent
#         Seismic Laboratory for Imaging and Modeling
#         Department of Earth & Ocean Sciences
#         The University of British Columbia
#         
# Date  : December, 07

from   rsfproj import *
try:
    import numpy   as     np
except:
    raise Exception('Cannot find numpy')
import os

########################################################################
# INITIALIZATION
########################################################################
matlab         = WhereIs('matlab')
if not matlab:
    raise Exception('Cannot find Matlab')

matROOT = '../matfcts'

# define problem parameters
N = 1024
np.random.seed(2007)
n = np.random.randint(np.round(N/5.),np.round(4./5.*N),1)
k = np.random.randint(1,np.round(N/5.),1)

# define plot parameters
min1 = 0
max1 = 150
min2 = -.1
max2 = 4

# define limited number of iterations
LimIters = 30

# construct problems
Flow('A',None,
     'sfspike n1=%(n)g n2=%(N)g | sfnoise rep=y seed=2008 | sfmath output="input/57" '%vars() )
pos = np.sort(np.round(np.random.rand(k)*N))
amp = np.random.randn(k)

Flow('x0',None,
     'sfspike n1=%d nsp=%d k1=%s mag=%s'
     %(N,k,string.join(map(str,pos),','),string.join(map(str,amp),',')) )

########################################################################
# COMPUTE PARETO CURVE
########################################################################
# number of points to compute the curve
nbpts = 100

matfun  = 'MakeParetoCurve'
Flow('res taus',['A','x0',os.path.join(matROOT,matfun+'.m')],
     matlab + '''
     -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun)s('$SOURCE','${SOURCES[1]}','$TARGET','${TARGETS[1]}',%(nbpts)d);quit"
     '''%vars(),stdin=0,stdout=-1)

Flow('pcurve','taus res','sfcmplx ${SOURCES[:-1]} | sftransp',stdin=0)
Plot('pcurve',
       '''
       sfgraph label1="\F15 7\F-1 x\F15 7\F-1 \_\s75 1" label2="\F15 7\F-1 y-Ax\F15 7\F-1 \_\s75 2" title=""
       min1=%(min1)g max1=%(max1)g min2=%(min2)g max2=%(max2)g symbol=+ symbolsz=6 parallel2=n
       '''%vars() )

suffix = ['SPGL1','IST','ISTc','IRLS']

########################################################################
# TRACK SPGL1
########################################################################
# define filename
matfun = 'Track'  +suffix[0]

for Iters,case in [[100,'_UnLim'],[LimIters,'_Lim']]:
    # define filenames
    xnorms = 'xnorms_'+suffix[0]+case
    rnorms = 'rnorms_'+suffix[0]+case
    lplot  = suffix[0]+case

    # track SPGL1
    Flow([xnorms,rnorms],['A','x0',os.path.join(matROOT,matfun+'.m')],
         matlab + '''
         -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun)s('$SOURCE','${SOURCES[1]}','$TARGET','${TARGETS[1]}',%(Iters)g);quit"
         '''%vars(),stdin=0,stdout=-1)
    Flow(lplot,[xnorms, rnorms],'sfcmplx ${SOURCES[:-1]}',stdin=0)

    # plot results
    Plot(lplot,
         '''
         sfgraph label1="\F15 7\F-1 x\F15 7\F-1 \_\s75 1" label2="\F15 7\F-1 y-Ax\F15 7\F-1 \_\s75 2" title=""
         min1=%(min1)g max1=%(max1)g min2=%(min2)g max2=%(max2)g
         plotcol=3 dash=5 parallel2=n
         '''%vars() )

########################################################################
# TRACK IST
########################################################################
# define filename
matfun = 'Track'  +suffix[1]

# define IST parameter (Lagrange multiplier)
lmult = 5e-4;

for Iters,case in [[2500,'_UnLim'],[LimIters,'_Lim']]:
    # define filenames
    xnorms = 'xnorms_'+suffix[1]+case
    rnorms = 'rnorms_'+suffix[1]+case
    lplot  = suffix[1]+case

    # track IST
    Flow([xnorms,rnorms],['A','x0',os.path.join(matROOT,matfun+'.m'),os.path.join(matROOT,'SolveIST.m')],
         matlab + '''
         -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun)s('$SOURCE','${SOURCES[1]}','$TARGET','${TARGETS[1]}',%(Iters)g,%(lmult)g);quit"
         '''%vars(),stdin=0,stdout=-1)
    Flow(lplot,[xnorms,rnorms],'sfcmplx ${SOURCES[:-1]} | sftransp',stdin=0)

    # plot results
    Plot(lplot,
         '''
         sfgraph label1="\F15 7\F-1 x\F15 7\F-1 \_\s75 1" label2="\F15 7\F-1 y-Ax\F15 7\F-1 \_\s75 2" title=""
         min1=%(min1)g max1=%(max1)g min2=%(min2)g max2=%(max2)g plotcol=4 dash=3 parallel2=n
         '''%vars() )

########################################################################
# TRACK ISTc
########################################################################
# define filename
matfun = 'Track'  +suffix[2]

for Iters,case,locInnerIters in [[500,'_UnLim',10],[LimIters,'_Lim',2]]:
    # define filenames
    xnorms = 'xnorms_'+suffix[2]+case
    rnorms = 'rnorms_'+suffix[2]+case
    lplot  = suffix[2]+case

    # track ISTc
    Flow([xnorms,rnorms,'lmults'+case],['A','x0',os.path.join(matROOT,matfun+'.m'),os.path.join(matROOT,'SolveISTc.m')],
         matlab + '''
         -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun)s('$SOURCE','${SOURCES[1]}','$TARGET','${TARGETS[1]}','${TARGETS[2]}',%(Iters)d,%(locInnerIters)d);quit"
         '''%vars(),stdin=0,stdout=-1)
    Flow(lplot,[xnorms,rnorms],'sfcmplx ${SOURCES[:-1]} | sftransp',stdin=0)

    # plot results
    Plot(lplot,
         '''
         sfgraph label1="\F15 7\F-1 x\F15 7\F-1 \_\s75 1" label2="\F15 7\F-1 y-Ax\F15 7\F-1 \_\s75 2" title=""
         min1=%(min1)g max1=%(max1)g min2=%(min2)g max2=%(max2)g parallel2=n
         plotcol=5
         '''%vars() )

########################################################################
# TRACK IRLS
########################################################################
# define IRLS parameters
a = .5

sigma = 0
damp = 1e-3

for InnerIters,Iters,case in [[1e2,5e3,'_UnLim'],[10,LimIters,'_Lim']]:
    # define filenames
    matfun = 'Track'  +suffix[3]
    xnorms = 'xnorms_'+suffix[3]+case
    rnorms = 'rnorms_'+suffix[3]+case
    lplot  = suffix[3]+case

    # track IRLS
    Flow([xnorms,rnorms],['A','x0',os.path.join(matROOT,matfun+'.m'),os.path.join(matROOT,'SolveIRLS.m')],
         matlab + '''
         -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun)s('$SOURCE','${SOURCES[1]}','$TARGET','${TARGETS[1]}',%(Iters)g,%(InnerIters)g,%(a)g,%(sigma)g,%(damp)g);quit"
         '''%vars(),stdin=0,stdout=-1)
    Flow(lplot,[xnorms,rnorms],'sfcmplx ${SOURCES[:-1]} | sftransp',stdin=0)

    # plot results
    Plot(lplot,
         '''
         sfgraph label1="\F15 7\F-1 x\F15 7\F-1 \_\s75 1" label2="\F15 7\F-1 y-Ax\F15 7\F-1 \_\s75 2" title=""
         min1=%(min1)g max1=%(max1)g min2=%(min2)g max2=%(max2)g plotcol=7 dash=2 parallel2=n
         '''%vars() )

########################################################################
# FINAL PLOTS
########################################################################
Result('plot','pcurve ISTc_UnLim IST_UnLim SPGL1_UnLim IRLS_UnLim','Overlay')
Result('plotLim','pcurve ISTc_Lim IST_Lim SPGL1_Lim IRLS_Lim','Overlay')

End()

########################################################################
# RUN TIME (iMac 2.4 GHz Intel Core 2 Duo, 2GB 667 MHz DDR2 SDRAM)
########################################################################
## time scons
## scons: done building targets.
##  real    1m24.173s
##  user    0m59.979s
##  sys     0m21.136s