#
# Convert a PyBaMM expression tree to a CasADi expression tree
#
from __future__ import annotations

import casadi
import numpy as np
from scipy import interpolate, special

import pybamm


class CasadiConverter:
    def __init__(self, casadi_symbols=None):
        self._casadi_symbols = casadi_symbols or {}

        pybamm.citations.register("Andersson2019")

    def convert(
        self,
        symbol: pybamm.Symbol,
        t: casadi.MX,
        y: casadi.MX,
        y_dot: casadi.MX,
        inputs: dict | None,
    ) -> casadi.MX:
        """
        This function recurses down the tree, converting the PyBaMM expression tree to
        a CasADi expression tree

        Parameters
        ----------
        symbol : :class:`pybamm.Symbol`
            The symbol to convert
        t : :class:`casadi.MX`
            A casadi symbol representing time
        y : :class:`casadi.MX`
            A casadi symbol representing state vectors
        y_dot : :class:`casadi.MX`
            A casadi symbol representing time derivatives of state vectors
        inputs : dict
            A dictionary of casadi symbols representing parameters

        Returns
        -------
        :class:`casadi.MX`
            The converted symbol
        """
        try:
            return self._casadi_symbols[symbol]
        except KeyError:
            # Change inputs to empty dictionary if it's None
            inputs = inputs or {}
            casadi_symbol = self._convert(symbol, t, y, y_dot, inputs)
            self._casadi_symbols[symbol] = casadi_symbol

            return casadi_symbol

    def _convert(self, symbol, t, y, y_dot, inputs):
        """See :meth:`CasadiConverter.convert()`."""
        if isinstance(
            symbol,
            pybamm.Scalar | pybamm.Array | pybamm.Time | pybamm.InputParameter,
        ):
            return casadi.MX(symbol.evaluate(t, y, y_dot, inputs))

        elif isinstance(symbol, pybamm.StateVector):
            if y is None:
                raise ValueError("Must provide a 'y' for converting state vectors")
            return casadi.vertcat(*[y[y_slice] for y_slice in symbol.y_slices])

        elif isinstance(symbol, pybamm.StateVectorDot):
            if y_dot is None:
                raise ValueError("Must provide a 'y_dot' for converting state vectors")
            return casadi.vertcat(*[y_dot[y_slice] for y_slice in symbol.y_slices])

        elif isinstance(symbol, pybamm.BinaryOperator):
            left, right = symbol.children
            # process children
            converted_left = self.convert(left, t, y, y_dot, inputs)
            converted_right = self.convert(right, t, y, y_dot, inputs)

            if isinstance(symbol, pybamm.Modulo):
                return casadi.fmod(converted_left, converted_right)
            if isinstance(symbol, pybamm.Minimum):
                return casadi.fmin(converted_left, converted_right)
            if isinstance(symbol, pybamm.Maximum):
                return casadi.fmax(converted_left, converted_right)
            if isinstance(symbol, pybamm.KroneckerProduct):
                return casadi.kron(converted_left, converted_right)

            # _binary_evaluate defined in derived classes for specific rules
            return symbol._binary_evaluate(converted_left, converted_right)

        elif isinstance(symbol, pybamm.UnaryOperator):
            converted_child = self.convert(symbol.child, t, y, y_dot, inputs)
            if isinstance(symbol, pybamm.Transpose):
                return converted_child.T
            if isinstance(symbol, pybamm.AbsoluteValue):
                return casadi.fabs(converted_child)
            if isinstance(symbol, pybamm.Floor):
                return casadi.floor(converted_child)
            if isinstance(symbol, pybamm.Ceiling):
                return casadi.ceil(converted_child)
            return symbol._unary_evaluate(converted_child)

        elif isinstance(symbol, pybamm.Function):
            converted_children = [
                self.convert(child, t, y, y_dot, inputs) for child in symbol.children
            ]
            # Special functions
            if symbol.function == np.min:
                return casadi.mmin(*converted_children)
            elif symbol.function == np.max:
                return casadi.mmax(*converted_children)
            elif symbol.function == np.abs:
                return casadi.fabs(*converted_children)
            elif symbol.function == np.sqrt:
                return casadi.sqrt(*converted_children)
            elif symbol.function == np.sin:
                return casadi.sin(*converted_children)
            elif symbol.function == np.arcsinh:
                return casadi.arcsinh(*converted_children)
            elif symbol.function == np.arccosh:
                return casadi.arccosh(*converted_children)
            elif symbol.function == np.tanh:
                return casadi.tanh(*converted_children)
            elif symbol.function == np.cosh:
                return casadi.cosh(*converted_children)
            elif symbol.function == np.sinh:
                return casadi.sinh(*converted_children)
            elif symbol.function == np.cos:
                return casadi.cos(*converted_children)
            elif symbol.function == np.exp:
                return casadi.exp(*converted_children)
            elif symbol.function == np.log:
                return casadi.log(*converted_children)
            elif symbol.function == np.sign:
                return casadi.sign(*converted_children)
            elif symbol.function == special.erf:
                return casadi.erf(*converted_children)
            elif isinstance(symbol, pybamm.Interpolant):
                if symbol.interpolator == "linear":
                    solver = "linear"
                elif symbol.interpolator == "cubic":
                    solver = "bspline"
                elif symbol.interpolator == "pchip":
                    x_np = np.array(symbol.x[0])
                    y_np = np.array(symbol.y)
                    pchip_interp = interpolate.PchipInterpolator(x_np, y_np)
                    d_np = pchip_interp.derivative()(x_np)
                    x = converted_children[0]

                    def hermite_poly(i):
                        x0 = x_np[i]
                        x1 = x_np[i + 1]
                        h_val = x1 - x0
                        h_val_mx = casadi.MX(h_val)
                        y0 = casadi.MX(y_np[i])
                        y1 = casadi.MX(y_np[i + 1])
                        d0 = casadi.MX(d_np[i])
                        d1 = casadi.MX(d_np[i + 1])
                        xn = (x - x0) / h_val_mx
                        h00 = 2 * xn**3 - 3 * xn**2 + 1
                        h10 = xn**3 - 2 * xn**2 + xn
                        h01 = -2 * xn**3 + 3 * xn**2
                        h11 = xn**3 - xn**2
                        return (
                            h00 * y0
                            + h10 * h_val_mx * d0
                            + h01 * y1
                            + h11 * h_val_mx * d1
                        )

                    # Build piecewise polynomial for points inside the domain.
                    inside = casadi.MX.zeros(x.shape)
                    for i in range(len(x_np) - 1):
                        cond = casadi.logic_and(x >= x_np[i], x <= x_np[i + 1])
                        inside = casadi.if_else(cond, hermite_poly(i), inside)

                    # Extrapolation:
                    left = hermite_poly(0)  # For x < x_np[0]
                    right = hermite_poly(len(x_np) - 2)  # For x > x_np[-1]

                    # if greater than the maximum, use right; otherwise, use the piecewise value.
                    result = casadi.if_else(
                        x < x_np[0], left, casadi.if_else(x > x_np[-1], right, inside)
                    )
                    return result
                else:  # pragma: no cover
                    raise NotImplementedError(
                        f"Unknown interpolator: {symbol.interpolator}"
                    )

                if len(converted_children) == 1:
                    if solver == "linear":
                        test = casadi.MX.interpn_linear(
                            symbol.x, symbol.y.flatten(), converted_children
                        )
                        if test.shape[0] == 1 and test.shape[1] > 1:
                            # for some reason, pybamm.Interpolant always returns a column vector, so match that
                            test = test.T
                        return test
                    elif solver == "bspline":
                        bspline = interpolate.make_interp_spline(
                            symbol.x[0], symbol.y, k=3
                        )
                        knots = [bspline.t]
                        coeffs = bspline.c.flatten()
                        degree = [bspline.k]
                        m = len(coeffs) // len(symbol.x[0])
                        f = casadi.Function.bspline(
                            symbol.name, knots, coeffs, degree, m
                        )
                        return f(converted_children[0])
                    else:
                        return casadi.interpolant(
                            "LUT", solver, symbol.x, symbol.y.flatten()
                        )(*converted_children)
                elif len(converted_children) in [2, 3]:
                    if solver == "linear":
                        return casadi.MX.interpn_linear(
                            symbol.x,
                            symbol.y.ravel(order="F"),
                            converted_children,
                        )
                    elif solver == "bspline" and len(converted_children) == 2:
                        bspline = interpolate.RectBivariateSpline(
                            symbol.x[0], symbol.x[1], symbol.y
                        )
                        [tx, ty, c] = bspline.tck
                        [kx, ky] = bspline.degrees
                        knots = [tx, ty]
                        coeffs = c
                        degree = [kx, ky]
                        m = 1
                        f = casadi.Function.bspline(
                            symbol.name, knots, coeffs, degree, m
                        )
                        return f(casadi.hcat(converted_children).T).T
                    else:
                        LUT = casadi.interpolant(
                            "LUT", solver, symbol.x, symbol.y.ravel(order="F")
                        )
                        res = LUT(casadi.hcat(converted_children).T).T
                        return res
                else:  # pragma: no cover
                    raise ValueError(
                        f"Invalid converted_children count: {len(converted_children)}"
                    )

            elif symbol.function.__name__.startswith("elementwise_grad_of_"):
                differentiating_child_idx = int(symbol.function.__name__[-1])
                # Create dummy symbolic variables in order to differentiate using CasADi
                dummy_vars = [
                    casadi.MX.sym("y_" + str(i)) for i in range(len(converted_children))
                ]
                func_diff = casadi.gradient(
                    symbol.differentiated_function(*dummy_vars),
                    dummy_vars[differentiating_child_idx],
                )
                # Create function and evaluate it using the children
                casadi_func_diff = casadi.Function("func_diff", dummy_vars, [func_diff])
                return casadi_func_diff(*converted_children)
            # Other functions
            else:
                return symbol._function_evaluate(converted_children)
        elif isinstance(symbol, pybamm.Concatenation):
            converted_children = [
                self.convert(child, t, y, y_dot, inputs) for child in symbol.children
            ]
            if isinstance(symbol, pybamm.NumpyConcatenation | pybamm.SparseStack):
                return casadi.vertcat(*converted_children)
            # DomainConcatenation specifies a particular ordering for the concatenation,
            # which we must follow
            elif isinstance(symbol, pybamm.DomainConcatenation):
                slice_starts = []
                all_child_vectors = []
                for i in range(symbol.secondary_dimensions_npts):
                    child_vectors = []
                    for child_var, slices in zip(
                        converted_children, symbol._children_slices, strict=True
                    ):
                        for child_dom, child_slice in slices.items():
                            slice_starts.append(symbol._slices[child_dom][i].start)
                            child_vectors.append(
                                child_var[child_slice[i].start : child_slice[i].stop]
                            )
                    all_child_vectors.extend(
                        [
                            v
                            for _, v in sorted(
                                zip(slice_starts, child_vectors, strict=False)
                            )
                        ]
                    )
                return casadi.vertcat(*all_child_vectors)

        else:
            raise TypeError(
                f"Cannot convert symbol of type '{type(symbol)}' to CasADi. Symbols must all be "
                "'linear algebra' at this stage."
            )
