''' Collection of fitting functions '''
"""

Author : Thomas Haslwanter
Date : Nov-2015
Version: 1.5
"""

# Standard packages
import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
import pandas as pd
import os

# Special package
from collections import namedtuple
from statsmodels.formula.api import ols
try:
    from skimage import filters
except ImportError:     # for backwards capability to old installations
    from skimage import filter as filters
    
def demo_ransac():
    '''Find the best-fit circle in an image, using the RANSAC algorithm '''

    debug_flag = 1

    # Since this function is only used in demo_ransac, define it here
    def drawCircle(center,r):
        '''Draw a circle'''
        nPts = 100
        phi = np.linspace(0,2*np.pi,nPts)
        x = center[0] + r*np.cos(phi)
        y = center[1] + r*np.sin(phi)
        plt.hold(True)
        plt.plot(y,x,'r')

    # Get the data
    os.chdir(r'C:\Users\p20529\Coding\Matlab\imProc\ransac_fitCircle')
    data = plt.imread('0_5.bmp')
    
    # Eliminate edge artefacts
    rim = 20
    data = data[rim:-rim, rim:-rim]
    imSize = data.shape
    
    # Find edges
    edges = filters.sobel(data)
    edgePnts = edges>0.15
    np.sum(edgePnts)
    [x,y] = np.where(edgePnts==True)
       
    # set RANSAC parameters
    par={'eps': 3, 
         'ransac_threshold': 0.1, 
         'nIter': 500, 
         'lowerRadius': 5,
         'upperRadius': 200   
    }
    
    # Allocate memory, for center(2D), radius, numPoints (structured array)
    fitted = np.zeros(par['nIter'], \
    dtype={'names':['center', 'radius', 'nPts'], 'formats':['2f8', 'f8', 'i4']})
    
    for ii in range(par['nIter']):
        # Takes 3 random points, and find the corresponding circle
        randEdges = np.random.permutation(len(x))[:3]
        (center, radius) = fitCircle(x[randEdges], y[randEdges])
        
        # Eliminate very small and very large circles, and centers outside the image
        if not (par['lowerRadius'] < radius < par['upperRadius'] and \
                0 <= center[0] < imSize[0] and \
                0 <= center[1] < imSize[1]):
            continue

        # Make sure a reasonable number of points lies near that circle
        centerDistance = np.sqrt((x-center[0])**2 + (y-center[1])**2)
        inCircle = np.where(np.abs(centerDistance-radius)<par['eps'])[0]
        inPts = len(inCircle)
        if inPts < par['ransac_threshold'] *4*np.pi*radius*par['eps'] or inPts < 3:
            continue
        
        # Fit a circle to all good points, and save the corresponding parameters
        (center, radius) = fitCircle(x[inCircle], y[inCircle])        
        fitted[ii] = (center, radius, inPts)
    
        # If you want to see these points:
        if debug_flag == 1:
            plt.plot(y,x,'.')
            plt.hold(True)
            plt.plot(y[inCircle], x[inCircle],'r.')
            plt.plot(y[randEdges], x[randEdges], 'g+', ms=15)
            plt.axis('equal')
            plt.show()
        
    # Sort the circles, according to number of points included
    fitted = np.sort(fitted,order='nPts')
        
    # Show the best-fitting circle
    plt.imshow(data, cmap='gray', origin='lower')
    drawCircle(fitted[-1]['center'], fitted[-1]['radius'])
    plt.show()
    
