"""Dual-head fine-tuning trainer for ModernBERT.

Supports the ablation grid: {base, +DAPT, +DAPT+TAPT} x {+/-class weighting} x {CE vs focal loss}.
Uses HuggingFace Trainer with custom loss computation for the dual-head architecture.
"""

import unsloth  # noqa: F401 — monkey-patches transformers for faster kernels

from dataclasses import dataclass
from pathlib import Path
from itertools import product

import numpy as np
import torch
from datasets import DatasetDict
from sklearn.metrics import f1_score, cohen_kappa_score
from transformers import (
    AutoModel,
    AutoTokenizer,
    Trainer,
    TrainingArguments,
    EvalPrediction,
)

from ..common.config import FinetuneConfig
from .data import (
    CATEGORIES,
    NUM_CATEGORIES,
    NUM_SPECIFICITY,
    load_finetune_data,
    compute_class_weights,
)
from .model import DualHeadModernBERT, ordinal_predict, softmax_predict


@dataclass
class DualHeadCollator:
    """Pads input_ids/attention_mask and stacks label columns into tensors."""

    tokenizer: object
    padding: str = "longest"

    def __call__(self, features: list[dict]) -> dict:
        # Pop labels before padding (they're scalars, not sequences)
        cat_labels = [f.pop("category_labels") for f in features]
        spec_labels = [f.pop("specificity_labels") for f in features]
        weights = [f.pop("sample_weight") for f in features]
        spec_masks = [f.pop("spec_mask") for f in features]

        # Pad input sequences
        batch = self.tokenizer.pad(
            features, padding=self.padding, return_tensors="pt"
        )

        # Add labels back as tensors
        batch["category_labels"] = torch.tensor(cat_labels, dtype=torch.long)
        batch["specificity_labels"] = torch.tensor(spec_labels, dtype=torch.long)
        batch["sample_weight"] = torch.tensor(weights, dtype=torch.float)
        batch["spec_mask"] = torch.tensor(spec_masks, dtype=torch.float)

        return batch


class DualHeadTrainer(Trainer):
    """HuggingFace Trainer adapted for dual-head model output."""

    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        outputs = model(**inputs)
        loss = outputs["loss"]
        return (loss, outputs) if return_outputs else loss

    def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
        inputs = self._prepare_inputs(inputs)
        with torch.no_grad():
            outputs = model(**inputs)

        loss = outputs["loss"]
        if prediction_loss_only:
            return (loss, None, None)

        cat_logits = outputs["category_logits"]
        spec_logits = outputs["specificity_logits"]
        cat_labels = inputs["category_labels"]
        spec_labels = inputs["specificity_labels"]

        return loss, (cat_logits, spec_logits), (cat_labels, spec_labels)


def build_compute_metrics(spec_head_type: str = "independent"):
    """Return a compute_metrics function for the dual-head model."""

    def compute_metrics(eval_pred: EvalPrediction) -> dict:
        (cat_logits, spec_logits) = eval_pred.predictions
        (cat_labels, spec_labels) = eval_pred.label_ids

        # Category predictions
        cat_preds = np.argmax(cat_logits, axis=1)
        cat_macro_f1 = f1_score(cat_labels, cat_preds, average="macro")
        cat_weighted_f1 = f1_score(cat_labels, cat_preds, average="weighted")
        cat_per_class = f1_score(
            cat_labels, cat_preds, average=None, labels=range(NUM_CATEGORIES)
        )

        # Specificity predictions
        spec_tensor = torch.tensor(spec_logits, dtype=torch.float32)
        if spec_head_type == "softmax":
            spec_preds = softmax_predict(spec_tensor).numpy()
        else:
            spec_preds = ordinal_predict(spec_tensor).numpy()
        spec_macro_f1 = f1_score(
            spec_labels, spec_preds, average="macro",
            labels=range(NUM_SPECIFICITY),
        )
        spec_weighted_f1 = f1_score(
            spec_labels, spec_preds, average="weighted",
            labels=range(NUM_SPECIFICITY),
        )
        spec_per_class = f1_score(
            spec_labels, spec_preds, average=None,
            labels=range(NUM_SPECIFICITY),
        )
        # Quadratic weighted kappa for ordinal agreement
        spec_qwk = cohen_kappa_score(
            spec_labels, spec_preds, weights="quadratic"
        )

        # Aggregate metric for model selection
        combined_f1 = (cat_macro_f1 + spec_macro_f1) / 2

        metrics = {
            "combined_macro_f1": combined_f1,
            "cat_macro_f1": cat_macro_f1,
            "cat_weighted_f1": cat_weighted_f1,
            "spec_macro_f1": spec_macro_f1,
            "spec_weighted_f1": spec_weighted_f1,
            "spec_qwk": spec_qwk,
        }

        # Per-class F1 for category
        for i, name in enumerate(CATEGORIES):
            short = name.replace(" ", "").replace("/", "")[:6]
            metrics[f"cat_f1_{short}"] = cat_per_class[i]

        # Per-class F1 for specificity
        for i in range(NUM_SPECIFICITY):
            metrics[f"spec_f1_L{i+1}"] = spec_per_class[i]

        return metrics

    return compute_metrics


