# coding=utf-8
import warnings
import numpy as np
import yafflib

__author__ = 'markusro'


class Function(object):


    def __init__(self):
        super(Function, self).__init__()
        self._id_string = None
        self._num_parameters = None
        self._pretty_name = ""
        self._function = None

    def get_id_string(cls):
        if cls._id_string is None:
            raise NotImplementedError("You need to set the id_string")
        return cls._id_string

    def set_id_string(cls, s):
        cls._id_string = s

    def get_num_paramters(cls):
        if cls._num_paramters is None:
            raise NotImplementedError("You need to set the num_paramters")
        return cls._num_paramters

    def set_num_paramters(cls, s):
        cls._num_paramters = s

    def get_function(cls):
        if cls._function is None:
            raise NotImplementedError("You need to set f")
        return cls._function

    def set_function(cls, f):
        cls._function = f

    def calculate(self, x, p):
        if self._function is None: raise NotImplementedError("You need to set f")
        data = self._function(x, p)
        return x,data[0],data[1]


class HavNegCmplx(Function):
    def __init__(self):
        super(HavNegCmplx, self).__init__()
        self.set_function(hn)
        self.set_id_string("hn")
        self.set_num_paramters(4)


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


class ConductivityCmplx(Function):
    def __init__(self):
        super(ConductivityCmplx, self).__init__()
        self.set_num_paramters(3)
        self.set_function(cond_cmplx)
        self.set_id_string("conductivity")


def cond_cmplx(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


class PowerCmplx(Function):
    def __init__(self):
        super(PowerCmplx, self).__init__()
        self.set_function(power_cmplx)
        self.set_num_paramters(2)
        self.set_id_string("power_law")


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


class EpsInftyCmplx(Function):
    def __init__(self):
        super(EpsInftyCmplx, self).__init__()
        self.set_function(static_cmplx)
        self.set_id_string("e_inf")
        self.set_num_paramters(1)


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


class YaffCmplx(Function):
    def __init__(self):
        super(YaffCmplx, self).__init__()
        self.set_function(yaff)
        self.set_id_string("yaff")
        self.set_num_paramters(8)


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


class ModelFunction(object):
    def __init__(self):
        super(ModelFunction, self).__init__(self)
        self.model = None

    @classmethod
    def select_model(self, model):
        self._functions_avail = ", ".join( cls().get_id_string() for cls in Function.__subclasses__())
        for cls in Function.__subclasses__():
            if model == cls().get_id_string():
                self.model = cls()
                return True
        warnings.warn("Function not found: %s \n(available functions: %s)"%(model, self._functions_avail))
        return False