"""NL-LFR inference and learning, optimization, and instantiation."""
import equinox as eqx
import jax
import jax.numpy as jnp
import numpy as np
import optimistix as optx

from . import _misc
from ._config import PRINT_EVERY, SEED, SOLVER, DeviceLike
from ._data_manager import FrequencyData, InputOutputData
from ._model_structures import ModelBLA, ModelNonlinearLFR
from ._solve import SolveResult, solve
from .static._feature_maps import AbstractFeatureMap
from .static._nonlin_funcs import AbstractNonlinearFunction, BasisFunctionModel


# When changing these constants, also update the corresponding docstrings!
EPSILON = 1e-8
FIXED_POINT_ITERS = 10
LAMBDA_W = 1e-2
MAX_ITER_INFERENCE_AND_LEARNING = 1000
MAX_ITER_OPTIMIZATION = 100
SIGMA = 1e-4


class DecVarsInferenceLearning(eqx.Module):
    """Decision variables for inference and learning."""

    B_w_star: jnp.ndarray = eqx.field(converter=jnp.asarray)
    C_z_star: jnp.ndarray = eqx.field(converter=jnp.asarray)
    D_yw_star: jnp.ndarray = eqx.field(converter=jnp.asarray)
    D_zu_star: jnp.ndarray = eqx.field(converter=jnp.asarray)
    
        
class ArgsInferenceLearning(eqx.Module):
    """Static arguments for inference and learning optimization."""
    
    theta_uy: tuple  # (A, B_u, C_y)
    phi: AbstractFeatureMap
    lambda_w: float
    fixed_point_iters: int
    freq_data: tuple  # (frequencies, sampling frequency, U, Y, G_yu)
    Lambda: jnp.ndarray  # shape (N//2 + 1, ny, ny)
    Tz_inv: jnp.ndarray  # shape (nz, nz)
    epsilon: float
    N: int
    recompute_fixed_point: bool


class ArgsNonlinearOptimization(eqx.Module):
    """Static arguments for nonlinear LFR optimization."""
    
    theta_static: ModelNonlinearLFR
    u: jnp.ndarray  # shape (N, nu, R)
    Y: jnp.ndarray  # shape (N//2 + 1, ny, R)
    Lambda: jnp.ndarray  # shape (N//2 + 1, ny, ny)
    x0: jnp.ndarray  # shape (nx, R)
    offset: int


### Public functions ###


