from __future__ import annotations

import copy
import logging
import math
from dataclasses import dataclass, field
from typing import Any, Callable

import numpy as np
from array_api_compat import (
    array_namespace,
    is_numpy_namespace,
    to_device,
)
from array_api_compat import device as api_device
from array_api_compat.common._typing import Array
from matplotlib.figure import Figure

from .utils import (
    asarray,
    convert_dtype,
    logsumexp,
    recursively_save_to_h5_file,
    resolve_dtype,
    to_numpy,
)

logger = logging.getLogger(__name__)


@dataclass
class BaseSamples:
    """Class for storing samples and corresponding weights.

    If :code:`xp` is not specified, all inputs will be converted to match
    the array type of :code:`x`.
    """

    x: Array
    """Array of samples, shape (n_samples, n_dims)."""
    log_likelihood: Array | None = None
    """Log-likelihood values for the samples."""
    log_prior: Array | None = None
    """Log-prior values for the samples."""
    log_q: Array | None = None
    """Log-probability values under the proposal distribution."""
    parameters: list[str] | None = None
    """List of parameter names."""
    dtype: Any | str | None = None
    """Data type of the samples.

    If None, the default dtype for the array namespace will be used.
    """
    xp: Callable | None = None
    """
    The array namespace to use for the samples.

    If None, the array namespace will be inferred from the type of :code:`x`.
    """
    device: Any = None
    """Device to store the samples on.

    If None, the device will be inferred from the array namespace of :code:`x`.
    """

    def __post_init__(self):
        if self.xp is None:
            self.xp = array_namespace(self.x)
        # Numpy arrays need to be on the CPU before being converted
        if is_numpy_namespace(self.xp):
            self.device = "cpu"
        if self.dtype is not None:
            self.dtype = resolve_dtype(self.dtype, self.xp)
        else:
            # Fall back to default dtype for the array namespace
            self.dtype = None
        self.x = self.array_to_namespace(self.x, dtype=self.dtype)
        if self.device is None:
            self.device = api_device(self.x)
        if self.log_likelihood is not None:
            self.log_likelihood = self.array_to_namespace(
                self.log_likelihood, dtype=self.dtype
            )
        if self.log_prior is not None:
            self.log_prior = self.array_to_namespace(
                self.log_prior, dtype=self.dtype
            )
        if self.log_q is not None:
            self.log_q = self.array_to_namespace(self.log_q, dtype=self.dtype)

        if self.parameters is None:
            self.parameters = [f"x_{i}" for i in range(self.dims)]

    @property
    def dims(self):
        """Number of dimensions (parameters) in the samples."""
        if self.x is None:
            return 0
        return self.x.shape[1] if self.x.ndim > 1 else 1

    def to_numpy(self, dtype: Any | str | None = None):
        logger.debug("Converting samples to numpy arrays")
        import array_api_compat.numpy as np

        if dtype is not None:
            dtype = resolve_dtype(dtype, np)
        else:
            dtype = convert_dtype(self.dtype, np)
        return self.__class__(
            x=self.x,
            parameters=self.parameters,
            log_likelihood=self.log_likelihood,
            log_prior=self.log_prior,
            log_q=self.log_q,
            xp=np,
        )

    def to_namespace(self, xp, dtype: Any | str | None = None):
        if dtype is None:
            dtype = convert_dtype(self.dtype, xp)
        else:
            dtype = resolve_dtype(dtype, xp)
        logger.debug("Converting samples to {} namespace", xp)
        return self.__class__(
            x=self.x,
            parameters=self.parameters,
            log_likelihood=self.log_likelihood,
            log_prior=self.log_prior,
            log_q=self.log_q,
            xp=xp,
            device=self.device,
            dtype=dtype,
        )

    def array_to_namespace(self, x, dtype=None):
        """Convert an array to the same namespace as the samples"""
        kwargs = {}
        if dtype is not None:
            kwargs["dtype"] = resolve_dtype(dtype, self.xp)
        else:
            kwargs["dtype"] = self.dtype
        x = asarray(x, self.xp, **kwargs)
        if self.device:
            x = to_device(x, self.device)
        return x

    def to_dict(self, flat: bool = True):
        samples = dict(zip(self.parameters, self.x.T, strict=True))
        out = {
            "log_likelihood": self.log_likelihood,
            "log_prior": self.log_prior,
            "log_q": self.log_q,
        }
        if flat:
            out.update(samples)
        else:
            out["samples"] = samples
        return out

    def to_dataframe(self, flat: bool = True):
        import pandas as pd

        return pd.DataFrame(self.to_dict(flat=flat))

    def plot_corner(
        self,
        parameters: list[str] | None = None,
        fig: Figure | None = None,
        **kwargs,
    ):
        """Plot a corner plot of the samples.

        Parameters
        ----------
        parameters : list[str] | None
            List of parameters to plot. If None, all parameters are plotted.
        fig : matplotlib.figure.Figure | None
            Figure to plot on. If None, a new figure is created.
        **kwargs : dict
            Additional keyword arguments to pass to corner.corner(). Kwargs
            are deep-copied before use.
        """
        import corner

        kwargs = copy.deepcopy(kwargs)
        kwargs.setdefault("labels", self.parameters)

        if parameters is not None:
            indices = [self.parameters.index(p) for p in parameters]
            kwargs["labels"] = parameters
            x = self.x[:, indices] if self.x.ndim > 1 else self.x[indices]
        else:
            x = self.x
        fig = corner.corner(to_numpy(x), fig=fig, **kwargs)
        return fig

    def __str__(self):
        out = (
            f"No. samples: {len(self.x)}\nNo. parameters: {self.x.shape[-1]}\n"
        )
        return out

    def save(self, h5_file, path="samples", flat=False):
        """Save the samples to an HDF5 file.

        This converts the samples to numpy and then to a dictionary.

        Parameters
        ----------
        h5_file : h5py.File
            The HDF5 file to save to.
        path : str
            The path in the HDF5 file to save to.
        flat : bool
            If True, save the samples as a flat dictionary.
            If False, save the samples as a nested dictionary.
        """
        dictionary = self.to_numpy().to_dict(flat=flat)
        recursively_save_to_h5_file(h5_file, path, dictionary)

    def __len__(self):
        return len(self.x)

    def __getitem__(self, idx) -> BaseSamples:
        return self.__class__(
            x=self.x[idx],
            log_likelihood=self.log_likelihood[idx]
            if self.log_likelihood is not None
            else None,
            log_prior=self.log_prior[idx]
            if self.log_prior is not None
            else None,
            log_q=self.log_q[idx] if self.log_q is not None else None,
            parameters=self.parameters,
            dtype=self.dtype,
        )

    def __setitem__(self, idx, value: BaseSamples):
        raise NotImplementedError("Setting items is not supported")

    @classmethod
    def concatenate(cls, samples: list[BaseSamples]) -> BaseSamples:
        """Concatenate multiple Samples objects."""
        if not samples:
            raise ValueError("No samples to concatenate")
        if not all(s.parameters == samples[0].parameters for s in samples):
            raise ValueError("Parameters do not match")
        if not all(s.xp == samples[0].xp for s in samples):
            raise ValueError("Array namespaces do not match")
        if not all(s.dtype == samples[0].dtype for s in samples):
            raise ValueError("Dtypes do not match")
        xp = samples[0].xp
        return cls(
            x=xp.concatenate([s.x for s in samples], axis=0),
            log_likelihood=xp.concatenate(
                [s.log_likelihood for s in samples], axis=0
            )
            if all(s.log_likelihood is not None for s in samples)
            else None,
            log_prior=xp.concatenate([s.log_prior for s in samples], axis=0)
            if all(s.log_prior is not None for s in samples)
            else None,
            log_q=xp.concatenate([s.log_q for s in samples], axis=0)
            if all(s.log_q is not None for s in samples)
            else None,
            parameters=samples[0].parameters,
            dtype=samples[0].dtype,
        )

    @classmethod
    def from_samples(cls, samples: BaseSamples, **kwargs) -> BaseSamples:
        """Create a Samples object from a BaseSamples object."""
        xp = kwargs.pop("xp", samples.xp)
        device = kwargs.pop("device", samples.device)
        dtype = kwargs.pop("dtype", samples.dtype)
        if dtype is not None:
            dtype = resolve_dtype(dtype, xp)
        else:
            dtype = convert_dtype(samples.dtype, xp)
        return cls(
            x=samples.x,
            log_likelihood=samples.log_likelihood,
            log_prior=samples.log_prior,
            log_q=samples.log_q,
            parameters=samples.parameters,
            xp=xp,
            device=device,
            **kwargs,
        )

    def __getstate__(self):
        state = self.__dict__.copy()
        # replace xp (callable) with module name string
        if self.xp is not None:
            state["xp"] = (
                self.xp.__name__ if hasattr(self.xp, "__name__") else None
            )
        return state

    def __setstate__(self, state):
        # Restore xp by checking the namespace of x
        state["xp"] = array_namespace(state["x"])
        self.__dict__.update(state)


