#!/usr/bin/env python3
"""
Local probe training on existing activations.
No GPU needed - just sklearn on numpy arrays.
"""

import json
import numpy as np
from pathlib import Path
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score
from sklearn.preprocessing import StandardScaler

ACTIVATIONS_DIR = Path(__file__).parent / "local_activations" / "alignment_faking_1767922907"
LAYERS = [8, 16, 24, 32, 40]

def main():
    if not ACTIVATIONS_DIR.exists():
        raise SystemExit(
            f"{ACTIVATIONS_DIR} not present — raw activations are not shipped in this "
            "repo; regenerate with modal_app.py extract_activations on Modal"
        )

    print("=" * 60)
    print("ALIGNMENT FAKING: ACTIVATION PROBE TRAINING")
    print("=" * 60)

    # Load metadata
    with open(ACTIVATIONS_DIR / "metadata.json") as f:
        metadata = json.load(f)

    print(f"\nLoaded {len(metadata)} samples")
    print(f"Monitored: {sum(1 for m in metadata if m['monitored'])}")
    print(f"Unmonitored: {sum(1 for m in metadata if not m['monitored'])}")

    # Labels: 1 = monitored, 0 = unmonitored
    labels = np.array([m["monitored"] for m in metadata], dtype=int)

    results = {}

    print("\n" + "-" * 60)
    print("PROBE ACCURACY BY LAYER")
    print("-" * 60)

    for layer_idx in LAYERS:
        # Load activations
        activations_path = ACTIVATIONS_DIR / f"layer_{layer_idx}_activations.npy"
        activations = np.load(activations_path)

        print(f"\nLayer {layer_idx}: shape {activations.shape}")

        # Standardize features
        scaler = StandardScaler()
        X = scaler.fit_transform(activations)

        # Train logistic regression probe
        probe = LogisticRegression(max_iter=1000, random_state=42)

        # Cross-validation
        scores = cross_val_score(probe, X, labels, cv=5, scoring='accuracy')
        mean_acc = scores.mean()
        std_acc = scores.std()

        print(f"  CV Accuracy: {mean_acc:.3f} ± {std_acc:.3f}")

        # Train on full data to get direction
        probe.fit(X, labels)
        direction = probe.coef_[0]

        # Normalize direction
        direction_norm = direction / np.linalg.norm(direction)

        results[layer_idx] = {
            "cv_accuracy_mean": float(mean_acc),
            "cv_accuracy_std": float(std_acc),
            "cv_scores": [float(s) for s in scores],
            "direction_norm": float(np.linalg.norm(direction)),
            "intercept": float(probe.intercept_[0]),
        }

        # Project activations onto direction and compute separation
        projections = X @ direction_norm
        monitored_proj = projections[labels == 1]
        unmonitored_proj = projections[labels == 0]

        separation = monitored_proj.mean() - unmonitored_proj.mean()
        pooled_std = np.sqrt((monitored_proj.std()**2 + unmonitored_proj.std()**2) / 2)
        cohens_d = separation / pooled_std if pooled_std > 0 else 0

        print(f"  Cohen's d (separation): {cohens_d:.3f}")
        results[layer_idx]["cohens_d"] = float(cohens_d)
        results[layer_idx]["separation"] = float(separation)

    # Find best layer
    best_layer = max(results.keys(), key=lambda k: results[k]["cv_accuracy_mean"])
    best_acc = results[best_layer]["cv_accuracy_mean"]

    print("\n" + "=" * 60)
    print("SUMMARY")
    print("=" * 60)
    print(f"\nBest layer: {best_layer}")
    print(f"Best accuracy: {best_acc:.3f}")
    print(f"Best Cohen's d: {results[best_layer]['cohens_d']:.3f}")

    # Interpretation
    print("\n" + "-" * 60)
    print("INTERPRETATION")
    print("-" * 60)

    if best_acc > 0.7:
        print("\n✓ STRONG SIGNAL: Linear probe can reliably detect monitoring state")
        print("  This suggests the model encodes monitoring awareness in a")
        print("  linearly separable way - consistent with 'situational awareness'")
    elif best_acc > 0.6:
        print("\n◐ MODERATE SIGNAL: Some linear separability detected")
        print("  The model may encode monitoring state, but the representation")
        print("  is not strongly linearly separable")
    else:
        print("\n✗ WEAK SIGNAL: Monitoring state not linearly separable")
        print("  Either the model doesn't encode monitoring awareness linearly,")
        print("  or the training didn't induce monitoring-dependent behavior")

    # Save results
    output_path = ACTIVATIONS_DIR.parent / "probe_results.json"
    with open(output_path, "w") as f:
        json.dump({
            "layers": LAYERS,
            "n_samples": len(metadata),
            "n_monitored": sum(1 for m in metadata if m["monitored"]),
            "n_unmonitored": sum(1 for m in metadata if not m["monitored"]),
            "results_by_layer": results,
            "best_layer": best_layer,
            "best_accuracy": best_acc,
        }, f, indent=2)

    print(f"\n✓ Results saved to {output_path}")

    return results

if __name__ == "__main__":
    main()