def inference_and_learning(
    bla: ModelBLA,
    data: InputOutputData,
    phi: AbstractFeatureMap,
    nw: int,
    *,
    lambda_w: float = LAMBDA_W,
    fixed_point_iters: int = FIXED_POINT_ITERS,
    solver: optx.AbstractLeastSquaresSolver | optx.AbstractMinimiser = SOLVER,
    freq_weighting: bool = True,
    max_iter: int = MAX_ITER_INFERENCE_AND_LEARNING,
    print_every: int = PRINT_EVERY,
    return_solve_details: bool = False,
    seed: int = SEED,
    epsilon: float = EPSILON,
    recompute_fixed_point: bool = True,
    device: DeviceLike = None
) -> ModelNonlinearLFR | tuple[ModelNonlinearLFR, SolveResult]:
    """Perform inference and learning.

    Parameters
    ----------
    bla : `ModelBLA`
        Parametric BLA model; is kept fixed during optimization.
    data : `InputOutputData`
        Measured input-output data object.
    phi : `AbstractFeatureMap`
        Any feature mapping that is linear in the parameters.
    nw : int
        Dimension of the latent signal `w`.
    lambda_w : float
        Regularization weight that controls the solution variance of `w`
        during inference. Defaults to `1e-2`.
    fixed_point_iters : int
        Number of fixed-point iterations for internal consistency. Defaults to `10`.
    solver : `optx.AbstractLeastSquaresSolver` or `optx.AbstractMinimiser`
        Any least-squares solver or general minimization solver from the
        Optimistix or Optax libraries. Defaults to 
        `optx.LevenbergMarquardt(rtol=1e-3, atol=1e-6)`.
    freq_weighting : bool
        Whether to use frequency weighting based on the inverse of the output noise 
        variance estimate. Defaults to `True`.
    max_iter : int
        Maximum number of optimization iterations. Defaults to `1000`.
    print_every : int
        Frequency of printing iteration information. If set to `0`, only a
        summary is printed. If set to `-1`, no printing is done. Defaults to `1`.
    return_solve_details : bool
        Whether to return detailed information about the optimization process. This is
        useful for e.g. plotting convergence curves. Defaults to `False`.
    seed : int
        Random seed for parameter initialization. Defaults to `42`.
    epsilon : float
        Numerical regularization constant for matrix inversion. Defaults to `1e-8`.
    recompute_fixed_point : bool
        Whether to recompute the fixed-point iterations with during the backward pass.
        Doing so increases computation time but reduces memory consumption, which can 
        be useful for large-scale problems. Defaults to `True`.
    device : `DeviceLike`, optional
        Device on which to perform the computations. Can be either a device
        name (`"cpu"`, `"gpu"`, or `"tpu"`) or a specific JAX device. If not
        provided, the default JAX device is used.
        
    Returns
    -------
    `ModelNonlinearLFR`
        Fully initialized NL-LFR model.
    `SolveResult`, optional
        More details about the optimization process, only returned if
        `return_solve_details` is `True`.

    """
    logging_enabled = print_every != -1
    if logging_enabled:
        header = " NL-LFR inference and learning "
        print(f"{header:=^72}")
        
    freq_weighting = _validate_weighting(
        freq_weighting, data.freq.Y_var_noise, logging_enabled
    )

    theta0, args = _prepare_inference_and_learning(
        bla, data, phi, nw, lambda_w, fixed_point_iters,
        seed, epsilon, recompute_fixed_point, freq_weighting
    )

    # Optimize the model parameters
    if logging_enabled:
        print("Starting iterative optimization...")
        if print_every > 0:
            bla_loss = _loss_bla(bla, data.freq, freq_weighting)
            print(f"    BLA loss: {bla_loss:.4e}")

    solve_result = solve(
        theta0, solver, args, _loss_inference_and_learning, 
        max_iter, print_every, device
    )

    scalar_loss, beta = solve_result.aux

    model = ModelNonlinearLFR(
        A=args.theta_uy[0],
        B_u=args.theta_uy[1],
        C_y=args.theta_uy[2],
        D_yu=bla.D_yu,
        B_w=solve_result.theta.B_w_star,
        C_z=args.Tz_inv @ solve_result.theta.C_z_star,
        D_yw=solve_result.theta.D_yw_star,
        D_zu=args.Tz_inv @ solve_result.theta.D_zu_star,
        func_static=_create_basis_function_model_given_beta(nw, phi, beta),
        ts=data.time.ts,
        norm=bla.norm
    )

    if logging_enabled:
        offset = data.time.u.shape[0]
        x_bla = _misc.compute_steady_state_bla_state(bla, data)
        _misc.evaluate_model_performance(
            model, data, x0=x_bla[-offset, :, :],
            offset=offset, solve_result=solve_result
        )
        
        # Compute recursive loss of the NL-LFR model to check consistency with the
        # inference and learning solution. We prepend a single period and start from
        # the corresponding steady-state BLA state.
        u_ext = _misc.extend_signal(data.time.u, offset)
        y_hat = model._simulate(u_ext, x0=x_bla[-offset, :, :])[0]
        y_hat = y_hat[offset:, ...]  # discard offset samples
        Y_hat = jnp.fft.rfft(y_hat, axis=0)
        loss_recursive = _loss_output_spectrum(data.freq.Y, Y_hat, args.Lambda)
        scalar_loss_recursive = _misc.scalar_valued(loss_recursive)

        print("NL-LFR loss consistency check:")
        print(f"    Inference & learning loss = {scalar_loss:.4e}")
        print(f"    Recursive simulation loss = {scalar_loss_recursive:.4e}\n")

    if return_solve_details:
        return model, solve_result
    return model


