"""Config-driven fine-tuning of the CEFR classifier (runs on the local GPU box).

Usage (from the repo root):
    uv run --group train python training/train_cefr.py --config <arm>.toml --dry-run
    uv run --group train python training/train_cefr.py --config training/configs/en_only.toml

--dry-run builds the dataset, prints the split summary and runs the leakage
check without importing torch — use it to validate a config cheaply.

Comparability across experiment arms: the split is seeded per (corpus|level)
stratum, so the English test documents are identical whether or not the
multilingual subsets are present (encoded as a unit test in
tests/test_cefr_splitting.py).
"""

import argparse
import json
import math
import tomllib
from collections import Counter
from pathlib import Path

from data_loading import load_passages
from evaluation import evaluate_views, lang_filtered, predict_probs

from tutor.ml.cefr.metrics import classification_report
from tutor.ml.cefr.preprocessing import CANONICAL_LEVELS, Passage
from tutor.ml.cefr.splitting import SPLITS, assign_splits

MODELS_DIR = Path("models/cefr")


def load_config(path: Path) -> dict:
    with path.open("rb") as handle:
        return tomllib.load(handle)


def build_parts(config: dict) -> dict[str, list[Passage]]:
    data = config["data"]
    passages = load_passages(
        data["subsets"],
        chunking=data.get("chunking", True),
        target_words=data.get("target_words", 200),
        max_words=data.get("max_words", 300),
    )
    doc_strata = {p.doc_id: f"{p.corpus}|{p.level}" for p in passages}
    assignment = assign_splits(
        doc_strata,
        ratios=tuple(data.get("ratios", [0.8, 0.1, 0.1])),
        seed=config["run"].get("seed", 13),
    )
    parts = {split: [p for p in passages if assignment[p.doc_id] == split] for split in SPLITS}

    # Generalization-audit arms (ADR 0003 arm 5): the excluded corpora vanish
    # from train/val but keep their canonical test docs, so numbers stay comparable.
    excluded = set(data.get("exclude_corpora_from_train", []))
    if excluded:
        for split in ("train", "val"):
            parts[split] = [p for p in parts[split] if p.corpus not in excluded]

    doc_ids = {split: {p.doc_id for p in parts[split]} for split in SPLITS}
    for first in SPLITS:
        for second in SPLITS:
            if first < second and doc_ids[first] & doc_ids[second]:
                msg = f"document leakage between {first} and {second}"
                raise AssertionError(msg)
    return parts


def summarize(parts: dict[str, list[Passage]]) -> str:
    lines = ["", "split | passages | docs | per-level passage counts", "---|---|---|---"]
    for split in SPLITS:
        passages = parts[split]
        levels = Counter(p.level for p in passages)
        level_cells = ", ".join(f"{lvl}:{levels.get(lvl, 0)}" for lvl in CANONICAL_LEVELS)
        lines.append(
            f"{split} | {len(passages)} | {len({p.doc_id for p in passages})} | {level_cells}"
        )
    corpora = Counter((p.corpus, parts_split) for parts_split in SPLITS for p in parts[parts_split])
    lines += ["", "corpus x split:"]
    for corpus in sorted({c for c, _ in corpora}):
        cells = ", ".join(f"{split}:{corpora.get((corpus, split), 0)}" for split in SPLITS)
        lines.append(f"  {corpus}: {cells}")
    return "\n".join(lines)


def class_weight_values(train_passages: list[Passage]) -> list[float]:
    """Inverse-frequency weights, normalised to mean 1 (absent classes get weight 0)."""
    counts = Counter(p.level for p in train_passages)
    raw = [1.0 / counts[lvl] if counts.get(lvl) else 0.0 for lvl in CANONICAL_LEVELS]
    present = [w for w in raw if w > 0]
    mean = sum(present) / len(present)
    return [w / mean for w in raw]


