# -*- encoding: utf8 -*-
from PyQt4.QtGui import QColor
import numpy as np
import pyqtgraph as pg
from PyQt4.QtCore import *

import CustomWidgets

from mathlib import id_to_color, FitFunctionCreator, Functions
import libyaff


class Data:
    def __init__(self, frequency=np.zeros(1), die_real=np.zeros(1), die_imag=np.zeros(1)):
        self.frequency = frequency
        self.epsilon = die_real + 1j * die_imag
        self.frequency_fit = frequency
        self.epsilon_fit = die_real*0 + 1j * die_imag*0
        myPen_imag = pg.mkPen(width=3, color=(255,255,127))
        myPen_real = pg.mkPen(width=3, color=(51,255,127))

        self.data_curve_imag = pg.PlotDataItem(x=[np.nan], y=[np.nan],pen=QColor(0,0,0,0), symbol='o',
                                          symbolBrush=(255,127,0,127))
        self.data_curve_real = pg.PlotDataItem(x=[np.nan], y=[np.nan],pen=QColor(0,0,0,0), symbol='s',
                                          symbolBrush=(119,202,92,127))
        self.fitted_curve_imag = pg.PlotDataItem(np.array([np.nan]), np.array([np.nan]), pen=myPen_imag)
        self.fitted_curve_real = pg.PlotDataItem(np.array([np.nan]), np.array([np.nan]), pen=myPen_real)
        self.length = len(frequency)
        self.meta = dict()
        self.fit_limits = [frequency.min(), frequency.max(), die_imag.min(), die_imag.max()]

        self.fit_param = None
        self.fit_funcs = None # list of fit functions
        self.hide_funcs = None # remove these func from the data

    def set_fit(self, param, funcs):
        self.fit_funcs = funcs
        self.hide_funcs = []
        self.fit_param = param
        fit_real, fit_imag = FitFunctionCreator().fitfcn(param, self.frequency_fit, *funcs)
        self.epsilon_fit =  fit_real + 1j*fit_imag

    def set_data(self,f,e_real,e_imag):
        self.frequency = f
        self.epsilon = e_real + 1j*e_imag
        self.epsilon_fit  =  0*e_real + 1j*e_imag*0
        self.fit_limits = [f.min(), f.max(), e_imag.min(), e_imag.max()]
        self.data_curve_imag.setData(f,e_imag)
        self.data_curve_real.setData(f,e_real)

    def set_fit_xlimits(self, xmin, xmax):
        self.fit_limits[0] = xmin
        self.fit_limits[1] = xmax
        self.frequency_fit = self.frequency[(self.frequency <= xmax) & (self.frequency >= xmin)]

    def set_fit_ylimits(self, ymin, ymax):
        self.fit_limits[2] = ymin
        self.fit_limits[3] = ymax

    def get_data(self):
        #mask = np.ones(len(self.frequency), dtype='bool')
        mask = (self.frequency > self.fit_limits[0]) & (self.frequency < self.fit_limits[1])
        #mask &= (self.epsilon.imag > self.fit_limits[2]) & (self.epsilon.imag < self.fit_limits[3])
        return self.frequency[mask], self.epsilon[mask]

    def remove_curves(self):
        print "remove data_curve"
        #if self.data_curve is not None: self.data_curve.remove()
        print "remove fitted_curve"
        #if self.fitted_curve is not None: self.fitted_curve.remove()

