# Copyright 2020- The Blackjax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Step size adaptation"""
from typing import Callable, NamedTuple

import jax
import jax.numpy as jnp

from blackjax.mcmc.hmc import HMCState
from blackjax.optimizers.dual_averaging import dual_averaging
from blackjax.types import PRNGKey

__all__ = [
    "DualAveragingAdaptationState",
    "dual_averaging_adaptation",
    "find_reasonable_step_size",
]


# -------------------------------------------------------------------
#                        DUAL AVERAGING
# -------------------------------------------------------------------


class DualAveragingAdaptationState(NamedTuple):
    """State carried through the dual averaging procedure.

    log_step_size
        The logarithm of the current value of the step size.
    log_step_size_avg
        The time-weighted average of the values that the logarithm of the step
        size has taken so far.
    step
        The current iteration step.
    avg_err
        The time average of the value of the quantity :math:`H_t`, the
        difference between the target acceptance rate and the current
        acceptance rate.
    mu
        Arbitrary point the values of log_step_size are shrunk towards. Chose
        to be :math:`\\log(10 \\epsilon_0)` where :math:`\\epsilon_0` is chosen
        in this context to be the step size given by the
        `find_reasonable_step_size` procedure.

    """

    log_step_size: float
    log_step_size_avg: float
    step: int
    avg_error: float
    mu: float


def dual_averaging_adaptation(
    target: float, t0: int = 10, gamma: float = 0.05, kappa: float = 0.75
) -> tuple[Callable, Callable, Callable]:
    """Tune the step size in order to achieve a desired target acceptance rate.

    Let us note :math:`\\epsilon` the current step size, :math:`\\alpha_t` the
    metropolis acceptance rate at time :math:`t` and :math:`\\delta` the desired
    aceptance rate. We define:

    .. math:
        H_t = \\delta - \\alpha_t

    the error at time t. We would like to find a procedure that adapts the
    value of :math:`\\epsilon` such that :math:`h(x) =\\mathbb{E}\\left[H_t|\\epsilon\\right] = 0`

    Following :cite:p:`nesterov2009primal`, the authors of :cite:p:`hoffman2014no` proposed the following update scheme. If
    we note :math:`x = \\log \\epsilon` we follow:

    .. math:
        x_{t+1} \\LongLeftArrow \\mu - \\frac{\\sqrt{t}}{\\gamma} \\frac{1}{t+t_0} \\sum_{i=1}^t H_i
        \\overline{x}_{t+1} \\LongLeftArrow x_{t+1}\\, t^{-\\kappa}  + \\left(1-t^\\kappa\\right)\\overline{x}_t

    :math:`\\overline{x}_{t}` is guaranteed to converge to a value such that
    :math:`h(\\overline{x}_t)` converges to 0, i.e. the Metropolis acceptance
    rate converges to the desired rate.

    See reference :cite:p:`hoffman2014no` (section 3.2.1) for a detailed discussion.

    Parameters
    ----------
    t0: float >= 0
        Free parameter that stabilizes the initial iterations of the algorithm.
        Large values may slow down convergence. Introduced in :cite:p:`hoffman2014no` with a default
        value of 10.
    gamma:
        Controls the speed of convergence of the scheme. The authors of :cite:p:`hoffman2014no` recommend
        a value of 0.05.
    kappa: float in [0.5, 1]
        Controls the weights of past steps in the current update. The scheme will
        quickly forget earlier step for a small value of `kappa`. Introduced
        in :cite:p:`hoffman2014no`, with a recommended value of .75
    target:
        Target acceptance rate.

    Returns
    -------
    init
        A function that initializes the state of the dual averaging scheme.
    update
        A function that updates the state of the dual averaging scheme.

    """
    da_init, da_update, da_final = dual_averaging(t0, gamma, kappa)

    def init(inital_step_size: float) -> DualAveragingAdaptationState:
        """Initialize the state of the dual averaging scheme.

        The parameter :math:`\\mu` is set to :math:`\\log(10 \\epsilon_1)`
        where :math:`\\epsilon_1` is the initial value of the step size.
        """
        return DualAveragingAdaptationState(*da_init(inital_step_size))

    def update(
        da_state: DualAveragingAdaptationState, acceptance_rate: float
    ) -> DualAveragingAdaptationState:
        """Update the state of the Dual Averaging adaptive algorithm.

        Parameters
        ----------
        da_state:
            The current state of the dual averaging algorithm.
        acceptance_rate: float in [0, 1]
            The current metropolis acceptance rate.

        Returns
        -------
        The updated state of the dual averaging algorithm.

        """
        gradient = target - acceptance_rate
        return DualAveragingAdaptationState(*da_update(da_state, gradient))

    def final(da_state: DualAveragingAdaptationState) -> float:
        return jnp.exp(da_state.log_step_size_avg)

    return init, update, final


# -------------------------------------------------------------------
#                 REASONABLE FIRST STEP SIZE
# -------------------------------------------------------------------


