import numpy as np

from scipy.integrate import solve_ivp

from typing import Sequence, Callable, Literal
from ..sensor_utils import TimeFunc
from ..exceptions import RydiquleError

def scipy_solve(eoms_base: np.ndarray, const: np.ndarray,
                eom_time_r: np.ndarray, const_r: np.ndarray,
                eom_time_i: np.ndarray, const_i: np.ndarray,
                time_inputs: Sequence[TimeFunc],
                t_eval: np.ndarray, init_cond: np.ndarray,
                eqns: Literal["loop", "comp"] = "loop",
                **kwargs
                ) -> np.ndarray:
    """
    Solve a set of Optical Bloch Equations (OBEs) with rydiqule's time solving convention
    using scipy's `solve_ivp`. 

    Uses matrix components of the equations of motion provided by the methods of a :meth:`~.Sensor`.
    Designed to be used as a wrapped function within :func:`~.timesolvers.solve_time`.
    Builds and solves equations of motion according rydiqule's time solving conventions.
    Sets up and solves dx/dt = A(t)x + b(t)

    Args
    ----
    eoms_base: numpy.ndarray
        The matrix of shape `(*l,n,n)` representing the non time-varying portion of the matrix A
        in the equations of motion.
    const: numpy.ndarray
        The array of shape `(*l, n)` representing the non time-varying portion of the vector b in the
        equations of motion.
    eoms_time_r: numpy.ndarray
        The matrix of shape `(n_t, *l, n, n)` representing the real time-varying portion of the matrix A,
        where n_t is the length of `time_inputs`.
        The ith slice along the first axis should be multiplied by the real part
        of the ith entry in `time_inputs`.
    const_r: numpy.nd_array
        The matrix of shape `(n_t, *l, n)` representing the real time-varying portion of the vector b,
        where n_t is the length of `time_inputs`.
        The ith slice along the first axis should be multiplied by the real part
        of the ith entry in `time_inputs`.
    eoms_time_i: numpy.ndarray
        The matrix of shape `(n_t, *l, n, n)` representing the imaginary time-varying portion of the matrix A,
        where n_t is the length of `time_inputs`.
        The ith slice along the first axis should be multiplied by the imaginary part
        of the ith entry in `time_inputs`.
    const_i: numpy.nd_array
        The matrix of shape `(n_t, *l, n)` representing the imaginary time-varying portion of the vector b,
        where n_t is the length of `time_inputs`.
        The ith slice along the first axis should be multiplied by the imaginary part
        of the ith entry in `time_inputs`.
    time_inputs: list(callable)
        List of callable functions of length `n_t`.
        The functions should take a single floating point
        as an input representing the time in microseconds,
        and return a real or complex floating point value represent an
        electric field in V/m at that time.
    t_eval: numpy.ndarray
        Array of times to sample the integration at.
    init_cond: (numpy.ndarray)
        Matrix of shape `(*l, n)` representing the initial state of the system.
    eqns: {"loop", "comp"}
        Function used of generating equations of motion. One of "loop" or "comp", corresponding
        to defining time-dependent equations of motion as a loop over time-dependent components
        or with a list comprehension. List comprehensions are preferred for longer solves and
        loops are preferred for shorter solves.
    **kwargs: dict
        Additional keyword arguments passed to the nbkode solver constructor.

    Returns
    -------
    numpy.ndarray
        The matrix solution of shape `(*l,n,n_t)`
        representing the density matrix of the system at each time t.
    """
    
    _derEqns = {"loop": _derEqns_loop, "comp": _derEqn_comp}

    sol_shape = eoms_base.shape[:-1]

    try:
        equations = _derEqns[eqns](eoms_base, const,
                                   eom_time_r, const_r,
                                   eom_time_i, const_i,
                                   time_inputs)
    except KeyError as err:
        raise RydiquleError("\'eqns\' must be one of \'loop\' or \'comp\'.") from err
    
    init_cond = init_cond.ravel()

    method = kwargs.pop("method", "RK45")
    rtol = kwargs.pop("rtol", 1e-6)

    sol_flat = solve_ivp(equations, (t_eval[0], t_eval[-1]),
                         init_cond, t_eval=t_eval,
                         method=method, rtol=rtol,
                         **kwargs).y

    return sol_flat.reshape(sol_shape+(sol_flat.shape[-1],))


def _derEqns_loop(obes_base: np.ndarray, const_base: np.ndarray,
             obes_time_r: np.ndarray, const_r: np.ndarray,
             obes_time_i: np.ndarray,  const_i: np.ndarray,
             time_inputs: Sequence[TimeFunc],
             ) -> Callable[[float, np.ndarray], np.ndarray]:
    """
    Function to build the callable passed to scipy's solve_ivp in :func:`~.scipy_solve`.

    Uses the base and time matrix components of the eoms to build a function of
    vector and scalar time that has the expected input/output of functions passed
    to `scipy.integrate.solve_ivp()`
    """

    t_func_num = obes_time_r.shape[0]
    input_shape = obes_base.shape[:-1]

    def func(t: float, A_flat: np.ndarray) -> np.ndarray:

        # create OBEs at time t
        obe_total = obes_base.copy()
        const_total = const_base.copy()

        for idx in range(t_func_num):
            ti = time_inputs[idx](t)
            obe_total += ti.real*obes_time_r[idx] + ti.imag*obes_time_i[idx]
            const_total += ti.real*const_r[idx] + ti.imag*const_i[idx]

        # Unflatten A for matrix broadcasting
        A_stack = A_flat.reshape(input_shape+(1,))
        result = np.matmul(obe_total, A_stack).squeeze(-1) + const_total
        # Flatten result to match input
        return np.ravel(result)

    return func


def _derEqn_comp(obes_base: np.ndarray, const: np.ndarray,
            obes_time_r: np.ndarray, const_r: np.ndarray,
            obes_time_i: np.ndarray,  const_i: np.ndarray,
            time_inputs: Sequence[TimeFunc],
            ) -> Callable[[float, np.ndarray], np.ndarray]:
    """
    Function to build the callable passed to scipy's solve_ivp in :func:`~.scipy_solve`.

    Uses the base and time matrix components of the eoms to build a function of vector and
    scalar time that has the expected input/output of functions passed to
    `scipy.integrate.solve_ivp()`
    """
    input_shape = obes_base.shape[:-1]
    l = obes_time_r.shape[0]

    def func(t: float, A_flat: np.ndarray) -> np.ndarray:

        obe_total = obes_base + np.sum([time_inputs[idx](t).real*obes_time_r[idx]
                                        + time_inputs[idx](t).imag*obes_time_i[idx]
                                        for idx in range(l)], axis=0)
        const_total = const + np.sum([time_inputs[idx](t).real*const_r[idx]
                                      + time_inputs[idx](t).imag*const_i[idx]
                                      for idx in range(l)], axis=0)

        # Unflatten A for matrix broadcasting
        A_stack = A_flat.reshape(input_shape+(1,))
        result = np.matmul(obe_total, A_stack).squeeze(-1) + const_total
        # Flatten result to match input
        return np.ravel(result)
 
    return func