def optimize(
    model: ModelNonlinearLFR,
    data: InputOutputData,
    *,
    solver: optx.AbstractLeastSquaresSolver | optx.AbstractMinimiser = SOLVER,
    freq_weighting: bool = True,
    max_iter: int = MAX_ITER_OPTIMIZATION,
    print_every: int = PRINT_EVERY,
    return_solve_details: bool = False,
    offset: int | None = None,
    device: DeviceLike = None
) -> ModelNonlinearLFR | tuple[ModelNonlinearLFR, SolveResult]:
    """Refine the parameters of an NL-LFR model using time-domain simulations.

    Parameters
    ----------
    model : `ModelNonlinearLFR`
        Initial NL-LFR model to optimize.
    data : `InputOutputData`
        Measured input-output data object.
    solver : `optx.AbstractLeastSquaresSolver` or `optx.AbstractMinimiser`
        Any least-squares solver or general minimization solver from the
        Optimistix or Optax libraries. Defaults to 
        `optx.LevenbergMarquardt(rtol=1e-3, atol=1e-6)`.
    freq_weighting : bool
        Whether to use frequency weighting based on the inverse of the output noise 
        variance estimate. Defaults to `True`.
    max_iter : int
        Maximum number of optimization iterations. Defaults to `100`.
    print_every : int
        Frequency of printing iteration information. If set to `0`, only a
        summary is printed. If set to `-1`, no printing is done. Defaults to `1`.
    return_solve_details : bool
        Whether to return detailed information about the optimization process. This is
        useful for e.g. plotting convergence curves. Defaults to `False`.
    offset : int, optional
        A non-negative integer which is used to select the unknown initial state of 
        the time-domain simulations. Specifically, the initial state is selected from 
        the steady-state BLA state simulations as `x0 = x_bla[-offset, :, :]`. The 
        number of input samples are prepended accordingly. This warm-up should ensure
        the system reaches steady-state before the main period begins, allowing for 
        leakage-free DFT computations (the prepended samples are later discarded). 
        Defaults to 10% of the data length if not provided.
    device : `DeviceLike`, optional
        Device on which to perform the computations. Can be either a device
        name (`"cpu"`, `"gpu"`, or `"tpu"`) or a specific JAX device. If not
        provided, the default JAX device is used.

    Returns
    -------
    `ModelNonlinearLFR`
        NL-LFR model with optimized parameters.
    `SolveResult`, optional
        More details about the optimization process, only returned if
        `return_solve_details` is `True`.

    """
    logging_enabled = print_every != -1
    if logging_enabled:
        header = " NL-LFR optimization "
        print(f"{header:=^72}")
        
    freq_weighting = _validate_weighting(
        freq_weighting, data.freq.Y_var_noise, logging_enabled
    )

    if offset is None:  # we start 10% "ahead of time"
        offset = int(np.ceil(0.1 * data.time.u.shape[0]))

    theta0, args = _prepare_nonlin_optimization(data, model, offset, freq_weighting)
 
    # Optimize the model parameters
    if logging_enabled:
        print("Starting iterative optimization...")
        if print_every > 0:
            bla_loss = _loss_bla(model._bla, data.freq, freq_weighting)
            print(f"    BLA loss: {bla_loss:.4e}")

    solve_result = solve(
        theta0, solver, args, _loss_nonlin_optimization, max_iter, print_every, device
    )

    # Combine optimized dynamic parameters with static ones
    model_opti = eqx.combine(solve_result.theta, args.theta_static)
    
    # Add original BLA to optimized model
    model_opti = eqx.tree_at(lambda tree: tree._bla, model_opti, replace=model._bla)
    
    if logging_enabled:
        _misc.evaluate_model_performance(
            model_opti, data, x0=args.x0, offset=offset, solve_result=solve_result
        )

    if return_solve_details:
        return model_opti, solve_result
    return model_opti


