"""Classification optimization command - uses cached activations for speed.

This optimizer extracts activations ONCE for all layers, then tests different
configurations (layer, aggregation, threshold, classifier) by slicing cached
activations and fitting classifiers - no repeated model forward passes.
"""

import sys
import json
import time
from typing import List, Dict, Any, Optional
import os
import torch
import numpy as np


def execute_optimize_classification(args):
    """
    Execute classification optimization with cached activations.

    Strategy:
    1. Generate contrastive pairs ONCE
    2. Extract activations for ALL layers ONCE (full sequence)
    3. For each configuration, slice cached activations and fit classifier

    This is O(1 forward pass + N classifier fits) instead of O(N forward passes).
    """
    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.activations.prompt_construction_strategy import PromptConstructionStrategy
    from wisent.core.contrastive_pairs.lm_eval_pairs.lm_task_pairs_generation import build_contrastive_pairs
    from wisent.core.classifiers.classifiers.models.logistic import LogisticClassifier
    from wisent.core.classifiers.classifiers.models.mlp import MLPClassifier
    from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score

    print(f"\n{'='*80}")
    print(f"🔍 CLASSIFICATION PARAMETER OPTIMIZATION (CACHED ACTIVATIONS)")
    print(f"{'='*80}")
    print(f"   Model: {args.model}")
    print(f"   Limit per task: {args.limit}")
    print(f"   Device: {args.device or 'auto'}")
    print(f"{'='*80}\n")

    # 1. Load model
    print(f"📦 Loading model...")
    model = WisentModel(args.model, device=args.device)
    total_layers = model.num_layers
    print(f"   ✓ Model has {total_layers} layers\n")

    # 2. Determine layer range
    if args.layer_range:
        start, end = map(int, args.layer_range.split('-'))
        layers_to_test = list(range(start, end + 1))
    else:
        start_layer = total_layers // 3
        end_layer = (2 * total_layers) // 3
        layers_to_test = list(range(start_layer, end_layer + 1))

    # Configuration options
    classifier_types = ['logistic', 'mlp']

    # Map string strategies to enums
    aggregation_map = {
        'average': ActivationAggregationStrategy.MEAN_POOLING,
        'final': ActivationAggregationStrategy.LAST_TOKEN,
        'first': ActivationAggregationStrategy.FIRST_TOKEN,
        'max': ActivationAggregationStrategy.MAX_POOLING,
        'min': ActivationAggregationStrategy.MEAN_POOLING,  # No min, use mean as fallback
        'continuation': ActivationAggregationStrategy.CONTINUATION_TOKEN,
        'choice': ActivationAggregationStrategy.CHOICE_TOKEN,
        'last_token': ActivationAggregationStrategy.LAST_TOKEN,
        'first_token': ActivationAggregationStrategy.FIRST_TOKEN,
        'mean_pooling': ActivationAggregationStrategy.MEAN_POOLING,
    }

    prompt_strategy_map = {
        'multiple_choice': PromptConstructionStrategy.MULTIPLE_CHOICE,
        'role_playing': PromptConstructionStrategy.ROLE_PLAYING,
        'direct_completion': PromptConstructionStrategy.DIRECT_COMPLETION,
        'instruction_following': PromptConstructionStrategy.INSTRUCTION_FOLLOWING,
        'chat_template': PromptConstructionStrategy.CHAT_TEMPLATE,
    }

    prompt_strategies = ['chat_template', 'direct_completion']  # Reduced for speed

    print(f"🎯 Testing layers: {layers_to_test[0]} to {layers_to_test[-1]} ({len(layers_to_test)} layers)")
    print(f"🤖 Classifier types: {', '.join(classifier_types)}")
    print(f"🔄 Aggregation methods: {', '.join(args.aggregation_methods)}")
    print(f"📝 Prompt strategies: {', '.join(prompt_strategies)}")
    print(f"📊 Thresholds: {args.threshold_range}\n")

    # 3. Get list of tasks
    default_tasks = ["arc_easy", "arc_challenge", "hellaswag", "winogrande", "gsm8k"]
    task_list = args.tasks if args.tasks else default_tasks

    print(f"📋 Optimizing {len(task_list)} tasks: {', '.join(task_list)}\n")

    # Create activation collector
    collector = ActivationCollector(model=model, store_device="cpu", dtype=torch.float32)

    all_results = {}

    # 4. Process each task
    for task_idx, task_name in enumerate(task_list, 1):
        print(f"\n{'='*80}")
        print(f"Task {task_idx}/{len(task_list)}: {task_name}")
        print(f"{'='*80}")

        task_start_time = time.time()

        try:
            # STEP 1: Generate contrastive pairs ONCE
            print(f"\n  📝 Generating contrastive pairs...")
            pairs = build_contrastive_pairs(
                task_name=task_name,
                limit=args.limit,
                seed=42
            )
            print(f"     ✓ Generated {len(pairs)} pairs")

            if not pairs:
                print(f"     ⚠️ No pairs generated, skipping task")
                continue

            # Split into train/test
            split_idx = int(len(pairs) * 0.8)
            train_pairs = pairs[:split_idx]
            test_pairs = pairs[split_idx:]
            print(f"     Train: {len(train_pairs)}, Test: {len(test_pairs)}")

            # STEP 2: Extract activations for ALL layers ONCE per prompt strategy
            # We cache activations per prompt strategy since they differ
            cached_activations = {}  # {prompt_strategy: {'train': [...], 'test': [...]}}

            for prompt_strategy in prompt_strategies:
                print(f"\n  🧠 Extracting activations (strategy: {prompt_strategy})...")
                strategy_enum = prompt_strategy_map.get(prompt_strategy, PromptConstructionStrategy.CHAT_TEMPLATE)

                train_activations = []
                test_activations = []

                # Extract for train pairs
                for i, pair in enumerate(train_pairs):
                    if i % 10 == 0:
                        print(f"     Train: {i+1}/{len(train_pairs)}", end='\r')
                    updated_pair = collector.collect_for_pair(
                        pair,
                        layers=None,  # ALL layers
                        aggregation=ActivationAggregationStrategy.CONTINUATION_TOKEN,  # Will be re-aggregated later
                        return_full_sequence=True,  # Keep full sequence for flexible aggregation
                        prompt_strategy=strategy_enum,
                    )
                    train_activations.append(updated_pair)
                print(f"     ✓ Train activations extracted ({len(train_activations)} pairs)")

                # Extract for test pairs
                for i, pair in enumerate(test_pairs):
                    if i % 10 == 0:
                        print(f"     Test: {i+1}/{len(test_pairs)}", end='\r')
                    updated_pair = collector.collect_for_pair(
                        pair,
                        layers=None,
                        aggregation=ActivationAggregationStrategy.CONTINUATION_TOKEN,
                        return_full_sequence=True,
                        prompt_strategy=strategy_enum,
                    )
                    test_activations.append(updated_pair)
                print(f"     ✓ Test activations extracted ({len(test_activations)} pairs)")

                cached_activations[prompt_strategy] = {
                    'train': train_activations,
                    'test': test_activations
                }

            # STEP 3: Grid search over configurations using cached activations
            print(f"\n  🔍 Testing configurations...")

            total_combinations = (
                len(layers_to_test) * len(classifier_types) *
                len(args.aggregation_methods) * len(args.threshold_range) *
                len(prompt_strategies)
            )
            print(f"     Total combinations: {total_combinations}")

            best_score = -1
            best_config = None
            combinations_tested = 0

            for prompt_strategy in prompt_strategies:
                cached = cached_activations[prompt_strategy]
                train_pairs_with_acts = cached['train']
                test_pairs_with_acts = cached['test']

                for layer in layers_to_test:
                    layer_name = str(layer)

                    for agg_method in args.aggregation_methods:
                        # Aggregate activations from cached full sequences
                        agg_enum = aggregation_map.get(agg_method, ActivationAggregationStrategy.MEAN_POOLING)

                        # Build feature matrices
                        X_train_pos, X_train_neg = [], []
                        X_test_pos, X_test_neg = [], []

                        for pair in train_pairs_with_acts:
                            pos_act = _aggregate_cached(pair.positive_response.layers_activations, layer_name, agg_enum)
                            neg_act = _aggregate_cached(pair.negative_response.layers_activations, layer_name, agg_enum)
                            if pos_act is not None and neg_act is not None:
                                X_train_pos.append(pos_act)
                                X_train_neg.append(neg_act)

                        for pair in test_pairs_with_acts:
                            pos_act = _aggregate_cached(pair.positive_response.layers_activations, layer_name, agg_enum)
                            neg_act = _aggregate_cached(pair.negative_response.layers_activations, layer_name, agg_enum)
                            if pos_act is not None and neg_act is not None:
                                X_test_pos.append(pos_act)
                                X_test_neg.append(neg_act)

                        if not X_train_pos or not X_test_pos:
                            continue

                        # Create training data: positive=1, negative=0
                        X_train = np.vstack(X_train_pos + X_train_neg)
                        y_train = np.array([1] * len(X_train_pos) + [0] * len(X_train_neg))

                        X_test = np.vstack(X_test_pos + X_test_neg)
                        y_test = np.array([1] * len(X_test_pos) + [0] * len(X_test_neg))

                        for classifier_type in classifier_types:
                            # Train classifier
                            if classifier_type == 'logistic':
                                clf = LogisticClassifier()
                            else:
                                clf = MLPClassifier()

                            try:
                                clf.fit(X_train, y_train)
                            except Exception:
                                continue

                            # Get probabilities
                            try:
                                probs = clf.predict_proba(X_test)[:, 1]
                            except Exception:
                                probs = clf.predict(X_test).astype(float)

                            for threshold in args.threshold_range:
                                combinations_tested += 1

                                # Apply threshold
                                y_pred = (probs >= threshold).astype(int)

                                # Calculate metrics
                                acc = accuracy_score(y_test, y_pred)
                                f1 = f1_score(y_test, y_pred, zero_division=0)
                                prec = precision_score(y_test, y_pred, zero_division=0)
                                rec = recall_score(y_test, y_pred, zero_division=0)

                                # Select metric to optimize
                                metric_map = {'f1': f1, 'accuracy': acc, 'precision': prec, 'recall': rec}
                                score = metric_map.get(args.optimization_metric, f1)

                                if score > best_score:
                                    best_score = score
                                    best_config = {
                                        'layer': layer,
                                        'classifier_type': classifier_type,
                                        'aggregation': agg_method,
                                        'threshold': threshold,
                                        'prompt_construction_strategy': prompt_strategy,
                                        'token_targeting_strategy': agg_method,
                                        'accuracy': acc,
                                        'f1_score': f1,
                                        'precision': prec,
                                        'recall': rec,
                                        'generation_count': len(X_test),
                                    }

                                if combinations_tested % 50 == 0:
                                    print(f"     Progress: {combinations_tested}/{total_combinations}, best {args.optimization_metric}: {best_score:.4f}", end='\r')

            print(f"\n\n  ✅ Best config for {task_name}:")
            if best_config:
                print(f"      Layer: {best_config['layer']}")
                print(f"      Classifier: {best_config['classifier_type']}")
                print(f"      Aggregation: {best_config['aggregation']}")
                print(f"      Threshold: {best_config['threshold']:.2f}")
                print(f"      Prompt Strategy: {best_config['prompt_construction_strategy']}")
                print(f"      Performance metrics:")
                print(f"        • Accuracy:  {best_config['accuracy']:.4f}")
                print(f"        • F1 Score:  {best_config['f1_score']:.4f}")
                print(f"        • Precision: {best_config['precision']:.4f}")
                print(f"        • Recall:    {best_config['recall']:.4f}")

                all_results[task_name] = {
                    'best_config': best_config,
                    'optimization_metric': args.optimization_metric,
                    'best_score': best_score,
                    'combinations_tested': combinations_tested
                }
            else:
                print(f"      ⚠️ No valid configuration found")

            task_time = time.time() - task_start_time
            print(f"\n  ⏱️  Task completed in {task_time:.1f}s")

        except Exception as e:
            print(f"\n❌ Task '{task_name}' optimization failed:")
            print(f"   Error: {e}")
            import traceback
            traceback.print_exc()
            raise

    # 5. Save results
    results_dir = getattr(args, 'classifiers_dir', None) or './optimization_results'
    os.makedirs(results_dir, exist_ok=True)

    model_name_safe = args.model.replace('/', '_')
    results_file = os.path.join(results_dir, f'classification_optimization_{model_name_safe}.json')

    with open(results_file, 'w') as f:
        json.dump({
            'model': args.model,
            'optimization_metric': args.optimization_metric,
            'results': all_results
        }, f, indent=2)

    print(f"\n{'='*80}")
    print(f"📊 OPTIMIZATION COMPLETE")
    print(f"{'='*80}")
    print(f"✅ Results saved to: {results_file}\n")

    # Print summary
    if all_results:
        print(f"📋 SUMMARY BY TASK:")
        print(f"-" * 140)
        print(f"{'Task':<20} | {'Layer':>5} | {'Classifier':<10} | {'Agg':<12} | {'Prompt':<20} | {'Thresh':>6} | {'F1':>6} | {'Acc':>6}")
        print(f"-" * 140)
        for task_name, result in all_results.items():
            config = result['best_config']
            print(f"{task_name:<20} | {config['layer']:>5} | {config['classifier_type']:<10} | "
                  f"{config['aggregation']:<12} | {config['prompt_construction_strategy']:<20} | "
                  f"{config['threshold']:>6.2f} | {config['f1_score']:>6.4f} | {config['accuracy']:>6.4f}")
        print(f"-" * 140)
        print()

    # Handle --show-comparisons and --save-comparisons
    show_comparisons = getattr(args, 'show_comparisons', 0)
    save_comparisons = getattr(args, 'save_comparisons', None)

    if (show_comparisons > 0 or save_comparisons) and all_results:
        print("\n📊 Generating comparison data (optimized vs default config)...")

        all_comparisons = []
        for task_name, result in all_results.items():
            best_config = result['best_config']

            default_config = {
                'layer': total_layers // 2,
                'aggregation': 'average',
                'threshold': 0.5,
                'classifier_type': 'logistic',
                'prompt_construction_strategy': 'chat_template',
            }

            all_comparisons.append({
                'task': task_name,
                'default_config': default_config,
                'optimized_config': {
                    'layer': best_config['layer'],
                    'aggregation': best_config['aggregation'],
                    'threshold': best_config['threshold'],
                    'classifier_type': best_config['classifier_type'],
                    'prompt_construction_strategy': best_config['prompt_construction_strategy'],
                },
                'optimized_metrics': {
                    'f1': best_config['f1_score'],
                    'accuracy': best_config['accuracy'],
                    'precision': best_config['precision'],
                    'recall': best_config['recall'],
                },
                'improvements': {
                    'layer_change': best_config['layer'] - default_config['layer'],
                    'aggregation_change': default_config['aggregation'] != best_config['aggregation'],
                    'threshold_change': best_config['threshold'] - default_config['threshold'],
                },
            })

        if save_comparisons:
            os.makedirs(os.path.dirname(save_comparisons) if os.path.dirname(save_comparisons) else ".", exist_ok=True)
            with open(save_comparisons, 'w') as f:
                json.dump({
                    'model': args.model,
                    'optimization_metric': args.optimization_metric,
                    'comparisons': all_comparisons,
                }, f, indent=2)
            print(f"💾 Saved comparisons to: {save_comparisons}")

        if show_comparisons > 0:
            print(f"\n📊 Configuration Comparisons (showing {min(show_comparisons, len(all_comparisons))} tasks):\n")
            for i, comp in enumerate(all_comparisons[:show_comparisons]):
                print(f"{'─'*80}")
                print(f"Task: {comp['task']}")
                print(f"{'─'*80}")
                print(f"DEFAULT CONFIG:")
                print(f"  Layer: {comp['default_config']['layer']}, Agg: {comp['default_config']['aggregation']}, "
                      f"Threshold: {comp['default_config']['threshold']}")
                print(f"OPTIMIZED CONFIG:")
                print(f"  Layer: {comp['optimized_config']['layer']}, Agg: {comp['optimized_config']['aggregation']}, "
                      f"Threshold: {comp['optimized_config']['threshold']:.2f}")
                print(f"  Classifier: {comp['optimized_config']['classifier_type']}, "
                      f"Prompt: {comp['optimized_config']['prompt_construction_strategy']}")
                print(f"METRICS:")
                print(f"  F1: {comp['optimized_metrics']['f1']:.4f}, Accuracy: {comp['optimized_metrics']['accuracy']:.4f}")
                print()


