import logging

from collections.abc import Callable
from typing import Literal, cast

import numpy as np
import pymc as pm

from better_optimize import basinhopping, minimize
from better_optimize.constants import minimize_method
from pymc.blocking import DictToArrayBijection, RaveledVars
from pymc.initial_point import make_initial_point_fn
from pymc.model.transform.optimization import freeze_dims_and_data
from pymc.util import get_default_varnames
from pytensor.tensor import TensorVariable
from scipy.optimize import OptimizeResult

from pymc_extras.inference.laplace_approx.idata import (
    add_data_to_inference_data,
    add_fit_to_inference_data,
    add_optimizer_result_to_inference_data,
    map_results_to_inference_data,
)
from pymc_extras.inference.laplace_approx.scipy_interface import (
    GradientBackend,
    scipy_optimize_funcs_from_loss,
    set_optimizer_function_defaults,
)

_log = logging.getLogger(__name__)


def get_nearest_psd(A: np.ndarray) -> np.ndarray:
    """
    Compute the nearest positive semi-definite matrix to a given matrix.

    This function takes a square matrix and returns the nearest positive semi-definite matrix using
    eigenvalue decomposition. It ensures all eigenvalues are non-negative. The "nearest" matrix is defined in terms
    of the Frobenius norm.

    Parameters
    ----------
    A : np.ndarray
        Input square matrix.

    Returns
    -------
    np.ndarray
        The nearest positive semi-definite matrix to the input matrix.
    """
    C = (A + A.T) / 2
    eigval, eigvec = np.linalg.eigh(C)
    eigval[eigval < 0] = 0

    return eigvec @ np.diag(eigval) @ eigvec.T


def _make_initial_point(model, initvals=None, random_seed=None, jitter_rvs=None):
    jitter_rvs = [] if jitter_rvs is None else jitter_rvs

    ipfn = make_initial_point_fn(
        model=model,
        jitter_rvs=set(jitter_rvs),
        return_transformed=True,
        overrides=initvals,
    )

    start_dict = ipfn(random_seed)
    vars_dict = {var.name: var for var in model.continuous_value_vars}
    initial_params = DictToArrayBijection.map(
        {var_name: value for var_name, value in start_dict.items() if var_name in vars_dict}
    )

    return initial_params


def _compute_inverse_hessian(
    optimizer_result: OptimizeResult | None,
    optimal_point: np.ndarray | None,
    f_fused: Callable | None,
    f_hessp: Callable | None,
    use_hess: bool,
    method: minimize_method | Literal["BFGS", "L-BFGS-B"],
):
    """
    Compute the Hessian matrix or its inverse based on the optimization result and the method used.

    Downstream functions (e.g. laplace approximation) will need the inverse Hessian matrix. This function computes it
    in the cheapest way possible, depending on the optimization method used and the available compiled functions.

    Parameters
    ----------
    optimizer_result: OptimizeResult, optional
        The result of the optimization, containing the optimized parameters and possibly an approximate inverse Hessian.
    optimal_point: np.ndarray, optional
        The optimal point found by the optimizer, used to compute the Hessian if necessary. If not provided, it will be
        extracted from the optimizer result.
    f_fused: callable, optional
        The compiled function representing the loss and possibly its gradient and Hessian.
    f_hessp: callable, optional
        The compiled function for Hessian-vector products, if available.
    use_hess: bool
        Whether the Hessian matrix was used in the optimization.
    method: minimize_method
        The optimization method used, which determines how the Hessian is computed.

    Returns
    -------
    H_inv: np.ndarray
        The inverse Hessian matrix, computed based on the optimization method and available functions.
    """
    if optimal_point is None and optimizer_result is None:
        raise ValueError("At least one of `optimal_point` or `optimizer_result` must be provided.")

    x_star = optimizer_result.x if optimizer_result is not None else optimal_point
    n_vars = len(x_star)

    if method == "BFGS" and optimizer_result is not None:
        # If we used BFGS, the optimizer result will contain the inverse Hessian -- we can just use that rather than
        # re-computing something
        if hasattr(optimizer_result, "lowest_optimization_result"):
            # We did basinhopping, need to get the inner optimizer results
            H_inv = getattr(optimizer_result.lowest_optimization_result, "hess_inv", None)
        else:
            H_inv = getattr(optimizer_result, "hess_inv", None)

    elif method == "L-BFGS-B" and optimizer_result is not None:
        # Here we will have a LinearOperator representing the inverse Hessian-Vector product.
        if hasattr(optimizer_result, "lowest_optimization_result"):
            # We did basinhopping, need to get the inner optimizer results
            f_hessp_inv = getattr(optimizer_result.lowest_optimization_result, "hess_inv", None)
        else:
            f_hessp_inv = getattr(optimizer_result, "hess_inv", None)

        if f_hessp_inv is not None:
            basis = np.eye(n_vars)
            H_inv = np.stack([f_hessp_inv(basis[:, i]) for i in range(n_vars)], axis=-1)
        else:
            H_inv = None

    elif f_hessp is not None:
        # In the case that hessp was used, the results object will not save the inverse Hessian, so we can compute it from
        # the hessp function, using euclidian basis vector.
        basis = np.eye(n_vars)
        H = np.stack([f_hessp(x_star, basis[:, i]) for i in range(n_vars)], axis=-1)
        H_inv = np.linalg.inv(get_nearest_psd(H))

    elif use_hess and f_fused is not None:
        # If we compiled a hessian function, just use it
        _, _, H = f_fused(x_star)
        H_inv = np.linalg.inv(get_nearest_psd(H))

    else:
        H_inv = None

    return H_inv


