"""
modelling.py
------------
Wraps the intfloat/multilingual-e5-large-instruct SentenceTransformer with LoRA
fine-tuning support via PEFT, and exposes a simple .train() interface.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Any, Literal

import torch
from datasets import Dataset
from peft import LoraConfig, TaskType, get_peft_model
from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, SentenceTransformerTrainingArguments
from sentence_transformers.evaluation import InformationRetrievalEvaluator
from sentence_transformers.losses import CachedMultipleNegativesRankingLoss
from sentence_transformers.util import cos_sim
from transformers import EarlyStoppingCallback

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Configuration dataclasses – give users a clean surface to customise runs
# ---------------------------------------------------------------------------

@dataclass
class LoraSettings:
    """LoRA adapter hyper-parameters."""

    r: int = 8
    lora_alpha: int = 8
    lora_dropout: float = 0.0
    bias: Literal["none", "all", "lora_only"] = "none"
    target_modules: list[str] = field(
        default_factory=lambda: ["query", "key", "value", "dense"]
    )


@dataclass
class TrainingSettings:
    """Training hyper-parameters exposed to the caller."""

    max_steps: int = 100
    per_device_train_batch_size: int = 4
    per_device_eval_batch_size: int = 8
    learning_rate: float = 1e-5
    lr_scheduler_type: str = "cosine"
    optim: str = "adafactor"
    eval_steps: int = 8
    fp16: bool = True
    early_stopping_patience: int = 2
    mini_batch_size: int = 128
    evaluator_batch_size: int = 32


@dataclass
class TrainResult:
    """Holds evaluation metrics before and after fine-tuning."""

    baseline: dict[str, Any]
    final_metrics: dict[str, Any]

    def improvement(self) -> dict[str, float]:
        """Return per-metric delta (final − baseline)."""
        return {
            k: self.final_metrics.get(k, 0) - self.baseline.get(k, 0)
            for k in self.final_metrics
        }


# ---------------------------------------------------------------------------
# Main modelling class
# ---------------------------------------------------------------------------

class MultilingualE5LargeInstructModelling:
    """
    Fine-tunes *intfloat/multilingual-e5-large-instruct* with LoRA adapters.

    Parameters
    ----------
    lora_settings:
        LoRA adapter configuration.  Defaults to :class:`LoraSettings`.
    training_settings:
        Training hyper-parameters.  Defaults to :class:`TrainingSettings`.

    Example
    -------
    >>> from fiesta import MultilingualE5LargeInstructModelling
    >>> modeller = MultilingualE5LargeInstructModelling()
    >>> results = modeller.train(train_ds, test_ds, evaluator_data, save_dir="./model")
    """

    MODEL_ID: str = "intfloat/multilingual-e5-large-instruct"

    def __init__(
        self,
        lora_settings: LoraSettings | None = None,
        training_settings: TrainingSettings | None = None,
    ) -> None:
        self.lora_settings = lora_settings or LoraSettings()
        self.training_settings = training_settings or TrainingSettings()

        device = "cuda" if torch.cuda.is_available() else "cpu"
        if self.training_settings.fp16 and device == "cpu":
            logger.warning(
                "fp16=True is set but no CUDA GPU was detected. "
                "fp16 training on CPU is not supported and will be disabled automatically by the trainer. "
                "Consider setting TrainingSettings(fp16=False) for CPU runs."
            )
        logger.info("Loading base model %s on %s ...", self.MODEL_ID, device)

        self.model = SentenceTransformer(self.MODEL_ID, device=device)
        self.model[0].auto_model.enable_input_require_grads()

        lora_cfg = LoraConfig(
            r=self.lora_settings.r,
            lora_alpha=self.lora_settings.lora_alpha,
            lora_dropout=self.lora_settings.lora_dropout,
            bias=self.lora_settings.bias,
            target_modules=self.lora_settings.target_modules,
            task_type=TaskType.FEATURE_EXTRACTION,
        )
        self.model[0].auto_model = get_peft_model(self.model[0].auto_model, lora_cfg)
        self.model[0].auto_model.print_trainable_parameters()

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _build_training_args(self, save_dir: str) -> SentenceTransformerTrainingArguments:
        ts = self.training_settings
        return SentenceTransformerTrainingArguments(
            output_dir=save_dir,
            max_steps=ts.max_steps,
            per_device_train_batch_size=ts.per_device_train_batch_size,
            per_device_eval_batch_size=ts.per_device_eval_batch_size,
            gradient_checkpointing=True,
            learning_rate=ts.learning_rate,
            lr_scheduler_type=ts.lr_scheduler_type,
            optim=ts.optim,
            eval_steps=ts.eval_steps,
            eval_strategy="steps",
            logging_strategy="steps",
            report_to="none",
            logging_steps=ts.eval_steps,
            save_strategy="steps",
            save_steps=ts.eval_steps,
            save_only_model=True,
            save_total_limit=1,
            load_best_model_at_end=True,
            greater_is_better=True,
            metric_for_best_model="eval_cosine_accuracy@10",
            disable_tqdm=False,
            fp16=ts.fp16,
        )

    def _build_evaluator(
        self,
        queries: dict[str, str],
        corpus: dict[str, str],
        relevant_docs: dict[str, list[str]],
    ) -> InformationRetrievalEvaluator:
        # InformationRetrievalEvaluator requires sets for relevant_docs
        relevant_docs_sets: dict[str, set[str]] = {
            k: set(v) for k, v in relevant_docs.items()
        }
        return InformationRetrievalEvaluator(
            queries=queries,
            corpus=corpus,
            relevant_docs=relevant_docs_sets,
            score_functions={"cosine": cos_sim},
            batch_size=self.training_settings.evaluator_batch_size,
        )

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def train(
        self,
        train_dataset: Dataset,
        test_dataset: Dataset,
        evaluator_data: dict[str, Any],
        save_dir: str,
    ) -> TrainResult:
        """
        Fine-tune the model and return baseline + final evaluation metrics.

        Parameters
        ----------
        train_dataset:
            HuggingFace Dataset with at least ``anchor`` and ``positive`` columns.
        test_dataset:
            Same schema as ``train_dataset``, used for evaluation.
        evaluator_data:
            Dict produced by :meth:`MultilingualE5LargeInstructPreProcessing.get_evaluator_data`
            with keys ``test_queries``, ``train_test_chunks``,
            ``testQueryId_to_chunkId``.
        save_dir:
            Directory where the best checkpoint will be saved.

        Returns
        -------
        TrainResult
            A dataclass with ``baseline`` and ``final_metrics`` dicts, plus
            an ``improvement()`` helper method.
        """
        args = self._build_training_args(save_dir)
        evaluator = self._build_evaluator(
            queries=evaluator_data["test_queries"],
            corpus=evaluator_data["train_test_chunks"],
            relevant_docs=evaluator_data["testQueryId_to_chunkId"],
        )

        logger.info("Evaluating baseline ...")
        baseline = evaluator(self.model)

        loss_fn = CachedMultipleNegativesRankingLoss(
            self.model,
            mini_batch_size=self.training_settings.mini_batch_size,
        )

        trainer = SentenceTransformerTrainer(
            model=self.model,
            args=args,
            train_dataset=train_dataset.select_columns(["anchor", "positive"]),
            eval_dataset=test_dataset.select_columns(["anchor", "positive"]),
            loss=loss_fn,
            evaluator=evaluator,
            callbacks=[
                EarlyStoppingCallback(
                    self.training_settings.early_stopping_patience, 0.0
                )
            ],
        )

        logger.info("Starting fine-tuning ...")
        trainer.train()

        logger.info("Evaluating fine-tuned model ...")
        final_metrics = evaluator(self.model)

        trainer.save_model()
        logger.info("Model saved to %s", save_dir)

        return TrainResult(baseline=baseline, final_metrics=final_metrics)
