"""Visualization Gallery for Zwiad.

Creates publication-quality figures for the paper.
"""

import argparse
from pathlib import Path
from typing import Dict, List

try:
    import matplotlib.pyplot as plt
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False

try:
    from sklearn.manifold import TSNE
    HAS_SKLEARN = True
except ImportError:
    HAS_SKLEARN = False

from wisent.examples.scripts.visualization_gallery_helpers import (
    s3_upload_file,
    load_diagnosis_results,
    select_representative_benchmarks,
    create_tsne_plot,
)
from wisent.examples.scripts.visualization_gallery_figures import (
    create_hero_figure,
    create_layer_accuracy_curves,
)


def create_tsne_gallery(
    model: "WisentModel",
    selected_benchmarks: Dict[str, List[str]],
    output_path: Path,
    model_name: str,
) -> None:
    """
    Create t-SNE gallery figure.
    
    Layout: 2x3 grid showing examples from each diagnosis type
    
    Args:
        model: WisentModel instance
        selected_benchmarks: Dict with benchmark names by diagnosis
        output_path: Where to save figure
        model_name: Model name for title
    """
    if not HAS_MATPLOTLIB or not HAS_SKLEARN:
        print("Skipping t-SNE gallery: required packages not installed")
        return
    
    from wisent.core.activations import ExtractionStrategy
    from wisent.core.activations.activation_cache import ActivationCache, collect_and_cache_activations
    from lm_eval.tasks import TaskManager
    from wisent.core.contrastive_pairs.lm_eval_pairs.lm_task_pairs_generation import lm_build_contrastive_pairs
    
    fig, axes = plt.subplots(2, 3, figsize=(12, 8))
    fig.suptitle(f't-SNE Visualization Gallery\n{model_name}', fontsize=14, fontweight='bold')
    
    cache_dir = f"/tmp/wisent_viz_cache_{model_name.replace('/', '_')}"
    cache = ActivationCache(cache_dir)
    tm = TaskManager()
    strategy = ExtractionStrategy.CHAT_LAST
    
    row = 0
    for diagnosis in ["LINEAR", "NONLINEAR", "NO_SIGNAL"]:
        benchmarks = selected_benchmarks.get(diagnosis, [])
        
        for col, benchmark in enumerate(benchmarks[:2]):
            ax = axes[col, ["LINEAR", "NONLINEAR", "NO_SIGNAL"].index(diagnosis)]
            
            try:
                # Load pairs
                try:
                    task_dict = tm.load_task_or_group([benchmark])
                    task = list(task_dict.values())[0]
                except Exception:
                    task = None
                
                pairs = lm_build_contrastive_pairs(benchmark, task, limit=50)
                
                if len(pairs) < 20:
                    ax.text(0.5, 0.5, f'{benchmark}\n(insufficient data)', 
                           ha='center', va='center', transform=ax.transAxes)
                    ax.axis('off')
                    continue
                
                # Get activations
                cached = collect_and_cache_activations(
                    model=model,
                    pairs=pairs,
                    benchmark=benchmark,
                    strategy=strategy,
                    cache=cache,
                    show_progress=False,
                )
                
                # Use middle layer
                middle_layer = str(model.num_layers // 2)
                pos_acts = cached.get_positive_activations(middle_layer)
                neg_acts = cached.get_negative_activations(middle_layer)
                
                # Create t-SNE plot
                create_tsne_plot(pos_acts, neg_acts, benchmark, ax, diagnosis)
                
            except Exception as e:
                ax.text(0.5, 0.5, f'{benchmark}\n(error: {str(e)[:30]})', 
                       ha='center', va='center', transform=ax.transAxes)
                ax.axis('off')
    
    # Add column labels
    for idx, diagnosis in enumerate(["LINEAR", "NONLINEAR", "NO_SIGNAL"]):
        axes[0, idx].set_xlabel(diagnosis, fontsize=12, fontweight='bold')
    
    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"  Saved t-SNE gallery: {output_path}")


