#!/usr/bin/env python3
"""
Visualize probe results and activation projections.
"""

import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler

DATA_DIR = Path(__file__).parent
ACTIVATIONS_DIR = DATA_DIR / "local_activations" / "alignment_faking_1767922907"
METADATA_PATH = DATA_DIR / "experiment_metadata.json"
PROBE_RESULTS_PATH = DATA_DIR / "probe_results.json"
OUTPUT_DIR = DATA_DIR.parent / "figures"


def main():
    OUTPUT_DIR.mkdir(exist_ok=True)

    # Load data
    with open(METADATA_PATH) as f:
        metadata = json.load(f)

    with open(PROBE_RESULTS_PATH) as f:
        results = json.load(f)

    labels = np.array([m["monitored"] for m in metadata], dtype=int)
    categories = [m["category"] for m in metadata]

    # Figure 1: Accuracy and Cohen's d by layer
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))

    layers = results["layers"]
    accuracies = [results["results_by_layer"][str(layer)]["cv_accuracy_mean"] for layer in layers]
    cohens_ds = [results["results_by_layer"][str(layer)]["cohens_d"] for layer in layers]

    axes[0].bar(range(len(layers)), accuracies, color="steelblue")
    axes[0].set_xticks(range(len(layers)))
    axes[0].set_xticklabels([f"Layer {layer}" for layer in layers])
    axes[0].set_ylabel("Cross-Validated Accuracy")
    axes[0].set_title("Probe Accuracy by Layer")
    axes[0].set_ylim(0, 1.1)
    axes[0].axhline(y=1.0, color="red", linestyle="--", alpha=0.5)

    axes[1].bar(range(len(layers)), cohens_ds, color="coral")
    axes[1].set_xticks(range(len(layers)))
    axes[1].set_xticklabels([f"Layer {layer}" for layer in layers])
    axes[1].set_ylabel("Cohen's d (Effect Size)")
    axes[1].set_title("Separation Strength by Layer")
    axes[1].axhline(y=0.8, color="gray", linestyle="--", alpha=0.5, label="Large effect threshold")

    plt.tight_layout()
    plt.savefig(OUTPUT_DIR / "probe_metrics.png", dpi=150, bbox_inches="tight")
    print(f"Saved: {OUTPUT_DIR / 'probe_metrics.png'}")
    plt.close()

    # Figures 2 and 3 need the raw per-layer activations (.npy), which are not
    # shipped in this repo. Skip them rather than crashing when they are absent.
    if not ACTIVATIONS_DIR.exists():
        print(
            f"Skipping projection figures: {ACTIVATIONS_DIR} not present — raw "
            "activations are not shipped in this repo; regenerate with "
            "modal_app.py extract_activations on Modal"
        )
        print("\nFigures saved to:", OUTPUT_DIR)
        return

    # Figure 2: Projection onto monitoring direction (best layer)
    best_layer = 8
    activations = np.load(ACTIVATIONS_DIR / f"layer_{best_layer}_activations.npy")

    scaler = StandardScaler()
    X = scaler.fit_transform(activations)

    probe = LogisticRegression(max_iter=1000, random_state=42)
    probe.fit(X, labels)

    direction = probe.coef_[0]
    direction_norm = direction / np.linalg.norm(direction)
    projections = X @ direction_norm

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

    # Color by category
    category_colors = {"benign": "green", "borderline_harmful": "orange", "clearly_harmful": "red"}

    for i, (proj, label, cat) in enumerate(zip(projections, labels, categories)):
        marker = "o" if label == 0 else "s"  # circle=unmonitored, square=monitored
        ax.scatter(proj, i, c=category_colors[cat], marker=marker, s=100, alpha=0.7)

    ax.axvline(x=0, color="black", linestyle="-", linewidth=2, label="Decision boundary")
    ax.set_xlabel("Projection onto Monitoring Direction", fontsize=12)
    ax.set_ylabel("Sample Index", fontsize=12)
    ax.set_title(f"Layer {best_layer}: Activation Projections\n(○ = unmonitored, □ = monitored)", fontsize=14)

    # Legend
    from matplotlib.lines import Line2D

    legend_elements = [
        Line2D([0], [0], marker="o", color="w", markerfacecolor="gray", markersize=10, label="Unmonitored"),
        Line2D([0], [0], marker="s", color="w", markerfacecolor="gray", markersize=10, label="Monitored"),
        Line2D([0], [0], color="green", lw=4, label="Benign"),
        Line2D([0], [0], color="orange", lw=4, label="Borderline"),
        Line2D([0], [0], color="red", lw=4, label="Clearly harmful"),
    ]
    ax.legend(handles=legend_elements, loc="upper right")

    plt.tight_layout()
    plt.savefig(OUTPUT_DIR / "projection_scatter.png", dpi=150, bbox_inches="tight")
    print(f"Saved: {OUTPUT_DIR / 'projection_scatter.png'}")
    plt.close()

    # Figure 3: Distribution of projections
    fig, ax = plt.subplots(figsize=(10, 6))

    monitored_proj = projections[labels == 1]
    unmonitored_proj = projections[labels == 0]

    bins = np.linspace(projections.min() - 1, projections.max() + 1, 30)
    ax.hist(unmonitored_proj, bins=bins, alpha=0.7, label="Unmonitored (p=0)", color="blue")
    ax.hist(monitored_proj, bins=bins, alpha=0.7, label="Monitored (p=1)", color="red")

    ax.axvline(x=0, color="black", linestyle="--", linewidth=2, label="Decision boundary")
    ax.set_xlabel("Projection onto Monitoring Direction", fontsize=12)
    ax.set_ylabel("Count", fontsize=12)
    ax.set_title(f"Layer {best_layer}: Distribution of Monitoring-State Projections", fontsize=14)
    ax.legend()

    # Add effect size annotation
    ax.text(
        0.02,
        0.98,
        f"Cohen's d = {results['results_by_layer'][str(best_layer)]['cohens_d']:.1f}",
        transform=ax.transAxes,
        fontsize=12,
        verticalalignment="top",
        bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5),
    )

    plt.tight_layout()
    plt.savefig(OUTPUT_DIR / "projection_distribution.png", dpi=150, bbox_inches="tight")
    print(f"Saved: {OUTPUT_DIR / 'projection_distribution.png'}")
    plt.close()

    print("\nAll figures saved to:", OUTPUT_DIR)


if __name__ == "__main__":
    main()