# -*- encoding: utf-8 -*-
__author__ = 'markusro'

from PyQt4.QtGui import QColor
from PyQt4.QtCore import QObject, pyqtSignal, QThread, pyqtSlot

import numpy as np
from scipy import optimize as opt, odr

import libyaff


def id_to_color( id ):
    colors = [
        QColor(255, 255, 255),
        QColor(168, 149, 17),
        QColor(45, 142, 15),
        QColor(160, 16, 36),
        QColor(54, 22, 115),
        QColor(36, 10, 85),
        QColor(118, 8, 23),
        QColor(31, 105, 7),
        QColor(124, 109, 8),
    ]
    chosen_color = colors[id%len(colors)]
    return chosen_color


class FitFunctionCreator(QObject):
    new_data = pyqtSignal(np.ndarray, np.ndarray)

    def __init__( self ):
        super(FitFunctionCreator, self).__init__()
        self.data = None
        self.functions = Functions()


    def fitfcn( self, p0, x, *funcs ):
        if x.ndim == 2:
            self.data = np.zeros(x.shape)
        else:
            self.data = np.zeros((2, x.size))
        ndx = 0
        for fn in funcs:  # loop over functions and add the results
            f, num_p = fn.function, fn.param_number
            p = p0[ndx:ndx+num_p]
            if x.ndim == 2:
                x = x[0]
            result = f(p, x)
            # fn.widget.updateTable(p)
            self.data += result  # fit functions take only 1-dim x
            ndx += num_p
        self.new_data.emit(x, self.data)
        return self.data

    def fitfcn_imag( self, p0, x, *funcs ):
        if x.ndim == 2:
            self.data = np.zeros(x.shape)
        else:
            self.data = np.zeros((2, x.size))
        ndx = 0
        for fn in funcs:  # loop over functions and add the results
            f, num_p = fn.function, fn.param_number
            p = p0[ndx:ndx+num_p]
            if x.ndim == 2:
                x = x[0]
            result = f(p, x)
            self.data += result  # fit functions take only 1-dim x
            ndx += num_p
        self.new_data.emit(x, self.data)
        return self.data[1]


class FitRoutine(QObject):
    finished_fit = pyqtSignal()
    data_ready = pyqtSignal(np.ndarray, np.ndarray)

    def __init__( self ):
        super(FitRoutine, self).__init__()
        self.f = FitFunctionCreator()
        self.f.new_data.connect(self.data_ready.emit)
        self._fitter = self.fit_odr_cmplx
        self._odr_fit = None
        self._start_parameter = None

    @property
    def start_parameter( self ):
        return self._start_parameter

    @start_parameter.setter
    def start_paramter( self, p0 ):
        self._start_parameter = p0

    @property
    def fitter( self ):
        return self._fitter

    @fitter.setter
    def fitter( self, f ):
        self._fitter = f

    def fit_odr_cmplx( self, x, y, p0, fixed, fcns ):
        self._start_parameter = p0
        if np.iscomplexobj(y) and y.ndim == 1:
            weights = 1/np.abs(y)**2
            we = np.resize(weights, (2, weights.size))
            # we = 1/N.array([y.real**2, y.imag**2])
            y = np.array([y.real, y.imag])
        else:
            raise NotImplementedError, "need complex input for now"
        dat = odr.Data(x, y, we=we)
        mod = odr.Model(self.f.fitfcn, extra_args=fcns)
        self._odr_fit = odr.ODR(dat, mod, p0, ifixx=(0,), ifixb=fixed, maxit=800)

    def fit_odr_imag( self, x, y, p0, fixed, fcns ):
        self._start_parameter = p0
        if np.iscomplexobj(y) and y.ndim == 1:
            we = 1/np.imag(y)**2
        else:
            raise NotImplementedError, "need complex input for now"
        dat = odr.Data(x, y.imag, we=we)
        mod = odr.Model(self.f.fitfcn_imag, extra_args=fcns)
        self._odr_fit = odr.ODR(dat, mod, p0, ifixx=(0,), ifixb=fixed, maxit=800)

    @pyqtSlot()
    def fit( self ):
        try:
            self._odr_fit.run()
        except RuntimeError:
            print "muh"
        self.finished_fit.emit()

    def result( self ):
        if self._odr_fit.output is None:
            self._odr_fit.output = odr.Output([self.start_parameter, None, None])
            self._odr_fit.output.stopreason = ["Aborted by user"]
        return self._odr_fit.output