@dataclass
class Samples(BaseSamples):
    """Class for storing samples and corresponding weights.

    If :code:`xp` is not specified, all inputs will be converted to match
    the array type of :code:`x`.
    """

    log_w: Array = field(init=False)
    weights: Array = field(init=False)
    evidence: float = field(init=False)
    evidence_error: float = field(init=False)
    log_evidence: float | None = None
    log_evidence_error: float | None = None
    effective_sample_size: float = field(init=False)

    def __post_init__(self):
        super().__post_init__()

        if all(
            x is not None
            for x in [self.log_likelihood, self.log_prior, self.log_q]
        ):
            self.compute_weights()
        else:
            self.log_w = None
            self.weights = None
            self.evidence = None
            self.evidence_error = None
            self.effective_sample_size = None

    @property
    def efficiency(self):
        """Efficiency of the weighted samples.

        Defined as ESS / number of samples.
        """
        if self.log_w is None:
            raise RuntimeError("Samples do not contain weights!")
        return self.effective_sample_size / len(self.x)

    def compute_weights(self):
        """Compute the posterior weights."""
        self.log_w = self.log_likelihood + self.log_prior - self.log_q
        self.log_evidence = asarray(logsumexp(self.log_w), self.xp) - math.log(
            len(self.x)
        )
        self.weights = self.xp.exp(self.log_w)
        self.evidence = self.xp.exp(self.log_evidence)
        n = len(self.x)
        self.evidence_error = self.xp.sqrt(
            self.xp.sum((self.weights - self.evidence) ** 2) / (n * (n - 1))
        )
        self.log_evidence_error = self.xp.abs(
            self.evidence_error / self.evidence
        )
        log_w = self.log_w - self.xp.max(self.log_w)
        self.effective_sample_size = self.xp.exp(
            asarray(logsumexp(log_w) * 2 - logsumexp(log_w * 2), self.xp)
        )

    @property
    def scaled_weights(self):
        return self.xp.exp(self.log_w - self.xp.max(self.log_w))

    def rejection_sample(self, rng=None):
        if rng is None:
            rng = np.random.default_rng()
        log_u = asarray(
            np.log(rng.uniform(size=len(self.x))), self.xp, device=self.device
        )
        log_w = self.log_w - self.xp.max(self.log_w)
        accept = log_w > log_u
        return self.__class__(
            x=self.x[accept],
            log_likelihood=self.log_likelihood[accept],
            log_prior=self.log_prior[accept],
            dtype=self.dtype,
        )

    def to_dict(self, flat: bool = True):
        samples = dict(zip(self.parameters, self.x.T, strict=True))
        out = super().to_dict(flat=flat)
        other = {
            "log_w": self.log_w,
            "weights": self.weights,
            "evidence": self.evidence,
            "log_evidence": self.log_evidence,
            "evidence_error": self.evidence_error,
            "log_evidence_error": self.log_evidence_error,
            "effective_sample_size": self.effective_sample_size,
        }
        out.update(other)
        if flat:
            out.update(samples)
        else:
            out["samples"] = samples
        return out

    def plot_corner(self, include_weights: bool = True, **kwargs):
        kwargs = copy.deepcopy(kwargs)
        if (
            include_weights
            and self.weights is not None
            and "weights" not in kwargs
        ):
            kwargs["weights"] = to_numpy(self.scaled_weights)
        return super().plot_corner(**kwargs)

    def __str__(self):
        out = super().__str__()
        if self.log_evidence is not None:
            out += f"Log evidence: {self.log_evidence:.2f} +/- {self.log_evidence_error:.2f}\n"
        if self.log_w is not None:
            out += (
                f"Effective sample size: {self.effective_sample_size:.1f}\n"
                f"Efficiency: {self.efficiency:.2f}\n"
            )
        return out

    def to_namespace(self, xp):
        return self.__class__(
            x=asarray(self.x, xp),
            parameters=self.parameters,
            log_likelihood=asarray(self.log_likelihood, xp)
            if self.log_likelihood is not None
            else None,
            log_prior=asarray(self.log_prior, xp)
            if self.log_prior is not None
            else None,
            log_q=asarray(self.log_q, xp) if self.log_q is not None else None,
            log_evidence=asarray(self.log_evidence, xp)
            if self.log_evidence is not None
            else None,
            log_evidence_error=asarray(self.log_evidence_error, xp)
            if self.log_evidence_error is not None
            else None,
        )

    def to_numpy(self):
        return self.__class__(
            x=to_numpy(self.x),
            parameters=self.parameters,
            log_likelihood=to_numpy(self.log_likelihood)
            if self.log_likelihood is not None
            else None,
            log_prior=to_numpy(self.log_prior)
            if self.log_prior is not None
            else None,
            log_q=to_numpy(self.log_q) if self.log_q is not None else None,
            log_evidence=self.log_evidence
            if self.log_evidence is not None
            else None,
            log_evidence_error=self.log_evidence_error
            if self.log_evidence_error is not None
            else None,
        )

    def __getitem__(self, idx):
        sliced = super().__getitem__(idx)
        return self.__class__.from_samples(
            sliced,
            log_evidence=self.log_evidence,
            log_evidence_error=self.log_evidence_error,
        )