def run_training(parts: dict[str, list[Passage]], config: dict, config_path: Path) -> None:
    # Heavy imports kept local so --dry-run works without torch installed.
    import mlflow
    import torch
    from datasets import Dataset
    from transformers import (
        AutoModelForSequenceClassification,
        AutoTokenizer,
        DataCollatorWithPadding,
        Trainer,
        TrainingArguments,
    )

    run_cfg, model_cfg, train_cfg = config["run"], config["model"], config["train"]
    tracking = config.get("tracking", {})
    run_name = run_cfg["name"]
    out_dir = MODELS_DIR / run_name
    max_length = model_cfg.get("max_length", 512)

    tokenizer = AutoTokenizer.from_pretrained(model_cfg["base"])
    model = AutoModelForSequenceClassification.from_pretrained(
        model_cfg["base"],
        num_labels=len(CANONICAL_LEVELS),
        id2label=dict(enumerate(CANONICAL_LEVELS)),
        label2id={lvl: i for i, lvl in enumerate(CANONICAL_LEVELS)},
    )

    rank = {lvl: i for i, lvl in enumerate(CANONICAL_LEVELS)}

    def to_dataset(passages: list[Passage]) -> Dataset:
        dataset = Dataset.from_dict(
            {"text": [p.text for p in passages], "labels": [rank[p.level] for p in passages]}
        )
        return dataset.map(
            lambda batch: tokenizer(batch["text"], truncation=True, max_length=max_length),
            batched=True,
            remove_columns=["text"],
        )

    weights = None
    if train_cfg.get("class_weights", True):
        weights = torch.tensor(class_weight_values(parts["train"]), dtype=torch.float)

    class WeightedTrainer(Trainer):
        def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
            labels = inputs.pop("labels")
            outputs = model(**inputs)
            loss = torch.nn.functional.cross_entropy(
                outputs.logits,
                labels,
                weight=weights.to(outputs.logits.device) if weights is not None else None,
            )
            return (loss, outputs) if return_outputs else loss

    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        report = classification_report(list(labels), list(logits.argmax(axis=-1)))
        return {
            key: value
            for key, value in report.items()
            if key != "n" and not key.startswith("support_")
        }

    effective_batch = train_cfg.get("batch_size", 8) * train_cfg.get("grad_accum", 2)
    steps_per_epoch = math.ceil(len(parts["train"]) / effective_batch)
    warmup_steps = int(
        steps_per_epoch * train_cfg.get("epochs", 4) * train_cfg.get("warmup_ratio", 0.1)
    )

    arguments = TrainingArguments(
        output_dir=str(out_dir / "checkpoints"),
        seed=run_cfg.get("seed", 13),
        learning_rate=train_cfg.get("lr", 2e-5),
        num_train_epochs=train_cfg.get("epochs", 4),
        per_device_train_batch_size=train_cfg.get("batch_size", 8),
        per_device_eval_batch_size=train_cfg.get("batch_size", 8) * 2,
        gradient_accumulation_steps=train_cfg.get("grad_accum", 2),
        warmup_steps=warmup_steps,
        weight_decay=train_cfg.get("weight_decay", 0.01),
        fp16=train_cfg.get("fp16", True),
        eval_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=1,
        load_best_model_at_end=True,
        metric_for_best_model="macro_f1",
        greater_is_better=True,
        logging_steps=50,
        report_to="none",
    )
    trainer = WeightedTrainer(
        model=model,
        args=arguments,
        train_dataset=to_dataset(parts["train"]),
        eval_dataset=to_dataset(parts["val"]),
        data_collator=DataCollatorWithPadding(tokenizer),
        compute_metrics=compute_metrics,
    )

    mlflow.set_tracking_uri(tracking.get("uri", "sqlite:///mlflow.db"))
    mlflow.set_experiment(tracking.get("experiment", "cefr-classifier"))
    with mlflow.start_run(run_name=run_name):
        for section, values in config.items():
            for key, value in values.items():
                mlflow.log_param(f"{section}.{key}", str(value))
        for split in SPLITS:
            mlflow.log_param(f"n_{split}_passages", len(parts[split]))
        if weights is not None:
            mlflow.log_param("class_weight_values", [round(w, 3) for w in weights.tolist()])

        trainer.train()

        probs = predict_probs(trainer.model, tokenizer, parts["test"], max_length=max_length)
        en_passages, en_probs = lang_filtered(parts["test"], probs, "en")
        results = {
            "en": evaluate_views(en_passages, en_probs),  # headline (the product target)
            "all": evaluate_views(parts["test"], probs),
        }
        for scope, views in results.items():
            for view, report in views.items():
                mlflow.log_metrics({f"test_{scope}_{view}_{k}": v for k, v in report.items()})

        out_dir.mkdir(parents=True, exist_ok=True)
        trainer.save_model(str(out_dir / "model"))
        tokenizer.save_pretrained(str(out_dir / "model"))
        meta = {
            "levels": list(CANONICAL_LEVELS),
            "max_length": max_length,
            "chunking": config["data"].get("chunking", True),
            "target_words": config["data"].get("target_words", 200),
            "max_words": config["data"].get("max_words", 300),
        }
        (out_dir / "preprocessing.json").write_text(json.dumps(meta, indent=2))
        mlflow.log_artifact(str(config_path))
        mlflow.log_dict(results, "test_metrics.json")

        print(json.dumps(results, indent=2))
        print(f"\nModel saved to {out_dir / 'model'}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--dry-run", action="store_true", help="build + split + checks, no GPU")
    args = parser.parse_args()

    config = load_config(args.config)
    parts = build_parts(config)
    print(summarize(parts))
    print("\nLeakage check: OK (document ids are disjoint across splits)")
    if args.dry_run:
        return
    run_training(parts, config, args.config)


if __name__ == "__main__":
    main()