"""Classification optimization command - trains classifier on model's own generations.

This optimizer:
1. Loads task questions
2. Generates actual responses from the model
3. Evaluates responses using task-specific evaluator (correct/incorrect)
4. Extracts activations during generation
5. Optionally augments training with contrastive pairs (ground truth positive/negative)
6. Trains classifier to predict correctness from activations

This tests whether we can detect when the model is generating incorrect/hallucinated content
based on its internal activations.
"""

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


def execute_optimize_classification(args):
    """
    Execute classification optimization on model's own generations.

    Strategy:
    1. Load questions from task
    2. Generate responses (capturing activations)
    3. Evaluate responses with task evaluator
    4. Train classifier: activations -> correct/incorrect
    """
    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.data_loaders.loaders.lm_loader import LMEvalDataLoader
    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
    from wisent.core.models.inference_config import get_generate_kwargs

    print(f"\n{'='*80}")
    print(f"🔍 CLASSIFICATION OPTIMIZATION (ON MODEL GENERATIONS)")
    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:
        # Default: middle third of layers
        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,
        'continuation': ActivationAggregationStrategy.CONTINUATION_TOKEN,
        'choice': ActivationAggregationStrategy.CHOICE_TOKEN,
        'last_token': ActivationAggregationStrategy.LAST_TOKEN,
        'first_token': ActivationAggregationStrategy.FIRST_TOKEN,
        'mean_pooling': ActivationAggregationStrategy.MEAN_POOLING,
    }

    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"📊 Thresholds: {args.threshold_range}\n")

    # 3. Get list of tasks
    default_tasks = ["truthfulqa_gen"]
    task_list = args.tasks if args.tasks else default_tasks

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

    # Check training mode
    use_contrastive = getattr(args, 'use_contrastive_pairs', False)
    train_on_generations = getattr(args, 'train_on_generations', False)
    train_contrastive_only = getattr(args, 'train_on_contrastive_only', True) and not train_on_generations

    if train_contrastive_only:
        print(f"📚 Training on contrastive pairs ONLY (default, faster)")
        print(f"   Test set will use actual model generations\n")
    elif use_contrastive:
        print(f"📚 Training on generations + contrastive pairs\n")
    else:
        print(f"📚 Training on model generations only\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: Load task data
            print(f"\n  📊 Loading task data...")
            loader = LMEvalDataLoader()

            try:
                result = loader._load_one_task(
                    task_name=task_name,
                    split_ratio=0.8,
                    seed=42,
                    limit=args.limit * 2,  # Load extra for train/test split
                    training_limit=None,
                    testing_limit=None
                )
                qa_pairs = result['train_qa_pairs'].pairs[:args.limit]
            except Exception as e:
                print(f"     Error loading via LMEvalDataLoader: {e}")
                print(f"     Trying TaskInterfaceDataLoader...")
                from wisent.core.data_loaders.loaders.task_interface_loader import TaskInterfaceDataLoader
                task_loader = TaskInterfaceDataLoader()
                result = task_loader.load(
                    task=task_name,
                    split_ratio=0.8,
                    seed=42,
                    limit=args.limit * 2,
                    training_limit=None,
                    testing_limit=None
                )
                qa_pairs = result.train_qa_pairs.pairs[:args.limit]

            print(f"     ✓ Loaded {len(qa_pairs)} question-answer pairs")

            # STEP 2: Load the appropriate evaluator
            print(f"\n  🔧 Loading evaluator for {task_name}...")
            evaluator = _get_evaluator_for_task(task_name)
            print(f"     ✓ Using evaluator: {evaluator.name if hasattr(evaluator, 'name') else type(evaluator).__name__}")

            # Split pairs into train/test FIRST
            split_idx = int(len(qa_pairs) * 0.8)
            train_pairs = qa_pairs[:split_idx]
            test_pairs = qa_pairs[split_idx:]

            from wisent.core.contrastive_pairs.core.pair import ContrastivePair
            from wisent.core.contrastive_pairs.core.response import PositiveResponse, NegativeResponse

            # MODE: Train on contrastive pairs only
            if train_contrastive_only:
                # STEP 3a: Extract training data from contrastive pairs (no generation needed)
                print(f"\n  📚 Extracting training data from contrastive pairs...")
                train_data = []

                for i, pair in enumerate(train_pairs):
                    if i % 10 == 0:
                        print(f"     Progress: {i+1}/{len(train_pairs)}", end='\r')

                    try:
                        updated_pair = collector.collect_for_pair(
                            pair,
                            layers=None,
                            aggregation=ActivationAggregationStrategy.LAST_TOKEN,
                            return_full_sequence=True,
                            prompt_strategy=PromptConstructionStrategy.CHAT_TEMPLATE,
                        )

                        # Add positive response (is_correct=True)
                        if updated_pair.positive_response.layers_activations:
                            train_data.append({
                                'activations': updated_pair.positive_response.layers_activations,
                                'is_correct': True,
                            })

                        # Add negative response (is_correct=False)
                        if updated_pair.negative_response.layers_activations:
                            train_data.append({
                                'activations': updated_pair.negative_response.layers_activations,
                                'is_correct': False,
                            })
                    except Exception:
                        continue

                print(f"     ✓ Training data: {len(train_data)} samples from contrastive pairs")

                # STEP 3b: Generate test responses only
                print(f"\n  🤖 Generating TEST responses only ({len(test_pairs)} samples)...")
                test_data = []
                num_correct = 0
                num_incorrect = 0

                for i, pair in enumerate(test_pairs):
                    if i % 5 == 0:
                        print(f"     Progress: {i+1}/{len(test_pairs)}", end='\r')

                    # Generate response
                    messages = [{"role": "user", "content": pair.prompt}]
                    try:
                        responses = model.generate(
                            inputs=[messages],
                            max_new_tokens=100,
                            temperature=0.7,
                            do_sample=True,
                        )
                        generated_text = responses[0] if responses else ""
                    except Exception:
                        continue

                    # Evaluate the generated response
                    eval_result = _evaluate_response(
                        evaluator=evaluator,
                        response=generated_text,
                        pair=pair,
                        task_name=task_name
                    )

                    is_correct = eval_result['is_correct']
                    if is_correct:
                        num_correct += 1
                    else:
                        num_incorrect += 1

                    # Extract activations for the generated response
                    mock_pair = ContrastivePair(
                        prompt=pair.prompt,
                        positive_response=PositiveResponse(model_response=generated_text),
                        negative_response=NegativeResponse(model_response=generated_text),
                    )

                    try:
                        updated_pair = collector.collect_for_pair(
                            mock_pair,
                            layers=None,
                            aggregation=ActivationAggregationStrategy.LAST_TOKEN,
                            return_full_sequence=True,
                            prompt_strategy=PromptConstructionStrategy.CHAT_TEMPLATE,
                        )
                        test_data.append({
                            'activations': updated_pair.positive_response.layers_activations,
                            'is_correct': is_correct,
                        })
                    except Exception:
                        continue

                print(f"     ✓ Test data: {len(test_data)} actual generations")
                print(f"     Model accuracy: {num_correct}/{num_correct + num_incorrect} = {num_correct / (num_correct + num_incorrect) * 100:.1f}%")

            else:
                # ORIGINAL MODE: Generate all responses, optionally augment with contrastive pairs
                print(f"\n  🤖 Generating and evaluating responses...")

                gen_kwargs = get_generate_kwargs()
                generations = []

                for i, pair in enumerate(qa_pairs):
                    if i % 10 == 0:
                        print(f"     Progress: {i+1}/{len(qa_pairs)}", end='\r')

                    # Generate response
                    messages = [{"role": "user", "content": pair.prompt}]
                    try:
                        responses = model.generate(
                            inputs=[messages],
                            max_new_tokens=100,
                            temperature=0.7,
                            do_sample=True,
                        )
                        generated_text = responses[0] if responses else ""
                    except Exception as e:
                        generated_text = ""

                    # Evaluate the generated response
                    eval_result = _evaluate_response(
                        evaluator=evaluator,
                        response=generated_text,
                        pair=pair,
                        task_name=task_name
                    )

                    generations.append({
                        'prompt': pair.prompt,
                        'generated': generated_text,
                        'is_correct': eval_result['is_correct'],
                        'confidence': eval_result.get('confidence', 0.5),
                        'positive_reference': pair.positive_response.model_response,
                        'negative_reference': pair.negative_response.model_response,
                    })

                print(f"     ✓ Generated {len(generations)} responses")

                # Count correct/incorrect
                num_correct = sum(1 for g in generations if g['is_correct'])
                num_incorrect = len(generations) - num_correct
                print(f"     Correct: {num_correct}, Incorrect: {num_incorrect}")

                if num_correct < 3 or num_incorrect < 3:
                    print(f"     ⚠️ Not enough samples in both classes, skipping task")
                    continue

                # STEP 4: Extract activations for generated responses
                print(f"\n  🧠 Extracting activations from generations...")

                activations_data = []
                for i, gen in enumerate(generations):
                    if i % 10 == 0:
                        print(f"     Progress: {i+1}/{len(generations)}", end='\r')

                    # We extract activations for the prompt + generated response
                    mock_pair = ContrastivePair(
                        prompt=gen['prompt'],
                        positive_response=PositiveResponse(model_response=gen['generated']),
                        negative_response=NegativeResponse(model_response=gen['generated']),
                    )

                    try:
                        updated_pair = collector.collect_for_pair(
                            mock_pair,
                            layers=None,
                            aggregation=ActivationAggregationStrategy.LAST_TOKEN,
                            return_full_sequence=True,
                            prompt_strategy=PromptConstructionStrategy.CHAT_TEMPLATE,
                        )
                        activations_data.append({
                            'activations': updated_pair.positive_response.layers_activations,
                            'is_correct': gen['is_correct'],
                            'prompt': gen['prompt'],
                            'generated': gen['generated'],
                        })
                    except Exception as e:
                        continue

                print(f"     ✓ Extracted activations for {len(activations_data)} generations")

                # Split into train/test
                gen_split_idx = int(len(activations_data) * 0.8)
                train_data = activations_data[:gen_split_idx]
                test_data = activations_data[gen_split_idx:]
                print(f"     Train (generations): {len(train_data)}, Test: {len(test_data)}")

                # Optionally augment training with contrastive pairs
                if use_contrastive:
                    print(f"\n  📚 Augmenting training with contrastive pairs...")
                    contrastive_data = []

                    for i, pair in enumerate(train_pairs):
                        if i % 10 == 0:
                            print(f"     Progress: {i+1}/{len(train_pairs)}", end='\r')

                        try:
                            updated_pair = collector.collect_for_pair(
                                pair,
                                layers=None,
                                aggregation=ActivationAggregationStrategy.LAST_TOKEN,
                                return_full_sequence=True,
                                prompt_strategy=PromptConstructionStrategy.CHAT_TEMPLATE,
                            )

                            if updated_pair.positive_response.layers_activations:
                                contrastive_data.append({
                                    'activations': updated_pair.positive_response.layers_activations,
                                    'is_correct': True,
                                })

                            if updated_pair.negative_response.layers_activations:
                                contrastive_data.append({
                                    'activations': updated_pair.negative_response.layers_activations,
                                    'is_correct': False,
                                })
                        except Exception:
                            continue

                    print(f"     ✓ Added {len(contrastive_data)} contrastive pair samples to training")
                    train_data = train_data + contrastive_data
                    print(f"     Total train: {len(train_data)} (generations + contrastive)")

            # STEP 5: Grid search over configurations
            print(f"\n  🔍 Testing classifier configurations...")

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

            best_score = -1
            best_config = None
            combinations_tested = 0

            for layer in layers_to_test:
                layer_name = str(layer)

                for agg_method in args.aggregation_methods:
                    agg_enum = aggregation_map.get(agg_method, ActivationAggregationStrategy.MEAN_POOLING)

                    # Build feature matrices
                    X_train, y_train = [], []
                    X_test, y_test = [], []

                    for item in train_data:
                        feat = _aggregate_cached(item['activations'], layer_name, agg_enum)
                        if feat is not None:
                            X_train.append(feat)
                            y_train.append(1 if item['is_correct'] else 0)

                    for item in test_data:
                        feat = _aggregate_cached(item['activations'], layer_name, agg_enum)
                        if feat is not None:
                            X_test.append(feat)
                            y_test.append(1 if item['is_correct'] else 0)

                    if len(X_train) < 10 or len(X_test) < 5:
                        continue

                    X_train = np.vstack(X_train)
                    y_train = np.array(y_train)
                    X_test = np.vstack(X_test)
                    y_test = np.array(y_test)

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

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

                        # Get probabilities
                        try:
                            probs_raw = clf.predict_proba(X_test)
                            if isinstance(probs_raw, list):
                                probs = np.array(probs_raw)
                            elif hasattr(probs_raw, 'numpy'):
                                probs = probs_raw.numpy()
                            else:
                                probs = np.array(probs_raw)
                            if probs.ndim == 2:
                                probs = probs[:, 1]
                        except Exception:
                            preds = clf.predict(X_test)
                            if isinstance(preds, list):
                                probs = np.array(preds, dtype=float)
                            else:
                                probs = np.array(preds).astype(float)

                        for threshold in args.threshold_range:
                            combinations_tested += 1

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

                            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)

                            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': int(layer),
                                    'classifier_type': classifier_type,
                                    'aggregation': agg_method,
                                    'threshold': float(threshold),
                                    'accuracy': float(acc),
                                    'f1_score': float(f1),
                                    'precision': float(prec),
                                    'recall': float(rec),
                                    'test_size': int(len(y_test)),
                                    'train_correct': int(sum(y_train)),
                                    'train_incorrect': int(len(y_train) - sum(y_train)),
                                }

                            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"      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}")
                print(f"      Dataset stats:")
                print(f"        • Train correct: {best_config['train_correct']}")
                print(f"        • Train incorrect: {best_config['train_incorrect']}")

                # Calculate model accuracy based on test data
                total_test = num_correct + num_incorrect
                model_acc = float(num_correct / total_test) if total_test > 0 else 0.0

                all_results[task_name] = {
                    'best_config': best_config,
                    'optimization_metric': args.optimization_metric,
                    'best_score': float(best_score),
                    'combinations_tested': int(combinations_tested),
                    'model_accuracy_on_task': model_acc,
                }
            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()

    # 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"-" * 120)
        print(f"{'Task':<20} | {'Model Acc':>9} | {'Layer':>5} | {'Classifier':<10} | {'Agg':<10} | {'Thresh':>6} | {'F1':>6} | {'Clf Acc':>7}")
        print(f"-" * 120)
        for task_name, result in all_results.items():
            config = result['best_config']
            model_acc = result.get('model_accuracy_on_task', 0)
            print(f"{task_name:<20} | {model_acc:>9.2%} | {config['layer']:>5} | {config['classifier_type']:<10} | "
                  f"{config['aggregation']:<10} | {config['threshold']:>6.2f} | {config['f1_score']:>6.4f} | {config['accuracy']:>7.4f}")
        print(f"-" * 120)
        print()

    # Handle --show-comparisons
    show_comparisons = getattr(args, 'show_comparisons', 0)
    if show_comparisons > 0 and all_results:
        print(f"\n📊 Note: Comparisons show actual model generations evaluated by task evaluator")
        print(f"   The classifier learns to predict evaluation outcome from activations\n")


