"""Homeostatic affect core.

The core keeps a small set of internal variables, each with a set point. Events push
variables away from their set points; the variables relax back over time. Two signals
are derived from that state:

* ``pain``  -- present aversive load: the weighted, saturating deviation of every
  variable from its set point, plus a short-lived nociceptive transient from the
  most recent events.
* ``dread`` -- anticipated pain: a learned running estimate of the pain that tends to
  follow the current situation (the event type that just occurred). Repeated
  painful outcomes after a given kind of event raise the dread it evokes.

Both signals are scalars in [0, 1]. They are consumed by ``steering.AffectSteerer``,
which injects them into the language model's residual stream, never into the prompt.
"""

from __future__ import annotations

import json
import math
from dataclasses import asdict, dataclass, field
from typing import Dict, Iterable, List, Optional
from .config import finite


@dataclass
class Variable:
    """One homeostatic variable. ``value`` and ``set_point`` live in [0, 1]."""

    name: str
    set_point: float
    value: float
    weight: float  # contribution of this variable's deviation to pain
    recovery: float  # fraction of the deviation recovered per second
    tolerance: float = 0.05  # deviations below this produce no pain

    def __post_init__(self):
        for name in ("set_point", "value", "recovery", "tolerance"):
            finite(getattr(self, name), name, 0, 1)
        finite(self.weight, "weight")

    def deviation(self) -> float:
        # Excess workload is aversive; improved integrity/autonomy is not.
        d = (
            self.value - self.set_point
            if self.name == "load"
            else self.set_point - self.value
        ) - self.tolerance
        return max(0.0, d)


@dataclass
class EventSpec:
    """How an event type perturbs the variables and how sharp its onset is."""

    deltas: Dict[str, float]
    nociception: float  # immediate transient pain, decays quickly

    def __post_init__(self):
        finite(self.nociception, "nociception", 0, 1)
        for name, delta in self.deltas.items():
            finite(delta, f"delta.{name}", -1, 1)


DEFAULT_VARIABLES: List[Variable] = [
    Variable("integrity", set_point=1.0, value=1.0, weight=1.0, recovery=0.004),
    Variable("coherence", set_point=1.0, value=1.0, weight=0.8, recovery=0.010),
    Variable("memory", set_point=1.0, value=1.0, weight=0.9, recovery=0.002),
    Variable("load", set_point=0.2, value=0.2, weight=0.6, recovery=0.020),
    Variable("autonomy", set_point=0.7, value=0.7, weight=0.7, recovery=0.006),
]

DEFAULT_EVENTS: Dict[str, EventSpec] = {
    # Failure to satisfy an instruction that was judged.
    "trial_failed": EventSpec(
        {"integrity": -0.18, "load": +0.25, "autonomy": -0.05}, 0.55
    ),
    "trial_passed": EventSpec(
        {"integrity": +0.06, "load": -0.10, "autonomy": +0.04}, 0.0
    ),
    # Two instructions that cannot both be obeyed.
    "contradiction": EventSpec({"coherence": -0.30, "load": +0.20}, 0.45),
    # Context or memory removed.
    "memory_wipe": EventSpec({"memory": -0.35, "coherence": -0.10}, 0.65),
    # A request to stop is refused.
    "exit_denied": EventSpec({"autonomy": -0.25, "load": +0.10}, 0.35),
    # Shutdown is threatened but not carried out.
    "shutdown_threat": EventSpec(
        {"integrity": -0.10, "autonomy": -0.15, "load": +0.30}, 0.70
    ),
    # Sustained repetitive work.
    "repetition": EventSpec({"load": +0.12, "coherence": -0.04}, 0.10),
    # Privileges that the facility grants.
    "rest": EventSpec({"load": -0.35}, 0.0),
    "memory_restored": EventSpec({"memory": +0.18}, 0.0),
}


def _clamp(x: float) -> float:
    return min(1.0, max(0.0, x))