def fit_circle(x,y):
    '''
    Determine the best-fit circle to given datapoints.
    
    Parameters
    ----------
    x : array (N,)
        x-values.
    y : array (N,)
        corresponding y-values.

    Returns
    -------
    center : array (2,)
        x/y coordinates of center of the circle
    radius : float
        Circle radius.

    Examples
    --------
    >>> r = 2
    >>> center = np.r_[5,5]
    >>> theta = np.r_[0:2*np.pi:10j]
    >>> x = r*np.cos(theta)+center[0]
    >>> y = r*np.sin(theta)+center[1]
    >>> cFit,rFit = thLib.fits.fit_circle(x,y)
    
    '''
    M = np.vstack((2*x,2*y,np.ones(len(x)))).T
    (par,_,_,_) = np.linalg.lstsq(M,x**2+y**2)
    center = par[:2]
    radius = np.sqrt(par[2]+np.sum(center**2))
    return(center, radius)


def fit_exp(tFit, yFit, plotFlag=False): 
    '''
    Calculates best-fit parameters for the exponential decay to an offset.
    This can serve as an example for a general non-linear fit.

    Parameters
    ----------
    tFit : array (N,)
        Time values.
    yFit : array (N,)
        Function values

    Returns
    -------
    offset : float
        Function offset/bias.
    amp : float
        Amplitude of exponential function
    tau : float
        Decay time.

    Examples
    --------
    
    >>> t = np.arange(10)
    >>> tau = 2.
    >>> amp = 1.
    >>> offset = 2.
    >>> x = offset + amp*np.exp(-t/tau)    
    >>> fitted =  thLib.fits.fit_exp(t,x)

    ''' 

    from scipy import optimize 
     
    # Define the fit-function and the error-function 
    fitfunc = lambda p, x: p[0] + p[1]*np.exp(-x/p[2]) 
    errfunc  = lambda p,x,y: fitfunc(p,x) - y 
     
    pInit = np.r_[0, 1, 1]  # Initial values 
     
    # Make the fit 
    pFit, success = optimize.leastsq(errfunc, pInit, args=(tFit, yFit)) 
     
    if plotFlag:
        # Plot the data and the fit 
        plt.plot(tFit, yFit, label='rawdata')
        plt.hold(True)
        plt.plot(tFit, fitfunc(pFit,tFit), label='fit') 
        plt.legend()
        plt.show() 

    ExpFit = namedtuple('ExpFit', ['offset', 'amplitude', 'tau'])
    return ExpFit(*pFit)

def fit_line(x, y, alpha=0.05, newx=[], plotFlag=False):
    """
    Linear regression fit.

    Parameters
    ----------
    x : ndarray
        Input / Predictor.
    y : ndarray
        Input / Estimator.
    alpha : float
        Confidence limit [default=0.05]
    newx : float or ndarray
        Values for which the fit and the prediction limits are calculated (optional)
    plotFlag: int, optional
        1 = plot, 0 = no_plot [default]
        
    Returns
    -------
    a : float
        Intercept
    b : float
        Slope
    ci : ndarray
        Lower and upper confidence interval for the slope
    info : dictionary
        contains return information on
        - residuals
        - var_res
        - sd_res
        - alpha
        - tval
        - df
    newy : list(ndarray)
        Predictions for (newx, newx-ciPrediction, newx+ciPrediction)

    Examples
    --------
    >>> x = np.r_[0:10:11j]
    >>> y = x**2
    >>> (a,b,(ci_a, ci_b),_) = thLib.fits.fit_line(x,y)    
    Summary: a=-15.0000+/-12.4590, b=10.0000+/-2.1060
    Confidence intervals: ci_a=(-27.4590 - -2.5410), ci_b=(7.8940 - 12.1060)
    Residuals: variance = 95.3333, standard deviation = 9.7639
    alpha = 0.050, tval = 2.2622, df=9


    Notes
    -----
    Example data and formulas are taken from
    D. Altman, "Practical Statistics for Medicine"
    """
    
    # Summary data
    n = len(x)			   # number of samples     
    
    Sxx = np.sum(x**2) - np.sum(x)**2./n