def connect(
    bla: ModelBLA,
    func_static: AbstractNonlinearFunction,
    sigma: float = SIGMA
) -> ModelNonlinearLFR:
    """Instantiate an NL-LFR model given the BLA and a static nonlinear function.

    Parameters
    ----------
    bla : `ModelBLA`
        Parametric BLA model.
    func_static : `AbstractNonlinearFunction`
        Nonlinear function object that defines the static relation between the latent 
        signals `z` and `w`. Note that the `seed` attribute of `func_static` is used
        to randomly initialize the matrices `B_w`, `C_z`, `D_yw`, and `D_zu`.
    sigma : float
        Standard deviation of the random values used to initialize `B_w` and `D_yw`.  
        Choosing a small value ensures the initial NL-LFR loss remains close to the  
        corresponding BLA loss. However, setting it too small reduces gradient 
        magnitude and can slow down or stall learning. Defaults to `1e-4`.

    Returns
    -------
    `ModelNonlinearLFR`
        NL-LFR model with (partly random) initial parameters.

    """
    nw, nz = func_static.nw, func_static.nz
    ny, nu = bla.D_yu.shape
    nx = bla.A.shape[0]

    key = _misc.get_key(func_static.seed, "connect")

    key_B_w, key_C_z, key_D_yw, key_D_zu = jax.random.split(key, 4)
    return ModelNonlinearLFR(
        A=bla.A,
        B_u=bla.B_u,
        C_y=bla.C_y,
        D_yu=bla.D_yu,
        B_w=sigma * jax.random.normal(key_B_w, (nx, nw)),
        C_z=jax.random.normal(key_C_z, (nz, nx)),
        D_yw=sigma * jax.random.normal(key_D_yw, (ny, nw)),
        D_zu=jax.random.normal(key_D_zu, (nz, nu)),
        func_static=func_static,
        ts=bla.ts,
        norm=bla.norm
    )


### Internal helpers ###


def _loss_output_spectrum(
    Y: jnp.ndarray,
    Y_hat: jnp.ndarray,
    Lambda: jnp.ndarray
) -> jnp.ndarray:
    """Compute the weighted loss between the measured and simulated output spectra.

    Parameters
    ----------
    Y : jnp.ndarray, shape (N//2 + 1, ny, R)
        Measured output spectrum, averaged over periods.
    Y_hat : jnp.ndarray, shape (N//2 + 1, ny, R)
        Simulated output spectrum.
    Lambda : jnp.ndarray, shape (N//2 + 1, ny, ny)
        Weight matrix.

    """
    F_tot, _, R = Y.shape
    return jnp.sqrt(Lambda / (F_tot * R)) @ (Y - Y_hat)


def _prepare_inference_and_learning(
    bla: ModelBLA,
    data: InputOutputData,
    phi: AbstractFeatureMap,
    nw: int,
    lambda_w: float,
    fixed_point_iters: int,
    seed: int,
    epsilon: float,
    recompute_fixed_point: bool,
    freq_weighting: bool
) -> tuple[ModelNonlinearLFR, ArgsInferenceLearning]:
    """Prepare initial guess and function arguments for inference and learning."""
    nz = phi.nz
    ny, nx = bla.C_y.shape
    N, nu = data.time.u.shape[:2]

    freqs = data.freq.f
    U = data.freq.U
    Y = data.freq.Y

    # Initialize theta_wz
    key = jax.random.key(seed)
    key_B_w, key_C_z, key_D_yw, key_D_zu = jax.random.split(key, 4)

    B_w_star = jax.random.normal(key_B_w, (nx, nw))
    C_z_star = jax.random.normal(key_C_z, (nz, nx))
    D_zu_star = jax.random.normal(key_D_zu, (nz, nu))
    D_yw_star = jax.random.normal(key_D_yw, (ny, nw))

    theta_wz = DecVarsInferenceLearning(B_w_star, C_z_star, D_yw_star, D_zu_star)
    theta_uy = (bla.A, bla.B_u, bla.C_y)

    # Compute z_star normalization
    G_zu = ModelBLA(  # parametric u->z frequency response; not the true BLA
        A=bla.A, B_u=bla.B_u, C_y=C_z_star, D_yu=D_zu_star, ts=bla.ts, norm=bla.norm
    )._frequency_response(freqs)

    Z_star = G_zu @ U
    z_star = np.fft.irfft(Z_star, axis=0)
    z_star_min, z_star_max = z_star.min(axis=(0, 2)), z_star.max(axis=(0, 2))
    Tz_inv = jnp.diag(2 / (z_star_max - z_star_min))

    G_yu = bla._frequency_response(freqs)
    f_data = (freqs, 1 / data.time.ts, U, Y, G_yu)

    args = ArgsInferenceLearning(
        theta_uy=theta_uy,
        phi=phi,
        lambda_w=lambda_w,
        fixed_point_iters=fixed_point_iters,
        freq_data=f_data,
        Lambda=_compute_weighting_matrix(data.freq, freq_weighting),
        Tz_inv=Tz_inv,
        epsilon=epsilon,
        N=N,
        recompute_fixed_point=recompute_fixed_point,
    )
    return theta_wz, args


