import copy
import logging
from typing import Any, Callable

import array_api_compat.numpy as np

from ...flows.base import Flow
from ...history import SMCHistory
from ...samples import SMCSamples
from ...utils import (
    asarray,
    effective_sample_size,
    track_calls,
    update_at_indices,
)
from ..mcmc import MCMCSampler

logger = logging.getLogger(__name__)


class SMCSampler(MCMCSampler):
    """Base class for Sequential Monte Carlo samplers."""

    def __init__(
        self,
        log_likelihood: Callable,
        log_prior: Callable,
        dims: int,
        prior_flow: Flow,
        xp: Callable,
        dtype: Any | str | None = None,
        parameters: list[str] | None = None,
        rng: np.random.Generator | None = None,
        preconditioning_transform: Callable | None = None,
    ):
        super().__init__(
            log_likelihood=log_likelihood,
            log_prior=log_prior,
            dims=dims,
            prior_flow=prior_flow,
            xp=xp,
            dtype=dtype,
            parameters=parameters,
            preconditioning_transform=preconditioning_transform,
        )
        self.rng = rng or np.random.default_rng()
        self._adapative_target_efficiency = False

    @property
    def target_efficiency(self):
        return self._target_efficiency

    @target_efficiency.setter
    def target_efficiency(self, value: float | tuple):
        """Set the target efficiency.

        Parameters
        ----------
        value : float or tuple
            If a float, the target efficiency to use for all iterations.
            If a tuple of two floats, the target efficiency will adapt from
            the first value to the second value over the course of the SMC
            iterations. See `target_efficiency_rate` for details.
        """
        if isinstance(value, float):
            if not (0 < value < 1):
                raise ValueError("target_efficiency must be in (0, 1)")
            self._target_efficiency = value
            self._adapative_target_efficiency = False
        elif len(value) != 2:
            raise ValueError(
                "target_efficiency must be a float or tuple of two floats"
            )
        else:
            value = tuple(map(float, value))
            if not (0 < value[0] < value[1] < 1):
                raise ValueError(
                    "target_efficiency tuple must be in (0, 1) and increasing"
                )
            self._target_efficiency = value
            self._adapative_target_efficiency = True

    def current_target_efficiency(self, beta: float) -> float:
        """Get the current target efficiency based on beta."""
        if self._adapative_target_efficiency:
            return self._target_efficiency[0] + (
                self._target_efficiency[1] - self._target_efficiency[0]
            ) * (beta**self.target_efficiency_rate)
        else:
            return self._target_efficiency

    def determine_beta(
        self,
        samples: SMCSamples,
        beta: float,
        beta_step: float,
        min_step: float,
        beta_tolerance: float = 1e-6,
    ) -> tuple[float, float]:
        """Determine the next beta value.

        Parameters
        ----------
        samples : SMCSamples
            The current samples.
        beta : float
            The current beta value.
        beta_step : float
            The fixed beta step size if not adaptive.
        min_step : float
            The minimum beta step size.
        beta_tolerance : float
            Tolerance when checking for beta convergence.

        Returns
        -------
        beta : float
            The new beta value.
        min_step : float
            The new minimum step size if adaptive_min_step is True.
        """
        if not self.adaptive:
            beta += beta_step
            if beta >= 1.0:
                beta = 1.0
        else:
            beta_prev = beta
            beta_min = beta_prev
            beta_max = 1.0
            eff_beta_max = effective_sample_size(
                samples.log_weights(beta_max)
            ) / len(samples)
            if eff_beta_max >= self.current_target_efficiency(beta_prev):
                beta_min = 1.0
            target_eff = self.current_target_efficiency(beta_prev)
            while beta_max - beta_min > beta_tolerance:
                beta_try = 0.5 * (beta_max + beta_min)
                eff = effective_sample_size(
                    samples.log_weights(beta_try)
                ) / len(samples)
                if eff >= target_eff:
                    beta_min = beta_try
                else:
                    beta_max = beta_try
            beta_star = beta_min

            if self.adaptive_min_step:
                min_step = min_step * (1 - beta_prev) / (1 - beta_star)
            beta = max(beta_star, beta_prev + min_step)
            beta = min(beta, 1.0)
        return beta, min_step

    @track_calls
    def sample(
        self,
        n_samples: int,
        n_steps: int | None = None,
        adaptive: bool = True,
        min_step: float | None = None,
        max_n_steps: int | None = None,
        target_efficiency: float = 0.5,
        target_efficiency_rate: float = 1.0,
        n_final_samples: int | None = None,
        checkpoint_callback: Callable[[dict], None] | None = None,
        checkpoint_every: int | None = None,
        checkpoint_file_path: str | None = None,
        resume_from: str | bytes | dict | None = None,
        store_sample_history: bool = True,
        beta_tolerance: float = 1e-6,
    ) -> SMCSamples:
        """Sample using the SMC sampler.

        Parameters
        ----------
        n_samples : int
            The number of samples (particles) to use in the SMC sampler.
        n_steps : int, optional
            The number of SMC iterations to perform. Must be specified if
            :code:`adaptive=False`. Default is None.
        adaptive : bool, optional
            Whether to adaptively determine the beta schedule. Default is True.
        min_step : float, optional
            The minimum beta step size when using adaptive beta. Default is None,
            which means no minimum step size.
        max_n_steps : int, optional
            The maximum number of SMC iterations to perform when using adaptive
            beta. Default is None, which means no maximum.
        target_efficiency : float or tuple, optional
            The target sample efficiency (ESS / n_samples) to aim for at each
            SMC iteration. Can be a single float in (0, 1) or a tuple of two
            floats specifying a range to adapt between. Default is 0.5.
        target_efficiency_rate : float, optional
            When using a tuple for target_efficiency, this controls the rate at
            which the target efficiency adapts from the first value to the
            second value as beta increases. Default is 1.0 (linear adaptation).
        n_final_samples : int, optional
            If specified, the number of final samples to produce after the SMC
            iterations. If not specified, the number of final samples will be
            the same as n_samples. Default is None.
        checkpoint_callback : callable, optional
            A callback function to call with a checkpoint dictionary at regular
            intervals during sampling. Default is None (no checkpointing).
        checkpoint_every : int, optional
            The number of iterations between checkpoints when using checkpoint
            callback.
            Default is None, which means no regular checkpointing.
        checkpoint_file_path : str, optional
            If using checkpoint_callback, this can be used to specify a file
            path to save checkpoints to. Default is None.
        resume_from : str, bytes, or dict, optional
            If specified, this can be used to resume sampling from a previous
            checkpoint. Can be a file path, bytes object, or checkpoint
            dictionary. Default is None (start from scratch).
        store_sample_history : bool, optional
            Whether to store the history of samples at each iteration in
            :code:`self.history.sample_history`. Default is True.
        beta_tolerance : float, optional
            Tolerance for determining convergence of beta when using adaptive
            beta. Default is 1e-6.

        Returns
        -------
        final_samples : SMCSamples
            The final samples after running the SMC sampler, with log evidence
            and log evidence error estimates.

        Raises
        ------
        ValueError
            If both n_steps is None and adaptive is False, or if
            target_efficiency is not in (0, 1) when a float, or if
            target_efficiency tuple is not valid, or if log probabilities
            contain NaN values.
        """
        resumed = resume_from is not None
        if resumed:
            samples, beta, iterations = self.restore_from_checkpoint(
                resume_from
            )
        else:
            samples = self.draw_initial_samples(n_samples)
            samples = SMCSamples.from_samples(
                samples, xp=self.xp, beta=0.0, dtype=self.dtype
            )
            beta = 0.0
            iterations = 0
            self.history = SMCHistory()
        self.fit_preconditioning_transform(samples.x)

        if store_sample_history:
            self.history.sample_history.append(samples)

        if self.xp.isnan(samples.log_q).any():
            raise ValueError("Log proposal contains NaN values")
        if self.xp.isnan(samples.log_prior).any():
            raise ValueError("Log prior contains NaN values")
        if self.xp.isnan(samples.log_likelihood).any():
            raise ValueError("Log likelihood contains NaN values")

        logger.debug(f"Initial sample summary: {samples}")

        # Remove the n_final_steps from sampler_kwargs if present
        self.sampler_kwargs = self.sampler_kwargs or {}
        n_final_steps = self.sampler_kwargs.pop("n_final_steps", None)

        self.target_efficiency = target_efficiency
        self.target_efficiency_rate = target_efficiency_rate

        if n_steps is not None:
            beta_step = 1 / n_steps
        elif not adaptive:
            raise ValueError("Either n_steps or adaptive=True must be set")
        else:
            beta_step = np.nan
        self.adaptive = adaptive

        if min_step is None:
            if max_n_steps is None:
                min_step = 0.0
                self.adaptive_min_step = False
            else:
                min_step = 1 / max_n_steps
                self.adaptive_min_step = True
        else:
            self.adaptive_min_step = False

        iterations = iterations or 0
        if checkpoint_callback is None and checkpoint_every is not None:
            checkpoint_callback = self.default_file_checkpoint_callback(
                checkpoint_file_path
            )
        if checkpoint_callback is not None and checkpoint_every is None:
            checkpoint_every = 1

        run_smc_loop = True
        if resumed:
            last_beta = self.history.beta[-1] if self.history.beta else beta
            if last_beta >= 1.0:
                run_smc_loop = False

        def maybe_checkpoint(force: bool = False):
            if checkpoint_callback is None:
                return
            should_checkpoint = force or (
                checkpoint_every is not None
                and checkpoint_every > 0
                and iterations % checkpoint_every == 0
            )
            if not should_checkpoint:
                return
            state = self.build_checkpoint_state(samples, iterations, beta)
            checkpoint_callback(state)

        if run_smc_loop:
            while True:
                iterations += 1

                beta, min_step = self.determine_beta(
                    samples,
                    beta,
                    beta_step,
                    min_step,
                    beta_tolerance=beta_tolerance,
                )
                self.history.eff_target.append(
                    self.current_target_efficiency(beta)
                )

                logger.info(f"it {iterations} - beta: {beta}")
                self.history.beta.append(beta)

                ess = effective_sample_size(samples.log_weights(beta))
                eff = ess / len(samples)
                if eff < 0.1:
                    logger.warning(
                        f"it {iterations} - Low sample efficiency: {eff:.2f}"
                    )
                self.history.ess.append(ess)
                logger.info(
                    f"it {iterations} - ESS: {ess:.1f} ({eff:.2f} efficiency)"
                )
                self.history.ess_target.append(
                    effective_sample_size(samples.log_weights(1.0))
                )

                log_evidence_ratio = samples.log_evidence_ratio(beta)
                log_evidence_ratio_var = samples.log_evidence_ratio_variance(
                    beta
                )
                self.history.log_norm_ratio.append(log_evidence_ratio)
                self.history.log_norm_ratio_var.append(log_evidence_ratio_var)
                logger.info(
                    f"it {iterations} - Log evidence ratio: {log_evidence_ratio:.2f} +/- {np.sqrt(log_evidence_ratio_var):.2f}"
                )

                samples = samples.resample(beta, rng=self.rng)

                samples = self.mutate(samples, beta)
                if store_sample_history:
                    self.history.sample_history.append(samples)
                maybe_checkpoint()
                if beta == 1.0 or (
                    max_n_steps is not None and iterations >= max_n_steps
                ):
                    break

        # If n_final_samples is specified and differs, perform additional mutation steps
        if n_final_samples is not None and len(samples.x) != n_final_samples:
            logger.info(f"Generating {n_final_samples} final samples")
            final_samples = samples.resample(
                1.0, n_samples=n_final_samples, rng=self.rng
            )
            samples = self.mutate(final_samples, 1.0, n_steps=n_final_steps)

        samples.log_evidence = samples.xp.sum(
            asarray(self.history.log_norm_ratio, self.xp)
        )
        samples.log_evidence_error = samples.xp.sqrt(
            samples.xp.sum(asarray(self.history.log_norm_ratio_var, self.xp))
        )
        maybe_checkpoint(force=True)

        final_samples = samples.to_standard_samples()
        logger.info(
            f"Log evidence: {final_samples.log_evidence:.2f} +/- {final_samples.log_evidence_error:.2f}"
        )
        return final_samples

    def mutate(self, particles):
        raise NotImplementedError

    def log_prob(self, z, beta=None):
        x, log_abs_det_jacobian = self.preconditioning_transform.inverse(z)
        samples = SMCSamples(x, xp=self.xp, beta=beta, dtype=self.dtype)
        log_q = self.prior_flow.log_prob(samples.x)
        samples.log_q = samples.array_to_namespace(log_q)
        samples.log_prior = self.log_prior(samples)
        samples.log_likelihood = self.log_likelihood(samples)
        log_prob = samples.log_p_t(
            beta=beta
        ).flatten() + samples.array_to_namespace(log_abs_det_jacobian)

        log_prob = update_at_indices(
            log_prob, self.xp.isnan(log_prob), -self.xp.inf
        )
        return log_prob

    def build_checkpoint_state(
        self, samples: SMCSamples, iteration: int, beta: float
    ) -> dict:
        """Prepare a serializable checkpoint payload for the sampler state."""
        return super().build_checkpoint_state(
            samples,
            iteration,
            meta={"beta": beta},
        )

    def _checkpoint_extra_state(self) -> dict:
        history_copy = copy.deepcopy(self.history)
        rng_state = (
            self.rng.bit_generator.state
            if hasattr(self.rng, "bit_generator")
            else None
        )
        return {
            "history": history_copy,
            "rng_state": rng_state,
            "sampler_kwargs": getattr(self, "sampler_kwargs", None),
        }

    def restore_from_checkpoint(
        self, source: str | bytes | dict
    ) -> tuple[SMCSamples, float, int]:
        samples, state = super().restore_from_checkpoint(source)
        meta = state.get("meta", {}) if isinstance(state, dict) else {}
        beta = None
        if isinstance(meta, dict):
            beta = meta.get("beta", None)
        if beta is None:
            beta = state.get("beta", 0.0)
        iteration = state.get("iteration", 0)
        self.history = state.get("history", SMCHistory())
        rng_state = state.get("rng_state")
        if rng_state is not None and hasattr(self.rng, "bit_generator"):
            self.rng.bit_generator.state = rng_state
        samples = SMCSamples.from_samples(
            samples, xp=self.xp, beta=beta, dtype=self.dtype
        )
        return samples, beta, iteration


class NumpySMCSampler(SMCSampler):
    def __init__(
        self,
        log_likelihood,
        log_prior,
        dims,
        prior_flow,
        xp,
        dtype=None,
        parameters=None,
        preconditioning_transform=None,
    ):
        if preconditioning_transform is not None:
            preconditioning_transform = preconditioning_transform.new_instance(
                xp=np
            )
        super().__init__(
            log_likelihood,
            log_prior,
            dims,
            prior_flow=prior_flow,
            xp=xp,
            dtype=dtype,
            parameters=parameters,
            preconditioning_transform=preconditioning_transform,
        )
