import math
import struct
import hashlib
import random
from typing import List, Tuple


class BeliefSketch:
    """A tiny particle-filter based belief sketch for a single scalar variable.

    Purpose: maintain a compact, mission-tuned belief for a scalar safety variable
    (e.g., battery, obstacle distance) and emit a probability-of-violation (PoV)
    and an entropy summary. The implementation is intentionally small and
    deterministic when a seed is provided so CI/edge tests are reproducible.

    This is not a full particle filter for production use, but a compact
    lightweight sketch suitable for the GuardRail.Space MVP and unit tests.
    """

    def __init__(
        self,
        name: str,
        num_particles: int = 64,
        prior_mean: float = 0.0,
        prior_std: float = 1.0,
        seed: int = None,
    ):
        self.name = name
        self.num_particles = max(8, int(num_particles))
        self.rng = random.Random(seed)
        # initialize particles from a normal prior
        self.particles: List[float] = [
            self.rng.gauss(prior_mean, prior_std) for _ in range(self.num_particles)
        ]
        # equal weights (implicit)

    def update(self, observation: float, obs_std: float = 1.0) -> None:
        """Update belief with a scalar observation using likelihood weighting
        followed by multinomial resampling and small jitter to avoid particle
        degeneracy. Deterministic when seed provided at construction.
        """
        if obs_std <= 0:
            obs_std = 1e-6
        # quick check: if observation is far outside current particle spread,
        # reinitialize particles around the observation so the sketch remains
        # responsive to strong, previously-unseen signals.
        n = len(self.particles)
        mean0 = sum(self.particles) / n if n else 0.0
        var0 = sum((p - mean0) ** 2 for p in self.particles) / n if n else 0.0
        std0 = math.sqrt(max(var0, 1e-12))
        if abs(observation - mean0) > max(3.0 * std0, 3.0 * obs_std):
            self.particles = [self.rng.gauss(observation, max(1e-6, obs_std * 0.5)) for _ in range(self.num_particles)]
            return

        # compute weights proportional to Gaussian likelihood
        weights = []
        two_var = 2.0 * (obs_std ** 2)
        for p in self.particles:
            # gaussian likelihood unnormalized
            w = math.exp(-((observation - p) ** 2) / two_var)
            weights.append(w)
        total = sum(weights)
        # If likelihoods underflow to (near) zero, treat the observation as
        # dominant and reinitialize particles around the observation. This
        # keeps the sketch responsive to strong signals and avoids a degenerate
        # uniform-resample when all weights are effectively zero.
        if total <= 1e-12:
            self.particles = [self.rng.gauss(observation, max(1e-6, obs_std * 0.5)) for _ in range(self.num_particles)]
            return
        probs = [w / total for w in weights]

        # multinomial resampling
        cumulative = []
        c = 0.0
        for p in probs:
            c += p
            cumulative.append(c)

        new_particles = []
        for _ in range(self.num_particles):
            u = self.rng.random()
            # find first cumulative >= u
            for idx, c in enumerate(cumulative):
                if u <= c:
                    new_particles.append(self.particles[idx])
                    break
            else:
                new_particles.append(self.particles[-1])

        # jitter with small gaussian noise proportional to obs_std
        jitter_scale = max(1e-3, obs_std * 0.01)
        self.particles = [p + self.rng.gauss(0.0, jitter_scale) for p in new_particles]

    def pov(self, threshold: float, operator: str = ">") -> float:
        """Return probability-of-violation (PoV) that the scalar meets the
        violation predicate defined by (operator, threshold). Supported
        operators: '>' (default), '<', '>=', '<='.
        """
        if operator == ">":
            count = sum(1 for p in self.particles if p > threshold)
        elif operator == "<":
            count = sum(1 for p in self.particles if p < threshold)
        elif operator == ">=":
            count = sum(1 for p in self.particles if p >= threshold)
        elif operator == "<=":
            count = sum(1 for p in self.particles if p <= threshold)
        else:
            raise ValueError(f"unsupported operator: {operator}")
        return float(count) / float(len(self.particles))

    def entropy(self) -> float:
        """Return a simple continuous-entropy proxy (Gaussian entropy using
        empirical variance). This is cheap and stable for small sketches.
        """
        n = len(self.particles)
        if n == 0:
            return 0.0
        mean = sum(self.particles) / n
        var = sum((p - mean) ** 2 for p in self.particles) / n
        var = max(var, 1e-12)
        # differential entropy of Gaussian: 0.5*ln(2*pi*e*var)
        return 0.5 * math.log(2 * math.pi * math.e * var)

    def summarize(self, threshold: float, operator: str = ">") -> dict:
        """Return a compact summary dictionary with PoV and entropy.
        """
        return {"name": self.name, "pov": self.pov(threshold, operator), "entropy": self.entropy()}

    def serialize(self) -> bytes:
        """Produce a tiny deterministic byte fingerprint of the sketch. Format:
        8 bytes mean (double) + 8 bytes variance (double) + 8 bytes crc64-like
        truncated using sha256. This keeps the sketch small (<64 bytes).
        """
        n = len(self.particles)
        mean = sum(self.particles) / n if n else 0.0
        var = sum((p - mean) ** 2 for p in self.particles) / n if n else 0.0
        packed = struct.pack("!dd", mean, var)
        digest = hashlib.sha256(packed).digest()[:16]
        return packed + digest