from typing import Dict, List, Tuple, Union

import casadi as ca

import numpy as np

from pymoca.backends.casadi.model import Variable as _Variable


MATHEMATICAL_OPERATORS = [
    "__add__",
    "__radd__",
    "__sub__",
    "__rsub__",
    "__mul__",
    "__rmul__",
    "__floordiv__",
    "__truediv__",
    "__mod__",
    "__pow__",
    "__lshift__",
    "__rshift__",
    "__and__",
    "__xor__",
    "__or__",
    "__iadd__",
    "__isub__",
    "__imul__",
    "__idiv__",
    "__ifloordiv__",
    "__imod__",
    "__ipow__",
    "__ilshift__",
    "__irshift__",
    "__iand__",
    "__ixor__",
    "__ior__",
    "__neg__",
    "__pos__",
    "__abs__",
    "__invert__",
    "__complex__",
    "__int__",
    "__long__",
    "__float__",
    "__oct__",
    "__hex__",
    "__lt__",
    "__le__",
    "__eq__",
    "__ne__",
    "__ge__",
    "__gt__",
]


class BaseVariable(_Variable):
    _attr_set = {"value", "start", "min", "max", "nominal", "fixed"}

    def __new__(cls, *args, **kwargs):
        if not hasattr(cls, MATHEMATICAL_OPERATORS[0]):
            for attr in MATHEMATICAL_OPERATORS:

                def _f(self, *args, attr=attr, **kwargs):
                    return getattr(self.symbol, attr)(*args, **kwargs)

                setattr(cls, attr, _f)

        return super().__new__(cls)

    def __init__(self, name, dimensions=None, **kwargs):
        update_attrs = {}

        for k, v in kwargs.items():
            if k in Variable._attr_set:
                if "k" == "start":
                    raise NotImplementedError("Setting 'start' attribute is not supported yet")
                update_attrs[k] = v

        for k in update_attrs:
            kwargs.pop(k)

        if dimensions is None:
            dimensions = []

        super().__init__(ca.MX.sym(name, *dimensions), **kwargs)

        for k, v in update_attrs.items():
            setattr(self, k, v)

    def __MX__(self):  # noqa: N802
        return self.symbol

    def __getitem__(self, key):
        return self.symbol.__getitem__(key)


class Variable(BaseVariable):
    def der(self):
        try:
            return self._derivative
        except AttributeError:
            self._derivative = Variable(f"der({self.symbol.name()})")
            return self._derivative

    @property
    def has_derivative(self):
        return hasattr(self, "_derivative")


class ControlInput(BaseVariable):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.fixed = False


class ConstantInput(BaseVariable):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.fixed = True


class SymbolicParameter(BaseVariable):
    def __init__(self, name, *args, **kwargs):
        if args and "value" not in kwargs:
            kwargs["value"] = args[0]
        super().__init__(name, *args, **kwargs)

    @property
    def name(self):
        return self.symbol.name()


class Array:
    def __init__(self, type_, var_name, dimensions, **modifiers):
        self._array = np.empty(dimensions, dtype=object)
        self._names = np.empty(dimensions, dtype=object)

        for index in np.ndindex(dimensions):
            str_suffix = "[{}]".format(",".join(str(x + 1) for x in index))
            self._array[index] = type_(f"{var_name}{str_suffix}", **modifiers)
            self._names[index] = str_suffix

    def __getitem__(self, key):
        if isinstance(key, int):
            return self._array[key - 1]
        else:
            return self._array[tuple(x - 1 for x in key)]


