from Numeric import *
from LinearAlgebra import *
import time
__author__ = "V.A. Sole <sole@esrf.fr>"
# codes understood by the routine
CFREE       = 0
CPOSITIVE   = 1
CQUOTED     = 2
CFIXED      = 3
CFACTOR     = 4
CDELTA      = 5
CSUM        = 6
CIGNORED    = 7

ONED = 0

def LeastSquaresFit(model, parameters0, data0, maxiter = 100,constrains=[],
                        weightflag = 0,model_deriv=None):
    parameters = array(parameters0)
    if ONED:
      data = array(data0)
      x = data[0:2,0]
    #import SimplePlot
    #SimplePlot.plot([data[:,0],data[:,1]],yname='Received data')
    else:
        x=array(map(lambda y:y[0],data0))  
    if len(constrains) == 0:
        try:
            model(parameters,x)
            constrains = [[],[],[]]
            for i in range(len(parameters0)):
                constrains[0].append(0)
                constrains[1].append(0)
                constrains[2].append(0)
            return RestreinedLeastSquaresFit(model,parameters0,
                                    data0,maxiter,
                                    constrains,weightflag)
        except TypeError:
            import traceback
            import sys
            traceback.print_exception(sys.exc_type, sys.exc_value, sys.exc_traceback)
            
            print "Using Konrad Hinsen's method"
            print "You should reconsider how to write your function"
            from Scientific.Functions.LeastSquares import leastSquaresFit
            fittedpar, chi_sq = leastSquaresFit(model, parameters0, data0)
            return fittedpar, chi_sq, zeros(len(fittedpar),Float).tolist()
    else:
            return RestreinedLeastSquaresFit(model,parameters0,
                                    data0,maxiter,
                                    constrains,weightflag,model_deriv=model_deriv)