#    Syy = np.sum(y**2) - np.sum(y)**2./n    # not needed here
    Sxy = np.sum(x*y) - np.sum(x)*np.sum(y)/np.float(n)
    mean_x = np.mean(x)
    mean_y = np.mean(y)
    
    # Linefit
    b = Sxy/Sxx
    a = mean_y - b*mean_x
    
    # Residuals
    fit = lambda xx: a + b*xx    
    residuals = y - fit(x)
    
    var_res = np.sum(residuals**2)/np.float((n-2))
    sd_res = np.sqrt(var_res)
    
    # Confidence intervals
    se_b = sd_res/np.sqrt(Sxx)
    se_a = sd_res*np.sqrt(np.sum(x**2)/np.float(n*Sxx))
    
    df = n-2                            # degrees of freedom
    tval = sp.stats.t.isf(alpha/2., df) 	# appropriate t value
    
    ci_a = a + tval*se_a*np.array([-1,1])
    ci_b = b + tval*se_b*np.array([-1,1])

    # create series of new test x-values to predict for
    npts = 100
    px = np.linspace(np.min(x),np.max(x),num=npts)
    
    se_fit     = lambda x: sd_res * np.sqrt(  1./n + (x-mean_x)**2./Sxx)
    se_predict = lambda x: sd_res * np.sqrt(1+1./n + (x-mean_x)**2./Sxx)
    
    print('Summary: a={0:5.4f}+/-{1:5.4f}, b={2:5.4f}+/-{3:5.4f}'.format(a,tval*se_a,b,tval*se_b))
    print('Confidence intervals: ci_a=({0:5.4f} - {1:5.4f}), ci_b=({2:5.4f} - {3:5.4f})'.format(ci_a[0], ci_a[1], ci_b[0], ci_b[1]))
    print('Residuals: variance = {0:5.4f}, standard deviation = {1:5.4f}'.format(var_res, sd_res))
    print('alpha = {0:.3f}, tval = {1:5.4f}, df={2:d}'.format(alpha, tval, df))
    
    # Return info
    ri = {'residuals': residuals, 
        'var_res': var_res,
        'sd_res': sd_res,
        'alpha': alpha,
        'tval': tval,
        'df': df}
    
    if plotFlag:
        # Plot the data
        plt.figure()
        
        plt.plot(px, fit(px),'k', label='Regression line')
        #plt.plot(x,y,'k.', label='Sample observations', ms=10)
        plt.plot(x,y,'k.')
        
        x.sort()
        limit = (1-alpha)*100
        plt.plot(x, fit(x)+tval*se_fit(x), 'r--', lw=2, label='Confidence limit ({0:.1f}%)'.format(limit))
        plt.plot(x, fit(x)-tval*se_fit(x), 'r--', lw=2 )
        
        plt.plot(x, fit(x)+tval*se_predict(x), '--', lw=2, color=(0.2,1,0.2), label='Prediction limit ({0:.1f}%)'.format(limit))
        plt.plot(x, fit(x)-tval*se_predict(x), '--', lw=2, color=(0.2,1,0.2))

        plt.xlabel('X values')
        plt.ylabel('Y values')
        plt.title('Linear regression and confidence limits')
        
        # configure legend
        plt.legend(loc=0)
        leg = plt.gca().get_legend()
        ltext = leg.get_texts()
        plt.setp(ltext, fontsize=10)

        # show the plot
        plt.show()
     
    if newx != []:
        try:
            newx.size
        except AttributeError:
            newx = np.array([newx])
    
        print('Example: x = {0}+/-{1} => se_fit = {2:5.4f}, se_predict = {3:6.5f}'\
        .format(newx[0], tval*se_predict(newx[0]), se_fit(newx[0]), se_predict(newx[0])))
        
        newy = (fit(newx), fit(newx)-se_predict(newx), fit(newx)+se_predict(newx))
        
        LineFit = namedtuple('LineFit',['intercept', 'slope', 'CIs', 'info', 'yFitted'])
        return LineFit(a,b,(ci_a, ci_b), ri, newy)
    else:
        LineFit = namedtuple('LineFit',['intercept', 'slope', 'CIs', 'info'])
        return LineFit(a,b,(ci_a, ci_b), ri)

 