def _loss_inference_and_learning(
    theta: DecVarsInferenceLearning,
    args: ArgsInferenceLearning
) -> tuple:
    """Implement the inference and learning method and compute the loss."""
    freqs, fs, U, Y, G_yu = args.freq_data

    A = args.theta_uy[0]
    B_u = args.theta_uy[1]
    C_y = args.theta_uy[2].astype(complex)

    B_w = theta.B_w_star.astype(complex)
    C_z = (args.Tz_inv @ theta.C_z_star).astype(complex)
    D_yw = theta.D_yw_star
    D_zu = args.Tz_inv @ theta.D_zu_star

    nw = D_yw.shape[1]
    nz, nx = C_z.shape
    F_tot, nu, R = U.shape

    Theta = jnp.vstack((B_w, D_yw)).T @ jnp.vstack((B_w, D_yw))
    Theta += args.epsilon / args.lambda_w * jnp.eye(nw)

    z = 2 * jnp.pi * freqs / fs
    zj = jnp.exp(z * 1j)

    I_nx = jnp.eye(nx)

    def _compute_parametric_Gs(k):
        G_x = jnp.linalg.solve(zj[k] * I_nx - A, jnp.hstack((B_u, B_w)))
        return (
            C_y @ G_x[:, nu:] + D_yw,  # G_yw
            C_z @ G_x[:, :nu] + D_zu,  # G_zu
            C_z @ G_x[:, nu:],  # G_zw
        )

    G_yw, G_zu, G_zw = jax.vmap(_compute_parametric_Gs)(jnp.arange(F_tot))

    # 1) Nonparametric inference
    def _infer_nonparametric_signals(k):
        Psi = G_yw[k, ...].T @ args.Lambda[k, ...]
        Phi = Psi @ G_yw[k, ...] + args.lambda_w * Theta
        W_hat = jnp.linalg.solve(Phi, Psi @ (Y[k, ...] - G_yu[k, ...] @ U[k, ...]))
        Z_hat = G_zu[k, ...] @ U[k, ...] + G_zw[k, ...] @ W_hat
        return W_hat, Z_hat

    W_star, Z_star = jax.vmap(_infer_nonparametric_signals)(jnp.arange(F_tot))

    # 2) Parametric learning
    w_star = jnp.fft.irfft(W_star, n=args.N, axis=0)
    z_star = jnp.fft.irfft(Z_star, n=args.N, axis=0)

    w_star_stacked = jnp.transpose(w_star, (2, 0, 1)).reshape(args.N * R, nw)
    z_star_stacked = jnp.transpose(z_star, (2, 0, 1)).reshape(args.N * R, nz)

    # 2a) Compute beta
    phi_z_star = args.phi._compute_features(z_star_stacked)
    beta_hat = jnp.linalg.solve(
        phi_z_star.T @ phi_z_star, phi_z_star.T @ w_star_stacked
    )

    # 2b) Perform fixed-point iterations
    def _fixed_point_iteration(_, phi_z):
        w_stacked = phi_z @ beta_hat
        w = jnp.transpose(w_stacked.reshape(R, args.N, nw), (1, 2, 0))
        W = jnp.fft.rfft(w, axis=0)
        Z = G_zu @ U + G_zw @ W
        z = jnp.fft.irfft(Z, n=args.N, axis=0)
        z_stacked = jnp.transpose(z, (2, 0, 1)).reshape(args.N * R, nz)
        return args.phi._compute_features(z_stacked)
    
    fixed_point_fn = (
        jax.checkpoint(_fixed_point_iteration)
        if args.recompute_fixed_point
        else _fixed_point_iteration
    )

    phi_z = jax.lax.fori_loop(
        0,
        args.fixed_point_iters,
        fixed_point_fn,
        phi_z_star,
        unroll=True
    )

    w_hat_stacked = phi_z @ beta_hat
    w_hat = jnp.transpose(w_hat_stacked.reshape(R, args.N, nw), (1, 2, 0))
    W_beta = jnp.fft.rfft(w_hat, axis=0)

    # 2c) Compute loss
    Y_hat = G_yu @ U + G_yw @ W_beta
    loss = _loss_output_spectrum(Y, Y_hat, args.Lambda)
    return _misc.real_valued(loss), (_misc.scalar_valued(loss), beta_hat)


