from __future__ import annotations

from typing import TYPE_CHECKING, Any, Callable, Union

import numpy as np
from tqdm import trange

from ._typing import Array
from .step import Step
from .utils import ChainState, ChainStateHistory

if TYPE_CHECKING:
    try:
        from orng import ArrayRNG
    except ImportError:
        ArrayRNG = Any


class Sampler:
    """Class for running the MiniPCN sampler.

    Parameters
    ----------
    log_prob_fn : Callable
        Function to compute the log probability of the target distribution.
        It should take a single argument (the samples) and return the log
        probability.
    step_fn : Step | str
        Step object that defines the proposal distribution and the
        transformation to the target distribution. If a string is provided,
        it should be the name of a known step type (e.g., "pCN" or "tpCN").
    rng : np.random.Generator | ArrayRNG
        Random number generator for reproducibility.
    dims : int
        Number of dimensions of the target distribution.
    target_acceptance_rate : float, optional
        Target acceptance rate for the sampler. Default is 0.234.
    xp : Any, optional
        Array namespace to use (e.g., numpy, jax.numpy, torch). Default is numpy.
    **kwargs
        Additional keyword arguments to pass to the step function if `step_fn`
        is provided as a string.
    """

    def __init__(
        self,
        log_prob_fn: Callable,
        step_fn: Union[Step, str],
        rng: "ArrayRNG" | np.random.Generator,
        dims: int,
        target_acceptance_rate: float = 0.234,
        xp: Any = np,
        **kwargs,
    ) -> None:
        self.log_prob_fn = log_prob_fn

        if isinstance(step_fn, str):
            from .step import step_factory

            step_fn = step_factory(step_fn, dims, rng, xp, **kwargs)

        self.step_fn = step_fn
        self.rng = rng
        self.dims = dims
        self.target_acceptance_rate = target_acceptance_rate
        self.xp = xp

    def sample(
        self,
        x_init: Array,
        n_steps: int,
        verbose: bool = True,
    ) -> tuple[Array, ChainStateHistory]:
        """Run the MiniCrank sampler.

        Parameters
        ----------
        x_init : Array
            Initial samples for the sampler. Shape: (N, D), where N is the number
            of samples and D is the number of dimensions.
        n_steps : int
            Number of steps to run the sampler.
        verbose : bool, optional
            If True, display a progress bar. Default is True.
            If False, no progress bar is displayed.

        Returns
        -------
        chain : Array
            Samples generated by the sampler. Shape: (n_steps + 1, N, D).
        state_history : ChainStateHistory
            History of the chain states during the sampling process.
        """
        x = self.xp.atleast_2d(x_init)
        self.step_fn.initialise(x)
        log_prob_x = self.log_prob_fn(x)  # Shape: (N,)
        # Accumulate states functionally to avoid in-place updates (e.g. JAX)
        chain_states: list[Array] = [x]
        states = []
        with trange(
            n_steps, desc="Sampling", unit="step", disable=not verbose
        ) as pbar:
            for i in pbar:
                x_new, log_alpha_step = self.step_fn(x)
                log_prob_x_new = self.log_prob_fn(x_new)
                log_alpha = log_prob_x_new - log_prob_x + log_alpha_step
                alpha = self.xp.exp(
                    self.xp.minimum(self.xp.asarray(0.0), log_alpha)
                )  # Shape: (N,)

                accept = self.rng.uniform(size=len(x_new)) < alpha
                x = self.xp.where(accept[:, None], x_new, x)  # Shape: (N, D)
                log_prob_x = self.xp.where(accept, log_prob_x_new, log_prob_x)
                chain_states.append(x)

                state = ChainState(
                    it=i,
                    acceptance_rate=float(self.xp.sum(accept) / len(accept)),
                    target_acceptance_rate=self.target_acceptance_rate,
                )
                self.step_fn.update(state, x)
                state = self.step_fn.update_state(state)
                states.append(state)
                pbar_dict = {
                    "acceptance_rate": state.acceptance_rate,
                }
                pbar_dict.update(state.extra_stats)
                pbar.set_postfix(pbar_dict)

        chain = self.xp.stack(chain_states, axis=0)
        state_history = ChainStateHistory.from_chain_states(states)
        return chain, state_history