def RestreinedLeastSquaresFit(model0,parameters0,data0,maxiter,
                constrains0,weightflag,model_deriv=None):
    #get the codes:
    # 0 = Free       1 = Positive     2 = Quoted
    # 3 = Fixed      4 = Factor       5 = Delta
    # 6 = Sum        7 = ignored
    constrains=[[],[],[]]
    for i in range(len(parameters0)):
        constrains[0].append(constrains0[0][i])
        constrains[1].append(constrains0[1][i])
        constrains[2].append(constrains0[2][i])
    for i in range(len(parameters0)):
        if type(constrains[0][i]) == type('string'):
            #get the number
            if   constrains[0][i] == "FREE":
                 constrains[0][i] = CFREE
            elif constrains[0][i] == "POSITIVE":
                 constrains[0][i] = CPOSITIVE 
            elif constrains[0][i] == "QUOTED":
                 constrains[0][i] = CQUOTED 
            elif constrains[0][i] == "FIXED":
                 constrains[0][i] = CFIXED 
            elif constrains[0][i] == "FACTOR":
                 constrains[0][i] = CFACTOR 
                 constrains[1][i] = int(constrains[1][i]) 
            elif constrains[0][i] == "DELTA":
                 constrains[0][i] = CDELTA 
                 constrains[1][i] = int(constrains[1][i]) 
            elif constrains[0][i] == "SUM":
                 constrains[0][i] = CSUM 
                 constrains[1][i] = int(constrains[1][i]) 
            elif constrains[0][i] == "IGNORED":
                 constrains[0][i] = CIGNORED 
            elif constrains[0][i] == "IGNORE":
                 constrains[0][i] = CIGNORED 
            else:
               #I should raise an exception
               constrains[0][i] = 0
    # make a local copy of the function for an easy speed up ...
    model = model0
    parameters = array(parameters0)
    if ONED:
        data = array(data0)
        x = data[1:2,0]
    fittedpar = parameters.__copy__()
    flambda = 0.001
    iter = maxiter
    niter = 0
    if ONED:
        selfx = data [:,0]
        selfy = data [:,1]        
    else:
        selfx = array(map(lambda x:x[0],data0))
        selfy = array(map(lambda x:x[1],data0))
    selfweight = ones(selfy.shape,Float)
    if ONED:
        nr0, nc = data.shape
    else:
        nr0 = len(selfy)
        nc =  len(data0[0])
    if weightflag == 1:
            if nc == 3:
                #dummy = abs(data[0:nr0:inc,2])
                if ONED:
                    dummy = abs(data [:,2])
                else:
                    dummy = abs(array(map(lambda x:x[2],data0))) 
                selfweight = 1.0 / (dummy + equal(dummy,0))
                selfweight = selfweight * selfweight
            else:
                selfweight = 1.0 / (abs(selfy) + equal(abs(selfy),0))            
    n_param = len(parameters)
    selfalphazeros = zeros((n_param, n_param),Float)
    selfbetazeros = zeros((1,n_param),Float)
    index = arange(0,nr0,2)
    while (iter > 0):
        niter += 1
        if niter < 3:
        #if 1:
        #if 0:
                x=take(selfx,index)
                y=take(selfy,index)
                weight=take(selfweight,index)
        else:
                x=selfx
                y=selfy
                weight = selfweight       
                
        chisq0, alpha0, beta,\
        n_free, free_index, noigno, fitparam, derivfactor  =ChisqAlphaBeta(
                                                 model,fittedpar,
                                                 x,y,weight,constrains,model_deriv=model_deriv)
        nr, nc = alpha0.shape
        flag = 0
        while flag == 0:
            newpar = parameters.__copy__()
            if(1):
                alpha = alpha0 + flambda * identity(nr) * alpha0
                deltapar = matrixmultiply(beta, inverse(alpha))
            else:
                #an attempt to increase accuracy
                #(it was unsuccessful)
                alphadiag=sqrt(diagonal(alpha0))
                npar = len(sqrt(diagonal(alpha0)))
                narray = zeros((npar,npar),Float)
                for i in range(npar):
                    for j in range(npar):
                        narray[i,j] = alpha0[i,j]/(alphadiag[i]*alphadiag[j])                
                narray = inverse(narray + flambda * identity(nr))
                for i in range(npar):
                    for j in range(npar):
                        narray[i,j] = narray[i,j]/(alphadiag[i]*alphadiag[j])                
                deltapar = matrixmultiply(beta, narray)
            pwork = zeros(deltapar.shape, Float)
            for i in range(n_free):
                if constrains [0] [free_index[i]] == CFREE:
                    pwork [0] [i] = fitparam [i] + deltapar [0] [i]
                elif constrains [0] [free_index[i]] == CPOSITIVE:
                    #abs method
                    pwork [0] [i] = fitparam [i] + deltapar [0] [i]
                    #square method
                    #pwork [0] [i] = (sqrt(fitparam [i]) + deltapar [0] [i]) * \
                    #                (sqrt(fitparam [i]) + deltapar [0] [i])
                elif constrains [0] [free_index[i]] == CQUOTED:
                    pmax=max(constrains[1] [free_index[i]],
                            constrains[2] [free_index[i]])            
                    pmin=min(constrains[1] [free_index[i]],
                            constrains[2] [free_index[i]])
                    A = 0.5 * (pmax + pmin)
                    B = 0.5 * (pmax - pmin)
                    if (B != 0):
                        pwork [0] [i] = A + \
                                    B * sin(arcsin((fitparam[i] - A)/B)+ \
                                    deltapar [0] [i])
                    else:
                        print "Error processing constrained fit"
                        print "Parameter limits are",pmin,' and ',pmax
                        print "A = ",A,"B = ",B                                
                newpar [free_index[i]] = pwork [0] [i]
            newpar=array(getparameters(newpar,constrains))
            workpar = take(newpar,noigno)
            #yfit = model(workpar.tolist(), x)
            yfit = model(workpar,x)
            chisq = sum( weight * (y-yfit) * (y-yfit))
            if chisq > chisq0:
                flambda = flambda * 10.0
                if flambda > 1000:
                    flag = 1
                    iter = 0
            else:
                flag = 1 
                fittedpar = newpar.__copy__()
                if ((chisq0-chisq)/(chisq0+(chisq0==0))) < 0.01:
                    iter = 0
                chisq0 = chisq
                flambda = flambda / 10.0
                #print "iter = ",iter,"chisq = ", chisq
            iter = iter -1
    sigma0 = sqrt(abs(diagonal(inverse(alpha0))))
    sigmapar = getsigmaparameters(fittedpar,sigma0,constrains)
    return fittedpar.tolist(), chisq/(len(yfit)-len(sigma0)), sigmapar.tolist()

