#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Wed Mar 15 11:16:40 2017

@author: ltlustos
"""
from scipy.optimize import curve_fit
import numpy as np
import pylab as pl
from scipy.signal import argrelmax, argrelmin
import scipy.ndimage  # @UnresolvedImport


def mov_median(data, span):
    return scipy.ndimage.median_filter(data, span)


def mov_mean(data, span):
    return np.convolve(data, np.ones((span,))/span, mode='same')


def moment1(x, y=None, nstd=None, warning=True):
    if y is None:
        y = x
        x = np.arange(len(y))
    if y.min() < 0:
        if warning:
            print('Warniong: negative weights: [%g, %g]' % (y.min(), y.max()))
#            raise Exception('neg weight')
    if x.dtype in (np.uint8, np.uint16, np.uint32, np.uint64):
        x = x.astype(np.double)
    if y.dtype in (np.uint8, np.uint16, np.uint32, np.uint64):
        y = y.astype(np.double)
#    offs = y.min()
#    y -= offs
    m0 = np.double(np.sum(y))
    m1 = np.sum(y * x)
    m2 = np.sum(y * x ** 2)
    amean = m1 / m0
    astd = np.sqrt((m2 - m1 ** 2 / m0) / m0)
    if nstd is not None:
        idxx = np.where(np.abs(x-amean) < nstd*astd)
        amean, astd = moment1(x[idxx], y[idxx])
    return [amean, astd]


def local_min(x, n=3):
    return argrelmin(x, order=n)


def local_max(x, n=3):
    return argrelmax(x, order=n)


def find_local_max(x, y, mincount=1, sigma_est=6):
    y /= y.max()
    y = mov_mean(mov_mean(y ** 3, sigma_est), sigma_est)

    x = x[0:-3]
    y = y[0:-3]

    idxx = np.where(y < mincount)[0]
    y[idxx] = 0
    return x[local_max(y)], y[local_max(y)]

###############################################################


class GOF(object):

    def __init__(self, y, fy, weights=None):
        self.__y = y
        self.__fy = fy
        if weights is not None:
            weights /= np.sum(weights)
        else:
            weights = 1.0

        ybar = np.mean(self.__y)
        # total sum of squares
        self.__SST = np.sum(weights * (self.__y - ybar) ** 2)
        # estimated sum of squares
        self.__SSE = np.sum(weights * (self.__fy - self.__y) ** 2)

        if self.__SST != 0:
            self.__R2 = 1 - self.__SSE / self.__SST
        elif (self.__SST == 0) & (self.__SSE == 0):
            self.__R2 = np.NaN
        else:
            if np.sqrt(np.abs(self.__SSE)) < np.sqrt(np.finfo(np.double).eps) * np.mean(abs(self.__y)):
                self.__R2 = np.NaN
            else:
                self.__R2 = np.Inf
        if np.isnan(self.__R2):
            self.__R2 = -1.0

    @property
    def R2(self):
        return self.__R2

    @property
    def ESS(self):
        return self.__SSE

    @property
    def TSS(self):
        return self.__SST

###############################################################


class Fit(object):  # dont foreget the "object" root class
    def __init__(self, func):
        self.__func__ = func
        self._coeff_ = None
        self._var_matrix_ = None

    def __eval_func__(self, x, *p):
        return eval(self.__func__)

    def eval_func(self, x, p):
        return self.__eval_func__(x, *p)

    def fit(self, x, y, p0, weights=None):
        # from scipy.optimize import curve_fit
        self.__x__ = x
        self.__y__ = y
        self.__weights__ = weights
        if weights is not None:
            self.__weights__ = 1.0 / weights
        self._coeff_, self._var_matrix_ = curve_fit(self.__eval_func__, x, y, p0=p0, sigma=self.__weights__,
                                                    absolute_sigma=weights is not None)
        self.GOF = GOF(y, self.__eval_func__(x, *self._coeff_), weights=weights)
        return self._coeff_, self._var_matrix_, self.GOF

    def refit_weighted(self, n_rep=1):
        for i in range(n_rep):
            self.fit(self.__x__, self.__y__, p0=self.p, weights=self.eval_func(self.__x__, self.p)**2)
        return self._coeff_, self._var_matrix_, self.GOF

    def plot(self, x=None, p=None, legend=True):
        if x is None:
            x = np.linspace(self.__x__.min(), self.__x__.max(), 5*len(self.__x__))
            # decrease spacing between fit point by factor 5
        if p is None:
            p = self._coeff_
        if legend:
            pl.plot(x, self.__eval_func__(x, *p), label='%s\n%s, R2=%g' % (self.__func__, p, self.GOF.R2))
            pl.plot(self.__x__, self.__y__, 'o', label='data')
        else:
            pl.plot(x, self.__eval_func__(x, *p))
            pl.plot(self.__x__, self.__y__, 'o')
        if legend:
            pl.legend(loc='1', fontsize='small')

    @property
    def p(self):
        return self._coeff_
###############################################################


def gauss(x, p):
    return p[0] * np.exp(-(x-p[1])**2/(2.*p[2]**2))


def erf(x, p):
    import scipy.special
    return p[0] * scipy.special.erf((x - p[1]) / p[2] / np.sqrt(2.0))


def erf_r(x, p0, p1, p2):
    import scipy.special
    return p0 * (1 - scipy.special.erf((x - p1) / p2 / np.sqrt(2.0))) / 2.0


def erf_l(x, p0, p1, p2):
    import scipy.special
    return p0 * (1 + scipy.special.erf((x - p1) / p2 / np.sqrt(2.0))) / 2.0


def tanh_r(x, p0, p1, p2):
    return p0 * (1 - np.tanh((x - p1) / p2 / np.sqrt(2.0))) / 2.0


def tanh_l(x, p0, p1, p2):
    return p0 * (1 + np.tanh((x - p1) / p2 / np.sqrt(2.0))) / 2.0


class FitGauss(Fit):
    def __init__(self):
        # p[0] : amplitude
        # p[1] : mu
        # p[2]: sigma
        super(FitGauss, self).__init__('gauss(x, p)')

    def fit(self, x, y, p0=None, weights=None):
        if p0 is None:
            mu, std = moment1(x, y)
            p0 = [y.max(), mu, std]
        super(FitGauss, self).fit(x, y, p0, weights=weights)
        self._coeff_[2] = np.abs(self._coeff_[2])
        return self._coeff_, self._var_matrix_, self.GOF


class FitGauss_Erf_right(Fit):
    def __init__(self):
        # p[0] : amplitude gauss
        # p[1] : mu
        # p[2]: sigma
        # p[3] : amplitude erf
        super(FitGauss_Erf_right, self).__init__('gauss(x, p[0:3]) + erf_r(x, p[3], p[1], p[2])')

    def fit(self, x, y, p0=None, weights=None):
        if p0 is None:
            mu, std = moment1(x, y)
            p0 = [y.max(), mu, std, y.max()/10, mu, std]
        super(FitGauss_Erf_right, self).fit(x, y, p0, weights=weights)
        self._coeff_[2] = np.abs(self._coeff_[2])
        return self._coeff_, self._var_matrix_, self.GOF

    def plot(self, x=None, p=None):
        super(FitGauss_Erf_right, self).plot(x, p)
        import pylab as pl
        if x is None:
            x = np.linspace(self.__x__.min(), self.__x__.max(), 5*len(self.__x__))
            # decrease spacing between fit point by factor 5
        if p is None:
            p = self._coeff_
        pl.plot(x, erf_r(x, p[3], p[1], p[2]), ':', zorder=0, color='0.75')
        pl.plot(x, gauss(x, p[0:3]), ':', zorder=0, color='0.75')


class FitGauss_tanh_right(Fit):
    def __init__(self):
        # p[0] : amplitude gauss
        # p[1] : mu
        # p[2]: sigma
        # p[3] : amplitude erf
        super(FitGauss_tanh_right, self).__init__('gauss(x, p[0:3]) + tanh_r(x, p[3], p[1], p[2])')

    def fit(self, x, y, p0=None, weights=None):
        if p0 is None:
            mu, std = moment1(x, y)
            p0 = [y.max()-y[0], mu, std, y[0], mu, std]
            print('p0=', p0)
        super(FitGauss_tanh_right, self).fit(x, y, p0, weights=weights)
        self._coeff_[2] = np.abs(self._coeff_[2])
        return self._coeff_, self._var_matrix_, self.GOF

    def plot(self, x=None, p=None):
        super(FitGauss_tanh_right, self).plot(x, p)
        import pylab as pl
        if x is None:
            x = np.linspace(self.__x__.min(), self.__x__.max(), 5*len(self.__x__))
            # decrease spacing between fit point by factor 5
        if p is None:
            p = self._coeff_
        pl.plot(x, tanh_r(x, p[3], p[1], p[2]), ':', zorder=0, color='0.75')
        pl.plot(x, gauss(x, p[0:3]), ':', zorder=0, color='0.75')


class Fit2Gauss_tanh_right(Fit):
    def __init__(self):
        # p[0] : amplitude gauss
        # p[1] : mu
        # p[2]: sigma
        # p[3] : amplitude erf
        super(Fit2Gauss_tanh_right, self).__init__('gauss(x, p[0:3]) + tanh_r(x, p[3], p[1], p[2]) + gauss(x, p[4:])')

    def fit(self, x, y, p0=None, weights=None):
        if p0 is None:
            mu, std = moment1(x, y)
            p0 = [y.max(), mu, std, y.max()/10, mu, std]
        super(Fit2Gauss_tanh_right, self).fit(x, y, p0, weights=weights)
        self._coeff_[2] = np.abs(self._coeff_[2])
        return self._coeff_, self._var_matrix_, self.GOF

    def plot(self, x=None, p=None):
        super(Fit2Gauss_tanh_right, self).plot(x, p)
        import pylab as pl
        if x is None:
            x = np.linspace(self.__x__.min(), self.__x__.max(), 5*len(self.__x__))
            # decrease spacing between fit point by factor 5
        if p is None:
            p = self._coeff_
        pl.plot(x, tanh_r(x, p[3], p[1], p[2]), ':', zorder=0, color='0.75')
        pl.plot(x, gauss(x, p[0:3]), ':', zorder=0, color='0.75')
        pl.plot(x, gauss(x, p[4:]), ':', zorder=0, color='0.75')


class Fit2Gauss(Fit):
    def __init__(self):
        # p[0] : amplitude
        # p[1] : mu
        # p[2]: sigma
        super(Fit2Gauss, self).__init__('gauss(x, p[0:3]) + gauss(x, p[3:])')

    def fit(self, x, y, p0=None, weights=None):
        if p0 is None:
            # to be done
            mu, std = moment1(x, y)
            p0 = [y.max()/2, mu/2, std/2, y.max(), mu, std]
        super(Fit2Gauss, self).fit(x, y, p0, weights=weights)
        self._coeff_[2::3] = np.abs(self._coeff_[2::3])
        return self._coeff_, self._var_matrix_, self.GOF


class Fit3Gauss(Fit):
    def __init__(self):
        # p[0] : amplitude
        # p[1] : mu
        # p[2]: sigma
        super(Fit3Gauss, self).__init__('gauss(x, p[0:3]) + gauss(x, p[3:6]) + gauss(x, p[6:9])')

    def fit(self, x, y, p0=None, weights=None):
        if p0 is None:
            # to be done
            mu, std = moment1(x, y)
            p0 = [y.max()/2, mu/2, std/2, y.max(), mu, std]
        super(Fit3Gauss, self).fit(x, y, p0, weights=weights)
        self._coeff_[2::3] = np.abs(self._coeff_[2::3])
        return self._coeff_, self._var_matrix_, self.GOF


class Fit4Gauss(Fit):
    def __init__(self):
        # p[0] : amplitude
        # p[1] : mu
        # p[2]: sigma
        super(Fit4Gauss, self).__init__('gauss(x, p[0:3]) + gauss(x, p[3:6]) + gauss(x, p[6:9]) + gauss(x, p[9:12])')

    def fit(self, x, y, p0=None, weights=None):
        if p0 is None:
            # to be done
            mu, std = moment1(x, y)
            p0 = [y.max()/2, mu/2, std/2, y.max(), mu, std]
        super(Fit3Gauss, self).fit(x, y, p0, weights=weights)
        self._coeff_[2::3] = np.abs(self._coeff_[2::3])
        return self._coeff_, self._var_matrix_, self.GOF

###############################################################


class Fit_ToT_keV(Fit):
    def __init__(self):
        super(Fit_ToT_keV, self).__init__('(p[0] * x + p[1] - p[2] / ( x - p[3])) * np.array((x > '
                                          '((p[0]*p[3] - p[1] - np.sqrt(p[0]**2*p[3]**2 + 2*p[0]*p[1]*p[3] '
                                          '+ 4*p[0]*p[2] + p[1]**2))/(2*p[0]))))')

    def __eval_func__(self, x, *p):
        # x0 = (p[0]*p[3] - p[1] + np.sqrt(p[0]**2*p[3]**2 + 2*p[0]*p[1]*p[3] + 4*p[0]*p[2] + p[1]**2))/(2*p[0])
        # print x0
        y = p[0] * x + p[1] - p[2] / (x - p[3])
        return y * np.array(y > 0)

    def fit(self, x, y, p0=None, weights=None):
        if p0 is None:
            A = np.polyfit(x, y, 1)
            p0 = [A[1], A[0], 10, 3]
        return super(Fit_ToT_keV, self).fit(x, y, p0, weights=weights)