#  Copyright (c) 2024 zfit

from __future__ import annotations

import warnings
from typing import TYPE_CHECKING

import jacobi

if TYPE_CHECKING:
    import zfit

from collections.abc import Callable

import logging
from functools import lru_cache, wraps

import numpy as np
import scipy.stats
import tensorflow as tf
from scipy import optimize

import zfit.z.numpy as znp

from .. import z
from ..core.interfaces import ZfitIndependentParameter
from ..core.parameter import assign_values, assign_values_jit
from ..util.container import convert_to_container
from ..util.deprecation import deprecated_args


class NewMinimum(Exception):
    """Exception class for cases where a new minimum is found."""

    pass


class FailEvalLossNaN(Exception):
    pass


class RootFound(Exception):
    """Exception class for cases where a root is found, since SciPy root solvers don't really respect tol or xtol on
    initial evaluation."""

    pass


@deprecated_args(None, "Use cl for confidence level instead.", "sigma")
def compute_errors(
    result: zfit.result.FitResult,
    params: list[ZfitIndependentParameter],
    cl: float | None = None,
    rtol: float | None = None,
    method: str | None = None,
    covariance_method: str | Callable | None = None,
    sigma: float = None,
) -> tuple[
    dict[ZfitIndependentParameter, dict[str, float]],
    zfit.result.FitResult | None,
]:
    """Compute asymmetric errors of parameters by profiling the loss function in the fit result.

    This method finds the value for a given parameter where the loss function is ``cl`` away: for example
    for a cl of 68.3%, this is one (multiplied by the errordef). The other parameters are also minimized and
    not fixed. This method is comparably computationally intensive and, if possible, ``hesse`` should be used.
    However, since ``hesse`` does not capture asymmetric or non-parabolic shaped profiles well, this method is
    preferable.

    Args:
        result: fit result to be used to compute the uncertainties.
        params: The parameters to calculate the
            error. If None, use all parameters.
        cl: Confidence Level of the parameter to be determined. Defaults to 68.3%.
        rtol: relative tol between the computed and the exact roots
        method: type of solver, ``method`` argument of :py:func:`scipy.optimize.root`. Defaults to "hybr".
        covariance_method: The method to use to calculate the correlation matrix, will be forwarded directly
            to :py:meth:`FitResult.covariance`. Valid choices are
            by default {'minuit_hesse', 'hesse_np'} (or any other method defined in the result)
            or a Callable.

    Returns:
        out:
            A ``dict`` containing as keys the parameter and as value a ``dict`` which
            contains two keys 'lower' and 'upper', holding the calculated errors.
            Example: result[par1]['upper'] -> the asymmetric upper error of 'par1'
        out: a fit result is returned when a new minimum is found during the loss scan
    """
    if rtol is None:
        rtol = 0.003
    method = "hybr" if method is None else method
    if cl is None:
        if sigma is None:
            sigma = 1

    else:
        if sigma is None:
            sigma = scipy.stats.chi2(1).ppf(cl) ** 0.5
        else:
            raise ValueError("Cannot specify both sigma and cl.")
    # TODO: how to name things, sigma or cl?

    params = convert_to_container(params)
    new_result = None

    all_params = list(result.params.keys())
    loss = result.loss
    errordef = loss.errordef
    fmin = result.fmin
    rtol *= errordef
    minimizer = result.minimizer

    covariance = result.covariance(method=covariance_method, as_dict=True)
    if covariance is None:
        covariance = result.covariance(method="hesse_np", as_dict=True)
    if covariance is None:
        raise RuntimeError(
            "Could not compute covariance matrix. Check if the minimum is valid."
        )
    param_errors = {param: covariance[(param, param)] ** 0.5 for param in all_params}
    param_scale = np.array(list(param_errors.values()))

    ncalls = 0
    loss_min_tol = minimizer.tol * errordef * 2  # 2 is just to be tolerant
    try:
        to_return = {}
        for param in params:
            assign_values(all_params, result)

            logging.info(f"profiling the parameter {param}")
            param_error = param_errors[param]
            param_value = result.params[param]["value"]

            initial_values = {"lower": [], "upper": []}
            direction = {"lower": -1, "upper": 1}

            for ap in all_params:
                ap_value = result.params[ap]["value"]
                error_factor = (
                    covariance[(param, ap)] * (2 * errordef / param_error**2) ** 0.5
                )
                for d in ["lower", "upper"]:
                    step = direction[d] * error_factor * sigma
                    for ntrial in range(50):
                        ap_value_init = ap_value + step
                        if ap_value_init < ap.lower or ap_value_init > ap.upper:
                            step *= 0.8
                        else:
                            break
                    else:
                        raise RuntimeError(
                            f"Could not find a valid initial value for {ap} in {d} direction after {ntrial + 1} trials."
                            f" step tried: {step}. This should not happes, the error probably looks weird. Maybe plot"
                            f" the loss function for different parameter values and check if it looks reasonable."
                        )

                    initial_values[d].append(ap_value_init)

            index_poi = all_params.index(param)  # remember the index
            _ = loss.value_gradient(
                params=all_params, full=False
            )  # to make sure the loss is compiled

            @z.function(wraps="gradient")
            def optimized_loss_gradient(values, index):
                assert isinstance(index, int)
                assign_values(all_params, values)
                loss_value, gradient = loss.value_gradient(
                    params=all_params, full=False
                )
                if isinstance(gradient, (tuple, list)):
                    gradient = znp.asarray(gradient)
                gradient = znp.concatenate(
                    [gradient[:index_poi], gradient[index_poi + 1 :]]
                )
                return loss_value, gradient

            # TODO: improvement, use jacobian?
            root = None
            ntol = 999  # if it's right in the beginning, we think it's fine

            # TODO: should we add a "robust" or similar option to not skip this?
            # or evaluate and then decide ,maybe use krylov as it doesn't do a lot of calls in the beginning, it
            # approximates the jacobian

            def func(values, args):
                nonlocal ncalls, root, ntol

                ncalls += 1
                swap_sign = args

                try:
                    loss_value, gradient = optimized_loss_gradient(values, index_poi)
                except tf.errors.InvalidArgumentError:
                    msg = (
                        f"The evaluation of the errors of {param.name} failed due to too many NaNs"
                        " being produced in the loss and/or its gradient. This is most probably"
                        " caused by negative values returned from the PDF."
                    )
                    loss_value = znp.array(999999.0)
                    gradient = z.random.normal(stddev=0.1, shape=(len(all_params) - 1,))
                    # raise FailEvalLossNaN(msg)
                zeroed_loss = loss_value.numpy() - fmin

                gradient = znp.array(gradient)

                if swap_sign(param):  # mirror at x-axis to remove second zero
                    zeroed_loss = -zeroed_loss
                    gradient = -gradient
                    logging.info("Swapping sign in error calculation 'zfit_error'")

                elif zeroed_loss < -loss_min_tol:
                    assign_values(all_params, values)  # set values to the new minimum
                    raise NewMinimum("A new minimum is found.")

                downward_shift = errordef * sigma**2
                shifted_loss = zeroed_loss - downward_shift

                if abs(shifted_loss) < rtol:
                    if ntol > 3:
                        root = values[index_poi]
                        raise RootFound()
                    ntol += 1
                else:
                    ntol = 0

                return znp.concatenate([[shifted_loss], gradient])

            to_return[param] = {}
            swap_sign = {
                "lower": lambda p: p > param_value,
                "upper": lambda p: p < param_value,
            }
            for d in ["lower", "upper"]:
                try:
                    root_result = optimize.root(
                        fun=func,
                        args=(swap_sign[d],),
                        x0=np.array(initial_values[d]),
                        tol=rtol * 0.1,  # we won't stop like this anyway
                        options={
                            "diag": 1 / param_scale,  # scale factor for variables
                        },
                        method=method,
                    )
                except RootFound:
                    assert root is not None, "Should be changed inside function."
                else:
                    warnings.warn(
                        f"The root finding did not converge below {rtol} but stopped by its own criteria."
                    )
                    root = root_result.x[index_poi]
                to_return[param][d] = root - param_value

        assign_values(all_params, result)

    except NewMinimum as e:
        from .. import settings

        if settings.get_verbosity() >= 5:
            print(e)
        minimizer = result.minimizer
        loss = result.loss
        new_found_fmin = loss.value(full=False)
        new_result = minimizer.minimize(loss=loss)
        if new_result.fmin >= new_found_fmin + loss_min_tol:
            raise RuntimeError(
                "A new minimum was discovered but the minimizer was not able to find this on himself. "
                "This behavior is currently an exception but will most likely change in the future."
            )
        to_return, new_result_ = compute_errors(
            result=new_result, params=params, cl=cl, rtol=rtol, method=method
        )
        if new_result_ is not None:
            new_result = new_result_
    return to_return, new_result


