# -*- encoding: utf-8 -*-
import matplotlib
import numpy as N
from scipy import optimize as opt, optimize, odr
#from QDS import mini_func, multi_hn
from PyQt4.QtGui import QColor
__author__ = 'markusro'


def fit_anneal(x, y, p0, fixed):
    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):
    # 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


def fit_odr(x, y, p0, fixed):
    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 * N.pi * nu
    Phi = N.arctan((om * tau) ** a * N.sin(N.pi * a / 2.) / (1. + (om * tau) ** a * N.cos(N.pi * a / 2.)))
    e_loss = delta_eps * (1 + 2 * (om * tau) ** a * N.cos(N.pi * a / 2.) + (om * tau) ** (2. * a) ) ** (
    -b / 2.) * N.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 2 * e_loss


def id_to_color(id):
    """

    """
    colors = ['b', 'r', 'g', 'c', 'm', 'y', 'k']
    colors = [QColor(54,22,115),
              QColor(160,16,36),
              QColor(45,142,15),
              QColor(168,149,17),
              QColor(36,10,85),
              QColor(118,8,23),
              QColor(31,105,7),
              QColor(124,109,8),
              ]
    return colors[id % len(colors)]


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


def multi_hn(p, nu):
    conductivity = p[1]
    cond_beta = p[2]
    om = 2 * N.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 = N.arctan((om * tau) ** a * N.sin(N.pi * a / 2.) / (1. + (om * tau) ** a * N.cos(N.pi * a / 2.)))
        e_loss += 2 * delta_eps * (1 + 2 * (om * tau) ** a * N.cos(N.pi * a / 2.) + (om * tau) ** (2. * a) ) ** (
        -b / 2.) * N.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 = (N.sin(N.pi * a / 2. / (b + 1)) / N.sin(N.pi * a * b / 2. / (b + 1))) ** (1 / a)
    tau /= 2 * N.pi * f
    return tau


### define funcs here
class Functions:
    def __init__(self):
        self.list = {
            "hn":(self.hn_cmplx,4),
            "conductivity":(self.cond_cmplx,1),
            "power":(self.power_cmplx,2),
            "static":(self.static_cmplx,1),
        }

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

    def cond_cmplx(self, p, x):
        om = 2*N.pi*x
        sgma = p[0]
        cond = sgma/(1j*om)
        cplx = N.array([cond.real, -cond.imag])
        return cplx

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

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

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


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

    def fitfcn(self, p0, x, *funcs):
        self.data = N.zeros( x.shape )
        ndx = 0
        for fn in funcs: # loop over functions and add the results up
            f,num_p = self.functions.get(fn)
            p = p0[ndx:ndx+num_p]
            self.data += f(p, x[0]) # fit functions take only 1-dim x
            ndx += num_p
        return self.data

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