"""Improved training script with class balancing and GoEmotions dataset."""

import numpy as np
from pathlib import Path
from collections import Counter
import pandas as pd
from datasets import load_dataset, DatasetDict, Dataset
from sklearn.model_selection import train_test_split
from transformers import (
    Trainer,
    TrainingArguments,
    AutoTokenizer,
    AutoModelForSequenceClassification,
    EarlyStoppingCallback
)
from sklearn.metrics import accuracy_score, f1_score
from sklearn.utils.class_weight import compute_class_weight
import torch

from .config import config


# Mapping from GoEmotions 28 labels to our 8 core emotions
GOEMOTIONS_MAPPING = {
    # Joy/Happiness cluster
    'joy': 'joy',
    'amusement': 'joy',
    'excitement': 'joy',
    'optimism': 'joy',
    'pride': 'joy',
    'relief': 'joy',
    'admiration': 'joy',

    # Sadness cluster
    'sadness': 'sadness',
    'grief': 'sadness',
    'disappointment': 'sadness',
    'remorse': 'sadness',

    # Anger cluster
    'anger': 'anger',
    'annoyance': 'anger',
    'disapproval': 'anger',
    'disgust': 'anger',

    # Fear cluster
    'fear': 'fear',
    'nervousness': 'fear',

    # Love cluster
    'love': 'love',
    'caring': 'love',
    'desire': 'love',
    'gratitude': 'love',

    # Surprise cluster
    'surprise': 'surprise',
    'realization': 'surprise',
    'confusion': 'surprise',
    'curiosity': 'surprise',

    # Neutral (skip or map to neutral)
    'neutral': 'neutral',
    'approval': 'neutral',
    'embarrassment': 'neutral',
}

# Our target labels
TARGET_LABELS = ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise', 'neutral', 'sarcasm']


def compute_metrics(eval_pred):
    """Compute metrics for evaluation."""
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)

    return {
        "accuracy": round(accuracy_score(labels, predictions), 4),
        "f1_macro": round(f1_score(labels, predictions, average="macro", zero_division=0), 4),
        "f1_weighted": round(f1_score(labels, predictions, average="weighted", zero_division=0), 4),
    }


