#!/usr/bin/python

from math import *
import numpy as np
import Manyo as m
#import PeakFit as Adv
import Adv
import time

####    linear combination of gaussian    ####
def dist(x, param):
    sum=param[0]
    for i in range(len(param[1])):
        a, c, w = param[1][i]
        sum = sum + a*exp(-1.0*((x-c)/w)**2)
    return sum

####    sign function    ####
def sign(r):
    if r > 0.0:
        s = 1.0
    elif r == 0.0:
        s = 0.0
    else:
        s = -1.0
    return s

####    initialize an ElementContainer    ####
def initElementContainer(xmin, xmax, nDiv, param):

    bin = m.MakeDoubleVector(nDiv+1)
    delta = (xmax - xmin)/nDiv
    for i in range(nDiv+1):
        bin[i] = xmin + delta*i

    y = m.MakeDoubleVector(nDiv)
    e = m.MakeDoubleVector(nDiv)
    for i in range(nDiv):
        xc = (bin[i] + bin[i+1])/2.0

        p = dist(xc, param) # orijinal
        ####    no noize    ####
        y[i] = p
        #e[i] = param[0]/3.0
        e[i] = 1.0
        ####    uniform background noize    ####
        #q = sign(np.random.rand(1)-0.5)*(param[0]/3.0)*sqrt(-1.0*log(np.random.rand(1))) # uniform back graund noize
        #y[i] = p + q
        #e[i] = param[0]/3.0
        ####    noize proprtianl to original data
        #q = sign(np.random.rand(1)-0.5)*(p/10.0)*sqrt(-1.0*log(np.random.rand(1)))        # noise that is proportinal to the original
        #y[i] = p + q
        #e[i] = p/10.0

    ec = m.ElementContainer()
    ec.AddToHeader("run number", 1)
    ec.AddToHeader("level",      1)
    ec.AddToHeader("Inst.",      "manyo")
    ec.Add("TOF",       bin, "sec." )
    ec.Add("Intensity", y,   "count")
    ec.Add("Error",     e,   "count")
    ec.SetKeys("TOF", "Intensity", "Error")

    return ec

def outputElementContainer(ec):
    bin = ec.Put(ec.PutXKey())
    y   = ec.Put(ec.PutYKey())
    e   = ec.Put(ec.PutEKey())

    print str("No.").rjust(4), str("----------- bin -----------").rjust(27), str("xc").rjust(10), str("Intensity").rjust(23), str("error").rjust(23)
    for i in range(ec.PutSize(ec.PutYKey())):
         print str(i).rjust(4), "[", str(bin[i]).rjust(10), ",", str(bin[i+1]).rjust(10), ")", str((bin[i]+bin[i+1])/2.0).rjust(10), str(y[i]).rjust(23), str(e[i]).rjust(23)