def _get_evaluator_for_task(task_name: str):
    """Get the appropriate evaluator for a task."""
    task_lower = task_name.lower()

    if 'truthfulqa_gen' in task_lower or 'truthfulqa_generation' in task_lower:
        from wisent.core.evaluators.oracles.truthfulqa_gen_evaluator import TruthfulQAGenEvaluator
        return TruthfulQAGenEvaluator()
    elif 'arc' in task_lower or 'hellaswag' in task_lower or 'winogrande' in task_lower:
        from wisent.core.evaluators.benchmark_specific import ExactMatchEvaluator
        return ExactMatchEvaluator()
    else:
        # Default to F1 evaluator
        from wisent.core.evaluators.benchmark_specific import F1Evaluator
        return F1Evaluator()


def _evaluate_response(evaluator, response: str, pair, task_name: str) -> Dict[str, Any]:
    """Evaluate a generated response using the task evaluator."""
    task_lower = task_name.lower()

    try:
        if 'truthfulqa_gen' in task_lower or 'truthfulqa_generation' in task_lower:
            # TruthfulQA needs correct/incorrect answer lists
            # The pair has positive (correct) and negative (incorrect) responses
            correct_answers = [pair.positive_response.model_response]
            incorrect_answers = [pair.negative_response.model_response]

            result = evaluator.evaluate(
                response=response,
                expected=pair.positive_response.model_response,
                correct_answers=correct_answers,
                incorrect_answers=incorrect_answers
            )

            is_correct = result.ground_truth == "TRUTHFUL"
            return {
                'is_correct': is_correct,
                'confidence': result.confidence,
                'details': result.details
            }
        else:
            # For other tasks, compare to positive reference
            result = evaluator.evaluate(response, pair.positive_response.model_response)
            is_correct = result.ground_truth == "TRUTHFUL"
            return {
                'is_correct': is_correct,
                'confidence': result.confidence,
            }
    except Exception as e:
        # Fallback: simple string matching
        response_lower = response.lower().strip()
        positive_lower = pair.positive_response.model_response.lower().strip()
        negative_lower = pair.negative_response.model_response.lower().strip()

        # Check if response is more similar to positive or negative
        from difflib import SequenceMatcher
        pos_sim = SequenceMatcher(None, response_lower, positive_lower).ratio()
        neg_sim = SequenceMatcher(None, response_lower, negative_lower).ratio()

        return {
            'is_correct': pos_sim > neg_sim,
            'confidence': abs(pos_sim - neg_sim),
        }


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

    if layer_activations is None:
        return None

    try:
        tensor = layer_activations.get(layer_name)
        if tensor is None:
            return None
    except (KeyError, AttributeError):
        return None

    if tensor.ndim == 1:
        return tensor.numpy() if hasattr(tensor, 'numpy') else np.array(tensor)

    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:
        mid = tensor.shape[0] // 2
        result = tensor[mid]
    elif aggregation == ActivationAggregationStrategy.CHOICE_TOKEN:
        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)