class WeightedTrainer(Trainer):
    """Trainer with class weights for imbalanced datasets."""

    def __init__(self, class_weights=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.class_weights = class_weights

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        # By popping labels we prevent Deberta from internally calculating its own non-weighted loss.
        # This keeps the computational graph clean.
        labels = inputs.pop("labels", None)
        outputs = model(**inputs)
        logits = outputs.logits

        if labels is not None:
            if self.class_weights is not None:
                # Force weights and logits to float32 to prevent FP16/BF16 gradient convergence issues
                weight = torch.tensor(self.class_weights, device=logits.device, dtype=torch.float32)
                loss_fn = torch.nn.CrossEntropyLoss(weight=weight)
            else:
                loss_fn = torch.nn.CrossEntropyLoss()
                
            loss = loss_fn(logits.to(torch.float32), labels)
        else:
            loss = outputs.loss if hasattr(outputs, "loss") else None

        return (loss, outputs) if return_outputs else loss


def prepare_goemotions_dataset():
    """Load and prepare GoEmotions dataset with mapped labels."""
    print("📦 Loading GoEmotions dataset...")
    raw_dataset = load_dataset('google-research-datasets/go_emotions', 'simplified')

    # Get original label names
    original_labels = raw_dataset['train'].features['labels'].feature.names

    # Create label mapping
    label2id = {label: i for i, label in enumerate(TARGET_LABELS)}

    def map_emotions(example):
        """Map GoEmotions labels to our target labels."""
        mapped_labels = []
        for label_id in example['labels']:
            original_label = original_labels[label_id]
            if original_label in GOEMOTIONS_MAPPING:
                target_label = GOEMOTIONS_MAPPING[original_label]
                if target_label in label2id:
                    mapped_labels.append(label2id[target_label])

        # If no valid mapping, skip (return None which we filter later)
        if not mapped_labels:
            return {'label': -1, 'text': example['text']}

        # Use the first mapped label (most confident)
        return {'label': mapped_labels[0], 'text': example['text']}

    print("🔄 Mapping emotions to target labels...")
    mapped_dataset = {}
    for split in ['train', 'validation', 'test']:
        mapped = raw_dataset[split].map(map_emotions, remove_columns=['labels', 'id'])
        # Filter out unmapped samples
        mapped = mapped.filter(lambda x: x['label'] != -1)
        mapped_dataset[split] = mapped

    return DatasetDict(mapped_dataset)


def prepare_sarcasm_dataset():
    """Load and prepare Kaggle SARC dataset for sarcasm."""
    # Assuming user downloads the dataset to data/train-balanced-sarcasm.csv
    csv_path = config.data_dir / "train-balanced-sarcasm.csv"
    if not csv_path.exists():
        print(f"⚠️ Sarcasm dataset not found at {csv_path}. Please download it from:")
        print("   https://www.kaggle.com/datasets/danofer/sarcasm")
        print("   Skipping sarcasm data...")
        return None

    print(f"📦 Loading Kaggle Sarcasm dataset from {csv_path}...")
    # Load dataset, taking only sarcastic rows, and drop NA
    df = pd.read_csv(csv_path)
    df = df.dropna(subset=['comment', 'label'])
    
    # We only take actual sarcastic entries (label == 1), 
    # to avoid muddying neutral/normal text
    df_sarcastic = df[df['label'] == 1].copy()
    
    # Cap it so we don't flood the model with 500k sarcasm samples compared to 59k normal
    # We will sample 25,000 sarcastic texts (roughly 30% of total dataset)
    if len(df_sarcastic) > 25000:
        df_sarcastic = df_sarcastic.sample(n=25000, random_state=42)

    # Convert to our schema (sarcasm is label index 7)
    # Strategy 4: Leverage [SEP] for context windows
    if 'parent_comment' in df_sarcastic.columns:
        df_sarcastic['text'] = df_sarcastic['parent_comment'].astype(str) + " [SEP] " + df_sarcastic['comment'].astype(str)
    else:
        df_sarcastic['text'] = df_sarcastic['comment'].astype(str)
        
    df_sarcastic['label'] = 7  # 'sarcasm' index in TARGET_LABELS

    # Split into train/val/test
    train_texts, temp_texts, train_labels, temp_labels = train_test_split(
        df_sarcastic['text'], df_sarcastic['label'], test_size=0.2, random_state=42
    )
    val_texts, test_texts, val_labels, test_labels = train_test_split(
        temp_texts, temp_labels, test_size=0.5, random_state=42
    )

    # Strategy 1: Contrast Data Augmentation
    contrast_data = [
        "I love it when my tire pops on the highway",
        "I love spending 5 hours in traffic.",
        "Oh, fantastic! The server is down again.",
        "Great, another mandatory team-building exercise.",
        "I'm absolutely thrilled that my flight was canceled.",
        "Wow, you really outdid yourself this time. Breaking the production server on a Friday takes true talent.",
        "What a wonderful surprise, taking a pay cut.",
        "I just adore getting completely ignored.",
        "Best day ever, everything went wrong.",
        "I love getting stuck in the rain without an umbrella.",
        "I love it when people talk over me.",
        "It's just fantastic when my coffee spills all over my keyboard.",
        "Absolutely amazing how you managed to mess that up.",
        "Great job breaking the build.",
        "I'm so happy my alarm didn't go off today."
    ]
    
    train_texts_list = train_texts.tolist() + contrast_data
    train_labels_list = train_labels.tolist() + [7] * len(contrast_data)

    sarc_datasets = {
        'train': Dataset.from_dict({'text': train_texts_list, 'label': train_labels_list}),
        'validation': Dataset.from_dict({'text': val_texts.tolist(), 'label': val_labels.tolist()}),
        'test': Dataset.from_dict({'text': test_texts.tolist(), 'label': test_labels.tolist()}),
    }
    
    return DatasetDict(sarc_datasets)


def prepare_combined_dataset():
    """Combine dair-ai/emotion with GoEmotions for better coverage."""
    print("📦 Loading and combining datasets...")

    # Load dair-ai/emotion
    dair_dataset = load_dataset('dair-ai/emotion')
    dair_labels = ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']

    # Load GoEmotions
    go_dataset = prepare_goemotions_dataset()

    # dair-ai labels are already 0-5 for sadness, joy, love, anger, fear, surprise
    # which matches our TARGET_LABELS[0:6], so no remapping needed

    print("🔗 Combining datasets...")
    
    # Check if user has downloaded SARC
    try:
        sarc_dataset = prepare_sarcasm_dataset()
    except Exception as e:
        print(f"Error loading sarcasm dataset: {e}")
        sarc_dataset = None

    combined = {}
    for split in ['train', 'validation', 'test']:
        # Directly combine text and labels from both datasets
        # Convert to lists explicitly (newer datasets versions return Column objects)
        combined_text = list(dair_dataset[split]['text']) + list(go_dataset[split]['text'])
        combined_label = list(dair_dataset[split]['label']) + list(go_dataset[split]['label'])
        
        if sarc_dataset and split in sarc_dataset:
            combined_text += list(sarc_dataset[split]['text'])
            combined_label += list(sarc_dataset[split]['label'])
            
        combined_data = {
            'text': combined_text,
            'label': combined_label
        }
        combined[split] = Dataset.from_dict(combined_data).shuffle(seed=42)

    return DatasetDict(combined)


def train(
    output_dir: Path = None,
    use_sample: bool = False,
    num_train_samples: int = None,
    use_goemotions: bool = True,
    combine_datasets: bool = False,
    use_class_weights: bool = True,
) -> str:
    """
    Train the emotion classifier with improvements.

    Args:
        output_dir: Where to save the model
        use_sample: Use subset for quick testing
        num_train_samples: Limit training samples
        use_goemotions: Use GoEmotions dataset (larger, more diverse)
        combine_datasets: Combine GoEmotions with dair-ai/emotion
        use_class_weights: Apply class weights for imbalanced data

    Returns:
        Path to saved model
    """
    from datetime import datetime
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    # Save to unique dir if user doesn't pass one manually
    if not output_dir:
        if use_sample:
            output_dir = config.model_dir / "sample_models" / f"emotion_classifier_sample_{timestamp}"
        else:
            output_dir = config.model_dir / f"emotion_classifier_{timestamp}"
    
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Load dataset
    if combine_datasets:
        dataset = prepare_combined_dataset()
        num_labels = 8  # Including neutral and sarcasm
        label_names = TARGET_LABELS
    elif use_goemotions:
        dataset = prepare_goemotions_dataset()
        num_labels = 8  # Including neutral and sarcasm
        label_names = TARGET_LABELS
    else:
        print(f"📦 Loading dataset: {config.hf_dataset_name}")
        dataset = load_dataset(config.hf_dataset_name)
        num_labels = 6
        label_names = list(config.label_names)[:6]

    # Subset for quick testing
    if use_sample:
        print("⚡ Using sample subset for quick testing...")
        dataset["train"] = dataset["train"].select(range(min(2000, len(dataset["train"]))))
        dataset["validation"] = dataset["validation"].select(range(min(500, len(dataset["validation"]))))
        dataset["test"] = dataset["test"].select(range(min(500, len(dataset["test"]))))

    if num_train_samples and num_train_samples < len(dataset["train"]):
        dataset["train"] = dataset["train"].select(range(num_train_samples))

    print(f"\n📊 Dataset sizes:")
    print(f"   Train: {len(dataset['train'])}")
    print(f"   Validation: {len(dataset['validation'])}")
    print(f"   Test: {len(dataset['test'])}")

    # Show class distribution
    train_labels = dataset['train']['label']
    label_dist = Counter(train_labels)
    print(f"\n📈 Class distribution (train):")
    for label_id, count in sorted(label_dist.items()):
        if label_id < len(label_names):
            print(f"   {label_names[label_id]:>10}: {count:>5} ({count/len(train_labels)*100:.1f}%)")

    sarcasm_train_prior = None
    if "sarcasm" in label_names:
        sarcasm_idx = label_names.index("sarcasm")
        sarcasm_train_prior = label_dist.get(sarcasm_idx, 0) / len(train_labels)
        print(f"\n🎯 Sarcasm training prior: {sarcasm_train_prior:.2%}")

    # Compute class weights
    class_weights = None
    if use_class_weights:
        print("\n⚖️ Computing class weights for balancing...")
        unique_labels = sorted(set(train_labels))
        class_weights = compute_class_weight(
            class_weight='balanced',
            classes=np.array(unique_labels),
            y=np.array(train_labels)
        )
                
        print(f"   Weights: {dict(zip([label_names[i] for i in unique_labels], class_weights.round(2)))}")

    # Load tokenizer
    print(f"\n🔧 Loading tokenizer: {config.model_name}")
    tokenizer = AutoTokenizer.from_pretrained(config.model_name)

    # Tokenize with longer max_length for better context
    max_length = 256  # Increased from 128 for longer texts

    def tokenize_fn(examples):
        return tokenizer(
            examples["text"],
            padding="max_length",
            truncation=True,
            max_length=max_length
        )

    print(f"🔄 Tokenizing dataset (max_length={max_length})...")
    
    # Strategy 3: Trigger Word Masking
    import random
    import re
    TRIGGER_WORDS = ["love", "great", "amazing", "fantastic", "wonderful", "thrilled", "joy", "happy", "best"]
    
    def apply_masking(examples):
        # We only apply this to training text to force context over keywords
        masked_texts = []
        for text in examples["text"]:
            for word in TRIGGER_WORDS:
                if random.random() < 0.15:
                    text = re.sub(rf'\b{word}\b', "[MASK]", text, flags=re.IGNORECASE)
            masked_texts.append(text)
        examples["text"] = masked_texts
        return examples

    print("🎭 Applying trigger word masking to training set...")
    dataset["train"] = dataset["train"].map(apply_masking, batched=True)

    tokenized = dataset.map(tokenize_fn, batched=True, remove_columns=["text"])
    tokenized.set_format("torch")

    # Load model
    print(f"🧠 Loading model: {config.model_name}")
    model = AutoModelForSequenceClassification.from_pretrained(
        config.model_name,
        num_labels=num_labels,
        id2label={i: label for i, label in enumerate(label_names)},
        label2id={label: i for i, label in enumerate(label_names)},
    )

    if sarcasm_train_prior is not None and 0.0 < sarcasm_train_prior < 1.0:
        model.config.sarcasm_train_prior = float(sarcasm_train_prior)

    import torch
    
    # Strategy 2: Gradual Unfreezing / Differential Learning Rates
    print("🧠 Setting up differential learning rates...")
    head_params, body_params = [], []
    for name, param in model.named_parameters():
        if "classifier" in name or "pooler" in name:
            head_params.append(param)
        else:
            body_params.append(param)

    # DeBERTa transformer base gets 1e-6 to protect its understanding of language.
    # The new linear classifier gets 5e-5 to map those embeddings to our 8 specific labels quickly.
    optimizer_grouped_parameters = [
        {"params": body_params, "lr": 1e-6},
        {"params": head_params, "lr": 5e-5},
    ]
    custom_optimizer = torch.optim.AdamW(optimizer_grouped_parameters, weight_decay=0.01, eps=1e-6)

    # Training arguments
    training_args = TrainingArguments(
        output_dir=str(output_dir),
        eval_strategy="epoch",
        save_strategy="epoch",
        learning_rate=2e-5, # Back down slightly
        per_device_train_batch_size=4, # Dropped from 16 to 4 to prevent VRAM spill over to system RAM under pure FP32
        gradient_accumulation_steps=4, # 4x4 = 16 effective batch size mathematically identical to original
        per_device_eval_batch_size=8,
        num_train_epochs=5,
        warmup_ratio=0.1,
        weight_decay=0.01,
        adam_epsilon=1e-6, # CRITICAL: DeBERTa V3 AdamW requires this to avoid division by zero during early normalization
        load_best_model_at_end=True,
        metric_for_best_model="f1_macro",
        greater_is_better=True,
        logging_steps=10,
        report_to="none",
        fp16=False,
        bf16=False, # Disable BF16 entirely to guarantee purely stable FP32 gradients
        max_grad_norm=1.0, 
    )

    # Initialize trainer (with or without class weights)
    if use_class_weights and class_weights is not None:
        trainer = WeightedTrainer(
            class_weights=list(class_weights),
            model=model,
            args=training_args,
            train_dataset=tokenized["train"],
            eval_dataset=tokenized["validation"],
            compute_metrics=compute_metrics,
            callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
            optimizers=(custom_optimizer, None)
        )
    else:
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=tokenized["train"],
            eval_dataset=tokenized["validation"],
            compute_metrics=compute_metrics,
            callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
            optimizers=(custom_optimizer, None)
        )

    # Train
    print("\n🚀 Starting training with class balancing...")
    print("=" * 50)
    
    # Starting a fully fresh training run
    trainer.train()

    # Evaluate
    print("\n📊 Evaluating on test set...")
    test_results = trainer.evaluate(tokenized["test"])
    print(f"\n✅ Test Results:")
    print(f"   Accuracy: {test_results['eval_accuracy']:.2%}")
    print(f"   F1 (macro): {test_results['eval_f1_macro']:.2%}")
    print(f"   F1 (weighted): {test_results['eval_f1_weighted']:.2%}")

    # Save
    final_path = output_dir / "final"
    model.save_pretrained(final_path)
    tokenizer.save_pretrained(final_path)

    print(f"\n💾 Model saved to: {final_path}")

    return str(final_path)


if __name__ == "__main__":
    train(use_goemotions=True, use_class_weights=True)