# 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.
# """Public API for Late Adjusted Parallel Sampler (LAPS)"""


from typing import Any, NamedTuple

import jax
import jax.numpy as jnp

from blackjax.adaptation import laps_burn_in
from blackjax.adaptation.laps_burn_in import (
    equipartition_diagonal,
    equipartition_diagonal_loss,
)
from blackjax.adaptation.step_size import bisection_monotonic_fn
from blackjax.mcmc.adjusted_mclmc import build_kernel as build_kernel_mclmc
from blackjax.mcmc.hmc import HMCState
from blackjax.mcmc.integrators import (
    generate_isokinetic_integrator,
    mclachlan_coefficients,
    omelyan_coefficients,
)
from blackjax.util import run_eca


class AdaptationState(NamedTuple):
    steps_per_sample: float
    step_size: float
    stepsize_adaptation_state: Any  # the state of the bisection algorithm to find a stepsize
    iteration: int


class Adaptation:
    def __init__(
        self,
        adaptation_state,
        num_adaptation_samples,
        steps_per_sample=15,
        acc_prob_target=0.8,
        observables=lambda x: 0.0,
        observables_for_bias=lambda x: 0.0,
        contract=lambda x: 0.0,
    ):
        self.num_adaptation_samples = num_adaptation_samples
        self.observables = observables
        self.observables_for_bias = observables_for_bias
        self.contract = contract

        """
        num_adaptation_samples: amount of tuning in the adjusted phase before fixing params
        steps_per_sample: number of steps per sample
        acc_prob_target: target acceptance probability
        observables: function to compute observables, for diagnostics
        observables_for_bias: function to compute observables for bias, for diagnostics
        contract: function to contract observables, for diagnostics
        """

        # Determine the initial hyperparameters #

        step_size = adaptation_state.step_size

        # Initialize the bisection for finding the step size
        self.epsadap_update = bisection_monotonic_fn(acc_prob_target)
        stepsize_adaptation_state = (jnp.array([-jnp.inf, jnp.inf]), False)

        self.initial_state = AdaptationState(
            steps_per_sample, step_size, stepsize_adaptation_state, 0
        )

    def summary_statistics_fn(self, state, info, rng_key):
        return {
            "acceptance_probability": info.acceptance_rate,
            "equipartition_diagonal": equipartition_diagonal(state),
            "observables": self.observables(state.position),
            "observables_for_bias": self.observables_for_bias(state.position),
        }

    def update(self, adaptation_state, Etheta):
        acc_prob = Etheta["acceptance_probability"]
        equi_diag = equipartition_diagonal_loss(Etheta["equipartition_diagonal"])
        true_bias = self.contract(Etheta["observables_for_bias"])

        info_to_be_stored = {
            "L": adaptation_state.step_size * adaptation_state.steps_per_sample,
            "steps_per_sample": adaptation_state.steps_per_sample,
            "step_size": adaptation_state.step_size,
            "acc_prob": acc_prob,
            "equi_diag": equi_diag,
            "bias": true_bias,
            "observables": Etheta["observables"],
        }

        # Bisection to find step size
        stepsize_adaptation_state, step_size = self.epsadap_update(
            adaptation_state.stepsize_adaptation_state,
            adaptation_state.step_size,
            acc_prob,
        )

        return (
            AdaptationState(
                adaptation_state.steps_per_sample,
                step_size,
                stepsize_adaptation_state,
                adaptation_state.iteration + 1,
            ),
            info_to_be_stored,
        )


def bias(model):
    """should be transfered to benchmarks/"""

    def observables(position):
        return jnp.square(model.transform(position))

    def contract(sampler_E_x2):
        bsq = jnp.square(sampler_E_x2 - model.E_x2) / model.Var_x2
        return jnp.array([jnp.max(bsq), jnp.average(bsq)])

    return observables, contract


def while_steps_num(cond):
    if jnp.all(cond):
        return len(cond)
    else:
        return jnp.argmin(cond) + 1


