from typing import Callable, Union

import numpy as np
from tqdm import trange

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


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
        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.
    """

    def __init__(
        self,
        log_prob_fn: Callable,
        step_fn: Union[Step, str],
        rng: np.random.Generator,
        dims: int,
        target_acceptance_rate: float = 0.234,
        **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, **kwargs)

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

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

        Parameters
        ----------
        x_init : np.ndarray
            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 : np.ndarray
            Samples generated by the sampler. Shape: (n_steps + 1, N, D).
        state_history : ChainStateHistory
            History of the chain states during the sampling process.
        """
        x = x_init
        self.step_fn.initialise(x)
        log_prob_x = self.log_prob_fn(x)  # Shape: (N,)
        chain = np.empty((n_steps + 1, x.shape[0], x.shape[1]), dtype=x.dtype)
        chain[0] = 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 = np.exp(np.minimum(0, log_alpha))  # Shape: (N,)

                accept = self.rng.uniform(size=len(x_new)) < alpha
                x = np.where(accept[:, None], x_new, x)  # Shape: (N, D)
                log_prob_x = np.where(accept, log_prob_x_new, log_prob_x)
                chain[i + 1] = x

                state = ChainState(
                    it=i,
                    acceptance_rate=float(np.mean(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)

        state_history = ChainStateHistory.from_chain_states(states)
        return chain, state_history