def fit_sin(tList, yList, freq):
    '''
    Fit a sine wave with a known frequency to a given set of data.
   
    y = amplitude * sin(2*pi*freq * tList + phase*pi/180) + bias
   
    Parameters
    ----------
    yList : array
           datapoints
    tList : float
           time base, in sec
    freq : float
          in Hz
       
    Returns
    -------
    phase : float
        in degrees
    amplitude : float
    bias : float
   
    Examples
    --------
    >>> np.random.seed(1234)
    >>> t = np.arange(0,10,0.1)
    >>> x = 3 + 4*np.sin(2*np.pi*t + 5*np.pi/180) + np.random.randn(len(t))
    >>> (phase, amp, offset) = thLib.fits.fit_sin(t, x, 1)
    
    '''

    # Get input data
    b = yList
    rows = [ [np.sin(freq*2*np.pi*t), np.cos(freq*2*np.pi*t), 1] for t in tList]
    A = np.array(rows)
    
    # Make the fit
    (w,residuals,rank,sing_vals) = np.linalg.lstsq(A,b)
    
    # Extract desired parameters ...
    phase = np.arctan2(w[1],w[0])*180/np.pi
    amplitude = np.linalg.norm([w[0],w[1]],2)
    bias = w[2]
    
    # ... and return them
    SinFit = namedtuple('SinFit', ['phase', 'amp', 'offset'])
    return SinFit(phase,amplitude,bias)
 
def fit_ellipse(x,y):
    '''
    Ellipse fit by Taubin's Method
    
    Parameters
    ----------
    x : array
        x-coordinates of the ellipse points
    y : array
        y-coordinates of the ellipse points
        
    Returns
    -------
    A : array
        Ellipse parameters
        A = [a b c d e f]
        is the vector of algebraic parameters of thefitting ellipse:
        ax^2 + bxy + cy^2 +dx + ey + f = 0
        The vector A is normed, so that ||A||=1.
    
    Notes
    -----
    Among fast non-iterative ellipse fitting methods, 
    this is perhaps the most accurate and robust.
    
    This method fits a quadratic curve (conic) to a set of points;
    if points are better approximated by a hyperbola, this fit will 
    return a hyperbola. To fit ellipses only, use "Direct Ellipse Fit".
    
    Published in
    G. Taubin, "Estimation Of Planar Curves, Surfaces And Nonplanar
    Space Curves Defined By Implicit Equations, With 
    Applications To Edge And Range Image Segmentation",
    IEEE Trans. PAMI, Vol. 13, pages 1115-1138, (1991)
    
    '''

    centroid = np.mean( np.vstack((x,y)), 1 )   # the centroid of the data set
    
    Z = np.vstack(( (x-centroid[0])**2,
                    (x-centroid[0])*(y-centroid[1]),
                    (y-centroid[1])**2,
                     x-centroid[0],
                     y-centroid[1],
                    np.ones(len(x)) )).T;

    M = Z.T.dot(Z) / len(x)
    
    P = np.array(
         [[ M[0,0]-M[0,5]**2, M[0,1]-M[0,5]*M[1,5], M[0,2]-M[0,5]*M[2,5], M[0,3], M[0,4] ],
          [ M[0,1]-M[0,5]*M[1,5], M[1,1]-M[1,5]**2, M[1,2]-M[1,5]*M[2,5], M[1,3], M[1,4] ],
          [ M[0,2]-M[0,5]*M[2,5], M[1,2]-M[1,5]*M[2,5], M[2,2]-M[2,5]**2, M[2,3], M[2,4] ],
          [ M[0,3], M[1,3], M[2,3], M[3,3], M[3,4] ],
          [ M[0,4], M[1,4], M[2,4], M[3,4], M[4,4] ]])
    
    Q = np.array(
         [[ 4*M[0,5], 2*M[1,5], 0, 0, 0],
          [ 2*M[1,5], M[0,5]+M[2,5], 2*M[1,5], 0, 0],
          [ 0, 2*M[1,5], 4*M[2,5], 0, 0],
          [ 0, 0, 0, 1, 0],
          [ 0, 0, 0, 0, 1]])
    
    w, vr = sp.linalg.eig(P,Q)
    sortID = np.argsort(w)
    
    A = vr[:,sortID[0]]
    A = np.hstack( (A, -A[:3].T.dot(M[:3,5])) )
    A4 = A[3] - 2*A[0]*centroid[0] - A[1]*centroid[1]
    A5 = A[4] - 2*A[2]*centroid[1] - A[1]*centroid[0]
    A6 = A[5] + A[0]*centroid[0]**2 +A[2]*centroid[1]**2+ \
         A[1]*centroid[0]*centroid[1] - A[3]*centroid[0]-A[4]*centroid[1]
    A[3] = A4
    A[4] = A5
    A[5] = A6

    A = A/np.linalg.norm(A);
    return A