def numerical_pdf_jacobian(func, params):  # TODO: jit?
    params = list(params.values())
    return znp.asarray(jacobi.jacobi(func, params)[0].T)


# @z.function(wraps="autodiff")
def autodiff_pdf_jacobian(func, params):
    params = list(params.values())
    # TODO(WrappedVariable): this is needed if we want to use wrapped Variables
    # params = z.math._extract_tfparams(params)

    # the below fails for some cases (i.e. CB) with an internal error
    # ValueError: Internal error: Tried to take gradients (or similar) of a variable without handle data:
    # Tensor("Placeholder_1:0", shape=(), dtype=resource)

    # we didn't report that yet, it's too hard to reproduce with a minimal example currently.

    # with tf.GradientTape(watch_accessed_variables=False) as t2:
    #     t2.watch(params)
    #     with tf.GradientTape(watch_accessed_variables=False) as t1:
    #         t1.watch(params)
    #         values = func()
    #
    #     grad = t1.gradient(values, params)
    #     grad = tf.convert_to_tensor(grad)
    # jacobian = t2.jacobian(grad, params)

    columns = []

    for p in params:
        vector = np.zeros(len(params))
        vector[params.index(p)] = 1.0
        with tf.autodiff.ForwardAccumulator(params, list(vector)) as acc:
            values = func()
        columns.append(acc.jvp(values))

    jacobian = z.convert_to_tensor(columns)

    return jacobian


