import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import seaborn as sns
from torch.utils.data import Dataset
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.model_selection import train_test_split
from transformers import(
    DebertaV2Tokenizer,
    DebertaV2ForSequenceClassification,
    TrainingArguments,
    Trainer,
    EarlyStoppingCallback
)
# Set seed for reproducibility
SEED = 42
torch.manual_seed(SEED)
np.random.seed(SEED)

class WeightedTrainer(Trainer):
    """Custom trainer that uses weighted loss function for imbaalance data."""

    def __init__(self, *args, class_weights=None, **kwargs):
        """Initialize trainer with class weights.
        Args:
        class_weights: Tensor of weights for each class
        *args, **kwargs: Arguments for parent Trainer class
        """
        super().__init__(*args, **kwargs)
        self.class_weights = class_weights.to(self.model.device)
        self.loss_fn = nn.CrossEntropyLoss(
                        weight=self.class_weights,
                        label_smoothing=0.1
                    )

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        """Compute loss using the weighted loss function

        args:
            model: Model to train
            inputs: Input to batch
            return_outputs: whether to return output along with loss

        Returns:
            Loss value or tuple of (loss, outputs)
        """

        labels = inputs.get("labels")
        outputs = model(**inputs)
        logits = outputs.get("logits")
        loss = self.loss_fn(logits, labels)
        return (loss, outputs) if return_outputs else loss

class EmailClassification(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length):
        """Intialize dataset with texts, labels and tokenizer settings.
        Args:
            texts: List of email Text Strings
            labels: List of corresponding Labels
            tokenizer: Pretrained tokenizer
            max_length: Maximum sequence length
        """
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        """Return the number of sample in the dataset."""
        return len(self.texts)

    def __getitem__(self, idx):
        """Get a single item from the dataset

        Args:
            idx: Index of the Item
        Returns:
            Dict containing input_ids, attention_mask and labels
        """
        text = str(self.texts[idx])
        label = self.labels[idx]

        encoding = self.tokenizer(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt',
            padding='max_length'
        )

        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': torch.tensor(label, dtype=torch.long)
        }


def compute_class_weights(labels):
    """Compute class weights inversly proportional to class frequencies.
    Args:
        labels: Array of labels

    return:
        Tensor of Class weights
    """
    class_counts = np.bincount(labels)
    total_sample = len(labels)
    class_weights = total_sample / (len(class_counts) * class_counts)
    return torch.tensor(class_weights, dtype=torch.float)



def compute_metrics(pred):
    """Compute metrics for evaluation.

    Args:
        pred: Prediction object from trainer

    Returns:
        Dict of metrics
    """
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)

    # Calculate metrics
    report = classification_report(
        labels,
        preds,
        target_names=list(label2id.keys()),
        output_dict=True
    )

    # Extract Metrics
    results = {
        'accuracy': report['accuracy'],
        'f1_macro': report['macro avg']['f1-score'],
        'f1_weighted': report['weighted avg']['f1-score'],
        'precision_weighted': report['weighted avg']['precision'],
        'recall_weighted': report['weighted avg']['recall']
    }

    # Add class specific f1 scores
    for cls_name, cls_id in label2id.items():
        results[f'f1_{cls_name}'] = report[cls_name]['f1-score']

    return results

