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

import subprocess
import sys

# Install visualization dependencies
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 transformers import AutoModelForCausalLM, AutoTokenizer
from wisent.core.activations.core.collector import ActivationCollector
from wisent.core.contrastive_pairs.generation.lm_eval_pairs import generate_lm_eval_contrastive_pairs


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}...")
    model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.float16, device_map="auto")
    tokenizer = AutoTokenizer.from_pretrained(args.model)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    print(f"Generating {args.num_pairs} contrastive pairs from {args.task}...")
    pairs = generate_lm_eval_contrastive_pairs(
        task_name=args.task,
        model=model,
        tokenizer=tokenizer,
        num_pairs=args.num_pairs,
        prompt_strategy="chat_template",
    )

    print(f"Collecting activations from layer {args.layer}...")
    collector = ActivationCollector(model, tokenizer)
    
    pos_activations = []
    neg_activations = []
    
    for pair in pairs:
        pos_act = collector.collect_single(pair["positive"], layers=[args.layer])
        neg_act = collector.collect_single(pair["negative"], layers=[args.layer])
        
        pos_vec = pos_act[args.layer][-1, :].cpu().numpy()  # last token
        neg_vec = neg_act[args.layer][-1, :].cpu().numpy()
        
        pos_activations.append(pos_vec)
        neg_activations.append(neg_vec)

    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()
