#!/usr/bin/python -itt

import numpy as np

from scipy import optimize
from numpy import *

class Parameter:
    def __init__(self, value):
            self.value = value

    def set(self, value):
            self.value = value

    def __call__(self):
            return self.value

def fit(function, parameters, y, x = None):
    def f(params):
        i = 0
        for p in parameters:
            p.set(params[i])
            i += 1
        return y - function(x)

    if x is None: x = arange(y.shape[0])
    p = [param() for param in parameters]
    optimize.leastsq(f, p)


class gapdpulse( object ):

    bsl = -1.0
    height = 18.0
    start  = 50.
    rising = 1.2
    tau1 = 7.
    tau2 = 25.
    noise = 0.0001
    sampling = False

    def __call__(self, x):
        if type(x) != type(np.empty(1)):
            x = np.array(x)
        val = np.zeros(x.shape)
        zeros = np.zeros(x.shape)
        
        #shortcuts
        bsl = self.bsl()
        height = self.height()
        start = self.start()
        rising = self.rising()
        tau1 = self.tau1()
        tau2 = self.tau2()
        
        val += bsl
        
        
        
        # this is not really the maximum, but about it.... so I call it max
        max = start + rising * tau1

        e1 = height * (1 - np.exp(-(x-start)/tau1 ) )
        e2 = -1 * height * (1 - np.exp(-(x-max)/tau2 ) )
        val += np.where( x>start , e1, zeros)
        val += np.where( x>max , e2, zeros)

        val += np.random.normal(0, self.noise, x.shape)
        
        #sampling
        self.presampling = val.copy()
        if self.sampling == True:
            val = np.round(val*2)/2.0
        self.postsampling = val.copy()
        return val 
        






# Target function, 6  parameters
def fitfunc( p, x):

    # give names to the parameters
    bsl = p[0]
    height = p[1]
    start = p[2]
    rising = p[3]
    tau1 = p[4]
    tau2 = p[5]
    
    # these are just helper variables
    max = start + rising * tau1
    e1 = height * (1 - np.exp(-(x-start)/tau1 ) )   # the rising edge
    e2 = -1 * height * (1 - np.exp(-(x-max)/tau2 ) )# the falling edge 
    
    y = bsl
    if x > start:
        y += e1
    if x > max:
        y += e2
    return y

def fitfunc2(p, x):
    # give names to the parameters
    bsl = p[0]      # just the baseline
    height = p[1]   # this is not the pulse height, but the maximum of the discharging edge 
                    # it *would* be the pulse height only in case the 
                    # recharging process would set in later...
                    
    start = p[2]    # start of avalanche aka start of discharge of diode capacitance
    stop = p[3]     # stop of avalanche, or start of recharging of diode
    tau1 = p[4]     # time constant of discharging process
    tau2 = p[5]     # time constant of recharging process
    
    # these are just helper variables
    
    e1 = height * (1 - np.exp(-(x-start)/tau1 ) )
    e2 = -1 * height * (1 - np.exp(-(x-stop)/tau2 ) )
    zeros = np.zeros(x.shape)
    
    y = np.zeros(x.shape)
    y += bsl
    y += np.where( x>start , e1, zeros)
    y += np.where( x>stop , e2, zeros)
    return y



# Distance to the Targetfunction
def errfunc( p, x , y):
    return fitfunc2(p, x) - y



if __name__ == '__main__':
    import matplotlib.pyplot as plt
    plt.ion()
    from gen_TH1D import SignalGeneratorCSV
    
    g = SignalGeneratorCSV('EdgeOverlayTemplate_AllPixel_1440_0.csv')
    
    x = np.linspace(0,300,300,False)
    
    fig = plt.figure()
    plt.grid(True)
    
    plt.plot(x,g.csv_maxprob, label='max prob')

    pnames = ['bsl', 'height', 'start','stop','tau1','tau2']

    p0 = [-1., 10., 50., 70, 30., 30.] # Initial guess for the parameters
    
    # plot function before fitting
    #plt.plot( x, fitfunc2( p0, x) , '.:')
    
    p1, success = optimize.leastsq(errfunc, p0[:], args=(x[0:300], g.csv_maxprob[0:300]))
    
    # plot function after fitting
    plt.plot( x, fitfunc2( p1, x) , '.:')
    print 'Start values:'
    for par,name in zip(p0,pnames):
        print name, ':', par
    print
    print
    print '------------------'
    print 'Fit parameters'
    for par,name in zip(p1,pnames):
        print name, ':', par
    print '------------------'

    print 'Fit successfull?:',bool(success)