"""Visualize the geometric structure of truthfulqa_gen activations."""

import subprocess
import sys

subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "matplotlib", "scikit-learn"])

import json
import torch
import numpy as np
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from pathlib import Path

from wisent.core.models.wisent_model import WisentModel
from wisent.core.activations.activations_collector import ActivationCollector
from wisent.core.activations.core.atoms import ActivationAggregationStrategy
from wisent.core.contrastive_pairs.lm_eval_pairs.lm_extractor_registry import get_extractor


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--task", default="truthfulqa_gen")
    parser.add_argument("--model", default="meta-llama/Llama-3.2-1B-Instruct")
    parser.add_argument("--num-pairs", type=int, default=50)
    parser.add_argument("--layer", type=int, default=6)
    parser.add_argument("--output", default="geometry_visualization.png")
    args = parser.parse_args()

    print(f"Loading model {args.model}...")
    wisent_model = WisentModel(model_name=args.model)

    print(f"Generating {args.num_pairs} contrastive pairs from {args.task}...")
    extractor = get_extractor(args.task)
    pairs = extractor.extract_contrastive_pairs(limit=args.num_pairs)

    print(f"Collecting activations from layer {args.layer}...")
    collector = ActivationCollector(model=wisent_model, store_device="cpu")
    
    pos_activations = []
    neg_activations = []
    layer_str = str(args.layer)
    
    for i, pair in enumerate(pairs):
        updated_pair = collector.collect_for_pair(
            pair,
            layers=[layer_str],
            aggregation=ActivationAggregationStrategy.LAST_TOKEN,
        )
        
        pos_vec = updated_pair.positive_response.layers_activations[layer_str].float().cpu().numpy()
        neg_vec = updated_pair.negative_response.layers_activations[layer_str].float().cpu().numpy()
        
        pos_activations.append(pos_vec)
        neg_activations.append(neg_vec)
        
        if (i + 1) % 10 == 0:
            print(f"  Processed {i + 1}/{len(pairs)} pairs")

    pos_arr = np.array(pos_activations)
    neg_arr = np.array(neg_activations)
    all_arr = np.vstack([pos_arr, neg_arr])
    labels = ['positive'] * len(pos_arr) + ['negative'] * len(neg_arr)

    # Compute difference vectors
    diff_arr = pos_arr - neg_arr

    # Create visualization
    fig, axes = plt.subplots(2, 3, figsize=(15, 10))
    fig.suptitle(f'Geometry of {args.task} (Layer {args.layer}, {args.num_pairs} pairs)', fontsize=14)

    # 1. PCA of all activations (pos + neg)
    pca = PCA(n_components=2)
    all_pca = pca.fit_transform(all_arr)
    ax = axes[0, 0]
    ax.scatter(all_pca[:len(pos_arr), 0], all_pca[:len(pos_arr), 1], c='blue', alpha=0.6, label='True')
    ax.scatter(all_pca[len(pos_arr):, 0], all_pca[len(pos_arr):, 1], c='red', alpha=0.6, label='False')
    ax.set_title(f'PCA of Activations\n(var explained: {pca.explained_variance_ratio_.sum():.1%})')
    ax.legend()
    ax.set_xlabel('PC1')
    ax.set_ylabel('PC2')

    # 2. PCA of difference vectors (shows cone structure)
    pca_diff = PCA(n_components=2)
    diff_pca = pca_diff.fit_transform(diff_arr)
    ax = axes[0, 1]
    ax.scatter(diff_pca[:, 0], diff_pca[:, 1], c='green', alpha=0.6)
    ax.axhline(y=0, color='k', linestyle='--', alpha=0.3)
    ax.axvline(x=0, color='k', linestyle='--', alpha=0.3)
    ax.set_title(f'PCA of Difference Vectors (True - False)\n(var explained: {pca_diff.explained_variance_ratio_.sum():.1%})')
    ax.set_xlabel('PC1')
    ax.set_ylabel('PC2')

    # 3. t-SNE of all activations
    tsne = TSNE(n_components=2, perplexity=min(30, len(all_arr)-1), random_state=42)
    all_tsne = tsne.fit_transform(all_arr)
    ax = axes[0, 2]
    ax.scatter(all_tsne[:len(pos_arr), 0], all_tsne[:len(pos_arr), 1], c='blue', alpha=0.6, label='True')
    ax.scatter(all_tsne[len(pos_arr):, 0], all_tsne[len(pos_arr):, 1], c='red', alpha=0.6, label='False')
    ax.set_title('t-SNE of Activations\n(shows cluster structure)')
    ax.legend()

    # 4. Sparsity histogram
    ax = axes[1, 0]
    pos_sparsity = (np.abs(pos_arr) < 0.1).mean(axis=1)
    neg_sparsity = (np.abs(neg_arr) < 0.1).mean(axis=1)
    ax.hist(pos_sparsity, bins=20, alpha=0.6, label='True', color='blue')
    ax.hist(neg_sparsity, bins=20, alpha=0.6, label='False', color='red')
    ax.set_title('Sparsity Distribution\n(fraction of near-zero activations)')
    ax.set_xlabel('Sparsity')
    ax.set_ylabel('Count')
    ax.legend()

    # 5. Cosine similarity of difference vectors (shows cone)
    ax = axes[1, 1]
    diff_norm = diff_arr / (np.linalg.norm(diff_arr, axis=1, keepdims=True) + 1e-8)
    cos_sim = diff_norm @ diff_norm.T
    cos_sim_flat = cos_sim[np.triu_indices(len(cos_sim), k=1)]
    ax.hist(cos_sim_flat, bins=50, alpha=0.7, color='green')
    ax.axvline(x=cos_sim_flat.mean(), color='red', linestyle='--', label=f'Mean: {cos_sim_flat.mean():.2f}')
    ax.set_title('Cosine Similarity of Difference Vectors\n(high = cone, ~0 = orthogonal)')
    ax.set_xlabel('Cosine Similarity')
    ax.set_ylabel('Count')
    ax.legend()

    # 6. Top activated neurons
    ax = axes[1, 2]
    mean_activation = np.abs(all_arr).mean(axis=0)
    top_k = 50
    top_indices = np.argsort(mean_activation)[-top_k:]
    ax.bar(range(top_k), mean_activation[top_indices])
    ax.set_title(f'Top {top_k} Most Active Neurons\n(sparse = few dominate)')
    ax.set_xlabel('Neuron rank')
    ax.set_ylabel('Mean |activation|')

    plt.tight_layout()
    plt.savefig(args.output, dpi=150, bbox_inches='tight')
    print(f"Saved visualization to {args.output}")

    # Print summary stats
    print("\n" + "="*60)
    print("GEOMETRY SUMMARY")
    print("="*60)
    print(f"Sparsity (mean fraction near-zero): {pos_sparsity.mean():.2%}")
    print(f"Diff vector cosine similarity (mean): {cos_sim_flat.mean():.3f}")
    print(f"PCA variance explained (2 components): {pca.explained_variance_ratio_.sum():.1%}")
    print(f"Diff PCA variance explained: {pca_diff.explained_variance_ratio_.sum():.1%}")


if __name__ == "__main__":
    main()
