"""Structured telemetry utilities for the TNFR–Riemann program.

This module centralises the schema for experimental artifacts generated by
`benchmarks/riemann_program.py` and related notebooks.  Keeping a single
source of truth makes it easier to satisfy invariants #5 (Structural
Metrology) and #6 (Reproducible Dynamics) while preserving the academic
memo tone outlined in `theory/TNFR_RIEMANN_RESEARCH_NOTES.md`.
"""
from __future__ import annotations

import json
import math
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from statistics import fmean
from typing import Any, Iterable, Sequence

from ..physics.fields import (  # type: ignore
    compute_phase_curvature,
    compute_phase_gradient,
    compute_structural_potential,
    estimate_coherence_length,
)

__all__ = [
    "RiemannTelemetryRecord",
    "compute_field_aggregates",
    "write_telemetry_records",
]

@dataclass
class RiemannTelemetryRecord:
    """Minimal JSON-serialisable record for Riemann program experiments."""

    graph_size: int
    sigma: float
    min_eigenvalue: float
    eigenvalues: Sequence[float]
    estimate_sigma_critical: float | None = None
    seed: int | None = None
    notes: str | None = None
    phi_s_mean: float | None = None
    phi_s_max: float | None = None
    phase_gradient_mean: float | None = None
    phase_gradient_max: float | None = None
    phase_curvature_mean: float | None = None
    phase_curvature_max: float | None = None
    coherence_length: float | None = None
    created_at: str = field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )

    def to_dict(self) -> dict:
        rec = asdict(self)
        rec["eigenvalues"] = list(self.eigenvalues)
        return rec

def _aggregate(values: dict[Any, float]) -> tuple[float | None, float | None]:
    if not values:
        return None, None
    abs_values = [abs(v) for v in values.values()]
    if not abs_values:
        return None, None
    return float(fmean(abs_values)), float(max(abs_values, key=abs))

def compute_field_aggregates(G: Any) -> dict[str, float | None]:
    """Compute Φ_s, |∇φ|, K_φ, ξ_C aggregates for a TNFR graph."""

    phi_s = compute_structural_potential(G)
    grad = compute_phase_gradient(G)
    curvature = compute_phase_curvature(G)
    coherence = estimate_coherence_length(G)

    phi_mean, phi_max = _aggregate(phi_s)
    grad_mean, grad_max = _aggregate(grad)
    curv_mean, curv_max = _aggregate(curvature)

    coherence_clean = None if math.isnan(coherence) else float(coherence)

    return {
        "phi_s_mean": phi_mean,
        "phi_s_max": phi_max,
        "phase_gradient_mean": grad_mean,
        "phase_gradient_max": grad_max,
        "phase_curvature_mean": curv_mean,
        "phase_curvature_max": curv_max,
        "coherence_length": coherence_clean,
    }

def write_telemetry_records(
    records: Iterable[RiemannTelemetryRecord],
    export_path: str | Path,
) -> Path:
    """Write records to newline-delimited JSON and return the output path."""

    path = Path(export_path)
    path.parent.mkdir(parents=True, exist_ok=True)

    # Convert to list to avoid one-shot iterables when writing twice later.
    cached: list[RiemannTelemetryRecord] = list(records)
    if not cached:
        raise ValueError("No telemetry records provided")

    with path.open("w", encoding="utf-8") as fp:
        for record in cached:
            json.dump(record.to_dict(), fp)
            fp.write("\n")

    return path