class FunctionRegister:
    def __init__( self ):
        self.registry = { }

    def register_function( self, obj ):
        # print "FR:   Registering:",obj
        id_string = obj.id_label
        if self.registry.has_key(obj):
            raise AssertionError, "The object is already registered! This should NOT happen"
        self.registry[obj] = id_string
        #print "FR: ",self.registry

    def unregister_function( self, obj ):
        # print "FR: UnRegistering:",obj
        if self.registry.has_key(obj):
            self.registry.pop(obj)
        else:
            obj.deleteLater()
            raise AssertionError, "The object is not in the registry! This should NOT happen"
            #print "FR: ",self.registry

    def get_registered_functions( self ):
        return self.registry


# ############# deprecated #####################
def fit_odr_cmplx( x, y, p0, fixed, fcns ):
    f = FitFunctionCreator()
    #if x.ndim < 2:
    #    x = N.resize(x, (2,x.size))
    if np.iscomplexobj(y) and y.ndim == 1:
        weights = 1/np.abs(y)**2
        we = np.resize(weights, (2, weights.size))
        #we = 1/N.array([y.real**2, y.imag**2])
        y = np.array([y.real, y.imag])
    else:
        raise NotImplementedError, "need complex input for now"
    dat = odr.Data(x, y, we=we)
    mod = odr.Model(f.fitfcn, extra_args=fcns)
    fit = odr.ODR(dat, mod, p0, ifixx=(0,), ifixb=fixed, maxit=8000)
    fit.run()
    #print fit.output.pprint()
    return fit.output


### define funcs here
class Functions(QObject):
    def __init__( self ):
        super(Functions, self).__init__()
        self.list = {
            # provides functions: "id_string":(function, number_of_parameters)
            "hn": (self.hn_cmplx, 4),
            "conductivity": (self.cond_cmplx, 3),
            "power": (self.power_cmplx, 2),
            "static": (self.static_cmplx, 1),
            "yaff": (self.yaff, 8)
        }
        self.YAFF = libyaff.Yaff()

    def hn_cmplx( self, p, x ):
        om = 2*np.pi*x
        #hn = om*1j
        eps, t, a, b = p
        hn = eps/(1+(1j*om*t)**a)**b
        cplx = np.array([hn.real, -hn.imag])
        return cplx

    def cond_cmplx( self, p, x ):
        om = 2*np.pi*x
        sgma, isgma, n = p
        cond = sgma/(om**n)+isgma/(1j*om**n)  # Jonscher (Universal Dielectric Response: e",e' prop sigma/omega**n
        cplx = np.array([cond.real, -cond.imag])
        return cplx

    def power_cmplx( self, p, x ):
        om = 2*np.pi*x
        sgma, n = p
        power = sgma/(om*1j)**n
        cplx = np.array([power.real, -power.imag])
        return cplx

    def static_cmplx( self, p, x ):
        eps_inf = p[0]
        static = np.ones((2, x.size))*eps_inf
        static[1, :] *= 0  # set imag part zero
        #cplx = N.array([static.real, static.imag])
        return static

    def yaff( self, p, x ):
        ya = self.YAFF.yaff(p[:8], x)
        cplx = np.array([ya.imag, ya.real])
        return cplx

    def get( self, name ):
        return self.list[name]

    def get_function( self, name ):
        return self.list[name][0]