def train(df, model_name, label2id, id2label, text_column='email_processed',
                        max_length=256, batch_size=16, epochs=5,seed=42):

    """Train and evaluate model using train/val/test split.
    Args:
        df: DataFrame containing the dataset
        model_name: Name of the pretrained model
        label2id: Dictionary mapping label names to IDs
        id2label: Dictionary mapping IDs to label names
        text_column: Column name containing processed email text
        max_length: Maximum sequence length for tokenizer
        batch_size: Batch size for training
        epochs: Number of training epochs
        seed: Random seed for reproducibility

    Returns:
        Tuple of (trained model, tokenizer)
    """
    # Map labels to IDs
    df['label_id'] = df['type'].map(label2id)

    # Load tokenizer
    tokenizer = DebertaV2Tokenizer.from_pretrained(model_name)

    # Prepare Data
    texts = df[text_column].values
    labels = df['label_id'].values

    # Compute class weights
    class_weights = compute_class_weights(labels)
    print("\nClass weights for handling imbalance: ")
    for label, weight in zip(label2id.keys(), class_weights):
        print(f"{label}: {weight.item():.4f}")

    # First, split data into train and temp (val+test combined)
    train_texts, temp_texts, train_labels, temp_labels = train_test_split(
        texts,
        labels,
        test_size=0.3,  # 70% train, 30% for val+test
        stratify=labels,  # Maintain class distribution
        random_state=seed
    )

    # Further split temp into val and test
    val_texts, test_texts, val_labels, test_labels = train_test_split(
        temp_texts,
        temp_labels,
        test_size=0.5,  # 15% val, 15% test (of the total dataset)
        stratify=temp_labels,  # Maintain class distribution
        random_state=seed
    )

    # Create Datasets
    train_dataset = EmailClassification(
        train_texts, train_labels, tokenizer, max_length
    )
    val_dataset = EmailClassification(
        val_texts, val_labels, tokenizer, max_length
    )
    test_dataset = EmailClassification(
        test_texts, test_labels, tokenizer, max_length
    )

    # Load Model
    model = DebertaV2ForSequenceClassification.from_pretrained(
        model_name,
        num_labels=len(label2id),
        id2label=id2label,
        label2id=label2id,
    )
    model.gradient_checkpointing_enable()
    # Training arguments
    training_args = TrainingArguments(
        output_dir="./results",
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size,
        learning_rate=5e-5,  # Starting learning rate
        lr_scheduler_type="cosine",
        warmup_steps=120,
        weight_decay=0.01,
        logging_dir="./logs",
        logging_steps=100,
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="f1_weighted",
        greater_is_better=True,
        save_total_limit=1,
        seed=seed,
        report_to=None,
        fp16=True,  # Enable mixed precision training
        gradient_accumulation_steps=4,
    )

    # Define Trainer
    trainer = WeightedTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        compute_metrics=compute_metrics,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
        class_weights=class_weights,
    )

    # Train the model
    trainer.train()

    # Evaluate on validation set
    print("\nEvaluating on validation set:")
    val_results = trainer.evaluate()
    for metric_name, value in val_results.items():
        print(f"{metric_name}: {value:.4f}")

    # Evaluate on test set
    print("\nEvaluating on test set:")
    test_results = trainer.predict(test_dataset)
    test_metrics = compute_metrics(test_results)
    for metric_name, value in test_metrics.items():
        print(f"{metric_name}: {value:.4f}")

    # Get predictions for test set
    test_preds = np.argmax(test_results.predictions, axis=1)

    # Print classification report for test set
    print("\nTest Set Classification Report:")
    print(classification_report(
        test_labels,
        test_preds,
        target_names=list(label2id.keys())
    ))

    # Plot confusion matrix for test set
    plt.figure(figsize=(10, 8))
    cm = confusion_matrix(test_labels, test_preds)
    sns.heatmap(
        cm,
        annot=True,
        fmt='d',
        cmap='Blues',
        xticklabels=list(label2id.keys()),
        yticklabels=list(label2id.keys())
    )
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.title('Confusion Matrix - Test Set')
    plt.tight_layout()
    plt.savefig('confusion_matrix_test_set.png')
    plt.close()


    model.save_pretrained('./model')
    tokenizer.save_pretrained("./model")

    return model, tokenizer





if __name__ == "__main__":
    """Main function to run the complete pipeline."""
    print("Email Classification Training Pipeline")
    print(f"{'='*40}")
    # Label Mapping
    label2id = {
        'Incident': 0,
        'Request': 1,
        'Problem': 2,
        'Change': 3
    }
    id2label = {v: k for k, v in label2id.items()}
    # Model configuration
    MODEL_NAME = "microsoft/mdeberta-v3-base"
    MAX_LENGTH = 128
    BATCH_SIZE = 4
    EPOCHS = 5
    SEED = 42

    # Load your dataset (example)
    df = pd.read_csv("/content/processed_emails.csv")

    # Train and evaluate
    model, tokenizer = train(
        df=df,
        model_name=MODEL_NAME,
        label2id=label2id,
        id2label=id2label,
        text_column='email_processed',
        max_length=MAX_LENGTH,
        batch_size=BATCH_SIZE,
        epochs=EPOCHS,
        seed=SEED
    )