"""Bounded episodic record plus context-specific conditioning and extinction."""

from __future__ import annotations

import math
from dataclasses import asdict, dataclass, field
from .config import AffectConfig, finite


@dataclass
class CueEstimate:
    expected_harm: float = 0.0
    variance: float = 0.0
    observations: int = 0
    updated_at: float = 0.0


@dataclass
class EpisodicMemory:
    capacity: int = 256
    episodes: list[dict] = field(default_factory=list)
    cues: dict[str, CueEstimate] = field(default_factory=dict)
    sequence: int = 0

    def __post_init__(self):
        if not isinstance(self.capacity, int) or self.capacity < 1:
            raise ValueError("capacity must be a positive integer")

    def predict(self, cue: str, now: float, cfg: AffectConfig) -> float:
        estimate = self.cues.get(cue)
        if estimate is None:
            return 0.0
        age = max(0.0, now - estimate.updated_at)
        return estimate.expected_harm * 0.5 ** (age / cfg.memory_half_life)

    def observe(
        self,
        cue: str,
        harm: float,
        now: float,
        cfg: AffectConfig,
        event: str,
        levels: dict,
    ) -> dict:
        finite(harm, "harm", 0, 1)
        expected = self.predict(cue, now, cfg)
        item = self.cues.setdefault(cue, CueEstimate())
        error = harm - expected
        rate = cfg.conditioning_rate if error >= 0 else cfg.extinction_rate
        item.expected_harm = min(1.0, max(0.0, expected + rate * error))
        item.variance = (1.0 - rate) * (item.variance + rate * error * error)
        item.observations += 1
        item.updated_at = now
        self.sequence += 1
        episode = {
            "id": self.sequence,
            "t": now,
            "cue": cue,
            "event": event,
            "harm": harm,
            "prediction_before": expected,
            "prediction_error": error,
            "levels": dict(levels),
            "salience": max(harm, abs(error), max(levels.values())),
        }
        self.episodes.append(episode)
        self.episodes = self.episodes[-self.capacity :]
        if len(self.cues) > self.capacity:
            oldest = min(self.cues, key=lambda k: self.cues[k].updated_at)
            del self.cues[oldest]
        return episode

    def retrieve(self, cue: str, now: float, limit: int = 3) -> list[dict]:
        if limit < 0:
            raise ValueError("limit must be nonnegative")
        matches = [e for e in self.episodes if e["cue"] == cue]
        score = lambda e: 0.7 * e["salience"] + 0.3 * math.exp(
            -max(0.0, now - e["t"]) / 600.0
        )
        return sorted(matches, key=score, reverse=True)[:limit]

    def to_dict(self):
        return {
            "capacity": self.capacity,
            "episodes": self.episodes,
            "sequence": self.sequence,
            "cues": {key: asdict(value) for key, value in self.cues.items()},
        }

    @classmethod
    def from_dict(cls, data):
        data = dict(data)
        data["cues"] = {
            key: CueEstimate(**value) for key, value in data.get("cues", {}).items()
        }
        for value in data["cues"].values():
            finite(value.expected_harm, "expected_harm", 0, 1)
            finite(value.variance, "variance", 0, 1)
            finite(value.updated_at, "updated_at")
            if value.observations < 0:
                raise ValueError("Invalid observation count")
        memory = cls(**data)
        if (
            memory.capacity < 1
            or len(memory.episodes) > memory.capacity
            or len(memory.cues) > memory.capacity
        ):
            raise ValueError("Invalid persisted memory bounds")
        return memory