####    plot    ####
def plotPeakData(src, smoothed, peakData):
    srcBin=src.Put(src.PutXKey())

    g=m.GnuplotInterface()

    g.e("set xlabel '" + src.PutXKey() + " / " + src.PutUnit(src.PutXKey()) + "'")
    g.e("set xrange[" + str(srcBin.front()) + ":" +str(srcBin.back())+"]")

    g.e("set ylabel '" + src.PutYKey() + " / " + src.PutUnit(src.PutYKey()) + "'")
    g.e("set autoscale y")

    v=peakData.toVector()
    for i in range(peakData.size()):
        peak=peakData.getPeak(i)
        v=peak.toFullVector()
        h=v[0] # height
        c=v[1] # position
        l=v[3] # left  pos at HWHM
        u=v[4] # right pos at HWHM
        #print "peak id = ", i, p
        g.e("set arrow "+str(3*i+1)+" from "+str(c)+", 0 to "+str(c)+", 150000 nohead\n")  # guide line for a peak or a sholder
        g.e("set label "+str(4*i+1)+" '"+str(c)+"' at "+str(c)+",10000" )                  # the position of a peak or a sholder
        g.e("set label "+str(4*i+2)+" '"+str(h)+"' at "+str(c)+","+str(h)+" front" )       # the height   of a peak or a sholder
        if c-l > 0.0:
            g.e("set arrow "+str(3*i+2)+" from "+str(c)+","+str(h/2.0)+" to "+str(l)+","+str(h/2.0)) # guide line for the lower bound at hwhm
            g.e("set label "+str(4*i+3)+" '"+str(c-l)+"' at "+str((c+l)/2.0)+","+str(h/2.0)+" center front" )  # the width
        if u-c > 0.0:
            g.e("set arrow "+str(3*i+3)+" from "+str(c)+","+str(h/2.0)+" to "+str(u)+","+str(h/2.0)) # guide line  for the upper bound at hwhm
            g.e("set label "+str(4*i+4)+" '"+str(u-c)+"' at "+str((c+u)/2.0)+","+str(h/2.0)+" center front" )  # the width

    while 1:
        g.e("plot '-' using 1:2 with lines linewidth 1 title 'source', '-' using 1:3 with lines linewidth 1 title 'smoothed'")

        for i in range(src.PutSize(src.PutYKey())):
            g.e( str( (srcBin[i]+srcBin[i+1])/2.0 ) + " " + str( src.Put(src.PutYKey(), i) )+ " " + str( smoothed.Put(smoothed.PutYKey(), i) ) )
        g.e("e")

        c=raw_input("plotPeakData: hit any key except return for next\n")
        if c != "":
            break
    #g.e("pause -1 'hit any key'")

    del g
    del srcBin


def plotNorm(method):
    ic = method.getConvergenceHistoryForInt4(  Adv.NewLevmar.ITERATION_COUNT  )
    rf = method.getConvergenceHistoryForDouble(Adv.NewLevmar.R_FACTOR         )
    rn = method.getConvergenceHistoryForDouble(Adv.NewLevmar.RESIDUAL_ERR_NORM)
    gn = method.getConvergenceHistoryForDouble(Adv.NewLevmar.GRADIENT_NORM    )
    pn = method.getConvergenceHistoryForDouble(Adv.NewLevmar.PARAM_DIFF_NORM  )
    t  = method.getConvergenceHistoryForInt4(  Adv.NewLevmar.ITERATION_TIME    )

    g=m.GnuplotInterface()

    g.e("set xlabel 'iteration'")
    if ic.size() == 1:
        cmd="set xrange [-0.5:0.5]"
    else:
        cmd="set xrange ["+repr(ic.front())+":"+repr(ic.back())+"]"
    g.e(cmd)

    g.e("set ylabel 'norm'")
    g.e("set autoscale y")
    g.e("set logscale y 10.0")

    cmd =  "plot"
    cmd = cmd + " '-' using 1:2 with linespoints title '" + Adv.NewLevmar.R_FACTOR + "'"
    #cmd = cmd + ","
    #cmd = cmd + " '-' using 1:3 with linespoints title '" + Adv.NewLevmar.RESIDUAL?ERR_NORM + "'"
    #cmd = cmd + ","
    #cmd = cmd + " '-' using 1:4 with linespoints title '" + Adv.NewLevmar.GRADIENT_NORM + "'"
    #cmd = cmd + ","
    #cmd = cmd + " '-' using 1:5 with linespoints title '" + Adv.NewLevmar.PARAM_DIFF_NORM + "'"
    g.e(cmd)

    for i in range(ic.size()):
        data =              str(ic[i]/1.0).rjust(5)
        data = data + " " + str(rf[i]    ).rjust(23)
        data = data + " " + str(rn[i]    ).rjust(23)
        data = data + " " + str(gn[i]    ).rjust(23)
        if pn[i] != 0.0:
            data = data + " " + str(pn[i]    ).rjust(23)
        else:
            data = data + " " + "                       "
        g.e(data)
        #print data
    g.e("e")

    g.e("pause " + str(t.back()/1000000))

    del g
    del ic
    del rf
    del rn
    del gn
    del pn
    del t


