"""Metric collector for causal discovery performance tracking.

Records SHD, edge F1, orientation accuracy, and other metrics
across validation runs. Stdlib-only.
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Protocol


class _MetricStore(Protocol):
    def save_metric(self, metrics: DiscoveryMetrics) -> None: ...


@dataclass
class DiscoveryMetrics:
    """Metrics from a single causal discovery run."""

    run_id: str
    timestamp: float = field(default_factory=time.time)
    algorithm: str = ""
    dataset_size: int = 0
    num_variables: int = 0
    # Structure metrics
    shd: int = 0  # Structural Hamming Distance
    edge_precision: float = 0.0
    edge_recall: float = 0.0
    edge_f1: float = 0.0
    orientation_accuracy: float = 0.0
    # Performance metrics
    execution_time_ms: float = 0.0
    # Config
    config: dict = field(default_factory=dict)

    @property
    def edge_f1_computed(self) -> float:
        """Compute F1 from precision and recall."""
        if self.edge_precision + self.edge_recall == 0:
            return 0.0
        return (
            2 * (self.edge_precision * self.edge_recall) / (self.edge_precision + self.edge_recall)
        )

    def to_dict(self) -> dict:
        """Serialize to dictionary."""
        return {
            "run_id": self.run_id,
            "timestamp": self.timestamp,
            "algorithm": self.algorithm,
            "dataset_size": self.dataset_size,
            "num_variables": self.num_variables,
            "shd": self.shd,
            "edge_precision": round(self.edge_precision, 4),
            "edge_recall": round(self.edge_recall, 4),
            "edge_f1": round(self.edge_f1, 4),
            "orientation_accuracy": round(self.orientation_accuracy, 4),
            "execution_time_ms": round(self.execution_time_ms, 2),
            "config": self.config,
        }


class MetricCollector:
    """Collects and stores discovery metrics across runs."""

    def __init__(
        self,
        storage_path: str | Path | None = None,
        db: _MetricStore | None = None,
    ) -> None:
        self._runs: list[DiscoveryMetrics] = []
        self._storage_path = Path(storage_path) if storage_path else None
        self._db = db

    def record(self, metrics: DiscoveryMetrics) -> None:
        """Record a discovery run's metrics."""
        self._runs.append(metrics)
        if self._storage_path:
            self._persist()
        if self._db is not None:
            self._db.save_metric(metrics)

    def latest(self, n: int = 1) -> list[DiscoveryMetrics]:
        """Return the most recent n runs."""
        return self._runs[-n:]

    def by_algorithm(self, algorithm: str) -> list[DiscoveryMetrics]:
        """Filter runs by algorithm name."""
        return [r for r in self._runs if r.algorithm == algorithm]

    def best_run(self, metric: str = "edge_f1") -> DiscoveryMetrics | None:
        """Return the run with the best value for the given metric."""
        if not self._runs:
            return None
        return max(self._runs, key=lambda r: getattr(r, metric, 0))

    def summary(self) -> dict:
        """Aggregate summary across all recorded runs."""
        if not self._runs:
            return {
                "total_runs": 0,
                "algorithms": [],
                "best_shd": None,
                "best_f1": None,
            }
        algorithms = sorted({r.algorithm for r in self._runs})
        best_shd = min(self._runs, key=lambda r: r.shd)
        best_f1 = max(self._runs, key=lambda r: r.edge_f1)
        return {
            "total_runs": len(self._runs),
            "algorithms": algorithms,
            "best_shd": best_shd.shd,
            "best_f1": round(best_f1.edge_f1, 4),
            "avg_execution_ms": round(
                sum(r.execution_time_ms for r in self._runs) / len(self._runs),
                2,
            ),
        }

    def _persist(self) -> None:
        """Write all runs to storage."""
        if self._storage_path:
            self._storage_path.parent.mkdir(parents=True, exist_ok=True)
            data = [r.to_dict() for r in self._runs]
            self._storage_path.write_text(json.dumps(data, indent=2))

    def load(self, path: str | Path) -> None:
        """Load runs from a JSON file."""
        p = Path(path)
        if p.exists():
            data = json.loads(p.read_text())
            for d in data:
                config = d.pop("config", {})
                self._runs.append(DiscoveryMetrics(**d, config=config))

    @property
    def runs(self) -> list[DiscoveryMetrics]:
        """Return a copy of all recorded runs."""
        return list(self._runs)
