import numpy as np

from .model import Model
from .parameter import Parameters


class Data(object):
    def __init__(self, x, y, we=None, idx=None):
        self.x = np.asarray(x)
        self.y = np.asarray(y)
        if self.y.shape[0] != self.x.shape[0]:
            raise ValueError(f'x and y have different lengths {self.x.shape[0]} and {self.y.shape[0]}')

        self.we = self._calc_weights(we)
        self.idx = idx
        self.model = None
        self.minimizer = None
        self.parameter = Parameters()
        self.para_keys = None
        self.fun_kwargs = {}

    def __len__(self):
        return self.y.shape[0]

    def _calc_weights(self, we):
        if we is None:
            return 1.

        if isinstance(we, str):
            if we == 'y2':
                we_func = lambda yy: 1. / yy**2
            elif we.lower() == 'none':
                we_func = lambda yy: np.ones_like(len(yy))
            else:
                we_func = lambda yy: 1. / np.abs(yy)

            if np.iscomplexobj(self.y):
                weights = we_func(np.r_[self.y.real, self.y.imag])
            else:
                weights = we_func(self.y)

        else:
            we = 1. / np.asarray(we)
            if np.iscomplexobj(self.y):
                if np.iscomplexobj(we):
                    weights = np.r_[we.real, we.imag]
                else:
                    weights = np.tile(we)
            else:
                weights = we

        weights[weights == np.inf] = np.finfo(float).max

        return weights

    def set_model(self, func, *args, **kwargs):
        if isinstance(func, Model):
            self.model = func
        else:
            self.model = Model(func, *args)

        self.fun_kwargs.update(self.model.fun_kwargs)
        self.fun_kwargs.update(kwargs)
        self.parameter = Parameters()

        return self.model

    def get_model(self):
        return self.model

    def set_parameter(self, parameter, var=None, ub=None, lb=None,
                      default_bounds=False, fun_kwargs=None):
        """
        Creates parameter for this data.
        If no Model is available, it falls back to the model
        :param parameter: list of parameters
        :param var: list of boolean or boolean; False fixes parameter at given list index.
                    Single value is broadcast to all parameter
        :param ub: list of upper boundaries or float; Single value is broadcast to all parameter.
                       None means no bound.
        :param lb: list of lower bounds or float; Single value is broadcast to all parameter.
                       None means no bound.
        :param default_bounds: bool; If True, uses default_bounds of a Model if lb or ub are None
        :param fun_kwargs: dict; key-word arguments for model,
        :return: Parameters
        """
        model = self.model
        if model is None:
            # Data has no unique
            if self.minimizer is None:
                model = None
            else:
                model = self.minimizer.fit_model
                self.fun_kwargs.update(model.fun_kwargs)

        if model is None:
            raise ValueError('No model found, please set model before parameters')

        if default_bounds:
            if lb is None:
                lb = model.lb
            if ub is None:
                ub = model.ub

        self.para_keys = self.parameter.add_parameter(parameter, var=var, lb=lb, ub=ub)

        if fun_kwargs is not None:
            self.fun_kwargs.update(fun_kwargs)

        return self.para_keys

    def get_parameter(self, scaled: bool = False):
        """
        Returns list of parameters if set.
        :param scaled: If True, returns scaled values, otherwise unscaled
        :return: parameter array
        """
        if self.parameter is None:
            return

        if scaled:
            return [p.scaled_value for p in self.minimizer.parameters[self.parameter]]
        else:
            return [p.value for p in self.minimizer.parameters[self.parameter]]

    def cost(self, p):
        """
        Cost function :math:`y-f(p, x)`
        :param p: list of parameters
        :return:
        """
        y_pred = self.model.func(p, self.x, **self.fun_kwargs)
        resid = y_pred - self.y

        if np.iscomplexobj(resid):
            resid = np.r_[resid.real, resid.imag]

        resid *= self.we

        return resid

    def func(self, p, x):
        """
        Function :math:`f(p, x)`
        :param x:
        :param p: list of parameters
        :return:
        """
        return self.model.func(p, x, **self.fun_kwargs)