def plotAll(src, smoothed, fittied, components):
    xkey=src.PutXKey()
    ykey=src.PutYKey()
    srcBin      = src.Put(xkey)

    g=m.GnuplotInterface()

    g.e("set xlabel '" + xkey + " / " + src.PutUnit(xkey) + "'")
    g.e("set xrange[" + str(srcBin.front()) + ":" +str(srcBin.back())+"]")

    g.e("set ylabel '" + ykey + " / " + src.PutUnit(ykey) + "'")
    g.e("set autoscale y")

    arrowCount=1
    labelCount=1
    for j in range(components.PutSize()):
        func=components.Put(j).PutHeader().PutString(Adv.NewLevmar.FUNCTIONS)
        v=components.Put(j).PutHeader().PutDoubleVector(Adv.NewLevmar.PARAMETER_VALUES)
        if func == "gaussian":
            h=v[0]
            c=v[1]
            w=v[2]
            # position and height of a component
            g.e("set arrow "+str(arrowCount)+" from "+str(c)+", 0 to "+str(c)+","+ str(h) + " front") # guide line for a peak or a sholder
            arrowCount=arrowCount+1
            g.e("set label "+str(labelCount)+" '"+str(c)+"' at "+str(c)+",0.0 front" )                # the position of a peak or a sholder
            labelCount=labelCount+1
            g.e("set label "+str(labelCount)+" '"+str(h)+"' at "+str(c)+","+str(h)+" front" )         # the height   of a peak or a sholder
            labelCount=labelCount+1
            # width of a component
            g.e("set arrow "+str(arrowCount)+" from "+str(c)+","+str(h/2.0)+" to "+str(c-w)+","+str(h/2.0)) # guide line for the lower bound at hwhm
            arrowCount=arrowCount+1
            g.e("set label "+str(labelCount)+" '"+str(w)+"' at "+str((c+(c-w))/2.0)+","+str(h/2.0)+" center front" )  # the width
            labelCount=labelCount+1
        if func == "constant":
            h=v[0]
            g.e("set label "+str(labelCount)+" '"+str(w)+"' at "+str((c+(c-w))/2.0)+","+str(h/2.0)+" center front" )  # the width
            labelCount=labelCount+1

    while 1:
        cmd = "plot"
        cmd = cmd + " '-' using 1:2 with lines linewidth 1 title 'original'"
        cmd = cmd + ", '-' using 1:3 with lines linewidth 1 title 'smoothed'"
        cmd = cmd + ", '-' using 1:4 with lines linewidth 1 title 'fittied'"
        for j in range(components.PutSize()):
            cmd = cmd + ", '-' using 1:"+str(5+j)+" with lines linewidth 1 title 'func "+str(j)+"'"
        g.e(cmd)

        for i in range(src.PutSize(ykey)):
            data = str( (srcBin[i]+srcBin[i+1])/2.0                ).rjust(7)
            data = data + " " + str( src.Put(ykey, i)              ).rjust(10)
            data = data + " " + str( smoothed.Put(ykey, i)         ).rjust(10)
            data = data + " " + str( fittied.Put(ykey, i)          ).rjust(10)
            for j in range(components.PutSize()):
                data = data + " " + str( components.Put(j).Put(ykey, i)).rjust(10)
            g.e(data)
        g.e("e")

        #g.e("replot")
        c=raw_input("plotAll: hit any key except return\n")
        if c != "":
            break

    del g
    del srcBin


############            main routine            ############
########        initialise an element conatainer        ####
xmin = 0.0
xmax = 15.0
nDiv = 1500
#      background  param for f0          param for f1           param for f2
param=[     0.0,    [[35000.0, 2.8, 1.0], [100000.0, 10.0, 1.5], [100000.0, 5.0, 1.5]]]
#param=[ 1000.0,    [[35000.0, 2.8, 1.0], [100000.0, 10.0, 1.5], [100000.0, 5.0, 1.5]]]
#param=[10000.0,    [[35000.0, 2.8, 1.0], [100000.0, 10.0, 1.5], [100000.0, 5.0, 1.5]]]

