from functools import partial
from typing import Any

import numpy as np

from ...samples import SMCSamples
from ...utils import (
    asarray,
    determine_backend_name,
    track_calls,
)
from .base import SMCSampler


class MiniPCNSMC(SMCSampler):
    """MiniPCN SMC sampler."""

    rng = None

    def log_prob(self, x: SMCSamples, beta: float | None = None) -> Any:
        return super().log_prob(x, beta)

    @track_calls
    def sample(
        self,
        n_samples: int,
        n_steps: int = None,
        min_beta_step: float | None = None,
        max_beta_step: float | None = None,
        max_n_steps: int | None = None,
        adaptive: bool = True,
        target_efficiency: float = 0.5,
        target_efficiency_rate: float = 1.0,
        n_final_samples: int | None = None,
        sampler_kwargs: dict | None = None,
        rng: np.random.Generator | None = None,
        checkpoint_callback=None,
        checkpoint_every: int | None = None,
        checkpoint_file_path: str | None = None,
        resume_from: str | bytes | dict | None = None,
        beta_tolerance: float = 1e-6,
        store_sample_history: bool = True,
    ) -> SMCSamples:
        from orng import ArrayRNG

        self.sampler_kwargs = sampler_kwargs or {}
        self.sampler_kwargs.setdefault("n_steps", 5 * self.dims)
        self.sampler_kwargs.setdefault("target_acceptance_rate", 0.234)
        self.sampler_kwargs.setdefault("step_fn", "tpcn")
        self.sampler_kwargs.setdefault("verbose", True)
        self.backend_str = determine_backend_name(xp=self.xp)
        self.rng = rng or self.rng or ArrayRNG(backend=self.backend_str)
        return super().sample(
            n_samples,
            n_steps=n_steps,
            adaptive=adaptive,
            target_efficiency=target_efficiency,
            target_efficiency_rate=target_efficiency_rate,
            n_final_samples=n_final_samples,
            min_beta_step=min_beta_step,
            max_beta_step=max_beta_step,
            max_n_steps=max_n_steps,
            checkpoint_callback=checkpoint_callback,
            checkpoint_every=checkpoint_every,
            checkpoint_file_path=checkpoint_file_path,
            resume_from=resume_from,
            beta_tolerance=beta_tolerance,
            store_sample_history=store_sample_history,
        )

    def mutate(
        self, particles: SMCSamples, beta: float, n_steps: int | None = None
    ) -> SMCSamples:
        """Mutate particles using the MiniPCN sampler.

        Parameters
        ----------
        particles : SMCSamples
            The current particles to be mutated.
        beta : float
            The current inverse temperature.
        n_steps : int, optional
            The number of MCMC steps to take. If None, uses the default from
            :code:`sampler_kwargs`.

        Returns
        -------
        SMCSamples
            The mutated particles.
        """
        from minipcn import Sampler

        log_prob_fn = partial(self.log_prob, beta=beta)

        kwargs = self.sampler_kwargs.copy()
        n_steps = n_steps or kwargs.pop("n_steps")
        verbose = kwargs.pop("verbose")

        sampler = Sampler(
            log_prob_fn=log_prob_fn,
            rng=self.rng,
            dims=self.dims,
            xp=self.xp,
            **kwargs,
        )
        # Map to transformed dimension for sampling
        z = asarray(
            self.fit_preconditioning_transform(particles.x),
            xp=self.xp,
            dtype=self.dtype,
        )
        chain, history = sampler.sample(
            z,
            n_steps=n_steps or self.sampler_kwargs["n_steps"],
            verbose=verbose,
        )
        x = self.preconditioning_transform.inverse(chain[-1])[0]

        self.history.mcmc_acceptance.append(np.mean(history.acceptance_rate))

        samples = SMCSamples(
            x,
            xp=self.xp,
            beta=beta,
            dtype=self.dtype,
            parameters=self.parameters,
        )
        samples.log_q = samples.array_to_namespace(
            self.prior_flow.log_prob(samples.x)
        )
        samples.log_prior = samples.array_to_namespace(self.log_prior(samples))
        samples.log_likelihood = samples.array_to_namespace(
            self.log_likelihood(samples)
        )
        if samples.xp.isnan(samples.log_q).any():
            raise ValueError("Log proposal contains NaN values")
        return samples
