from __future__ import annotations

import numpy as np

from .model import Model
from .parameter import Parameters, Parameter


class Data:
    def __init__(self, x, y, we=None, idx=None, complex_type: int = 0):
        self.x = np.asarray(x)
        self.y = np.asarray(y)
        if self.x.size == 0 or self.y.size == 0:
            raise ValueError("Data is empty")

        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.we_string = self._calc_weights(we)
        self.idx = idx
        self.model = None
        self.minimizer = None
        self.parameter = Parameters()
        self.para_keys: list = []
        self.fun_kwargs = {}
        self.complex_type = complex_type

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

    def _calc_weights(self, we: str | np.ndarray | None) -> tuple[np.ndarray, str]:
        if isinstance(we, str) or we is None:
            if we is None or we.lower() == 'none':
                we_string = 'None'
                we_func = lambda yy: np.ones_like(len(yy))
            elif we == 'y2':
                we_string = we
                we_func = lambda yy: 1. / yy**2
            else:
                we_string = we
                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_string = 'yerr'  # This is pure speculation that array equals error
            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, we_string

    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,
                      values: list[float | Parameter],
                      *,
                      var: list[bool] = None,
                      ub: list[float] = None,
                      lb: list[float] = None,
                      default_bounds: bool = False,
                      fun_kwargs: dict = None
                      ):
        """
        Creates parameter for this data.
        If no Model is available, it falls back to the model
        :param values: 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 not None:
                model = self.minimizer.fit_model

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

        if len(values) != len(model.params):
            raise ValueError('Number of given parameter does not match number of model parameters')

        is_parameter = [isinstance(v, Parameter) for v in values]
        if all(is_parameter):
            for p_i in values:
                key = f"p{next(Parameters.parameter_counter)}"
                self.parameter.add_parameter(key, p_i)
        elif any(is_parameter):
            raise ValueError('list of parameter are not all float of Parameter')

        else:
            if var is None:
                var = [True] * len(values)

            if lb is None:
                if default_bounds:
                    lb = model.lb
                else:
                    lb = [None] * len(values)

            if ub is None:
                if default_bounds:
                    ub = model.ub
                else:
                    ub = [None] * len(values)

            arg_names = ['name', 'value', 'var', 'lb', 'ub']
            for parameter_arg in zip(model.params, values, var, lb, ub):
                self.parameter.add(**{arg_name: arg_value for arg_name, arg_value in zip(arg_names, parameter_arg)})

        self.para_keys = list(self.parameter.keys())

        self.fun_kwargs.update(model.fun_kwargs)
        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 replace_parameter(self, key: str, parameter: Parameter) -> None:
        tobereplaced = None
        for k, v in self.parameter.items():
            if v.name == parameter.name:
                tobereplaced = k
                break

        if tobereplaced is None:
            raise KeyError(f'Global parameter {key} not found in list of parameters')
        self.para_keys[self.para_keys.index(tobereplaced)] = key
        self.parameter.replace_parameter(tobereplaced, key, 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)
