"""Reproducible training pipeline for ACO specialist classifiers.

Trains three models (tier-router, tool-gater, verifier-gater) using
ModernBERT-base with focal loss and threshold calibration.

Loads data from narcolepticchicken/aco-traces (preprocessed source datasets)
and applies the same training recipe as the original v2 training.

Usage:
    # Train all three models
    uv run --with transformers,torch,datasets,scikit-learn,huggingface_hub,trackio train_aco.py

    # Train a single model
    uv run --with transformers,torch,datasets,scikit-learn,huggingface_hub train_aco.py --task tier_router

    # Via hf_jobs
    hf_jobs run --script train_aco.py --deps transformers,torch,datasets,scikit-learn,huggingface_hub --hardware a10g-large --timeout 8h
"""
import argparse, torch, numpy as np, json, os, sys
from datasets import load_dataset, Dataset
from transformers import (
    AutoTokenizer, AutoModelForSequenceClassification,
    TrainingArguments, Trainer, EarlyStoppingCallback
)
from sklearn.metrics import accuracy_score, f1_score, precision_recall_fscore_support
import torch.nn as nn
import torch.nn.functional as F

# ═══════════════════════════════════════════════════════════════
# Configuration — matches v2 training exactly
# ═══════════════════════════════════════════════════════════════

BASE_MODEL = "answerdotai/ModernBERT-base"
DATASET_REPO = "narcolepticchicken/aco-traces"
OUTPUT_REPO_TEMPLATE = "narcolepticchicken/aco-specialists-{task}"

TASK_CONFIG = {
    "tier_router": {
        "num_labels": 3,
        "id2label": {0: "easy", 1: "medium", 2: "hard"},
        "label2id": {"easy": 0, "medium": 1, "hard": 2},
    },
    "tool_gater": {
        "num_labels": 2,
        "id2label": {0: "no_tool", 1: "call_tool"},
        "label2id": {"no_tool": 0, "call_tool": 1},
    },
    "verifier_gater": {
        "num_labels": 2,
        "id2label": {0: "no_verify", 1: "verify"},
        "label2id": {"no_verify": 0, "verify": 1},
    },
}

TRAINING_ARGS = dict(
    output_dir="/tmp/aco_train",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=10,
    weight_decay=0.01,
    warmup_ratio=0.1,
    logging_strategy="steps",
    logging_steps=20,
    logging_first_step=True,
    eval_strategy="steps",
    eval_steps=50,
    save_strategy="steps",
    save_steps=200,
    load_best_model_at_end=True,
    metric_for_best_model="f1_macro",
    greater_is_better=True,
    save_total_limit=1,
    disable_tqdm=True,
    lr_scheduler_type="cosine",
    bf16=torch.cuda.is_available(),
    push_to_hub=False,  # we push manually after calibration
    report_to="none",
)

# ═══════════════════════════════════════════════════════════════
# Focal Loss
# ═══════════════════════════════════════════════════════════════

class FocalLoss(nn.Module):
    """Focal loss per Lin et al. (1708.02002)."""
    def __init__(self, gamma=2.0, alpha=None, reduction="mean"):
        super().__init__()
        self.gamma = gamma
        self.alpha = alpha
        self.reduction = reduction

    def forward(self, logits, labels):
        ce_loss = F.cross_entropy(logits, labels, reduction="none", weight=self.alpha)
        pt = torch.exp(-ce_loss)
        focal_loss = ((1 - pt) ** self.gamma) * ce_loss
        if self.reduction == "mean":
            return focal_loss.mean()
        elif self.reduction == "sum":
            return focal_loss.sum()
        return focal_loss

# ═══════════════════════════════════════════════════════════════
# Trainer with focal loss
# ═══════════════════════════════════════════════════════════════

class FocalTrainer(Trainer):
    def __init__(self, focal_gamma=2.0, class_alpha=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.focal_gamma = focal_gamma
        if class_alpha is not None:
            self.class_alpha = torch.tensor(class_alpha, dtype=torch.float)
        else:
            self.class_alpha = None

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.logits

        # Move alpha to correct device
        if self.class_alpha is not None:
            alpha = self.class_alpha.to(logits.device)
        else:
            alpha = None

        loss_fn = FocalLoss(gamma=self.focal_gamma, alpha=alpha)
        loss = loss_fn(logits, labels)

        return (loss, outputs) if return_outputs else loss

# ═══════════════════════════════════════════════════════════════
# Metrics
# ═══════════════════════════════════════════════════════════════

def compute_metrics(eval_pred, num_labels):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)
    acc = accuracy_score(labels, preds)
    f1 = f1_score(labels, preds, average="macro", zero_division=0)
    return {"accuracy": acc, "f1_macro": f1}

# ═══════════════════════════════════════════════════════════════
# Threshold calibration
# ═══════════════════════════════════════════════════════════════

def calibrate_threshold(model, tokenizer, eval_dataset, device, num_labels):
    """Grid search for best decision threshold on binary tasks."""
    if num_labels != 2:
        return 0.5

    model.eval()
    texts = eval_dataset["text"]
    labels_list = eval_dataset["labels"]
    encodings = tokenizer(texts, truncation=True, max_length=2048, padding=True)

    # Run inference
    input_ids = torch.tensor(encodings["input_ids"]).to(device)
    attention_mask = torch.tensor(encodings["attention_mask"]).to(device)

    all_probs = []
    bs = 32
    for i in range(0, len(input_ids), bs):
        with torch.no_grad():
            logits = model(
                input_ids=input_ids[i:i+bs],
                attention_mask=attention_mask[i:i+bs]
            ).logits
            probs = torch.softmax(logits, dim=-1).cpu().numpy()
            all_probs.append(probs)

    probs = np.vstack(all_probs)
    labels = np.array(labels_list)

    best_f1 = 0
    best_threshold = 0.5
    for t in np.linspace(0.05, 0.95, 19):
        preds = (probs[:, 1] >= t).astype(int)
        f1 = f1_score(labels, preds, average="macro", zero_division=0)
        if f1 > best_f1:
            best_f1 = f1
            best_threshold = t

    print(f"  Calibrated threshold: {best_threshold:.3f} (best F1={best_f1:.4f})")
    return float(best_threshold)

