#!/usr/bin/env python3
"""Calibrate the crystallization probe against Run 2 checkpoint data.

Runs the probe against known Run 2 snapshots to find the interaction
where warning/critical alerts would have fired, and tunes parameters.

Calibration targets (based on Run 2 analysis):
  - Phase 2 ecology: compression 0.05-0.10, unique tapes 400-800
  - Crystallization onset: unique tapes collapse below 100 (~6.7M)
  - Full monoculture: compression 0.018, unique tapes 13 (~7.0M)

The probe should:
  - Warning: fire BEFORE unique_tapes < 100 (pre-crystallization)
  - Critical: fire BEFORE compression < 0.03 (crystallization in progress)
  - No false positives during stable Phase 2 ecology
"""

import json
import sys
from pathlib import Path

import numpy as np

from abiogenesis.probe import CrystallizationProbe


def load_run2_data(data_dir: str = "data/bff_runs/"):
    """Load Run 2 checkpoint tapes in order."""
    data_path = Path(data_dir)

    final_meta_path = data_path / "final_10000000_meta.json"
    if not final_meta_path.exists():
        print(f"No final metadata found in {data_dir}")
        return None, None

    with open(final_meta_path) as f:
        meta = json.load(f)

    # Load tapes from each checkpoint (sorted by interaction count)
    checkpoint_files = sorted(data_path.glob("checkpoint_*.npz"))
    checkpoint_tapes = []
    for cp_path in checkpoint_files:
        cp_meta_path = cp_path.parent / (cp_path.stem + "_meta.json")
        if not cp_meta_path.exists():
            continue
        with open(cp_meta_path) as f:
            cp_meta = json.load(f)
        data = np.load(cp_path)
        checkpoint_tapes.append((cp_meta["interaction_count"], data["tapes"]))

    # Sort by interaction count
    checkpoint_tapes.sort(key=lambda x: x[0])

    return meta, checkpoint_tapes


def calibrate():
    """Sweep probe parameters against Run 2 data."""
    meta, checkpoint_tapes = load_run2_data()
    if meta is None:
        print("Cannot calibrate without Run 2 data.")
        return

    # Determine ground truth from probe readings with a reference probe
    ref_probe = CrystallizationProbe(window_size=3, critical_threshold=-0.01)
    print("Run 2 diversity timeline:")
    print(f"{'interaction':>12}  {'H':>6}  {'compression':>11}  {'unique':>6}")
    print("-" * 45)

    first_unique_below_100 = None
    first_compression_below_003 = None

    for interaction, tapes in checkpoint_tapes:
        reading = ref_probe.measure(tapes, interaction)
        if interaction % 1_000_000 == 0 or reading.unique_tapes < 100:
            print(f"{interaction:>12,}  {reading.shannon_entropy:>6.3f}  "
                  f"{reading.compression_ratio:>11.4f}  {reading.unique_tapes:>6}")
        if reading.unique_tapes < 100 and first_unique_below_100 is None:
            first_unique_below_100 = interaction
        if reading.compression_ratio < 0.03 and first_compression_below_003 is None:
            first_compression_below_003 = interaction

    print()
    print(f"Ground truth crystallization events:")
    print(f"  Unique tapes < 100 at: {first_unique_below_100:,}")
    print(f"  Compression < 0.03 at: {first_compression_below_003:,}")
    print()

    # Sweep parameters
    window_sizes = (3, 4, 5, 7)
    thresholds = (-0.001, -0.01, -0.05, -0.1, -0.5, -1.0)

    best_config = None
    best_score = float("inf")

    print(f"{'ws':>3}  {'thresh':>7}  {'warning@':>12}  {'critical@':>12}  "
          f"{'FP':>3}  {'warn<100':>8}  {'crit<.03':>8}  {'result':>6}")
    print("-" * 80)

    for ws in window_sizes:
        for thresh in thresholds:
            probe = CrystallizationProbe(window_size=ws, critical_threshold=thresh)

            first_warning = None
            first_critical = None
            false_positives = 0

            for interaction, tapes in checkpoint_tapes:
                reading = probe.measure(tapes, interaction)

                if reading.alert_level in ("warning", "critical") and first_warning is None:
                    first_warning = interaction
                if reading.alert_level == "critical" and first_critical is None:
                    first_critical = interaction

                # False positive: alert when compression > 0.5
                if reading.alert_level != "none" and reading.compression_ratio > 0.5:
                    false_positives += 1

            warning_ok = (first_warning is not None and
                         first_unique_below_100 is not None and
                         first_warning <= first_unique_below_100)
            critical_ok = (first_critical is not None and
                          first_compression_below_003 is not None and
                          first_critical <= first_compression_below_003)

            score = false_positives * 100
            if not warning_ok:
                score += 50
            if not critical_ok:
                score += 30

            ok = warning_ok and critical_ok and false_positives == 0
            if score < best_score:
                best_score = score
                best_config = {
                    "window_size": ws,
                    "critical_threshold": thresh,
                    "first_warning": first_warning,
                    "first_critical": first_critical,
                    "false_positives": false_positives,
                    "warning_ok": warning_ok,
                    "critical_ok": critical_ok,
                }

            w_str = f"{first_warning:>12,}" if first_warning else "       never"
            c_str = f"{first_critical:>12,}" if first_critical else "       never"
            print(
                f"{ws:>3}  {thresh:>7.3f}  {w_str}  {c_str}  "
                f"{false_positives:>3}  {str(warning_ok):>8}  {str(critical_ok):>8}  "
                f"{'OK' if ok else 'FAIL':>6}"
            )

    print()
    print("=" * 80)
    if best_config:
        print(f"Best configuration (score={best_score:.1f}):")
        for k, v in best_config.items():
            print(f"  {k}: {v}")

        # Show alert timeline for best config
        print()
        print("Alert timeline (best config):")
        probe = CrystallizationProbe(
            window_size=best_config["window_size"],
            critical_threshold=best_config["critical_threshold"],
        )
        for interaction, tapes in checkpoint_tapes:
            reading = probe.measure(tapes, interaction)
            if reading.alert_level != "none":
                print(
                    f"  [{interaction:>10,}] "
                    f"H={reading.shannon_entropy:>6.3f}  "
                    f"cr={reading.compression_ratio:.4f}  "
                    f"dH={reading.entropy_delta:>+7.3f}  "
                    f"unique={reading.unique_tapes:>4}  "
                    f"alert={reading.alert_level}"
                )


if __name__ == "__main__":
    calibrate()