def _aggregate_cached(layer_activations, layer_name: str, aggregation: "ActivationAggregationStrategy") -> Optional[np.ndarray]:
    """
    Aggregate cached full-sequence activations for a specific layer.

    Args:
        layer_activations: LayerActivations object with full sequences
        layer_name: Layer to extract (e.g., "10")
        aggregation: How to aggregate the sequence to a single vector

    Returns:
        Numpy array of shape (hidden_dim,) or None if layer not found
    """
    from wisent.core.activations.core.atoms import ActivationAggregationStrategy

    if layer_activations is None:
        return None

    # Get the tensor for this layer
    try:
        tensor = layer_activations.get(layer_name)
        if tensor is None:
            return None
    except (KeyError, AttributeError):
        return None

    # tensor shape: (seq_len, hidden_dim) or (hidden_dim,)
    if tensor.ndim == 1:
        # Already aggregated
        return tensor.numpy() if hasattr(tensor, 'numpy') else np.array(tensor)

    # Aggregate based on strategy
    if aggregation == ActivationAggregationStrategy.LAST_TOKEN:
        result = tensor[-1]
    elif aggregation == ActivationAggregationStrategy.FIRST_TOKEN:
        result = tensor[0]
    elif aggregation == ActivationAggregationStrategy.MEAN_POOLING:
        result = tensor.mean(dim=0)
    elif aggregation == ActivationAggregationStrategy.MAX_POOLING:
        result = tensor.max(dim=0).values
    elif aggregation == ActivationAggregationStrategy.CONTINUATION_TOKEN:
        # Use middle of sequence as approximation (proper continuation would need prompt_len)
        mid = tensor.shape[0] // 2
        result = tensor[mid]
    elif aggregation == ActivationAggregationStrategy.CHOICE_TOKEN:
        # Use token after middle
        mid = tensor.shape[0] // 2 + 1
        mid = min(mid, tensor.shape[0] - 1)
        result = tensor[mid]
    else:
        result = tensor[-1]

    return result.numpy() if hasattr(result, 'numpy') else np.array(result)
