"""
episode.py — defines the core episodic unit for HIEROMEM.
"""

from __future__ import annotations
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any
import numpy as np
import networkx as nx
import uuid


@dataclass
class Episode:
    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    e: np.ndarray = field(default_factory=lambda: np.zeros(1))
    G: nx.DiGraph = field(default_factory=nx.DiGraph)
    tocc: datetime = field(default_factory=datetime.now)
    meta: dict[str, Any] = field(default_factory=dict)
    parent: str | None = None
    trust: float = 1.0

    # ------------------------------------------------------------------ #
    # Validation and representation
    # ------------------------------------------------------------------ #
    def validate(self) -> None:
        if not (0.0 <= self.trust <= 1.0):
            raise ValueError(
                f"Invalid trust {self.trust} for episode {self.id}")
        if self.e is None or not isinstance(self.e, np.ndarray):
            raise TypeError("Embedding vector (e) must be a NumPy array.")
        if self.e.ndim != 1:
            raise ValueError("Embedding vector must be 1-dimensional.")
        if not isinstance(self.G, nx.DiGraph):
            raise TypeError("Scene graph (G) must be a NetworkX DiGraph.")

    def __repr__(self) -> str:
        return f"<Episode {self.id[:8]} | trust={self.trust:.2f}, nodes={len(self.G)}, dim={len(self.e)}>"

    # ------------------------------------------------------------------ #
    # Utility helpers
    # ------------------------------------------------------------------ #
    def age(self) -> float:
        """Return episode age in seconds."""
        return (datetime.now() - self.tocc).total_seconds()

    def clone(self, new_id: bool = True) -> Episode:
        """Return a deep copy.

        Preserves the original timestamp when keeping the same ID.
        """
        G_copy = self.G.copy(as_view=False)
        return Episode(
            id=str(uuid.uuid4()) if new_id else self.id,
            e=self.e.copy(),
            G=G_copy,
            tocc=datetime.now() if new_id else self.tocc,
            meta=self.meta.copy(),
            parent=self.parent,
            trust=self.trust,
        )

    def update_trust(self, used: bool, lam: float) -> None:
        """Update trust via an exponential moving average."""

        lam = float(min(max(lam, 1e-6), 0.999999))
        target = 1.0 if used else 0.0
        self.trust = float(lam * self.trust + (1.0 - lam) * target)
        self.trust = float(max(0.0, min(1.0, self.trust)))

    def similarity(self, other: Episode) -> float:
        """Compute cosine similarity."""
        a, b = self.e, other.e
        return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))

    # ------------------------------------------------------------------ #
    # Serialization
    # ------------------------------------------------------------------ #
    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "e": self.e.tolist(),
            "G": {
                "nodes": [
                    {"id": node_id, **dict(self.G.nodes[node_id])}
                    for node_id in self.G.nodes
                ],
                "edges": [
                    {"u": u, "v": v, **dict(self.G.edges[u, v])}
                    for u, v in self.G.edges
                ],
            },
            "tocc": self.tocc.isoformat(),
            "meta": self.meta,
            "parent": self.parent,
            "trust": self.trust,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Episode:
        e = np.array(data["e"], dtype=np.float32)
        G_data = data["G"]
        G = nx.DiGraph()
        for n in G_data.get("nodes", []):
            nid = n.pop("id")
            G.add_node(nid, **n)
        for edge in G_data.get("edges", []):
            u, v = edge.pop("u"), edge.pop("v")
            G.add_edge(u, v, **edge)
        return cls(
            id=data.get("id", str(uuid.uuid4())),
            e=e,
            G=G,
            tocc=datetime.fromisoformat(data["tocc"]),
            meta=data.get("meta", {}),
            parent=data.get("parent"),
            trust=float(data.get("trust", 1.0)),
        )