def laps(
    logdensity_fn,
    sample_init,
    ndims,
    num_steps1,
    num_steps2,
    num_chains,
    mesh,
    rng_key,
    microcanonical=True,
    alpha=1.9,
    save_frac=0.2,
    C=0.1,
    early_stop=True,
    r_end=0.01,
    bias_type=3,
    diagonal_preconditioning=True,
    integrator_coefficients=None,
    steps_per_sample=15,
    acc_prob=None,
    observables_for_bias=lambda x: x,
    all_chains_info=None,
    diagnostics=True,
    contract=lambda x: 0.0,
    superchain_size=1,
):
    """
    model: the target density object
    num_steps1: number of steps in the first phase
    num_steps2: number of steps in the second phase
    num_chains: number of chains
    mesh: the mesh object, used for distributing the computation across cpus and nodes
    rng_key: the random key
    alpha: L = sqrt{d} * alpha * variances
    save_frac: the fraction of samples used to estimate the fluctuation in the first phase
    C: constant in stage 1 that determines step size (eq (9) of EMAUS paper)
    early_stop: whether to stop the first phase early
    r_end
    diagonal_preconditioning: whether to use diagonal preconditioning
    integrator_coefficients: the coefficients of the integrator
    steps_per_sample: the number of steps per sample
    acc_prob: the acceptance probability
    observables: the observables (for diagnostic use)
    all_chains_info: summary statistics calculated and stored for all chain at each iteration so it can be memory intensive
    diagnostics: whether to return diagnostics
    """

    key_init, key1, key2 = jax.random.split(rng_key, 3)

    # initialize the chains
    initial_state = laps_burn_in.initialize(
        key_init,
        logdensity_fn,
        microcanonical,
        sample_init,
        num_chains,
        mesh,
        superchain_size,
    )

    # burn-in with the unadjusted method #
    kernel = laps_burn_in.build_kernel(logdensity_fn, ndims, microcanonical)
    save_num = (jnp.rint(save_frac * num_steps1)).astype(int)
    adap = laps_burn_in.Adaptation(
        ndims,
        microcanonical=microcanonical,
        alpha=alpha,
        bias_type=bias_type,
        save_num=save_num,
        C=C,
        r_end=r_end,
        observables_for_bias=observables_for_bias,
        contract=contract,
    )

    final_state, final_adaptation_state, info1 = run_eca(
        key1,
        initial_state,
        kernel,
        adap,
        num_steps1,
        num_chains,
        mesh,
        superchain_size,
        all_chains_info,
        early_stop=early_stop,
    )

    # refine the results with the adjusted method
    _acc_prob = acc_prob
    if integrator_coefficients is None:
        high_dims = ndims > 200
        _integrator_coefficients = (
            omelyan_coefficients if high_dims else mclachlan_coefficients
        )
        if acc_prob is None:
            _acc_prob = 0.9 if high_dims else 0.7

    else:
        _integrator_coefficients = integrator_coefficients
        if acc_prob is None:
            _acc_prob = 0.9

    gradient_calls_per_step = (
        len(_integrator_coefficients) // 2
    )  # The number of B updates in scheme = BABAB..AB is len(scheme)//2 + 1. The number of gradient calls is then len(scheme)//2, because the last B's gradient can be reused in the next step.

    if diagonal_preconditioning:
        inverse_mass_matrix = final_adaptation_state.inverse_mass_matrix

        # scale the stepsize so that it reflects averag scale change of the preconditioning
        average_scale_change = jnp.sqrt(jnp.average(inverse_mass_matrix))
        final_adaptation_state = final_adaptation_state._replace(
            step_size=final_adaptation_state.step_size / average_scale_change
        )

    else:
        inverse_mass_matrix = 1.0

    if microcanonical:
        integrator = generate_isokinetic_integrator(_integrator_coefficients)

        build_kernel = build_kernel_mclmc(
            integrator=integrator,
            logdensity_fn=logdensity_fn,
            inverse_mass_matrix=inverse_mass_matrix,
        )

        kernel = lambda key, state, adap: build_kernel(
            rng_key=key,
            state=state,
            step_size=adap.step_size,
            num_integration_steps=adap.steps_per_sample,
            L_proposal_factor=1.25,
        )

    else:
        raise ValueError("Only microcanonical mode is supported for LAPS.")

    initial_state = HMCState(
        final_state.position, final_state.logdensity, final_state.logdensity_grad
    )
    num_samples = num_steps2 // (gradient_calls_per_step * steps_per_sample)
    num_adaptation_samples = (
        num_samples // 2
    )  # number of samples after which the stepsize is fixed.

    final_adaptation_state = final_adaptation_state._replace(
        step_size=final_adaptation_state.step_size.item()
    )

    adap = Adaptation(
        final_adaptation_state,
        num_adaptation_samples,
        steps_per_sample,
        _acc_prob,
        contract=contract,
        observables_for_bias=observables_for_bias,
    )

    final_state, final_adaptation_state, info2 = run_eca(
        key2,
        initial_state,
        kernel,
        adap,
        num_samples,
        num_chains,
        mesh,
        superchain_size,
        all_chains_info,
    )

    if diagnostics:
        info = {"phase_1": info1, "phase_2": info2}
    else:
        info = None

    return info, gradient_calls_per_step, _acc_prob, final_state