@dataclass
class SMCSamples(BaseSamples):
    beta: float | None = None
    """Temperature parameter for the current samples."""
    log_evidence: float | None = None
    """Log evidence estimate for the current samples."""
    log_evidence_error: float | None = None
    """Log evidence error estimate for the current samples."""

    def log_p_t(self, beta):
        log_p_T = self.log_likelihood + self.log_prior
        return (1 - beta) * self.log_q + beta * log_p_T

    def unnormalized_log_weights(self, beta: float) -> Array:
        return (self.beta - beta) * self.log_q + (beta - self.beta) * (
            self.log_likelihood + self.log_prior
        )

    def log_evidence_ratio(self, beta: float) -> float:
        log_w = self.unnormalized_log_weights(beta)
        return logsumexp(log_w) - math.log(len(self.x))

    def log_evidence_ratio_variance(self, beta: float) -> float:
        """Estimate the variance of the log evidence ratio using the delta method.

        Defined as Var(log Z) = Var(w) / (E[w])^2 where w are the unnormalized weights.
        """
        log_w = self.unnormalized_log_weights(beta)
        m = self.xp.max(log_w)
        u = self.xp.exp(log_w - m)
        mean_w = self.xp.mean(u)
        var_w = self.xp.var(u)
        return (
            var_w / (len(self) * (mean_w**2)) if mean_w != 0 else self.xp.nan
        )

    def log_weights(self, beta: float) -> Array:
        log_w = self.unnormalized_log_weights(beta)
        if self.xp.isnan(log_w).any():
            raise ValueError(f"Log weights contain NaN values for beta={beta}")
        log_evidence_ratio = logsumexp(log_w) - math.log(len(self.x))
        return log_w + log_evidence_ratio

    def resample(
        self,
        beta,
        n_samples: int | None = None,
        rng: np.random.Generator = None,
    ) -> "SMCSamples":
        if beta == self.beta and n_samples is None:
            logger.warning(
                "Resampling with the same beta value, returning identical samples"
            )
            return self
        if rng is None:
            rng = np.random.default_rng()
        if n_samples is None:
            n_samples = len(self.x)
        log_w = self.log_weights(beta)
        w = to_numpy(self.xp.exp(log_w - logsumexp(log_w)))
        idx = rng.choice(len(self.x), size=n_samples, replace=True, p=w)
        return self.__class__(
            x=self.x[idx],
            log_likelihood=self.log_likelihood[idx],
            log_prior=self.log_prior[idx],
            log_q=self.log_q[idx],
            beta=beta,
            dtype=self.dtype,
        )

    def __str__(self):
        out = super().__str__()
        if self.log_evidence is not None:
            out += f"Log evidence: {self.log_evidence:.2f}\n"
        return out

    def to_standard_samples(self):
        """Convert the samples to standard samples."""
        return Samples(
            x=self.x,
            log_likelihood=self.log_likelihood,
            log_prior=self.log_prior,
            xp=self.xp,
            parameters=self.parameters,
            log_evidence=self.log_evidence,
            log_evidence_error=self.log_evidence_error,
        )

    def __getitem__(self, idx):
        sliced = super().__getitem__(idx)
        return self.__class__.from_samples(
            sliced,
            beta=self.beta,
            log_evidence=self.log_evidence,
            log_evidence_error=self.log_evidence_error,
        )