# ═══════════════════════════════════════════════════════════════
# Compute class alpha from training labels (focal loss setup)
# ═══════════════════════════════════════════════════════════════

def compute_alpha(train_dataset, num_labels):
    """Compute alpha weights as inverse class frequency."""
    labels = train_dataset["labels"]
    counts = {}
    for lb in labels:
        counts[lb] = counts.get(lb, 0) + 1
    total = len(labels)
    alpha = [total / (num_labels * max(counts.get(i, 1), 1)) for i in range(num_labels)]
    print(f"  Class counts: {counts}, alpha: {alpha}")
    return alpha

# ═══════════════════════════════════════════════════════════════
# Train one model
# ═══════════════════════════════════════════════════════════════

def train_model(task_name, push=True):
    print(f"\n{'='*60}")
    print(f"TRAINING: {task_name}")
    print(f"{'='*60}")

    cfg = TASK_CONFIG[task_name]
    num_labels = cfg["num_labels"]

    # Load data
    try:
        ds = load_dataset(DATASET_REPO, data_dir=f"data/{task_name}", split="train")
        ds_dict = ds.train_test_split(test_size=0.15, seed=42)
        train_ds = ds_dict["train"]
        eval_ds = ds_dict["test"]
    except Exception:
        # Fallback: load from pre-split files
        train_ds = Dataset.from_parquet(f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/data/{task_name}/train.parquet")
        eval_ds = Dataset.from_parquet(f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/data/{task_name}/test.parquet")

    print(f"  Train: {len(train_ds)}, Eval: {len(eval_ds)}")

    # Class alpha
    alpha = compute_alpha(train_ds, num_labels)

    # Tokenizer
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # Tokenize
    def tokenize_fn(examples):
        return tokenizer(examples["text"], truncation=True, max_length=2048, padding=False)

    train_tok = train_ds.map(tokenize_fn, batched=True)
    eval_tok = eval_ds.map(tokenize_fn, batched=True)

    # Model with explicit dropout
    model = AutoModelForSequenceClassification.from_pretrained(
        BASE_MODEL,
        num_labels=num_labels,
        id2label=cfg["id2label"],
        label2id=cfg["label2id"],
        attention_dropout=0.1,
        embedding_dropout=0.1,
        mlp_dropout=0.1,
        classifier_dropout=0.1,
        ignore_mismatched_sizes=True,
    )

    # Training args
    args = dict(TRAINING_ARGS)
    args["output_dir"] = f"/tmp/aco_train_{task_name}"
    training_args = TrainingArguments(**args)

    # Trainer
    trainer = FocalTrainer(
        focal_gamma=2.0,
        class_alpha=alpha,
        model=model,
        args=training_args,
        train_dataset=train_tok,
        eval_dataset=eval_tok,
        tokenizer=tokenizer,
        compute_metrics=lambda p: compute_metrics(p, num_labels),
        callbacks=[EarlyStoppingCallback(early_stopping_patience=4)],
    )

    # Train
    print("  Starting training...")
    trainer.train()

    # Calibrate threshold (binary tasks)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    threshold = calibrate_threshold(model, tokenizer, eval_ds, device, num_labels)

    # Save to model config
    model.config.threshold = threshold
    model.config.focal_gamma = 2.0
    model.config.alpha = alpha

    # Push to Hub
    if push:
        output_repo = OUTPUT_REPO_TEMPLATE.format(task=task_name)
        print(f"  Pushing to {output_repo}...")
        model.push_to_hub(output_repo)
        tokenizer.push_to_hub(output_repo)
        print(f"  Published: https://huggingface.co/{output_repo}")

    # Final eval
    from sklearn.metrics import classification_report
    # Quick eval on a subset for printing
    model.eval()
    texts = eval_ds["text"][:500]
    labels_list = eval_ds["labels"][:500]
    enc = tokenizer(texts, truncation=True, max_length=2048, padding=True, return_tensors="pt").to(device)
    with torch.no_grad():
        logits = model(**enc).logits
        probs = torch.softmax(logits, dim=-1).cpu().numpy()

    if num_labels == 2:
        preds = (probs[:, 1] >= threshold).astype(int)
    else:
        preds = np.argmax(probs, axis=-1)

    acc = accuracy_score(labels_list, preds)
    f1 = f1_score(labels_list, preds, average="macro", zero_division=0)
    print(f"  Final eval (first 500): acc={acc:.4f}, f1_macro={f1:.4f}")
    print(classification_report(labels_list, preds, zero_division=0))

    return model, tokenizer, threshold

# ═══════════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--task", choices=["tier_router", "tool_gater", "verifier_gater", "all"],
                        default="all")
    parser.add_argument("--no-push", action="store_true")
    args = parser.parse_args()

    tasks = ["tier_router", "tool_gater", "verifier_gater"] if args.task == "all" else [args.task]

    for task in tasks:
        try:
            train_model(task, push=not args.no_push)
        except Exception as e:
            print(f"  FAILED: {task}: {e}")
            import traceback; traceback.print_exc()

if __name__ == "__main__":
    main()