#!/usr/bin/python

from math import *
import numpy as np
import Manyo as m
import Levmar
import time

####    linear combination of gaussian    ####
#def dist(x, n, a, c, w):
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):
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

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

    #nPeak=1
    #a=m.MakeDoubleVector(nPeak)
    #c=m.MakeDoubleVector(nPeak)
    #w=m.MakeDoubleVector(nPeak)
    #a[0]=100000.0
    #c[0]=7.5
    #w[0]=2.5
    y = m.MakeDoubleVector(nDiv)
    e = m.MakeDoubleVector(nDiv)
    for i in range(nDiv):
        #p = dist(xc[i], nPeak, a, c, w)
        p = dist(xc[i], param)
        q = sign(np.random.rand(1)-0.5)*(p/10.0)*sqrt(-1.0*log(np.random.rand(1)))
        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 plot(src, dest):
    srcBin  = src.Put(src.PutXKey())
    srcY    = src.Put(src.PutYKey())
    destBin = dest.Put(dest.PutXKey())
    destY   = dest.Put(dest.PutYKey())

    count=0
    g=m.GnuplotInterface()
    g.e("set xrange[0: 15]")
    g.e("set yrange[0: 150000]")
    g.e("plot '-' using 1:2 with lines linewidth 1 title'%d', '-' using 1:3 with lines linewidth 4" %count)
    while 1:
        for i in range(srcY.size()):
            #print str( (srcBin[i]+srcBin[i+1])/2.0 ) + " " + str( srcY[i] )+ " " + str( destY[i] )
            g.e( str( (srcBin[i]+srcBin[i+1])/2.0 ) + " " + str( srcY[i] )+ " " + str( destY[i] ) )
        g.e("e")
        c=raw_input("")
        if c != "":
            break

####    main routine    ####
xmin = 0.0
xmax = 15.0
nDiv = 1500
param1=[1000.0,     [[100000.0, 5.0, 1.5], [100000.0, 10.0, 1.5], [35000.0, 2.8, 1.0] ]]

src = initElementContainer(xmin, xmax, nDiv, param1)
outputElementContainer(src)

domain = Levmar.Domain()
domain.setSource(src)
domain.setRange(xmin, xmax)
print "[", str(domain.getLowerBound()).rjust(10), ",", str(domain.getUpperBound()).rjust(10), "]"
print "[", str(domain.getLowerBoundID()).rjust(10), ",", str(domain.getUpperBoundID()).rjust(10), "]"
print
print "type of the domain", domain.getType()
#domain.setType(Domain.OO)
#print "type of the domain", domain.getType()

levmar=Levmar.Levmar()

paramSet = levmar.setDefaultParam(src);
paramSet.dump()

# initial values of fitting parameters
param=m.DoubleVector()
param.push_back(60683.0)
param.push_back(3.42)
param.push_back(1.14)
param.push_back(102955.0)
param.push_back(5.02)
param.push_back(1.66)
param.push_back(100978.0)
param.push_back(10.0)
param.push_back(1.27)

# box constrain for fitting parameters
lb=m.DoubleVector()
for i in range(param.size()):
    lb.push_back(param[i]*0.9)
ub=m.DoubleVector()
for i in range(param.size()):
    ub.push_back(param[i]*1.1)

paramSet.add(Levmar.Levmar.PARAMETER_VALUES, param)
paramSet.add(Levmar.Levmar.LOWER_BOUNDS,     lb)
paramSet.add(Levmar.Levmar.UPPER_BOUNDS,     ub)

#combinations of fitting function
funcStr=""
for i in range(len(param1[1])):
    funcStr=funcStr+" g"
   
parser=Levmar.FuncParser(funcStr) # a string expression for list of function names or symbols
funcList=parser.parse()
paramSet.add(Levmar.Levmar.FUNCTIONS, funcList)

paramSet.replace(Levmar.Levmar.CONSTRAIN,          Levmar.Levmar.NO_CONSTRAIN)  # default Levmar.Levmar.BOX
paramSet.replace(Levmar.Levmar.USE_NUMERICAL_DIFF, False)                             # default True
#paramSet.replace(Levmar.Levmar.DIFF_METHOD,        Levmar.Levmar.CENTRAL)       # default FOWARD
#paramSet.replace(Levmar.Levmar.USE_DATA_WEIGHTS,   False)                             # default True
paramSet.replace(Levmar.Levmar.MAX_ITERATIONS,     1000)                              # default 1000
paramSet.replace(Levmar.Levmar.OUTPUT_INTERVAL,    50)                               # default 50

paramSet.replace(Levmar.Levmar.SCALING_FACTOR,     Levmar.Levmar.DEFAULT_SCALING_FACTOR)     # default 0.001
paramSet.replace(Levmar.Levmar.TOLERANCE,          Levmar.Levmar.DEFAULT_TOLERANCE)          # default 1.0e-17
paramSet.replace(Levmar.Levmar.GRADIENT_TOLERANCE, Levmar.Levmar.DEFAULT_GRADIENT_TOLERANCE) # default 1.0e-17
paramSet.replace(Levmar.Levmar.RELATIVE_TOLERANCE, Levmar.Levmar.DEFAULT_RELATIVE_TOLERANCE) # default 1.0e-17
paramSet.dump()

print "parameter check: ", levmar.checkParam(src, domain, paramSet)  # parameter check, true if parametes are consistent

if  levmar.checkParam(src, domain, paramSet):
    levmar.toInnerForm(src, domain, paramSet)  # translate inner form
    levmar.fit()                               # fitting
    levmar.eval()                              # evaluate values of functions useing fitted parameter
    result=m.ElementContainer()
    levmar.toElementContainer(src, result)     # copy results to a element container
    plot(src, result)
    #outputElementContainer(result)

