telemetry.py

"""Telemetry module for RL-NPO.

Handles all logging, output writing, and visualization:
- Per-generation JSONL logging
- Run config JSON at start
- Final result JSON at end
- Reward curve PNG
"""

import json
import os
from datetime import datetime, timezone
from dataclasses import dataclass, field, asdict

import matplotlib
matplotlib.use("Agg")  # Non-interactive backend
import matplotlib.pyplot as plt
import numpy as np


@dataclass
class CandidateRecord:
    """A single scored candidate from one generation."""
    rank: int
    score: float
    delta: float
    text: str


@dataclass
class GenerationRecord:
    """Full record of one generation's results."""
    run_id: str
    generation: int
    timestamp: str
    parent_score: float
    candidates: list[dict]
    winner_text: str
    winner_score: float
    converged: bool


class Telemetry:
    """Manages telemetry output for an RL-NPO run.

    Parameters
    ----------
    run_id : str
        Unique run identifier.
    output_dir : str
        Root output directory (e.g., "runs").
    """

    def __init__(self, run_id: str, output_dir: str = "runs"):
        self.run_id = run_id
        self.run_dir = os.path.join(output_dir, run_id)
        os.makedirs(self.run_dir, exist_ok=True)

        self.jsonl_path = os.path.join(self.run_dir, "telemetry.jsonl")
        self.result_path = os.path.join(self.run_dir, "result.json")
        self.config_path = os.path.join(self.run_dir, "run_config.json")
        self.curve_path = os.path.join(self.run_dir, "reward_curve.png")

        self.reward_curve: list[float] = []

    def write_run_config(
        self,
        input_text: str,
        target_roi: str,
        target_vector_path: str,
        roi_mask_vertex_count: int,
        generations: int,
        mutations: int,
        openrouter_model: str,
    ) -> None:
        """Write run configuration JSON at the start of a run."""
        import torch

        # Detect hardware
        hardware = {"inference_mode": "text-only-cpu"}
        if torch.cuda.is_available():
            hardware = {
                "gpu": torch.cuda.get_device_name(0),
                "vram_gb": round(torch.cuda.get_device_properties(0).total_memory / 1e9, 1),
                "inference_mode": "text-only-gpu",
            }

        config = {
            "run_id": self.run_id,
            "started_at": datetime.now(timezone.utc).isoformat(),
            "input_text": input_text[:500] + ("..." if len(input_text) > 500 else ""),
            "target_roi": target_roi,
            "target_vector_path": target_vector_path,
            "roi_mask_vertex_count": roi_mask_vertex_count,
            "generations": generations,
            "mutations": mutations,
            "openrouter_model": openrouter_model,
            "tribe_model": "facebook/tribev2",
            "hardware": hardware,
        }
        with open(self.config_path, "w", encoding="utf-8") as f:
            json.dump(config, f, indent=2)

    def log_generation(
        self,
        generation: int,
        parent_score: float,
        candidates: list[dict],
        winner_text: str,
        winner_score: float,
        converged: bool,
    ) -> None:
        """Append one generation's results to the JSONL telemetry file.

        Parameters
        ----------
        generation : int
            Generation number (0-indexed).
        parent_score : float
            Score of the parent text entering this generation.
        candidates : list[dict]
            List of candidate dicts with keys: rank, score, delta, text.
        winner_text : str
            Text of the highest-scoring candidate.
        winner_score : float
            Score of the winning candidate.
        converged : bool
            Whether convergence was detected.
        """
        self.reward_curve.append(winner_score)

        record = {
            "run_id": self.run_id,
            "generation": generation,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "parent_score": round(parent_score, 6),
            "candidates": [
                {
                    "rank": c["rank"],
                    "score": round(c["score"], 6),
                    "delta": round(c["delta"], 6),
                    "text": c["text"][:300] + ("..." if len(c["text"]) > 300 else ""),
                }
                for c in candidates
            ],
            "winner_text": winner_text[:500] + ("..." if len(winner_text) > 500 else ""),
            "winner_score": round(winner_score, 6),
            "converged": converged,
        }

        with open(self.jsonl_path, "a", encoding="utf-8") as f:
            f.write(json.dumps(record) + "\n")

    def write_result(
        self,
        target_roi: str,
        input_text: str,
        output_text: str,
        score_initial: float,
        score_final: float,
        generations_run: int,
        converged: bool,
        convergence_generation: int | None,
        run_config: dict,
    ) -> None:
        """Write the final result JSON at the end of a run."""
        result = {
            "run_id": self.run_id,
            "target_roi": target_roi,
            "input_text": input_text,
            "output_text": output_text,
            "score_initial": round(score_initial, 6),
            "score_final": round(score_final, 6),
            "score_delta": round(score_final - score_initial, 6),
            "generations_run": generations_run,
            "converged": converged,
            "convergence_generation": convergence_generation,
            "reward_curve": [round(s, 6) for s in self.reward_curve],
            "reward_curve_png": self.curve_path,
            "telemetry_path": self.jsonl_path,
            "run_config": run_config,
        }

        with open(self.result_path, "w", encoding="utf-8") as f:
            json.dump(result, f, indent=2)

    def plot_reward_curve(self) -> str:
        """Generate and save a reward curve PNG.

        Returns
        -------
        str
            Path to the saved PNG.
        """
        if not self.reward_curve:
            return self.curve_path

        fig, ax = plt.subplots(figsize=(10, 5))

        generations = list(range(len(self.reward_curve)))
        scores = self.reward_curve

        # Main curve
        ax.plot(
            generations, scores,
            color="#6366f1", linewidth=2.5, marker="o",
            markersize=8, markerfacecolor="#818cf8",
            markeredgecolor="#4f46e5", markeredgewidth=1.5,
            zorder=3,
        )

        # Fill under curve
        ax.fill_between(
            generations, scores,
            alpha=0.15, color="#6366f1",
        )

        # Annotations
        if len(scores) > 0:
            ax.annotate(
                f"{scores[0]:.4f}",
                (0, scores[0]),
                textcoords="offset points",
                xytext=(10, 10),
                fontsize=9,
                color="#6b7280",
            )
            ax.annotate(
                f"{scores[-1]:.4f}",
                (len(scores) - 1, scores[-1]),
                textcoords="offset points",
                xytext=(10, -15),
                fontsize=9,
                fontweight="bold",
                color="#4f46e5",
            )

        delta = scores[-1] - scores[0] if len(scores) > 1 else 0
        ax.set_title(
            f"RL-NPO Reward Curve — {self.run_id}\n"
            f"Δ = {delta:+.4f}  |  {len(scores)} generations",
            fontsize=13, fontweight="bold", pad=15,
        )
        ax.set_xlabel("Generation", fontsize=11)
        ax.set_ylabel("Cosine Similarity (ROI)", fontsize=11)

        ax.grid(True, alpha=0.3, linestyle="--")
        ax.set_xlim(-0.2, len(scores) - 0.8)

        # Style
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)

        plt.tight_layout()
        plt.savefig(self.curve_path, dpi=150, bbox_inches="tight")
        plt.close(fig)

        return self.curve_path