#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Jul 23 15:02:12 2017

@author: Sean
"""
import numpy as np
from scipy.optimize import curve_fit

def cut(x,n):
    x[x<n] = 0
    return x

def d_thrsd(a,n = 0.15):
    a = np.array(a)
    x = a.shape
    if len(x) == 1:
        thrsd = a.max() - n
        a = cut(a,thrsd)
    else:
        for i in range(x[0]):
            thrsd = a[i,:].max() - n
            a[i,:] = cut(a[i,:],thrsd)
    return a

class para_fit:
    
    def __init__(self,thrsd = 0):
        self.thrsd = thrsd
        pass

    def para(self,x,a1, b1, c1):
        v = -c1*(x-a1)**2 + b1
        v = cut(v,self.thrsd)
        return v
    
    def lorentz(self,x,a1,c1):
        return (1/3.14)*c1*(1/((x-a1)**2 + c1))
    
    def Gauss(self,x,a1,b1,c1):
        return c1*np.exp(-(x-a1)**2/(2*b1**2))
    
    def dual_lrz(self,x,a1,c1,a2,c2,b1,b2):
        return b1 * self.lorentz(x,*[a1,c1]) + b2 * self.lorentz(x, *[a2,c2])
    
    def dual_para(self,x,a1,b1,a2,b2,c1,c2):
        return self.para(x,*[a1,b1,c1]) + self.para(x,*[a2,b2,c1])
    
    def dual_Gauss(self,x,a1,b1,c1,a2,b2,c2):
        return self.Gauss(x,*[a1,b1,c1]) + self.Gauss(x,*[a2,b2,c2])

    def fit(self,xdata, ydata, method = 'dual_lrz'):
        
        result = {
                'dual_lrz': curve_fit(self.dual_lrz, xdata, ydata, p0 = [0.1,1e-6,0.3,1e-6,0.5,0.5], bounds = (1e-7, [1,0.1,1,0.1,1,1])),
                
                'dual_para': curve_fit(self.dual_para, xdata, ydata,p0 = [0.1,0.1,0.4,0.1,1000,1000], bounds=(1e-7, [1,0.5,1,0.5,10000,10000])),
                    
                'dual_Gauss': curve_fit(self.dual_Gauss, xdata, ydata, p0 = [0.1,0.01,1,0.3,0.01,1], bounds = (1e-4, [1,0.2,10,1,0.2,10]))
                }[method]

        return result
    