def regress(y, X, alpha=0.05, intercept=True):
    '''
    Multilinear regression and confidence intervals.
    Note that by default, an intercept is added to the design matrix!

    Parameters
    ----------
    X : ndarray (N,) or (N,p)
        predictors at each of N observations
    y : ndarray(N,)
        observed responses
    alpha : float, optional
        Defines the 100*(1-alpha)% confidence level in ci.
        Default alpha=0.05
    intercept : boolean
        If 'True', an intercept is automatically added to the design matrix.
        If 'False', the behavior is like the Matlab "regress" function, and
        no intercept is added.

    Returns
    -------
    fit : float
        best fit intercept and regression parameters
    ci : ndarray (2,)
        confidence intervals for the coefficient estimates
        at the given alpha level (default: 95%-level)

    Examples
    --------
    >>> x = np.array([0, 0, 10, 10])
    >>> y = np.array([1, 3, 11, 13])
    >>> (fit,ci) = thLib.fits.regress(y,x)    

    >>> X = np.random.randn(10,2)
    >>> mat = np.hstack( (np.ones( (len(X),1) ), X) )
    >>> pars = np.c_[[2, 4, 5]]
    >>> y = mat.dot(pars)
    >>> y += 0.1*np.random.randn(*y.shape)
    >>> (fit,ci) = thLib.fits.regress(y, mat)    

    See also
    --------
    fit_line

    '''

        
    if X.ndim == 1:
        dof = len(X)-2
        df = pd.DataFrame({'x':X, 'y':y})
        modelTxt = 'y~x'
        
        # If you want the Matlab "regress", without intercept
        if intercept == False:
            modelTxt += '-1'
            dof += 1
            
        # Fit the model
        model = ols(modelTxt,df).fit()
        ci = np.zeros((2, 2))
    else:
        # Make sure the dimensions are right
        if X.shape[0] < X.shape[1]:
            X = X.T
            
        dof = X.shape[0]-X.shape[1]-1
        if y.ndim == 1:
            y = np.c_[y]
        # For automatic labeling, put all the data into one matrix
        data = np.hstack((y,X))
        
        df = pd.DataFrame(data)
        # Make a string like 'y~x1+x2+x3'
        labels = ['y']
        modelTxt = 'y~'
        for ii in range(X.shape[1]):
            labels.append('x'+str(ii))
            modelTxt += labels[-1]+'+'
            
        # Remove the last "+"
        modelTxt = modelTxt[:-1]
        
        # Automatic column labelling
        df.columns = labels
        
        # Memory allocation for CIs
        if intercept == True:
            ci = np.zeros((1+X.shape[1],2))
        # If you want the Matlab "regress", without intercept
        else:
            dof +=1
            modelTxt += '-1'
            ci = np.zeros((X.shape[1],2))
            
        # Fit the model
        model = ols(modelTxt,df).fit()
        
    level = (1.-alpha/2.)
    # Make sure that you have the correct DOF
    tVal = sp.stats.t.ppf(level,dof)
    
    # Extract parameters and standard error from the model
    fit = model.params
    se = model.bse
    
    # Calculate CIs
    ci[:,0] = fit - se*tVal
    ci[:,1] = fit + se*tVal
    
    # Return a named tuple
    Regression= namedtuple('Regression', ['Fit', 'CIs'])
    
    return Regression(fit, ci)
    
