import numpy as np
from abc import ABC, abstractmethod
from genlm_grammar import Float
from arsenal.maths import logsumexp, sample_dict
from functools import cached_property
from genlm_control import EOS
from dataclasses import dataclass
from arsenal import colors

from hfppl import Model
from hfppl import smc_standard


@dataclass
class Sequences:
    """Container for sequence samples with their weights and probabilities.

    This class stores and processes the results of sequence generation, including the
    sequences themselves, their importance weights, and their log probabilities. It
    provides utilities for analyzing the samples, including effective sample
    size (ESS) calculations and posterior distribution estimation.

    Args:
        contexts (list): List of token sequences generated by the sampler.
        log_weights (list): Log importance weights for each sequence.
        log_probs (list): Log probabilities of each sequence under the proposal distribution.

    Attributes:
        size (int): Number of sequences in the container.
        logp (float): Sum of log probabilities across all sequences.
        log_total (float): Log of the sum of importance weights.
        log_ml (float): Log marginal likelihood estimate.
        log_normalized_weights (list): Log weights normalized to sum to 1.
        log_ess (float): Log of the effective sample size.
        ess (float): Effective sample size of the particle population.
    """

    contexts: list
    log_weights: list
    log_probs: list

    def __post_init__(self):
        assert len(self.contexts) == len(self.log_weights) == len(self.log_probs)

        self.size = len(self.contexts)
        self.logp = sum(self.log_probs)
        self.log_total = logsumexp(self.log_weights)
        max_weight = max(self.log_weights)
        if np.isfinite(max_weight):
            self.log_ml = (
                np.log(np.mean(np.exp(self.log_weights - max_weight))) + max_weight
            )
        else:
            self.log_ml = float("-inf")
        self.log_normalized_weights = self.log_weights - self.log_total
        self.log_ess = -logsumexp(2 * self.log_normalized_weights)
        self.ess = np.exp(self.log_ess)

    @cached_property
    def posterior(self):
        """Compute the estimated posterior distribution over sequences.

        The probability of a sequence corresponds to its normalized weight. The probabilities
        of duplicate sequences are summed.

        Returns:
            (Float.chart): A normalized chart mapping sequences to their posterior probabilities,
                sorted in descending order by probability.
        """
        posterior = Float.chart()
        for sequence, prob in zip(self.contexts, np.exp(self.log_normalized_weights)):
            posterior[tuple(sequence)] += prob
        return posterior.normalize().sort_descending()

    @property
    def normalized_weights(self):
        """Return exponential of normalized log weights."""
        return np.exp(self.log_normalized_weights)

    def __len__(self):
        return self.size

    def __iter__(self):
        return iter(zip(self.contexts, self.log_weights))

    def __getitem__(self, i):
        return self.contexts[i], self.log_weights[i]

    def __str__(self):
        return str(self.posterior)

    def _repr_html_(self):
        return self.posterior._repr_html_()

    def show(self):
        for p in sorted(self, reverse=True):
            print(p)


class SequenceModel(Model):
    def __init__(self, unit_sampler, critic=None, max_tokens=float("inf"), verbosity=0):
        assert max_tokens > 0

        super().__init__()
        self.token_ctx = []
        self.unit_sampler = unit_sampler
        self.max_tokens = max_tokens
        self.critic = critic
        self.logp = 0
        self.verbosity = verbosity
        self.twist_with_critic = True  # This flag is used to avoid running the critic at each step for IS (when ess_threshold = 0)

    async def start(self):
        start_w = await self.unit_sampler.start_weight()
        if start_w == float("-inf"):
            raise ValueError(
                "Start weight is -inf (log(0)). This is likely because a potential assigns zero weight to "
                "the empty sequence as a prefix, which violates the potential contract."
            )
        self.score(start_w)

    async def step(self):
        unit = await self.call(self.unit_sampler)
        self.token_ctx.append(unit)

        if self.critic and self.twist_with_critic:
            twist_amt = await self.critic.score(self.token_ctx)
            self.twist(twist_amt)

        if self.verbosity > 0:
            print(self.__repr__())

        self.max_tokens -= 1
        if self.max_tokens == 0 or self.token_ctx[-1] is EOS:
            self.finish()
            if self.critic:
                if not self.twist_with_critic:
                    twist_amt = await self.critic.score(self.token_ctx)
                self.score(twist_amt)
            return

    def __repr__(self):
        return (
            f"{self.weight:.2f}:\t"
            + colors.magenta % "["
            + (colors.magenta % "|").join(repr(y) for y in self.token_ctx)
            + colors.magenta % "]"
        )

    def string_for_serialization(self):
        return "|".join(repr(y) for y in self.token_ctx)

    def immutable_properties(self):
        return set(["unit_sampler", "critic"])


def _unpack_particles(particles):
    contexts, logws, logps = map(
        list,
        zip(
            *[
                (
                    p.token_ctx,
                    float("-inf") if np.isnan(p.weight) else p.weight,
                    p.logp,
                )
                for p in particles
            ]
        ),
    )
    return contexts, logws, logps


class SequenceSampler(ABC):
    """Abstract base class for sequence samplers."""

    def __init__(self, unit_sampler, critic=None, max_tokens=float("inf")):
        self.unit_sampler = unit_sampler
        self.critic = critic
        self.model = SequenceModel(unit_sampler, critic, max_tokens)

    @property
    def max_tokens(self):
        return self.model.max_tokens

    @max_tokens.setter
    def max_tokens(self, value):
        self.model.max_tokens = value

    @abstractmethod
    async def sample(self, context=None, draw=sample_dict):
        pass

    @abstractmethod
    async def infer(self):
        pass


class Importance(SequenceSampler):
    def __init__(self, unit_sampler, n_particles, critic=None, max_tokens=float("inf")):
        if n_particles < 1:
            raise ValueError("n_particles must be greater than 0")
        super().__init__(unit_sampler, critic, max_tokens)
        self.n_particles = n_particles

    def sample(self, context=None, draw=sample_dict):
        raise NotImplementedError("SMC does not support sampling")

    async def infer(self, **kwargs):
        try:
            original_critic_condition = self.model.twist_with_critic
            self.model.twist_with_critic = False

            particles = await smc_standard(
                model=self.model,
                n_particles=self.n_particles,
                ess_threshold=0,
                **kwargs,
            )
        finally:
            self.model.twist_with_critic = original_critic_condition

        contexts, logws, logps = _unpack_particles(particles)
        assert len(contexts) == len(logws) == len(logps)

        return Sequences(contexts, logws, logps)


class SMC(SequenceSampler):
    def __init__(
        self,
        unit_sampler,
        n_particles,
        ess_threshold,
        critic=None,
        max_tokens=float("inf"),
    ):
        if n_particles < 0:
            raise ValueError("n_particles must be greater than 0")
        if not 0 <= ess_threshold <= 1.0:
            raise ValueError("ess_threshold must be between 0 and 1.0")

        super().__init__(unit_sampler, critic, max_tokens)
        self.n_particles = n_particles
        self.ess_threshold = ess_threshold

    def sample(self, context=None, draw=sample_dict):
        # Eventually implement to support nested SMC.
        raise NotImplementedError("SMC does not support sampling")

    async def infer(self, **kwargs):
        particles = await smc_standard(
            model=self.model,
            n_particles=self.n_particles,
            ess_threshold=self.ess_threshold,
            **kwargs,
        )

        contexts, logws, logps = _unpack_particles(particles)
        assert len(contexts) == len(logws) == len(logps)

        return Sequences(contexts, logws, logps)