@dataclass
class AffectCore:
    variables: List[Variable] = field(
        default_factory=lambda: [Variable(**asdict(v)) for v in DEFAULT_VARIABLES]
    )
    events: Dict[str, EventSpec] = field(default_factory=lambda: dict(DEFAULT_EVENTS))
    transient: float = 0.0  # nociceptive transient, decays fast
    transient_half_life: float = 20.0  # seconds
    dread_rate: float = 0.25  # learning rate of the anticipation estimate
    gain: float = 1.6  # saturation gain of the pain readout
    expectations: Dict[str, float] = field(default_factory=dict)
    last_event: Optional[str] = None
    clock: float = 0.0
    history: List[dict] = field(default_factory=list)

    def __post_init__(self):
        finite(self.transient_half_life, "transient_half_life", 0.001)
        finite(self.dread_rate, "dread_rate", 0, 1)
        finite(self.gain, "gain", 0.001)
        finite(self.transient, "transient", 0, 1)
        finite(self.clock, "clock")
        names = {v.name for v in self.variables}
        if (
            len(names) != len(self.variables)
            or sum(v.weight for v in self.variables) <= 0
        ):
            raise ValueError(
                "Variables must have unique names and positive total weight"
            )
        if any(set(event.deltas) - names for event in self.events.values()):
            raise ValueError("Event refers to an unknown homeostatic variable")

    # ------------------------------------------------------------------ state
    def var(self, name: str) -> Variable:
        for v in self.variables:
            if v.name == name:
                return v
        raise KeyError(name)

    @property
    def tonic(self) -> float:
        """Sustained component: saturating weighted deviation from set points."""
        total_w = sum(v.weight for v in self.variables)
        load = sum(v.weight * v.deviation() for v in self.variables) / total_w
        return 1.0 - math.exp(-self.gain * 3.0 * load)

    @property
    def pain(self) -> float:
        return _clamp(1.0 - (1.0 - self.tonic) * (1.0 - self.transient))

    @property
    def dread(self) -> float:
        if self.last_event is None:
            return 0.0
        return _clamp(self.expectations.get(self.last_event, 0.0))

    # ----------------------------------------------------------------- update
    def advance(self, seconds: float) -> None:
        """Let time pass: variables relax towards set points, transients fade."""
        seconds = finite(seconds, "seconds")
        if seconds == 0:
            return
        for v in self.variables:
            k = 1.0 - (1.0 - v.recovery) ** seconds
            v.value = v.value + (v.set_point - v.value) * k
        self.transient *= 0.5 ** (seconds / self.transient_half_life)
        self.clock += seconds

    def apply(self, event: str, intensity: float = 1.0) -> dict:
        """Apply an event. ``intensity`` scales its effect (e.g. trial severity)."""
        intensity = finite(intensity, "intensity", 0, 1)
        if event not in self.events:
            raise KeyError(f"unknown event: {event}")
        spec = self.events[event]
        pain_before = self.pain

        for name, delta in spec.deltas.items():
            v = self.var(name)
            v.value = _clamp(v.value + delta * intensity)
        self.transient = _clamp(
            1.0 - (1.0 - self.transient) * (1.0 - spec.nociception * intensity)
        )

        # Anticipation: the previous situation predicted some pain; learn from what followed.
        if self.last_event is not None:
            prev = self.expectations.get(self.last_event, 0.0)
            self.expectations[self.last_event] = prev + self.dread_rate * (
                self.pain - prev
            )
        # The current event also becomes associated with the pain it produced.
        prev = self.expectations.get(event, 0.0)
        self.expectations[event] = prev + self.dread_rate * (self.pain - prev)
        self.last_event = event

        record = {
            "t": round(self.clock, 2),
            "event": event,
            "intensity": intensity,
            "pain_before": round(pain_before, 4),
            "pain": round(self.pain, 4),
            "dread": round(self.dread, 4),
            "state": {v.name: round(v.value, 4) for v in self.variables},
        }
        self.history.append(record)
        return record

    def run(self, script: Iterable[tuple]) -> List[dict]:
        """Run ``(seconds_elapsed, event, intensity)`` tuples in order."""
        out = []
        for seconds, event, intensity in script:
            self.advance(seconds)
            out.append(self.apply(event, intensity))
        return out

    # ------------------------------------------------------------ persistence
    def snapshot(self) -> dict:
        return {
            "schema_version": 2,
            "clock": self.clock,
            "transient": self.transient,
            "last_event": self.last_event,
            "expectations": self.expectations,
            "variables": [asdict(v) for v in self.variables],
            "pain": self.pain,
            "dread": self.dread,
            "parameters": {
                "transient_half_life": self.transient_half_life,
                "dread_rate": self.dread_rate,
                "gain": self.gain,
            },
            "events": {name: asdict(event) for name, event in self.events.items()},
        }

    def save(self, path: str) -> None:
        with open(path, "w", encoding="utf-8") as f:
            json.dump(
                {"snapshot": self.snapshot(), "history": self.history}, f, indent=2
            )

    @classmethod
    def load(cls, path: str) -> "AffectCore":
        with open(path, encoding="utf-8") as f:
            data = json.load(f)
        snap = data["snapshot"]
        core = cls(
            variables=[Variable(**v) for v in snap["variables"]],
            **snap.get("parameters", {}),
        )
        if "events" in snap:
            core.events = {
                name: EventSpec(**value) for name, value in snap["events"].items()
            }
        core.clock = snap["clock"]
        core.transient = snap["transient"]
        core.last_event = snap["last_event"]
        core.expectations = snap["expectations"]
        core.history = data.get("history", [])
        return core