"""
Exhaustive layer combination analysis.

Tests all 2^N - 1 layer combinations to find optimal layer subsets
for geometry detection.
"""

import json
import os
import sys
import time
import torch
from datetime import datetime
from typing import Dict, List


def run_exhaustive_layer_analysis(
    task: str = "truthfulqa_gen",
    model: str = "meta-llama/Llama-3.2-1B-Instruct",
    num_pairs: int = 50,
    max_layers: int = 16,
    output_dir: str = "/home/ubuntu/output",
):
    """
    Run exhaustive layer combination analysis.
    
    Tests all 2^N - 1 layer combinations to find which layer subsets
    produce the strongest geometric structure detection.
    """
    from wisent.core.models.wisent_model import WisentModel
    from wisent.core.contrastive_pairs.diagnostics.control_vectors import (
        detect_geometry_exhaustive,
        ExhaustiveGeometryAnalysisResult,
    )
    from wisent.data.pairs.lm_eval_pairs import load_task_pairs
    from wisent.core.activations.aggregation import TokenAggregation
    
    print("=" * 80)
    print("EXHAUSTIVE LAYER COMBINATION ANALYSIS")
    print("=" * 80)
    print(f"Task: {task}")
    print(f"Model: {model}")
    print(f"Num pairs: {num_pairs}")
    print(f"Max layers: {max_layers}")
    print(f"Output dir: {output_dir}")
    print("=" * 80)
    
    # Calculate expected combinations
    expected_combos = 2 ** max_layers - 1
    print(f"\nExpected combinations to test: {expected_combos}")
    
    # Load model
    print("\n[1] Loading model...")
    start = time.time()
    wisent_model = WisentModel(model)
    num_layers = wisent_model.num_layers
    print(f"    Model loaded in {time.time() - start:.1f}s")
    print(f"    Model has {num_layers} layers")
    
    # Limit layers if needed
    layers_to_use = min(num_layers, max_layers)
    actual_combos = 2 ** layers_to_use - 1
    print(f"    Using {layers_to_use} layers -> {actual_combos} combinations")
    
    # Load pairs
    print(f"\n[2] Loading {num_pairs} pairs for {task}...")
    start = time.time()
    pairs = load_task_pairs(task, limit=num_pairs)
    print(f"    Loaded {len(pairs)} pairs in {time.time() - start:.1f}s")
    
    if not pairs:
        print("ERROR: No pairs loaded")
        return
    
    # Collect activations for all layers
    print(f"\n[3] Collecting activations for layers 1-{layers_to_use}...")
    layers_list = list(range(1, layers_to_use + 1))
    
    pos_by_layer: Dict[int, List[torch.Tensor]] = {l: [] for l in layers_list}
    neg_by_layer: Dict[int, List[torch.Tensor]] = {l: [] for l in layers_list}
    
    start = time.time()
    for i, pair in enumerate(pairs):
        if (i + 1) % 10 == 0:
            print(f"    Processing pair {i + 1}/{len(pairs)}...")
        
        # Get activations for positive response
        pos_acts = wisent_model.get_activations(
            pair.positive_response.text,
            layers=layers_list,
            token_aggregation=TokenAggregation.LAST,
        )
        
        # Get activations for negative response
        neg_acts = wisent_model.get_activations(
            pair.negative_response.text,
            layers=layers_list,
            token_aggregation=TokenAggregation.LAST,
        )
        
        for layer in layers_list:
            if layer in pos_acts and layer in neg_acts:
                pos_by_layer[layer].append(pos_acts[layer].float().cpu())
                neg_by_layer[layer].append(neg_acts[layer].float().cpu())
    
    print(f"    Collected activations in {time.time() - start:.1f}s")
    
    # Stack into tensors
    print("\n[4] Stacking activations...")
    pos_tensors = {}
    neg_tensors = {}
    for layer in layers_list:
        if pos_by_layer[layer] and neg_by_layer[layer]:
            pos_tensors[layer] = torch.stack(pos_by_layer[layer])
            neg_tensors[layer] = torch.stack(neg_by_layer[layer])
            print(f"    Layer {layer}: pos={pos_tensors[layer].shape}, neg={neg_tensors[layer].shape}")
    
    # Run exhaustive analysis
    print(f"\n[5] Running exhaustive analysis ({actual_combos} combinations)...")
    start = time.time()
    
    last_progress = 0
    def progress_callback(current: int, total: int):
        nonlocal last_progress
        pct = int(100 * current / total)
        if pct >= last_progress + 5:  # Report every 5%
            elapsed = time.time() - start
            rate = current / elapsed if elapsed > 0 else 0
            remaining = (total - current) / rate if rate > 0 else 0
            print(f"    Progress: {current}/{total} ({pct}%) - {rate:.1f} combos/sec - ETA: {remaining:.0f}s")
            last_progress = pct
    
    result = detect_geometry_exhaustive(
        pos_tensors,
        neg_tensors,
        max_layers=layers_to_use,
        combination_method="concat",
        progress_callback=progress_callback,
    )
    
    elapsed = time.time() - start
    print(f"\n    Completed in {elapsed:.1f}s ({actual_combos / elapsed:.1f} combos/sec)")
    
    # Print results
    print("\n" + "=" * 80)
    print("RESULTS")
    print("=" * 80)
    
    print(f"\nTotal combinations tested: {result.total_combinations}")
    print(f"\nBest combination: {result.best_combination}")
    print(f"Best score: {result.best_score:.4f}")
    print(f"Best structure: {result.best_structure.value}")
    
    print(f"\nBest single layer: L{result.single_layer_best}")
    print(f"Best single layer score: {result.single_layer_best_score:.4f}")
    print(f"Combination beats single: {result.combination_beats_single}")
    print(f"Improvement over single: {result.improvement_over_single:.4f}")
    
    print("\n--- Top 10 Combinations ---")
    for i, r in enumerate(result.top_10):
        layers_str = "+".join(f"L{l}" for l in r.layers)
        print(f"  {i+1}. {layers_str}: {r.best_structure.value} = {r.best_score:.4f}")
    
    print("\n--- Patterns ---")
    print(f"  Most important layers: {result.patterns.get('most_important_layers', [])}")
    print(f"  Optimal combination size: {result.patterns.get('optimal_combination_size', 1)}")
    print(f"  Dominant structure: {result.patterns.get('dominant_structure', 'unknown')}")
    print(f"  Best score by size: {result.patterns.get('best_score_by_size', {})}")
    print(f"  Early vs late ratio: {result.patterns.get('early_vs_late_ratio', 0):.2f}")
    
    print(f"\n--- Recommendation ---")
    print(f"  {result.recommendation}")
    
    # Save results
    os.makedirs(output_dir, exist_ok=True)
    output_file = os.path.join(output_dir, f"exhaustive_geometry_{task}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
    
    # Convert to serializable format
    results_json = {
        "task": task,
        "model": model,
        "num_pairs": num_pairs,
        "max_layers": layers_to_use,
        "total_combinations": result.total_combinations,
        "elapsed_seconds": elapsed,
        "best_combination": list(result.best_combination),
        "best_score": result.best_score,
        "best_structure": result.best_structure.value,
        "single_layer_best": result.single_layer_best,
        "single_layer_best_score": result.single_layer_best_score,
        "combination_beats_single": result.combination_beats_single,
        "improvement_over_single": result.improvement_over_single,
        "top_10": [
            {
                "layers": list(r.layers),
                "best_structure": r.best_structure.value,
                "best_score": r.best_score,
                "all_scores": r.all_scores,
            }
            for r in result.top_10
        ],
        "top_100": [
            {
                "layers": list(r.layers),
                "best_structure": r.best_structure.value,
                "best_score": r.best_score,
            }
            for r in result.all_results[:100]
        ],
        "patterns": {
            k: v if not isinstance(v, float) or not (v != v) else None  # Handle NaN
            for k, v in result.patterns.items()
        },
        "recommendation": result.recommendation,
    }
    
    with open(output_file, "w") as f:
        json.dump(results_json, f, indent=2)
    print(f"\nResults saved to: {output_file}")
    
    return result


if __name__ == "__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("--max-layers", type=int, default=16)
    parser.add_argument("--output-dir", default="/home/ubuntu/output")
    args = parser.parse_args()
    
    run_exhaustive_layer_analysis(
        task=args.task,
        model=args.model,
        num_pairs=args.num_pairs,
        max_layers=args.max_layers,
        output_dir=args.output_dir,
    )