def _prepare_nonlin_optimization(
    data: InputOutputData,
    model: ModelNonlinearLFR,
    offset: int,
    freq_weighting: bool = True
) -> tuple[ModelNonlinearLFR, ArgsNonlinearOptimization, ModelNonlinearLFR]:
    """Prepare initial guess and function arguments for nonlinear optimization."""    
    # Separate static and dynamic parameters
    theta0, theta_static = eqx.partition(model, eqx.is_inexact_array)
    
    # Ensure BLA is not part of the decision variables
    theta0 = eqx.tree_at(lambda tree: tree._bla, theta0, replace=None) 
    
    x_bla = _misc.compute_steady_state_bla_state(model._bla, data)

    args = ArgsNonlinearOptimization(
        theta_static=theta_static,
        u=_misc.extend_signal(data.time.u, offset) if offset > 0 else data.time.u, 
        Y=data.freq.Y,
        Lambda=_compute_weighting_matrix(data.freq, freq_weighting),
        x0=x_bla[-offset, :, :],
        offset=offset,
    )
    return theta0, args


def _loss_nonlin_optimization(
    theta: ModelNonlinearLFR, 
    args: ArgsNonlinearOptimization
) -> tuple:
    """Perform time-domain forward simulations and compute the loss."""
    theta = eqx.combine(theta, args.theta_static)

    # Simulate the model in time domain
    y_hat = theta._simulate(args.u, args.x0)[0]
    y_hat = y_hat[args.offset:, ...]  # discard offset samples
    Y_hat = jnp.fft.rfft(y_hat, axis=0)

    # Compute loss
    loss = _loss_output_spectrum(args.Y, Y_hat, args.Lambda)
    return _misc.real_valued(loss), (_misc.scalar_valued(loss),)

    
def _validate_weighting(
    freq_weighting: bool,
    var_noise: jnp.ndarray | None,
    print_warning: bool
) -> bool:
    """Check if weighting can be applied based on output noise variance availability."""
    if freq_weighting and var_noise is None:
        if print_warning:
            print(
                "Warning: Frequency weighting based on output noise variance "
                "requested, but such estimate is not available. Proceeding without "
                "weighting."
        )
        freq_weighting = False
    return freq_weighting


def _compute_weighting_matrix(
    freq_data: FrequencyData,
    freq_weighting: bool
) -> jnp.ndarray:
    """Compute weighting matrix containing the inverse output noise variances."""
    F_tot, ny = freq_data.Y.shape[:2]
    
    # Each (ny x ny) matrix is diagonal with elements equal to the inverse of the noise
    # variance for each output at each frequency. If this estimate is not available, no
    # weighting is applied (identity matrices are used).
    Lambda = np.zeros((F_tot, ny, ny))
    for k in range(F_tot):
        diag = np.diag(1 / freq_data.Y_var_noise[k]) if freq_weighting else np.eye(ny)
        Lambda[k, ...] = diag

    return jnp.asarray(Lambda)


def _loss_bla(bla: ModelBLA, freq_data: FrequencyData, freq_weighting: bool) -> float:
    """Compute the BLA performance on the loss function that is used in this module."""
    Lambda = _compute_weighting_matrix(freq_data, freq_weighting)
    Y = freq_data.Y
    U = freq_data.U

    Y_hat = bla._frequency_response(freq_data.f) @ U
    loss = _loss_output_spectrum(Y, Y_hat, Lambda)
    return _misc.scalar_valued(loss)


def _create_basis_function_model_given_beta(
    nw: int,
    phi: AbstractFeatureMap,
    beta: jnp.ndarray
) -> BasisFunctionModel:
    """Create a `BasisFunctionModel` instance given `beta`.

    The coefficients `beta` of a `BasisFunctionModel` are always initialized randomly 
    with a given seed. This function bypasses that behavior by immediately replacing 
    the randomly generated `beta` with the provided matrix obtained from inference and 
    learning. A dummy seed is still required for initialization, but it does not 
    affect the final result.
    """
    dummy_seed = 0
    return eqx.tree_at(
        where=lambda tree: tree.beta,
        pytree=BasisFunctionModel(nw, phi, dummy_seed),
        replace=beta,
    )