def train(config: FinetuneConfig) -> dict:
    """Run a single fine-tuning experiment. Returns final eval metrics."""
    print(f"\n{'='*60}")
    print(f"  SEC-cyBERT Fine-tuning")
    print(f"  Model:          {config.model.name_or_path}")
    print(f"  Loss:           {config.training.loss_type}")
    print(f"  Class weights:  {config.training.class_weighting}")
    print(f"  Output:         {config.training.output_dir}")
    print(f"{'='*60}\n")

    # Load tokenizer
    tokenizer_name = config.model.tokenizer or config.model.name_or_path
    tokenizer = AutoTokenizer.from_pretrained(
        tokenizer_name,
        trust_remote_code=config.model.trust_remote_code,
    )

    # Load and prepare data (shared cache across ablation runs)
    output_dir = Path(config.training.output_dir)
    # Use a stable cache location based on data paths, not output_dir
    cache_dir = Path(config.data.consensus_path).parent / ".finetune_data_cache"

    if cache_dir.exists():
        print(f"  Loading cached dataset from {cache_dir}...")
        split = DatasetDict.load_from_disk(str(cache_dir))
        print(f"  Train: {len(split['train']):,} | Val: {len(split['test']):,}\n")
    else:
        split = load_finetune_data(
            paragraphs_path=config.data.paragraphs_path,
            consensus_path=config.data.consensus_path,
            quality_path=config.data.quality_path,
            holdout_path=config.data.holdout_path,
            max_seq_length=config.data.max_seq_length,
            validation_split=config.data.validation_split,
            tokenizer=tokenizer,
            seed=config.training.seed,
        )
        cache_dir.mkdir(parents=True, exist_ok=True)
        split.save_to_disk(str(cache_dir))
        print(f"  Cached to {cache_dir}\n")

    # If not filtering spec confidence, override spec_mask to all 1.0
    if not config.training.filter_spec_confidence:
        split["train"] = split["train"].map(
            lambda x: {"spec_mask": 1.0}, desc="Disabling spec filter"
        )
        split["test"] = split["test"].map(
            lambda x: {"spec_mask": 1.0}, desc="Disabling spec filter"
        )

    # Compute class weights if needed
    category_weights = None
    specificity_weights = None
    if config.training.class_weighting:
        weights = compute_class_weights(split["train"])
        category_weights = torch.tensor(weights["category_weights"], dtype=torch.float32)
        specificity_weights = weights["specificity_weights"]
        print(f"  Category weights:    {[f'{w:.2f}' for w in weights['category_weights']]}")
        print(f"  Specificity weights: {[f'{w:.2f}' for w in weights['specificity_weights']]}\n")

    # Load backbone
    try:
        import flash_attn  # noqa: F401
        attn_impl = "flash_attention_2"
    except ImportError:
        attn_impl = "sdpa"

    backbone = AutoModel.from_pretrained(
        config.model.name_or_path,
        trust_remote_code=config.model.trust_remote_code,
        attn_implementation=attn_impl,
        dtype=torch.bfloat16 if config.training.bf16 else None,
    )
    hidden_size = backbone.config.hidden_size
    print(f"  Backbone parameters: {sum(p.numel() for p in backbone.parameters()) / 1e6:.0f}M")
    print(f"  Hidden size: {hidden_size}")
    print(f"  Attention: {attn_impl}\n")

    # Build dual-head model
    spec_head = config.training.specificity_head
    print(f"  Specificity head: {spec_head}")
    print(f"  Spec MLP dim: {config.training.spec_mlp_dim}")
    print(f"  Pooling: {config.training.pooling}")
    print(f"  Filter spec confidence: {config.training.filter_spec_confidence}")
    print(f"  Ordinal consistency: {config.training.ordinal_consistency_weight}\n")

    model = DualHeadModernBERT(
        backbone=backbone,
        hidden_size=hidden_size,
        num_categories=NUM_CATEGORIES,
        num_specificity=NUM_SPECIFICITY,
        dropout=config.training.dropout,
        loss_type=config.training.loss_type,
        focal_gamma=config.training.focal_gamma,
        category_weights=category_weights,
        specificity_weights=specificity_weights,
        category_loss_weight=config.training.category_loss_weight,
        specificity_loss_weight=config.training.specificity_loss_weight,
        specificity_head_type=spec_head,
        spec_mlp_dim=config.training.spec_mlp_dim,
        pooling=config.training.pooling,
        ordinal_consistency_weight=config.training.ordinal_consistency_weight,
    )

    # Data collator
    collator = DualHeadCollator(tokenizer=tokenizer)

    # Training arguments
    steps_per_epoch = len(split["train"]) // (
        config.training.per_device_train_batch_size
        * config.training.gradient_accumulation_steps
    )

    training_kwargs = dict(
        output_dir=str(output_dir),
        learning_rate=config.training.learning_rate,
        num_train_epochs=config.training.num_train_epochs,
        per_device_train_batch_size=config.training.per_device_train_batch_size,
        per_device_eval_batch_size=config.training.per_device_eval_batch_size,
        gradient_accumulation_steps=config.training.gradient_accumulation_steps,
        warmup_steps=int(config.training.warmup_ratio * steps_per_epoch * config.training.num_train_epochs),
        weight_decay=config.training.weight_decay,
        bf16=config.training.bf16,
        tf32=True,
        torch_compile=True,
        optim="adamw_torch_fused",
        gradient_checkpointing=config.training.gradient_checkpointing,
        dataloader_persistent_workers=True,
        dataloader_num_workers=config.training.dataloader_num_workers,
        logging_steps=config.training.logging_steps,
        eval_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=config.training.save_total_limit,
        seed=config.training.seed,
        remove_unused_columns=False,  # we have custom columns
        report_to="none",
        load_best_model_at_end=True,
        metric_for_best_model="spec_macro_f1",  # optimize for specificity (category already > 0.80)
        greater_is_better=True,
    )

    if config.training.gradient_checkpointing:
        training_kwargs["gradient_checkpointing_kwargs"] = {"use_reentrant": False}

    args = TrainingArguments(**training_kwargs)

    trainer = DualHeadTrainer(
        model=model,
        args=args,
        train_dataset=split["train"],
        eval_dataset=split["test"],
        data_collator=collator,
        compute_metrics=build_compute_metrics(spec_head_type=spec_head),
    )

    # Train
    resume = config.training.resume_from_checkpoint
    if resume is None and any(output_dir.glob("checkpoint-*")):
        resume = True
    trainer.train(resume_from_checkpoint=resume)

    # Save final model + tokenizer
    final_dir = output_dir / "final"
    print(f"\n  Saving final model to {final_dir}...")
    trainer.save_model(str(final_dir))
    tokenizer.save_pretrained(str(final_dir))

    # Final eval
    metrics = trainer.evaluate()
    print(f"\n  Final metrics:")
    for k, v in sorted(metrics.items()):
        if isinstance(v, float):
            print(f"    {k}: {v:.4f}")

    return metrics