class BaseObject(QObject):
    changedData = pyqtSignal()
    removeObj = pyqtSignal(QObject)

    def __init__(self, plt_real=None, plt_imag=None, limits=None):
        super(BaseObject, self).__init__()

        myPen = pg.mkPen( style=Qt.DotLine,
                          width=2.5)

        self.data_curve_real = pg.PlotDataItem(x=np.array([np.nan]),y=np.array([np.nan]), pen=myPen)
        self.plt_real = plt_real
        self.plt_real.addItem(self.data_curve_real)

        self.data_curve_imag = pg.PlotDataItem(x=np.array([np.nan]),y=np.array([np.nan]), pen=myPen)
        self.plt_imag = plt_imag
        self.plt_imag.addItem(self.data_curve_imag)

        self.limits = limits

        # private varaibles
        self.functions = Functions()
        self._color = QColor("white")
        self._id = None
        self._id_string = None
        self._widget = None
        self._frequency = np.logspace(np.log10(limits[0]), np.log10(limits[1]), 256)
        self._data = None
        self._func = None
        self._beta = None
        self._sd_beta = None
        self._selector_mask = None
        self._param_number = 0

    @property
    def param_number(self):
        return self._param_number

    @param_number.setter
    def param_number(self, num):
        self._param_number = num

    @property
    def id_string(self):
        return self._id_string

    @id_string.setter
    def id_string(self, id):
        #self._func = self.functions.get_function(id)
        self._id_string  = id


    @property
    def id_label(self):
        return self._id

    @id_label.setter
    def id_label(self, id):
        #self._func = self.functions.get_function(id)
        self._func = self.function
        self._id  = id

    @property
    def color(self):
        return self._color

    @color.setter
    def color(self, c):
        self._color = c
        self.data_curve_real.setPen(c)
        self.data_curve_imag.setPen(c)

    @property
    def widget(self):
        return self._widget

    @widget.setter
    def widget(self, wdgt):
        self._widget = wdgt
        self._widget.changedTable.connect(self.updateData) # TODO better to use self.setParameter
        self._widget.removeMe.connect(self.removeMe)

    def getParameter(self):
        p = self.widget.getTable() # TODO ugly ... should return self._beta etc ...?
        return p

    def getFixed(self):
        p = self.widget.fixedParameter()
        return p

    def setParameter(self, beta, sd_beta=None):
        self._beta    = beta
        self._sd_beta = sd_beta
        self.widget.updateTable(beta, sd_beta)
        self.updateData()

    def get_data(self):
        return self._frequency, self._data

    def removeMe(self):
        self.plt_imag.removeItem(self.data_curve_imag)
        self.plt_real.removeItem(self.data_curve_real)
        self.removeObj.emit(self)
        self.changedData.emit()

    def updateData(self):
        self._data = self._func(self.getParameter(), self._frequency)
        self.data_curve_real.setData(x=self._frequency, y=self._data[0])
        self.data_curve_imag.setData(x=self._frequency, y=self._data[1])
        self.changedData.emit()

    def resampleData(self, x):
        data = self._func(self.getParameter(), x)
        return np.array([x,data[0],data[1]]).T


    def clearData(self):
        self.data_curve_real.setData(x=[np.nan], y=[np.nan])
        self.data_curve_imag.setData(x=[np.nan], y=[np.nan])

    def function(self,p,x):
        raise NotImplementedError, "This needs to be implemented in your subclass"

class Conductivity(BaseObject):
    def __init__( self, plt_imag=None, plt_real=None, limits=None ):
        super(Conductivity, self).__init__(plt_real=plt_real, plt_imag=plt_imag, limits=limits)
        self.widget = CustomWidgets.ConductivityWidget()
        self.color = QColor("blue")
        self.id_label = "Cond."
        self.id_string = "cond"

        self.param_number = 3

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


class PowerComplex(BaseObject):
    def __init__( self, plt_real=None, plt_imag=None, limits=None ):
        super(PowerComplex, self).__init__(plt_real=plt_real, plt_imag=plt_imag, limits=limits)
        self.widget = CustomWidgets.PowerLawWidget()
        self.color = QColor("#ff44c4")
        self.id_label = 'Power Law'
        self.id_string = "pwr"
        self.param_number = 2

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


class Static(BaseObject):
    def __init__( self, plt_real=None, plt_imag=None, limits=None ):
        super(Static, self).__init__(plt_real=plt_real, plt_imag=plt_imag, limits=limits)
        self.widget = CustomWidgets.StaticWidget()
        self.color = QColor('#FF0F13')
        self.id_label = u'ε(∞)'
        self.id_string = "eps_infty"
        self.param_number = 1

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


class Peak(BaseObject):
    def __init__( self, id_num=None, plt_real=None, plt_imag=None, limits=None ):
        super(Peak, self).__init__(plt_real=plt_real, plt_imag=plt_imag, limits=limits)
        self.widget = CustomWidgets.PeakWidget()
        self.widget.setId(id_num)
        self.color = id_to_color(id_num)
        self.widget.setColor(self.color)
        self.id_num = id_num
        self.id_label = "Hav-Neg"
        self.id_string = "hn"
        self.param_number = 4

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

class YAFF(BaseObject):
    def __init__( self, plt_real=None, plt_imag=None, limits=None ):
        super(YAFF, self).__init__(plt_real=plt_real, plt_imag=plt_imag, limits=limits)
        self.widget = CustomWidgets.YaffWidget()
        self.widget.on_model_changed.connect(self.change_model)
        self.color = QColor(32, 120, 29, int(255*0.82))
        self._libyaff = libyaff.Yaff(self.widget.getYaffType())
        self.id_label = self._libyaff.label
        self.id_string = "yaff"
        self._param_number = self._libyaff.params
        self._selector_mask = self.widget.selector_mask

    @property
    def param_number(self):
        return self._param_number

    @param_number.setter
    def param_number(self, num=None):
        self._param_number = self._libyaff.params

    def change_model(self):
        self._libyaff = libyaff.Yaff(self.widget.getYaffType())
        self._selector_mask = self.widget.selector_mask
        self.id_label = self._libyaff.label
        self.param_number = self._libyaff.params
        self.updateData()


    def function( self, p, x ):
        ya = self._libyaff.loss( p, x)
        cplx = np.array([ya.imag, ya.real])
        return cplx