"""Classifier with LoRA and class-weighted loss."""

import torch
import torch.nn as nn
from transformers import AutoModelForSequenceClassification, Trainer
from peft import LoraConfig, get_peft_model, TaskType

from config import CONFIG


def load_model():
    """Load base model and wrap with LoRA adapter for efficient fine-tuning."""
    base_model = AutoModelForSequenceClassification.from_pretrained(
        CONFIG["model_name"],
        num_labels=CONFIG["num_labels"],
    )

    if CONFIG.get("gradient_checkpointing", False):
        base_model.gradient_checkpointing_enable()

    lora_config = LoraConfig(
        task_type=TaskType.SEQ_CLS,
        r=CONFIG.get("lora_r", 16),
        lora_alpha=CONFIG.get("lora_alpha", 32),
        lora_dropout=CONFIG.get("lora_dropout", 0.05),
        target_modules=CONFIG.get("lora_target_modules", ["query", "value"]),
        modules_to_save=["classifier"],  # train the classification head fully
    )

    model = get_peft_model(base_model, lora_config)
    model.print_trainable_parameters()

    return model


class WeightedTrainer(Trainer):
    """Trainer subclass that uses class-weighted CrossEntropyLoss.

    Compatible with transformers v4.x and v5.x.
    """

    def __init__(self, class_weights=None, **kwargs):
        super().__init__(**kwargs)
        if class_weights is not None:
            self._class_weights = torch.tensor(class_weights, dtype=torch.float32)
        else:
            self._class_weights = None

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        labels = inputs.get("labels")
        outputs = model(**inputs)
        logits = outputs.get("logits")

        if self._class_weights is not None:
            weight = self._class_weights.to(logits.device)
            loss_fn = nn.CrossEntropyLoss(weight=weight)
        else:
            loss_fn = nn.CrossEntropyLoss()

        loss = loss_fn(logits.view(-1, self.model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else loss