def ChisqAlphaBeta(model0, parameters, x,y,weight, constrains,model_deriv=None):
    model = model0
    #nr0, nc = data.shape
    n_param = len(parameters)
    n_free = 0
    fitparam=[]
    free_index=[]
    noigno = []
    derivfactor = []
    for i in range(n_param):
        if constrains[0] [i] != CIGNORED:
            noigno.append(i)
        if constrains[0] [i] == CFREE:
            fitparam.append(parameters [i])
            derivfactor.append(1.0)
            free_index.append(i)
            n_free += 1
        elif constrains[0] [i] == CPOSITIVE:
            fitparam.append(abs(parameters[i]))
            derivfactor.append(1.0)
            #fitparam.append(sqrt(abs(parameters[i])))
            #derivfactor.append(2.0*sqrt(abs(parameters[i])))
            free_index.append(i)
            n_free += 1
        elif constrains[0] [i] == CQUOTED:
            pmax=max(constrains[1] [i],constrains[2] [i])            
            pmin=min(constrains[1] [i],constrains[2] [i])
            if ((pmax-pmin) > 0) & \
               (parameters[i] < pmax) & \
               (parameters[i] > pmin):
                A = 0.5 * (pmax + pmin)
                B = 0.5 * (pmax - pmin)
                fitparam.append(parameters[i])
                #fitparam.append(arcsin((parameters[i] - A)/B))
                derivfactor.append(B*cos(arcsin((parameters[i] - A)/B)))
                free_index.append(i)
                n_free += 1                        
    fitparam = array(fitparam, Float)
    alpha = zeros((n_free, n_free),Float)
    beta = zeros((1,n_free),Float)
    delta = (fitparam + equal(fitparam,0.0)) * 0.00001
    nr  = x.shape[0]
    ##############
    # Prior to each call to the function one has to re-calculate the
    # parameters
    pwork = parameters.__copy__()
    for i in range(n_free):
        pwork [free_index[i]] = fitparam [i]
    newpar = getparameters(pwork.tolist(),constrains)
    newpar = take(newpar,noigno)    
    yfit = model(newpar, x)
    for i in range(n_free):
        if model_deriv is None:
            #pwork = parameters.__copy__()
            pwork [free_index[i]] = fitparam [i] + delta [i]
            newpar = getparameters(pwork.tolist(),constrains)
            newpar=take(newpar,noigno)
            f1 = model(newpar, x)
            pwork [free_index[i]] = fitparam [i] - delta [i]
            newpar = getparameters(pwork.tolist(),constrains)
            newpar=take(newpar,noigno)
            f2 = model(newpar, x)
            help0 = (f1-f2) / (2.0 * delta [i])
            help0 = help0 * derivfactor[i]        
            #removed I resize outside the loop:
            #help0 = resize(help0,(1,nr))
        else:
            newpar = getparameters(pwork.tolist(),constrains)
            help0=model_deriv(pwork,free_index[i],x)
            help0 = help0 * derivfactor[i]        
        
        if i == 0 :
            deriv = help0
        else:
            deriv = concatenate ((deriv,help0), 0)
    #line added to resize outside the loop
    deriv=resize(deriv,(n_free,nr))
    print y.shape
    print yfit.shape
    deltay = y - yfit
    help0 = weight * deltay
    for i in range(n_free):
        derivi = resize(deriv [i,:], (1,nr))
        help1 = resize(sum((help0 * derivi),1),(1,1))
        if i == 0:
            beta = help1
        else:
            beta = concatenate ((beta, help1), 1)
        help1 = innerproduct(deriv,weight*derivi)
        if i == 0:
            alpha = help1
        else:
            alpha = concatenate((alpha, help1),1)
    chisq = sum(help0 * deltay)
    return chisq, alpha, beta, \
           n_free, free_index, noigno, fitparam, derivfactor