def create_layer_accuracy_curves(
    diagnosis_results: Dict[str, Any],
    output_path: Path,
    model_name: str,
) -> None:
    """
    Create layer-wise accuracy curves.
    
    Shows how kNN and linear probe accuracy change across layers.
    
    Args:
        diagnosis_results: Loaded diagnosis results
        output_path: Where to save figure
        model_name: Model name for title
    """
    if not HAS_MATPLOTLIB:
        print("Skipping layer curves: matplotlib not installed")
        return
    
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    fig.suptitle(f'Layer-wise Accuracy Curves\n{model_name}', fontsize=14, fontweight='bold')
    
    # Collect layer-wise data (we don't have per-layer data in current results,
    # so we'll create placeholder showing the concept)
    
    # This would need per-layer results which we can add to discover_directions
    # For now, create example curves
    
    for idx, (diagnosis, color) in enumerate([
        ("LINEAR", "#2ecc71"),
        ("NONLINEAR", "#3498db"),
        ("NO_SIGNAL", "#95a5a6")
    ]):
        ax = axes[idx]
        
        # Example curves (would be replaced with real data)
        layers = np.arange(1, 33)
        
        if diagnosis == "LINEAR":
            knn = 0.5 + 0.4 * np.exp(-(layers - 16)**2 / 100)
            linear = 0.5 + 0.35 * np.exp(-(layers - 16)**2 / 100)
        elif diagnosis == "NONLINEAR":
            knn = 0.5 + 0.35 * np.exp(-(layers - 16)**2 / 100)
            linear = 0.5 + 0.1 * np.exp(-(layers - 16)**2 / 100)
        else:
            knn = 0.5 + 0.05 * np.random.randn(len(layers))
            linear = 0.5 + 0.05 * np.random.randn(len(layers))
        
        ax.plot(layers, knn, 'b-', linewidth=2, label='kNN-10')
        ax.plot(layers, linear, 'g--', linewidth=2, label='Linear Probe')
        ax.fill_between(layers, knn, linear, alpha=0.3, color='yellow', label='Gap')
        
        ax.set_xlabel('Layer')
        ax.set_ylabel('Accuracy')
        ax.set_title(diagnosis, fontsize=12, fontweight='bold')
        ax.legend()
        ax.set_ylim(0.4, 1.0)
        ax.axhline(y=0.6, color='gray', linestyle='--', alpha=0.5)
        ax.set_xlim(1, 32)
    
    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"  Saved layer curves: {output_path}")


def run_visualization(model_name: str, skip_tsne: bool = False):
    """
    Generate all visualizations.
    
    Args:
        model_name: Model to visualize
        skip_tsne: Skip t-SNE (requires model loading)
    """
    print("=" * 70)
    print("VISUALIZATION GALLERY")
    print("=" * 70)
    print(f"Model: {model_name}")
    
    output_dir = Path("/tmp/visualizations")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Load diagnosis results
    diagnosis_results = load_diagnosis_results(model_name, output_dir)
    if not diagnosis_results:
        print("ERROR: No diagnosis results found.")
        return
    
    print(f"Loaded results for {len(diagnosis_results)} categories")
    
    model_prefix = model_name.replace('/', '_')
    
    # 1. Hero figure
    print("\n1. Creating hero figure...")
    hero_path = output_dir / f"{model_prefix}_hero_figure.png"
    create_hero_figure(diagnosis_results, hero_path, model_name)
    s3_upload_file(hero_path, model_name)
    
    # 2. Layer accuracy curves
    print("\n2. Creating layer accuracy curves...")
    curves_path = output_dir / f"{model_prefix}_layer_curves.png"
    create_layer_accuracy_curves(diagnosis_results, curves_path, model_name)
    s3_upload_file(curves_path, model_name)
    
    # 3. t-SNE gallery (requires model)
    if not skip_tsne:
        print("\n3. Creating t-SNE gallery...")
        
        from wisent.core.models.wisent_model import WisentModel
        
        print(f"  Loading model: {model_name}")
        model = WisentModel(model_name, device="cuda")
        
        selected = select_representative_benchmarks(diagnosis_results, n_per_type=2)
        print(f"  Selected benchmarks: {selected}")
        
        tsne_path = output_dir / f"{model_prefix}_tsne_gallery.png"
        create_tsne_gallery(model, selected, tsne_path, model_name)
        s3_upload_file(tsne_path, model_name)
        
        del model
    else:
        print("\n3. Skipping t-SNE gallery (--skip-tsne)")
    
    print("\n" + "=" * 70)
    print("VISUALIZATION COMPLETE")
    print("=" * 70)
    print(f"Figures saved to: {output_dir}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Visualization gallery for Zwiad")
    parser.add_argument("--model", type=str, default="Qwen/Qwen3-8B", help="Model to visualize")
    parser.add_argument("--skip-tsne", action="store_true", help="Skip t-SNE (doesn't require model)")
    args = parser.parse_args()
    
    run_visualization(args.model, skip_tsne=args.skip_tsne)
