from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Callable, Generic, Sequence, TypeVar
import math
import random

S = TypeVar("S")
A = TypeVar("A")


def normalize_logweights(logw: Sequence[float]) -> list[float]:
    if not logw:
        return []
    finite = [w for w in logw if math.isfinite(w)]
    if not finite:
        return [1.0 / len(logw)] * len(logw)
    m = max(finite)
    ws = [math.exp(w - m) if math.isfinite(w) else 0.0 for w in logw]
    z = sum(ws)
    if not math.isfinite(z) or z <= 0.0:
        return [1.0 / len(logw)] * len(logw)
    return [w / z for w in ws]


def effective_sample_size(probs: Sequence[float]) -> float:
    s2 = sum(float(p) ** 2 for p in probs)
    return 0.0 if s2 <= 0.0 else 1.0 / s2


def systematic_resample(probs: Sequence[float], rng: random.Random) -> list[int]:
    n = len(probs)
    if n == 0:
        return []
    cdf: list[float] = []
    total = 0.0
    for p in probs:
        total += float(p)
        cdf.append(total)
    if total <= 0.0 or not math.isfinite(total):
        return list(range(n))
    cdf = [x / total for x in cdf]
    cdf[-1] = 1.0
    u0 = rng.random() / n
    out: list[int] = []
    j = 0
    for i in range(n):
        u = u0 + i / n
        while j < n - 1 and u > cdf[j]:
            j += 1
        out.append(j)
    return out


@dataclass(frozen=True)
class Proposal(Generic[A]):
    action: A
    probability: float


@dataclass
class Particle(Generic[S]):
    state: S
    logw: float = 0.0
    history: list[Any] = field(default_factory=list)


@dataclass(frozen=True)
class SMCConfig:
    n_particles: int = 32
    n_steps: int = 3
    ess_resample_fraction: float = 0.5
    seed: int = 0


class ParticleFilter(Generic[S, A]):
    """Importance-corrected SMC runtime.

    `propose(state, t)` returns a discrete proposal distribution q(a|state).
    `transition(state, action, t)` produces the next state.
    `log_target_increment(prev, action, nxt, t)` supplies the incremental
    log target density. The importance update is log w += log target - log q.
    Resampling is never performed after the final weighting step, so the final
    posterior weights are preserved.
    """

    def __init__(self, config: SMCConfig | None = None):
        self.config = config or SMCConfig()

    def run(
        self,
        initial_state: S,
        propose: Callable[[S, int], Sequence[Proposal[A]]],
        transition: Callable[[S, A, int], S],
        log_target_increment: Callable[[S, A, S, int], float],
    ) -> list[Particle[S]]:
        cfg = self.config
        rng = random.Random(cfg.seed)
        particles = [Particle(state=initial_state) for _ in range(cfg.n_particles)]
        for t in range(cfg.n_steps):
            nxt: list[Particle[S]] = []
            for p in particles:
                proposals = [x for x in propose(p.state, t) if x.probability > 0.0]
                if not proposals:
                    raise ValueError("proposal distribution is empty")
                z = sum(x.probability for x in proposals)
                probs = [x.probability / z for x in proposals]
                r = rng.random()
                acc = 0.0
                j = len(proposals) - 1
                for i, pr in enumerate(probs):
                    acc += pr
                    if r <= acc:
                        j = i
                        break
                selected = proposals[j]
                q = probs[j]
                new_state = transition(p.state, selected.action, t)
                inc = float(log_target_increment(p.state, selected.action, new_state, t))
                nxt.append(Particle(new_state, p.logw + inc - math.log(q), p.history + [selected.action]))
            particles = nxt
            probs_now = normalize_logweights([p.logw for p in particles])
            if t < cfg.n_steps - 1 and effective_sample_size(probs_now) < cfg.ess_resample_fraction * len(particles):
                ids = systematic_resample(probs_now, rng)
                particles = [Particle(particles[i].state, 0.0, list(particles[i].history)) for i in ids]
        final_probs = normalize_logweights([p.logw for p in particles])
        for p, pr in zip(particles, final_probs):
            setattr(p, "posterior", pr)
        return particles