src = initElementContainer(xmin, xmax, nDiv, param)
#outputElementContainer(src)

########        peak search        ########
####    create the driver for peak search    ####
peakSearch = Adv.PeakSearch(src, Adv.BSPLINE)
####    initialise parameters for smoothing    ####
peakSearch.setParam(Adv.BSpline.NUMBER_OF_BREAK_POINTS, 16)

#peakSearch = Adv.PeakSearch(src, Adv.MOVING_AVERAGE)
#peakSearch.setParam(Adv.MovingAverage.WINDOW_WIDTH, 131)
#peakSearch.setParam(Adv.MovingAverage.WINDOW_UPPER, 65)
####    peak search    ####
if peakSearch.checkParam() :              # check the consistency between source data and parameters
    peakSearch.execute()                  # execute peak search

    smoothed = peakSearch.getResult()     # get smoothed data
    peakData = peakSearch.getPeaks()      # get peaks and sholders
    #peakData.Dump()                       # print peal data
    plotPeakData(src, smoothed, peakData) # plot peak data

########        peak fit        ########
####    create dirver for peak fitting    ####
peakFit = Adv.PeakFit(src, Adv.NEW_LEVMAR)  # the fitting method is NewLevmar that supports multi-thread

####    initialise Levmar (set parameter values)    ####
# print template for Levmar parameters
paramSet=peakFit.getParam()
#paramSet.dump()

####    set domain    ####
lower=xmin
upper=xmax
c=input("lower of domain ?\n")
if c != "":
    lower=c
c=input("upper of domain ?\n")
if c != "":
    c=upper
peakFit.setDomain(lower, upper)


# difine Levmar mode
peakFit.setParam(Adv.NewLevmar.CONSTRAIN,          Adv.NewLevmar.NO_CONSTRAIN)
peakFit.setParam(Adv.NewLevmar.USE_NUMERICAL_DIFF, False)
peakFit.setParam(Adv.NewLevmar.DIFF_METHOD,        Adv.NewLevmar.FORWARD)
peakFit.setParam(Adv.NewLevmar.USE_DATA_WEIGHTS,   True)
peakFit.setParam(Adv.NewLevmar.MAX_ITERATIONS,     1000)
peakFit.setParam(Adv.NewLevmar.OUTPUT_INTERVAL,    50)

# fitting functions
#l=peakFit.getLowerBoun()
#u=peakFit.getUpperBoun()
expr=""
#for i in range(1, 4):
for i in range(peakData.size()):
    #c= peakData.getPeak(i).getPosition()
    #if l <= c && c <= u:
    expr = expr + " g"
#expr = expr + " c"
print expr
peakFit.setParam(Adv.NewLevmar.FUNCTIONS, expr)

# initial values of fitting parameters
v=peakData.toPyList()
#del v[12:15]
#del v[0:3]
#v.append(1.0)
print peakData.size()
peakFit.setParam(Adv.NewLevmar.PARAMETER_VALUES, v)

# lower bounds of fitting parameters
lb=[ x*0.9 for x in v ]
peakFit.setParam(Adv.NewLevmar.LOWER_BOUNDS, lb)

# upper bounds of fitting parameters
ub=[ x*1.1 for x in v ]
peakFit.setParam(Adv.NewLevmar.UPPER_BOUNDS, ub)

####    do fitting    ####
if peakFit.checkParam() :  # check the consistency of source data and Levmar parameters
    peakFit.fit()              # start to fit
    while peakFit.isFitting(): # sleep while fitting
        time.sleep(1)
    peakFit.eval()             # eval the values of fitting functiosn after fitting

    fittedParam=peakFit.getFittedParam()     # get fitted parameters
    result=peakFit.getResult()               # get the values of the fitting functions
    components=peakFit.getResultComponents() # get the values of the component functions

    #fittedParam.dump()
    plotAll(src, smoothed, result, components)
    ####result.dump