def getparameters(parameters,constrains):
    # 0 = Free       1 = Positive     2 = Quoted
    # 3 = Fixed      4 = Factor       5 = Delta
    newparam=[]
    for i in range(len(constrains [0])):
        if constrains[0][i] == CFREE:
            newparam.append(parameters[i])
        elif constrains[0][i] == CPOSITIVE:
            #newparam.append(parameters[i] * parameters[i])
            newparam.append(abs(parameters[i]))
        elif constrains[0][i] == CQUOTED:
            newparam.append(parameters[i])
        elif abs(constrains[0][i]) == CFIXED:
            newparam.append(parameters[i])
        elif constrains[0][i] == CFACTOR:
            newparam.append(constrains[2][i]*parameters[int(constrains[1][i])])
        elif constrains[0][i] == CDELTA:
            newparam.append(constrains[2][i]+parameters[int(constrains[1][i])])
        elif constrains[0][i] == CIGNORED:
            newparam.append(0)            
        elif constrains[0][i] == CSUM:
            newparam.append(constrains[2][i]-parameters[int(constrains[1][i])])            
    return newparam

def getsigmaparameters(parameters,sigma0,constrains):
    # 0 = Free       1 = Positive     2 = Quoted
    # 3 = Fixed      4 = Factor       5 = Delta
    n_free = 0
    sigma_par = zeros(parameters.shape,Float)
    for i in range(len(constrains [0])):
        if constrains[0][i] == CFREE:
            sigma_par [i] = sigma0[n_free]
            n_free += 1
        elif constrains[0][i] == CPOSITIVE:
            #sigma_par [i] = 2.0 * sigma0[n_free]
            sigma_par [i] = sigma0[n_free]
            n_free += 1
        elif constrains[0][i] == CQUOTED:
            pmax = max(constrains [1] [i], constrains [2] [i])
            pmin = min(constrains [1] [i], constrains [2] [i])
            A = 0.5 * (pmax + pmin)
            B = 0.5 * (pmax - pmin)
            if (B > 0) & (parameters [i] < pmax) & (parameters [i] > pmin):
                sigma_par [i] = abs(B) * cos(parameters[i]) * sigma0[n_free]
                n_free += 1
            else:
                sigma_par [i] = parameters[i]
        elif abs(constrains[0][i]) == CFIXED:
            sigma_par[i] = parameters[i]
    for i in range(len(constrains [0])):
        if constrains[0][i] == CFACTOR:
            sigma_par [i] = constrains[2][i]*sigma_par[int(constrains[1][i])]
        elif constrains[0][i] == CDELTA:
            sigma_par [i] = sigma_par[int(constrains[1][i])]
        elif constrains[0][i] == CSUM:
            sigma_par [i] = sigma_par[int(constrains[1][i])]
    return sigma_par

def fitpar2par(fitpar,constrains,free_index):
    newparam = []
    for i in range(len(constrains [0])):
        if constrains[0][free_index[i]] == CFREE:
            newparam.append(fitpar[i])
        elif constrains[0][free_index[i]] == CPOSITIVE:
            newparam.append(fitpar[i] * fitpar [i])
        elif abs(constrains[0][free_index[i]]) == CQUOTED:
            pmax=max(constrains[1] [free_index[i]],constrains[2] [free_index[i]])            
            pmin=min(constrains[1] [free_index[i]],constrains[2] [free_index[i]])
            A = 0.5 * (pmax + pmin)
            B = 0.5 * (pmax - pmin)
            newparam.append(A + B * sin(fitpar[i]))
    return newparam
    
def gauss(param, t):
    return param[0] + param[1] * t + SpecfitFuns.gauss(param[2:],t)

def test(npoints):
    xx = arange (npoints)
    xx=resize(xx,(npoints,1))
    #yy = 1000.0 * exp (- 0.5 * (xx * xx) /15)+ 2.0 * xx + 10.5
    yy = gauss([10.5,2,1000.0,20.,15],xx)
    yy=resize(yy,(npoints,1))
    sy = sqrt(abs(yy))
    sy=resize(sy,(npoints,1)) 
    data = concatenate((xx, yy, sy),1)
    parameters = [0.0,1.0,900.0, 25., 10]
    stime = time.time()
    fittedpar, chisq, sigmapar = LeastSquaresFit(gauss,parameters,data)
    etime = time.time()
    print "Took ",etime - stime, "seconds"
    print "chi square  = ",chisq
    print "Fitted pars = ",fittedpar
    print "Sigma pars  = ",sigmapar


if __name__ == "__main__":
  import profile
  import SpecfitFuns
  import Specfit
  profile.run('test(10000)',"test")
  import pstats
  p=pstats.Stats("test")
  p.strip_dirs().sort_stats(-1).print_stats()