def fit_anneal( x, y, p0, fixed, funcs ):
    raise NotImplementedError
    bounds = [(0, 1e14), (0, 1)]
    for i in xrange(len(p0[2:])/4):
        bounds.append((1e-4, 1e12))  # delta_eps
        bounds.append((1e-12, 1e3))  # tau
        bounds.append((0.1, 1))  # a
        bounds.append((0.1, 1))  # b
    ret = opt.anneal(mini_func, p0,
                     args=(x, y),
                     maxeval=20000,
                     maxiter=30000,
                     lower=[b[0] for b in bounds],
                     upper=[b[1] for b in bounds],
                     dwell=100,
                     full_output=1)
    #pmin, func_min, final_Temp, cooling_iters,accepted_tests, retval
    #retval : int
    #Flag indicating stopping condition::

    #       0 : Points no longer changing
    #       1 : Cooled to final temperature
    #       2 : Maximum function evaluations
    #       3 : Maximum cooling iterations reached
    #       4 : Maximum accepted query locations reached
    #       5 : Final point not the minimum amongst encountered points

    print "Stop reason", ret
    return ret[0]


def fit_lbfgsb( x, y, p0, fixed, funcs ):
    raise NotImplementedError
    # TODO fixed parameters…
    bounds = [(0, None), (0, 1)]
    for i in xrange(len(p0[3:])/4):
        bounds.append((1e-4, 1e12))  # delta_eps
        bounds.append((1e-12, 1e3))  # tau
        bounds.append((0.1, 1))  # a
        bounds.append((0.1, 1))  # b

    x, f_minvalue, info_dict = opt.fmin_l_bfgs_b(mini_func, p0,
                                                 fprime=None,
                                                 args=(x, y),
                                                 approx_grad=True,
                                                 bounds=bounds,
                                                 iprint=0,
                                                 maxfun=4000)
    if info_dict['warnflag'] != 0:
        print info_dict["task"]
    return x


# Replaced with fit_odr_cmplx
#
#def fit_odr(x, y, p0, fixed, funcs):
#    dat = odr.Data(x, y, 1.0 / y**2)
#    mod = odr.Model(multi_hn)
#    fit = odr.ODR(dat, mod, p0, ifixx=(0,), ifixb=fixed, maxit=2000)
#    fit.run()
#    return fit.output.beta


def hn( p, nu ):
    delta_eps, tau, a, b = p
    om = 2*np.pi*nu
    Phi = np.arctan((om*tau)**a*np.sin(np.pi*a/2.)/(1.+(om*tau)**a*np.cos(np.pi*a/2.)))
    e_loss = delta_eps*(1+2*(om*tau)**a*np.cos(np.pi*a/2.)+(om*tau)**(2.*a) )**(
        -b/2.)*np.sin(b*Phi)
    e_stor = delta_eps*(1+2*(om*tau)**a*np.cos(np.pi*a/2.)+(om*tau)**(2.*a) )**(
        -b/2.)*np.cos(b*Phi)
    return e_loss  # 2* oder nicht?


def mini_func( p, x, y ):
    res = y-multi_hn(p, x)
    # apply weights
    res /= 1/y
    return np.sqrt(np.dot(res, res))


def multi_hn( p, nu ):
    conductivity = p[1]
    cond_beta = p[2]
    om = 2*np.pi*nu
    e_loss = conductivity/om**cond_beta
    e_loss += p[0]
    #for key, igroup in groupby(p[3:], lambda x: x//4):
    for i in xrange(len(p[3:])/4):
        delta_eps, tau, a, b = p[3+i*4:3+(i+1)*4]
        #delta_eps, tau, a, b = list(igroup)
        #print delta_eps,tau,a,b
        #a = 0.5 *(1 + N.tanh(a))
        #b = 0.5 *(1 + N.tanh(b))
        Phi = np.arctan((om*tau)**a*np.sin(np.pi*a/2.)/(1.+(om*tau)**a*np.cos(np.pi*a/2.)))
        e_loss += 2*delta_eps*(1+2*(om*tau)**a*np.cos(np.pi*a/2.)+(om*tau)**(2.*a) )**(
            -b/2.)*np.sin(b*Phi)
        #e_stor = delta_eps * (1+ 2*(om*tau)**a * N.cos(N.pi*a/2.) + (om*tau)**(2.*a) )**(-b/2.)*N.cos(b*Phi)
    return e_loss


def tau_peak( f, a, b ):
    tau = (np.sin(np.pi*a/2./(b+1))/np.sin(np.pi*a*b/2./(b+1)))**(1/a)
    tau /= 2*np.pi*f
    return tau