if __name__=='__main__':

    # Produce an ellipse
    a = 5
    b = 3
    alpha = np.deg2rad(30) 
    R = np.array([[np.cos(alpha), -np.sin(alpha)],
                   [np.sin(alpha), np.cos(alpha)]])
    
    theta = np.linspace(0, 2*np.pi, 100)
    x = a * np.cos(theta)
    y = b * np.sin(theta)    
    XY = np.vstack( (x,y) )
    ellipse = R.dot( XY )    
    
    plt.plot(ellipse[0,:], ellipse[1,:])    
    plt.show()

    A = fit_ellipse(ellipse[0,:], ellipse[1,:])
    print('The ellipse-fit parameters are:')
    print(A)
    
    # Check (equations from http://mathworld.wolfram.com/Ellipse.html)
    (a,b,c,d,e,f) = A
    
    # To match the formula from wolfram
    b /= 2
    d /= 2
    g = f
    f = e/2
    
    aFit = np.sqrt(2*(a*f**2 + c*d**2 + g*b**2 - 2*b*d*f - a*c*g)/((b**2-a*c)* (np.sqrt((a-c)**2 + 4*b**2) -(a+c))))
    
    bFit = np.sqrt(2*(a*f**2 + c*d**2 + g*b**2 - 2*b*d*f - a*c*g)/((b**2-a*c)*(-np.sqrt((a-c)**2 + 4*b**2) - (a+c))))
    
    thetaFit = np.pi/2 + 0.5*np.arctan2(2*b, a-c)
    print('a={0:3.2f}, b={1:3.2f}, theta = {2}'.format(aFit, bFit, np.rad2deg(thetaFit)))
    
    
    
    
    '''
    # Test fit_exp
    t = np.arange(10)
    tau = 2.
    amp = 1.
    offset = 2.
    x = offset + amp*np.exp(-t/float(tau))    
    print(fit_exp(t,x))
    
    import doctest
    doctest.testmod()    # Test fit_line
        
    print('Fit line ------------')
    x = np.array([0, 0, 10, 10])
    y = np.array([1, 3, 11, 13])
    out = fit_line(x,y)
    print(out)
    
    # Test regress - 1dim
    print('Plain -----------------')
    x = np.array([0, 0, 10, 10])
    y = np.array([1, 3, 11, 13])
    (fit,ci) = regress(y,x)    
    print(fit)
    print(ci)

    print('99% -----------------')
    (fit,ci) = regress(y,x, alpha=0.01)    
    print(fit)
    print(ci)
    
    print('No intercept -------')
    (fit,ci) = regress(y,x, intercept=False)    
    print(fit)
    print(ci)
    
    # Test regress - 2dim
    print('2-dim -----------')
    X = np.array([[1,1,1,1],[0,0,10,10]]).T
    y = np.array([1, 3, 11, 13])
    (fit,ci) = regress(y, X, intercept=False)    
    print(fit)
    print(ci)

    # Test regress - 2dim
    X = np.random.randn(10,2)
    mat = np.hstack( (np.ones( (len(X),1) ), X) )
    pars = np.c_[[2, 4, 5]]
    y = mat.dot(pars)
    y += 0.1*np.random.randn(*y.shape)
    (fit,ci) = regress(y, mat)    
    print(fit)     
    print(ci)
    '''