# ── Ablation grid ────────────────────────────────────────────────────────────

ABLATION_CHECKPOINTS = {
    "base": "answerdotai/ModernBERT-large",
    "dapt": "../checkpoints/dapt/modernbert-large/final",
    "tapt": "../checkpoints/tapt/modernbert-large/final",
}


def ablate(config: FinetuneConfig) -> None:
    """Run the full ablation grid: 3 checkpoints x 2 weighting x 2 loss = 12 runs."""
    base_output = Path(config.training.output_dir)
    results: list[dict] = []

    checkpoints = list(ABLATION_CHECKPOINTS.items())
    weightings = [False, True]
    losses = ["ce", "focal"]

    combos = list(product(checkpoints, weightings, losses))
    print(f"\n  Ablation grid: {len(combos)} configurations")
    print(f"  Checkpoints: {[c[0] for c in checkpoints]}")
    print(f"  Weighting: {weightings}")
    print(f"  Losses: {losses}")
    print()

    for i, ((ckpt_name, ckpt_path), weighting, loss_type) in enumerate(combos):
        run_name = f"{ckpt_name}_{'weighted' if weighting else 'unweighted'}_{loss_type}"
        run_output = base_output / run_name

        print(f"\n{'#'*60}")
        print(f"  Ablation run {i+1}/{len(combos)}: {run_name}")
        print(f"{'#'*60}")

        # Override config for this run
        config.model.name_or_path = ckpt_path
        config.training.output_dir = str(run_output)
        config.training.class_weighting = weighting
        config.training.loss_type = loss_type

        try:
            metrics = train(config)
            results.append({"run": run_name, **metrics})
        except Exception as e:
            print(f"\n  FAILED: {run_name}: {e}")
            results.append({"run": run_name, "error": str(e)})

    # Summary table
    print(f"\n\n{'='*80}")
    print("  ABLATION RESULTS SUMMARY")
    print(f"{'='*80}")
    print(f"  {'Run':<45} {'Combined':>10} {'Cat F1':>10} {'Spec F1':>10} {'QWK':>10}")
    print(f"  {'-'*45} {'-'*10} {'-'*10} {'-'*10} {'-'*10}")

    for r in results:
        name = r["run"]
        if "error" in r:
            print(f"  {name:<45} {'FAILED':>10}")
        else:
            combined = r.get("eval_combined_macro_f1", 0)
            cat = r.get("eval_cat_macro_f1", 0)
            spec = r.get("eval_spec_macro_f1", 0)
            qwk = r.get("eval_spec_qwk", 0)
            print(f"  {name:<45} {combined:>10.4f} {cat:>10.4f} {spec:>10.4f} {qwk:>10.4f}")

    # Save results
    import json
    results_path = base_output / "ablation_results.json"
    results_path.parent.mkdir(parents=True, exist_ok=True)
    with open(results_path, "w") as f:
        json.dump(results, f, indent=2, default=str)
    print(f"\n  Results saved to {results_path}")