def find_MAP(
    method: minimize_method | Literal["basinhopping"] = "L-BFGS-B",
    *,
    model: pm.Model | None = None,
    use_grad: bool | None = None,
    use_hessp: bool | None = None,
    use_hess: bool | None = None,
    initvals: dict | None = None,
    random_seed: int | np.random.Generator | None = None,
    jitter_rvs: list[TensorVariable] | None = None,
    progressbar: bool = True,
    include_transformed: bool = True,
    freeze_model: bool = True,
    gradient_backend: GradientBackend = "pytensor",
    compile_kwargs: dict | None = None,
    compute_hessian: bool = False,
    **optimizer_kwargs,
) -> (
    dict[str, np.ndarray]
    | tuple[dict[str, np.ndarray], np.ndarray]
    | tuple[dict[str, np.ndarray], OptimizeResult]
    | tuple[dict[str, np.ndarray], OptimizeResult, np.ndarray]
):
    """
    Fit a PyMC model via maximum a posteriori (MAP) estimation using JAX and scipy.optimize.

    Parameters
    ----------
    model : pm.Model
        The PyMC model to be fit. If None, the current model context is used.
    method : str
        The optimization method to use. Valid choices are: Nelder-Mead, Powell, CG, BFGS, L-BFGS-B, TNC, SLSQP,
        trust-constr, dogleg, trust-ncg, trust-exact, trust-krylov, and basinhopping.

        See scipy.optimize.minimize documentation for details.
    use_grad : bool | None, optional
        Whether to use gradients in the optimization. Defaults to None, which determines this automatically based on
        the ``method``.
    use_hessp : bool | None, optional
        Whether to use Hessian-vector products in the optimization. Defaults to None, which determines this automatically based on
        the ``method``.
    use_hess : bool | None, optional
        Whether to use the Hessian matrix in the optimization. Defaults to None, which determines this automatically based on
        the ``method``.
    initvals : None | dict, optional
        Initial values for the model parameters, as str:ndarray key-value pairs. Partial initialization is permitted.
         If None, the model's default initial values are used.
    random_seed : None | int | np.random.Generator, optional
        Seed for the random number generator or a numpy Generator for reproducibility
    jitter_rvs : list of TensorVariables, optional
        Variables whose initial values should be jittered. If None, all variables are jittered.
    progressbar : bool, optional
        Whether to display a progress bar during optimization. Defaults to True.
    include_transformed: bool, optional
        Whether to include transformed variable values in the returned dictionary. Defaults to True.
    freeze_model: bool, optional
        If True, freeze_dims_and_data will be called on the model before compiling the loss functions. This is
        sometimes necessary for JAX, and can sometimes improve performance by allowing constant folding. Defaults to
        True.
    gradient_backend: str, default "pytensor"
        Which backend to use to compute gradients. Must be one of "pytensor" or "jax".
    compute_hessian: bool
        If True, the inverse Hessian matrix at the optimum will be computed and included in the returned
        InferenceData object. This is needed for the Laplace approximation, but can be computationally expensive for
        high-dimensional problems. Defaults to False.
    compile_kwargs: dict, optional
        Additional options to pass to the ``pytensor.function`` function when compiling loss functions.
    **optimizer_kwargs
        Additional keyword arguments to pass to the ``scipy.optimize`` function being used. Unless
        ``method = "basinhopping"``, ``scipy.optimize.minimize`` will be used. For ``basinhopping``,
        ``scipy.optimize.basinhopping`` will be used. See the documentation of these functions for details.

    Returns
    -------
    map_result: az.InferenceData
        Results of Maximum A Posteriori (MAP) estimation, including the optimized point, inverse Hessian, transformed
        latent variables, and optimizer results.
    """
    compile_kwargs = {} if compile_kwargs is None else compile_kwargs
    model = pm.modelcontext(model) if model is None else model

    if freeze_model:
        model = freeze_dims_and_data(model)

    initial_params = _make_initial_point(model, initvals, random_seed, jitter_rvs)

    do_basinhopping = method == "basinhopping"
    minimizer_kwargs = optimizer_kwargs.pop("minimizer_kwargs", {})

    if do_basinhopping:
        # For a nice API, we let the user set method="basinhopping", but if we're doing basinhopping we still need
        # another method for the inner optimizer. This will be set in the minimizer_kwargs, but also needs a default
        # if one isn't provided.

        method = minimizer_kwargs.pop("method", "L-BFGS-B")
        minimizer_kwargs["method"] = method

    use_grad, use_hess, use_hessp = set_optimizer_function_defaults(
        method, use_grad, use_hess, use_hessp
    )

    f_fused, f_hessp = scipy_optimize_funcs_from_loss(
        loss=-model.logp(jacobian=False),
        inputs=model.continuous_value_vars + model.discrete_value_vars,
        initial_point_dict=DictToArrayBijection.rmap(initial_params),
        use_grad=use_grad,
        use_hess=use_hess,
        use_hessp=use_hessp,
        gradient_backend=gradient_backend,
        compile_kwargs=compile_kwargs,
    )

    args = optimizer_kwargs.pop("args", ())

    # better_optimize.minimize will check if f_logp is a fused loss+grad Op, and automatically assign the jac argument
    # if so. That is why the jac argument is not passed here in either branch.

    if do_basinhopping:
        if "args" not in minimizer_kwargs:
            minimizer_kwargs["args"] = args
        if "hessp" not in minimizer_kwargs:
            minimizer_kwargs["hessp"] = f_hessp
        if "method" not in minimizer_kwargs:
            minimizer_kwargs["method"] = method

        optimizer_result = basinhopping(
            func=f_fused,
            x0=cast(np.ndarray[float], initial_params.data),
            progressbar=progressbar,
            minimizer_kwargs=minimizer_kwargs,
            **optimizer_kwargs,
        )

    else:
        optimizer_result = minimize(
            f=f_fused,
            x0=cast(np.ndarray[float], initial_params.data),
            args=args,
            hessp=f_hessp,
            progressbar=progressbar,
            method=method,
            **optimizer_kwargs,
        )

    if compute_hessian:
        H_inv = _compute_inverse_hessian(
            optimizer_result=optimizer_result,
            optimal_point=None,
            f_fused=f_fused,
            f_hessp=f_hessp,
            use_hess=use_hess,
            method=method,
        )
    else:
        H_inv = None

    raveled_optimized = RaveledVars(optimizer_result.x, initial_params.point_map_info)
    unobserved_vars = get_default_varnames(model.unobserved_value_vars, include_transformed=True)
    unobserved_vars_values = model.compile_fn(unobserved_vars, mode="FAST_COMPILE")(
        DictToArrayBijection.rmap(raveled_optimized)
    )

    optimized_point = {
        var.name: value for var, value in zip(unobserved_vars, unobserved_vars_values)
    }

    idata = map_results_to_inference_data(
        map_point=optimized_point, model=model, include_transformed=include_transformed
    )

    idata = add_fit_to_inference_data(idata=idata, mu=raveled_optimized, H_inv=H_inv, model=model)

    idata = add_optimizer_result_to_inference_data(
        idata=idata, result=optimizer_result, method=method, mu=raveled_optimized, model=model
    )

    idata = add_data_to_inference_data(
        idata=idata, progressbar=False, model=model, compile_kwargs=compile_kwargs
    )

    return idata