def covariance_with_weights(method, result, params):
    """Compute the covariance matrix of the parameters with weights."""
    from .. import run

    run.assert_executing_eagerly()
    loss = result.loss
    model = loss.model
    data = loss.data
    yields = None
    if loss.is_extended:
        yields = []
        for m in model:
            yields.append(m.get_yield())
    constraints = None
    if loss.constraints:
        constraints = loss.constraints

    old_vals = run(params)

    Hinv_dict = method(result=result, params=params)  # inverse of the hessian matrix
    if not Hinv_dict:
        return {}
    Hinv = dict_to_matrix(params, Hinv_dict)

    def func():
        values = []
        for i, (m, d) in enumerate(zip(model, data)):
            v = m.log_pdf(d)
            weights = d.weights
            if weights is not None:
                v *= weights
            values.append(v)
            if yields is not None:
                yi = yields[i]
                nevents_collected = (
                    tf.reduce_sum(weights) if weights is not None else d.nevents
                )
                term_new = tf.nn.log_poisson_loss(
                    nevents_collected, znp.log(yi), compute_full_loss=True
                )[
                    ..., None
                ]  # make it an array like the others
                values.append(term_new)
        if constraints is not None:
            for constraint in constraints:
                values.append(znp.reshape(constraint.value(), (-1,)))

        return znp.concatenate(values, axis=0)

    params_dict = {p.name: p for p in params}

    if run.get_autograd_mode():
        try:
            jacobian = autodiff_pdf_jacobian(func=func, params=params_dict)
        except ValueError as error:
            raise ValueError(
                "The autodiff could not be performed for the jacobian matrix. This can have a natural cause (see error above)"
                " or be a bug in the autodiff. In the latter case, you can try to use the numerical jacobian instead using:\n"
                "with zfit.run.set_autograd_mode(False):\n"
                "    # calculate here, i.e. result.errors()"
            ) from error
    else:

        def wrapped_func(values):
            assign_values(list(params_dict.values()), values)
            return np.array(func())

        jacobian = numerical_pdf_jacobian(func=wrapped_func, params=params_dict)

    C = znp.matmul(jacobian, znp.transpose(jacobian))
    covariance = np.asarray(znp.matmul(Hinv, znp.matmul(C, Hinv)))
    assign_values(params, old_vals)
    return matrix_to_dict(params, covariance)


def dict_to_matrix(params, matrix_dict):
    nparams = len(params)
    matrix = np.empty((nparams, nparams))
    if not matrix_dict:
        return None

    for i in range(nparams):
        pi = params[i]
        for j in range(i, nparams):
            pj = params[j]
            matrix[i, j] = matrix_dict[(pi, pj)]
            if i != j:
                matrix[j, i] = matrix_dict[(pi, pj)]

    return matrix


def matrix_to_dict(params, matrix):
    nparams = len(params)
    matrix_dict = {}
    if matrix is None:
        return matrix_dict

    for i in range(nparams):
        pi = params[i]
        for j in range(i, nparams):
            pj = params[j]
            matrix_dict[(pi, pj)] = matrix[i, j]
            if i != j:
                matrix_dict[(pj, pi)] = matrix[i, j]

    return matrix_dict


def np_cache(*args, **kwargs):
    """LRU cache implementation for functions whose FIRST parameter is a numpy array.

    >>> array = np.array([[1, 2, 3], [4, 5, 6]])
    >>> @np_cache(maxsize=256)
    ... def multiply(array, factor):
    ...     print("Calculating...")
    ...     return factor*array
    >>> multiply(array, 2)
    Calculating...
    array([[ 2,  4,  6],
           [ 8, 10, 12]])
    >>> multiply(array, 2)
    array([[ 2,  4,  6],
           [ 8, 10, 12]])
    >>> multiply.cache_info()
    CacheInfo(hits=1, misses=1, maxsize=256, currsize=1)
    """

    def decorator(function):
        @wraps(function)
        def wrapper(np_array, *args, **kwargs):
            hashable_array = array_to_tuple(np_array)
            return cached_wrapper(hashable_array, *args, **kwargs)

        @lru_cache(*args, **kwargs)
        def cached_wrapper(hashable_array, *args, **kwargs):
            array = np.array(hashable_array)
            return function(array, *args, **kwargs)

        def array_to_tuple(np_array):
            """Iterates recursively."""
            try:
                return tuple(array_to_tuple(_) for _ in np_array)
            except TypeError:
                return np_array

        # copy lru_cache attributes over too
        wrapper.cache_info = cached_wrapper.cache_info
        wrapper.cache_clear = cached_wrapper.cache_clear

        return wrapper

    return decorator
