# pre/_reprojection.py
"""Re-projection of trajectories for recovering intrusive models with
Operator Inference.
"""

__all__ = [
            "reproject_discrete",
            "reproject_continuous",
          ]

import numpy as np


# Reprojection schemes ========================================================
def reproject_discrete(f, basis, init, niters, inputs=None):
    """Sample re-projected trajectories of the discrete dynamical system

        q_{j+1} = f(q_{j}, u_{j}),  q_{0} = q0.

    Parameters
    ----------
    f : callable mapping (n,) ndarray (and (m,) ndarray) to (n,) ndarray
        Function defining the (full-order) discrete dynamical system. Accepts
        a full-order state vector and (optionally) an input vector and returns
        another full-order state vector.
    basis : (n,r) ndarray
        Basis for the low-dimensional linear subspace (e.g., POD basis).
    init : (n,) ndarray
        Initial condition for the iteration in the high-dimensional space.
    niters : int
        The number of iterations to do.
    inputs : (m,niters-1) or (niters-1) ndarray
        Control inputs, one for each iteration beyond the initial condition.

    Returns
    -------
    states_reprojected : (r,niters) ndarray
        Re-projected state trajectories in the projected low-dimensional space.
    """
    # Validate and extract dimensions.
    n,r = basis.shape
    if init.shape != (n,):
        raise ValueError("basis and initial condition not aligned")

    # Create the solution array and fill in the initial condition.
    states_ = np.empty((r,niters))
    states_[:,0] = basis.T @ init

    # Run the re-projection iteration.
    if inputs is None:
        for j in range(niters-1):
            states_[:,j+1] = basis.T @ f(basis @ states_[:,j])
    elif inputs.ndim == 1:
        for j in range(niters-1):
            states_[:,j+1] = basis.T @ f(basis @ states_[:,j], inputs[j])
    else:
        for j in range(niters-1):
            states_[:,j+1] = basis.T @ f(basis @ states_[:,j], inputs[:,j])

    return states_


def reproject_continuous(f, basis, states, inputs=None):
    """Sample re-projected trajectories of the continuous system of ODEs

        dq / dt = f(t, q(t), u(t)),     q(0) = q0.

    Parameters
    ----------
    f : callable mapping (n,) ndarray (and (m,) ndarray) to (n,) ndarray
        Function defining the (full-order) differential equation. Accepts a
        full-order state vector and (optionally) an input vector and returns
        another full-order state vector.
    basis : (n,r) ndarray
        Basis for the low-dimensional linear subspace.
    states : (n,k) ndarray
        State trajectories (training data).
    inputs : (m,k) or (k,) ndarray
        Control inputs corresponding to the state trajectories.

    Returns
    -------
    states_reprojected : (r,k) ndarray
        Re-projected state trajectories in the projected low-dimensional space.
    ddts_reprojected : (r,k) ndarray
        Re-projected velocities in the projected low-dimensional space.
    """
    # Validate and extract dimensions.
    if states.shape[0] != basis.shape[0]:
        raise ValueError("states and basis not aligned, first dimension "
                         f"{states.shape[0]} != {basis.shape[0]}")
    n,r = basis.shape
    _,k = states.shape

    # Create the solution arrays.
    states_ = basis.T @ states
    ddts_ = np.empty((r,k))

    # Run the re-projection iteration.
    if inputs is None:
        for j in range(k):
            ddts_[:,j] = basis.T @ f(basis @ states_[:,j])
    elif inputs.ndim == 1:
        for j in range(k):
            ddts_[:,j] = basis.T @ f(basis @ states_[:,j], inputs[j])
    else:
        for j in range(k):
            ddts_[:,j] = basis.T @ f(basis @ states_[:,j], inputs[:,j])

    return states_, ddts_
