"""
ml/model/stutter_trainer.py - Fine-tune wav2vec2 (+LoRA) stutter classifier
===========================================================================
Trains a sequence classifier on speech disfluency datasets (real & synthetic
lattice) using a wav2vec2-base encoder + LoRA adapters.

Enhancements:
  - Focal Loss (gamma=2.0) to focus gradients on hard disfluency boundaries
    and drive precision >90%.
  - Supports both binary detection-first (fluent vs stutter) and 4-way
    classification (fluent, repetition, prolongation, block).
  - LoRA on query/key/value projections with explicit persistence of both
    `projector` and `classifier` in `modules_to_save`.
  - Balanced inverse-frequency class weighting.

Outputs:
  ml/models/stutter/            Trainer checkpoint (LoRA adapters + head)
  ml/models/stutter/stutter_lora/  adapter_model.safetensors
  data/class_map.json           id<->label mapping used by inference
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Optional

import numpy as np
import os
import shutil
import tempfile
import torch
import torch.nn as nn
from datasets import Audio, load_from_disk
import datasets as _datasets

def clean_cache():
    """No-op kept for callers; caching is disabled in prepare_dataset."""
    return

from peft import LoraConfig, TaskType, get_peft_model
from sklearn.metrics import accuracy_score, precision_recall_fscore_support, f1_score
from transformers import (
    Trainer,
    TrainingArguments,
    Wav2Vec2FeatureExtractor,
    Wav2Vec2ForSequenceClassification,
)

SR = 16000
MAX_SECONDS = 8.0
MODEL_BASE = "facebook/wav2vec2-base"

# Full 4-way label space.
ID2LABEL = {
    0: "fluent_control",
    1: "stutter_repetition",
    2: "stutter_prolongation",
    3: "stutter_block",
}
LABEL2ID = {v: k for k, v in ID2LABEL.items()}

# Detection-first label space.
BIN_ID2LABEL = {0: "fluent", 1: "stutter"}
BIN_LABEL2ID = {v: k for k, v in BIN_ID2LABEL.items()}


def bin_label(label: str) -> int:
    """Map any canonical stutter subtype (or fluent) to the binary id."""
    return 0 if label == "fluent_control" or label == "fluent" else 1


class FocalLoss(nn.Module):
    """Multi-class Focal Loss to downweight easy negatives and emphasize hard boundaries."""
    def __init__(self, gamma: float = 2.0, alpha: Optional[torch.Tensor] = None):
        super().__init__()
        self.gamma = gamma
        self.alpha = alpha

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        weight = self.alpha.to(logits.device) if self.alpha is not None else None
        ce_loss = nn.functional.cross_entropy(logits, targets, reduction="none", weight=weight)
        pt = torch.exp(-ce_loss)
        focal = ((1.0 - pt) ** self.gamma) * ce_loss
        return focal.mean()


# --------------------------------------------------------------------------
# Preprocessing
# --------------------------------------------------------------------------
def _tokenize_batch(batch: dict, feat_extractor) -> dict:
    """audio_array (plain float32 column) -> input_values + attention_mask."""
    inp = []
    for arr in batch["audio_array"]:
        arr = np.asarray(arr, dtype=np.float32)
        if arr.ndim > 1:
            arr = arr.mean(axis=1)
        arr = arr[: int(SR * MAX_SECONDS)]  # truncate to max window
        inp.append(arr)
    fe = feat_extractor(
        inp,
        sampling_rate=SR,
        return_tensors="pt",
        padding="max_length",
        truncation=True,
        max_length=int(SR * MAX_SECONDS),
    )
    values = fe["input_values"]
    attention_mask = (values != 0).long()
    return {"input_values": values, "attention_mask": attention_mask}


def prepare_dataset(data_dir: str, feature_extractor, binary: bool = True,
                    balance_train: float = 0.0):
    _datasets.disable_caching()
    label_map = BIN_LABEL2ID if binary else LABEL2ID
    valid = set(label_map)
    
    if binary:
        def to_id(label: str) -> int:
            return BIN_LABEL2ID["stutter"] if label != "fluent_control" and label != "fluent" else BIN_LABEL2ID["fluent"]
    else:
        def to_id(label: str) -> int:
            if label in LABEL2ID:
                return LABEL2ID[label]
            if label == "stutter":
                return LABEL2ID["stutter_repetition"]
            return LABEL2ID["fluent_control"]

    if Path(data_dir).suffix == ".parquet":
        from datasets import Dataset
        ds = Dataset.from_parquet(data_dir)
    else:
        from datasets import load_from_disk
        ds = load_from_disk(data_dir)

    # Filter valid rows
    if binary:
        ds = ds.filter(lambda r: isinstance(r.get("label"), str))
    else:
        ds = ds.filter(lambda r: isinstance(r.get("label"), str))

    # Map labels
    ds = ds.map(lambda r: {"labels": to_id(r["label"])},
                remove_columns=["label"] if "label" in ds.column_names else None)

    # Tokenize waveforms into input tensors
    cols = {"audio", "audio_array", "text", "corpus", "id", "speaker_id"}
    ds = ds.map(lambda b: _tokenize_batch(b, feature_extractor), batched=True,
                remove_columns=list(cols & set(ds.column_names)))

    def _split(s: str):
        return ds.filter(lambda r: r.get("split") == s)

    tr, va, te = _split("train"), _split("val"), _split("test")

    # Optional rebalancing
    if balance_train > 0 and binary:
        labs = tr["labels"]
        counts = {0: sum(1 for x in labs if x == 0),
                  1: sum(1 for x in labs if x == 1)}
        if counts.get(1, 0) > 0 and counts.get(1, 0) < counts.get(0, 0):
            target = int(counts[0] * balance_train)
            fac = target // counts[1]
            if fac >= 1:
                keep = [i for i, l in enumerate(labs) if l == 1] * fac
                maj = [i for i, l in enumerate(labs) if l == 0]
                keep = (maj + keep)[: len(maj) + target]
                tr = tr.select(sorted(keep))
                print(f"[train] rebalanced train: fluent={counts[0]} "
                      f"stutter={target} (x{fac} oversample)")
    return tr, va, te


# --------------------------------------------------------------------------
# LoRA + trainer
# --------------------------------------------------------------------------
def _lora_config() -> LoraConfig:
    return LoraConfig(
        task_type=TaskType.SEQ_CLS,
        r=8,
        lora_alpha=16,
        lora_dropout=0.1,
        target_modules=["q_proj", "k_proj", "v_proj"],
        modules_to_save=["projector", "classifier"],
        bias="none",
    )


def _class_weights(dataset, n_classes: int) -> torch.Tensor:
    """Inverse-frequency weights normalised to sum to 1."""
    counts = dataset.to_pandas()["labels"].value_counts()
    n = n_classes
    w = torch.zeros(n)
    for i in range(n):
        c = int(counts.get(i, 0))
        w[i] = 1.0 / (c if c > 0 else 1.0)
    w = w / w.sum()
    return w


def compute_metrics(eval_pred):
    """Compute accuracy, precision, recall, and macro-F1."""
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=1)
    acc = float(accuracy_score(labels, preds))
    p, r, f, _ = precision_recall_fscore_support(labels, preds, average="macro", zero_division=0)
    return {
        "accuracy": acc,
        "precision": float(p),
        "recall": float(r),
        "macro_f1": float(f),
    }


def _save_class_map(dest: Path, binary: bool):
    id2l = BIN_ID2LABEL if binary else ID2LABEL
    dest.write_text(
        json.dumps({"label2id": {v: k for k, v in id2l.items()},
                    "id2label": id2l, "binary": binary}, indent=2),
        encoding="utf-8",
    )


def train(
    data_dir: str,
    out_dir: str = "ml/models/stutter",
    epochs: int = 5,
    lr: float = 3e-5,
    batch: int = 8,
    seed: int = 0,
    fp16: bool = True,
    binary: bool = True,
    balance_train: float = 0.0,
    focal_gamma: float = 2.0,
) -> None:
    rng = np.random.default_rng(seed)
    torch.manual_seed(seed)

    feat = Wav2Vec2FeatureExtractor(sampling_rate=SR)
    train_ds, val_ds, test_ds = prepare_dataset(
        data_dir, feat, binary=binary, balance_train=balance_train)

    class_map = BIN_LABEL2ID if binary else LABEL2ID
    n_classes = len(class_map)
    model = Wav2Vec2ForSequenceClassification.from_pretrained(
        MODEL_BASE,
        num_labels=n_classes,
        ignore_mismatched_sizes=True,
    )
    model = get_peft_model(model, _lora_config())
    try:
        model.print_trainable_parameters()
    except AttributeError:
        n = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print(f"[train] trainable params: {n:,}")

    weights = _class_weights(train_ds, n_classes)
    loss_fn = FocalLoss(gamma=focal_gamma, alpha=weights) if focal_gamma > 0 else None

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

    gpu_ok = torch.cuda.is_available()
    if not gpu_ok:
        print("[train] WARNING: no CUDA; falling back to CPU")
    elif not fp16:
        print("[train] fp32 (fp16 disabled)")

    args = TrainingArguments(
        output_dir=str(Path(out_dir) / "checkpoints"),
        num_train_epochs=epochs,
        per_device_train_batch_size=batch,
        per_device_eval_batch_size=batch,
        gradient_accumulation_steps=1,
        learning_rate=lr,
        lr_scheduler_type="cosine",
        warmup_steps=min(100, max(10, int(0.08 * len(train_ds) / (batch * 1)))),
        weight_decay=0.01,
        fp16=fp16 and gpu_ok,
        eval_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=2,
        logging_steps=25,
        seed=seed,
        report_to="none",
        dataloader_num_workers=0,
        remove_unused_columns=False,
    )
    trainer = CustomTrainer(
        model=model,
        args=args,
        train_dataset=train_ds,
        eval_dataset=val_ds,
        compute_metrics=compute_metrics,
    )
    trainer.train()
    trainer.save_model(str(Path(out_dir) / "stutter_lora"))

    test_metrics = trainer.predict(test_ds)
    test_scores = test_metrics.metrics
    summary = {
        "data_dir": data_dir,
        "out_dir": out_dir,
        "epochs": epochs,
        "seed": seed,
        "base_model": MODEL_BASE,
        "binary": binary,
        "focal_gamma": focal_gamma,
        "n_train": len(train_ds), "n_val": len(val_ds), "n_test": len(test_ds),
        "final_holdout": {k: v for k, v in test_scores.items()},
    }
    final_path = Path(out_dir) / "test_report.json"
    final_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")

    _save_class_map(Path(out_dir) / "class_map.json", binary)

    clean_cache()
    print(f"[train] done. Adapters + head at {out_dir}/stutter_lora")
    print(f"[train] held-out final scores -> {final_path}")


def _main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--data", default="data/metadata/hybrid_dataset")
    ap.add_argument("--out", default="ml/models/stutter")
    ap.add_argument("--epochs", type=int, default=5)
    ap.add_argument("--lr", type=float, default=3e-5)
    ap.add_argument("--batch", type=int, default=8)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--no-fp16", action="store_true", help="disable mixed precision")
    ap.add_argument("--binary", action="store_true", default=True)
    ap.add_argument("--no-binary", dest="binary", action="store_false")
    ap.add_argument("--focal-gamma", type=float, default=2.0, help="Focal Loss focusing factor")
    args = ap.parse_args()
    train(args.data, args.out, args.epochs, args.lr, args.batch, args.seed,
          fp16=not args.no_fp16, binary=args.binary, focal_gamma=args.focal_gamma)


if __name__ == "__main__":
    _main()