class ReasonableStepSizeState(NamedTuple):
    """State carried through the search for a reasonable first step size.

    step
        The current iteration step.
    direction: {-1, 1}
        Determines whether the step size should be increased or decreased during the
        previous step search. If direction = 1 it will be increased, otherwise decreased.
    previous_direction
        The previous direction. It is necessary to carry it because the choice of step
        size is made at the end of the search update.
    step_size
        The current step size in the search.

    """

    step: int
    direction: int
    previous_direction: int
    step_size: float


def find_reasonable_step_size(
    rng_key: PRNGKey,
    kernel_generator: Callable[[float], Callable],
    reference_state: HMCState,
    initial_step_size: float,
    target_accept: float = 0.65,
) -> float:
    """Find a reasonable initial step size during warmup.

    While the dual averaging scheme is guaranteed to converge to a reasonable
    value for the step size starting from any value, choosing a good first
    value can speed up the convergence. This heuristics doubles and halves the
    step size until the acceptance probability of the HMC proposal crosses the
    target value :cite:p:`hoffman2014no`.

    Parameters
    ----------
    rng_key
       Key used by JAX's random number generator.
    kernel_generator
        A function that takes a step size as an input and returns the corresponding
        sampling kernel.
    reference_hmc_state
        The location (HMC state) where this first step size must be found. This function
        never advances the chain.
    inverse_mass_matrix
        The inverse mass matrix relative to which the step size must be found.
    initial_step_size
        The first step size used to start the search.
    target_accept
        Once that value of the metropolis acceptance probability is reached we
        estimate that we have found a "reasonable" first step size.

    Returns
    -------
    float
        A reasonable first value for the step size.

    """
    fp_limit = jnp.finfo(jax.lax.dtype(initial_step_size))

    def do_continue(rss_state: ReasonableStepSizeState) -> bool:
        """Decides whether the search should continue.

        The search stops when it crosses the `target_accept` threshold, i.e.
        when the current direction is opposite to the previous direction.

        Note
        ----
        Per JAX's documentation :cite:p:`jax_finfo` the `jnp.finfo` object is cached so we do not
        occur any performance penalty when calling it repeatedly inside this
        function.

        """
        _, direction, previous_direction, step_size = rss_state

        not_too_large = (step_size < fp_limit.max) | (direction <= 0)
        not_too_small = (step_size > fp_limit.tiny) | (direction >= 0)
        is_step_size_not_extreme = not_too_large & not_too_small
        has_acceptance_rate_not_crossed_threshold = (previous_direction == 0) | (
            direction == previous_direction
        )
        return is_step_size_not_extreme & has_acceptance_rate_not_crossed_threshold

    def update(rss_state: ReasonableStepSizeState) -> ReasonableStepSizeState:
        """Perform one step of the step size search."""
        i, direction, _, step_size = rss_state
        subkey = jax.random.fold_in(rng_key, i)

        step_size = (2.0**direction) * step_size
        kernel = kernel_generator(step_size)
        _, info = kernel(subkey, reference_state)

        new_direction = jnp.where(target_accept < info.acceptance_rate, 1, -1)
        return ReasonableStepSizeState(i + 1, new_direction, direction, step_size)

    rss_state = ReasonableStepSizeState(0, 0, 0, initial_step_size)
    rss_state = jax.lax.while_loop(do_continue, update, rss_state)

    return rss_state.step_size


def bisection_monotonic_fn(acc_prob_wanted, reduce_shift=jnp.log(2.0), tolerance=0.03):
    """Bisection of a monotonically decreassing function, that doesn't require an initially bracketing interval."""

    def update(state, exp_x, acc_rate_new):
        bounds, terminated = state

        # update the bounds
        acc_high = acc_rate_new > acc_prob_wanted
        x = jnp.log(exp_x)

        def on_true(bounds):
            lower, upper = bounds
            # jax.debug.print("true: {x}", x=True)
            lower = jnp.max(jnp.array([lower, x]))
            return jnp.array([lower, upper]), lower + reduce_shift

        def on_false(bounds):
            lower, upper = bounds
            upper = jnp.min(jnp.array([upper, x]))
            return jnp.array([lower, upper]), upper - reduce_shift

        bounds_new, x_new = jax.lax.cond(acc_high, on_true, on_false, bounds)

        # if we have already found a bracketing interval, do bisection, otherwise further reduce or increase the bounds
        bracketing = jnp.all(jnp.isfinite(bounds_new))

        def reduce(bounds):
            return x_new

        def bisect(bounds):
            return jnp.average(bounds)

        x_new = jax.lax.cond(bracketing, bisect, reduce, bounds_new)

        stepsize = terminated * exp_x + (1 - terminated) * jnp.exp(x_new)

        terminated_new = (
            jnp.abs(acc_rate_new - acc_prob_wanted) < tolerance
        ) | terminated

        return (bounds_new, terminated_new), stepsize

    return update