class Model:
    _modifiers = {}
    _variables: Dict[str, Union[Array, BaseVariable, "Model"]] = {}
    _numeric_parameters: Dict[str, Union[float, int, bool, str]] = {}
    _derivatives: Dict[str, Variable] = {}
    _equations: List[ca.MX] = []
    _initial_equations: List[ca.MX] = []
    _inequalities: List[Tuple[ca.MX, float, float]] = []
    _initial_inequalities: List[Tuple[ca.MX, float, float]] = []
    _skip_variables = None

    def __init__(self, name, **modifiers):
        # Note that this method should be such that it's allowed to be called
        # multiple times, e.g. when a (modifiable) parameter is passed as a
        # modifier to its super class.
        self._variables = {}
        self._numeric_parameters = {}
        self._derivatives = {}
        self._equations = []
        self._initial_equations = []

        self.name = name
        # Value assignment can be done directly, but we move it to the value attribute to
        # make sure that all modifiers are a dictionary
        for k, v in modifiers.items():
            if not isinstance(v, dict):
                modifiers[k] = dict(value=v)

        self._modifiers = modifiers

        self.__prefix = "" if name is None else f"{self.name}."

        if Model._skip_variables is None:
            Model._skip_variables = dir(self)

    def add_variable(self, type_, var_name, *dimensions, **kwargs):
        if var_name in self._variables:
            raise Exception(f"Variable with name '{var_name}' already exists")

        if var_name in self._modifiers:
            kwargs = self.merge_modifiers(kwargs, self._modifiers.pop(var_name))
            # Explicit conversion to MX for our wrapper classes
            for k, v in kwargs.items():
                if isinstance(v, BaseVariable):
                    kwargs[k] = ca.MX(v)

        if dimensions:
            var = self._variables[var_name] = Array(
                type_, f"{self.__prefix}{var_name}", dimensions, **kwargs
            )
        else:
            var = self._variables[var_name] = type_(f"{self.__prefix}{var_name}", **kwargs)

        if isinstance(var, (Variable, ControlInput, ConstantInput)) and (
            isinstance(var.value, (ca.MX, BaseVariable)) or not np.isnan(var.value)
        ):
            # For states and algebraic states, we move the "value" part to an equation
            self.add_equation(var - var.value)
            var.value = np.nan

    def add_equation(self, equation, lb=None, ub=None):
        if lb is None and ub is None:
            self._equations.append(equation)
        elif lb is not None and ub is not None and lb == ub:
            self._equations.append(equation - lb)
        else:
            self._inequalities.append((equation, lb, ub))

    def add_initial_equation(self, equation, lb=None, ub=None):
        if lb is None and ub is None:
            self._initial_equations.append(equation)
        elif lb is not None and ub is not None and lb == ub:
            self._initial_equations.append(equation - lb)
        else:
            self._initial_inequalities.append((equation, lb, ub))

    def connect(self, a: "Connector", b: "Connector"):
        if not a.variables.keys() == b.variables.keys():
            raise Exception(
                f"Cannot connect port {a} of type {type(a)} to port {b} "
                f"of type {type(b)} as they have different variables."
            )

        self._equations.extend([a.variables[k] - b.variables[k] for k in a.variables.keys()])

    def der(self, var: Variable):
        return var.der()

    @property
    def variables(self):
        return self._variables.copy()

    @property
    def numeric_parameters(self):
        return self._numeric_parameters.copy()

    @property
    def equations(self):
        return self._equations.copy()

    @property
    def initial_equations(self):
        return self._initial_equations.copy()

    @property
    def inequalities(self):
        return self._inequalities.copy()

    @property
    def initial_inequalities(self):
        return self._initial_inequalities.copy()

    @staticmethod
    def merge_modifiers(a: dict, b: dict):
        """
        Recursive (not in place) merge of dictionaries.

        :param a: Base dictionary to merge.
        :param b: Dictionary to merge on top of base dictionary.
        :return: Merged dictionary
        """
        b = b.copy()

        for k, v in a.items():
            if isinstance(v, dict):
                b_node = b.setdefault(k, {})
                b[k] = Model.merge_modifiers(v, b_node)
            else:
                if k not in b:
                    b[k] = v

        return b

    def __MX__(self):  # noqa: N802
        return self.symbol

    def __getattr__(self, attr):
        try:
            return self._variables[attr]
        except KeyError:
            pass

        try:
            return self._numeric_parameters[attr]
        except KeyError:
            raise AttributeError(f"Attribute '{attr}' not found")

    def __setattr__(self, key, value):
        if self._skip_variables is None or key in self._skip_variables:
            super().__setattr__(key, value)
        else:
            try:
                value = self._modifiers.pop(key)["value"]
            except KeyError:
                pass
            self._numeric_parameters[key] = value

    def __repr__(self):
        return self.name

    def __str__(self):
        return self.name

    def flatten(self):
        if self._modifiers:
            raise Exception("Cannot flatten a model with remaining modifiers")

        m = FlattenedModel()

        all_variables = {}
        all_parameters = {}

        all_equations = []
        all_initial_equations = []

        all_inequalities = []
        all_initial_inequalities = []

        # First we expand arrays
        variables = {}
        for k, var in self._variables.items():
            if isinstance(var, Array):
                for el, suff in zip(var._array.ravel(), var._names.ravel()):
                    variables[f"{k}{suff}"] = el
            else:
                variables[k] = var

        # Move variables to flattened model
        for k, var in variables.items():
            if isinstance(var, Model):
                flatten_var = var.flatten()
                all_variables.update(flatten_var._variables)
                all_parameters.update(flatten_var._numeric_parameters)

                all_equations.extend(flatten_var._equations)
                all_initial_equations.extend(flatten_var._initial_equations)

                all_inequalities.extend(flatten_var._inequalities)
                all_initial_inequalities.extend(flatten_var._initial_inequalities)
            else:
                all_variables[f"{self.__prefix}{k}"] = var

        all_parameters.update(
            {f"{self.__prefix}{p}": v for p, v in self._numeric_parameters.items()}
        )
        all_equations.extend(self.equations)
        all_initial_equations.extend(self.initial_equations)

        all_inequalities.extend(self.inequalities)
        all_initial_inequalities.extend(self.initial_inequalities)

        m._variables = all_variables
        m._numeric_parameters = all_parameters

        m._equations = all_equations
        m._initial_equations = all_initial_equations

        m._inequalities = all_inequalities
        m._initial_inequalities = all_initial_inequalities

        return m


class FlattenedModel(Model):
    def __init__(self):
        super().__init__(None)


class Component(Model):
    pass


class Connector(Component):
    pass
