"""Fine-tune DistilBERT for binary controversy classification on Belnap corpus.

This script is the reproducible artifact for the model. Trains on the
110-record `barissozudogru/belnap-debate-corpus` to predict whether a
proposition is high-controversy (vs medium/low).

Honest caveats baked in:
  - 110 records is small; model favors memorization
  - Stratified 80/20 split; leave-one-out would be more rigorous
  - Reports accuracy + F1 alongside the always-predict-high baseline
"""
from __future__ import annotations

import json
import os
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from datasets import Dataset
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix, f1_score
from sklearn.model_selection import train_test_split
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    Trainer,
    TrainingArguments,
)

SEED = 42
BASE_MODEL = "distilbert-base-uncased"
DATA_CSV = "/Users/baris/Projects/belnap-corpus-dataset/propositions.csv"
OUT_DIR = Path("/Users/baris/Projects/belnap-controversy-classifier/model")
OUT_DIR.mkdir(parents=True, exist_ok=True)
METRICS_PATH = OUT_DIR.parent / "metrics.json"

torch.manual_seed(SEED)
np.random.seed(SEED)


def main():
    # 1. Load + binarize
    df = pd.read_csv(DATA_CSV)
    df["label"] = (df["controversy"] == "high").astype(int)
    pos_count = int(df["label"].sum())
    print(f"{len(df)} records, {pos_count} high ({pos_count/len(df):.1%}), {len(df) - pos_count} non-high")

    # 2. Stratified split
    train_df, eval_df = train_test_split(
        df,
        test_size=0.20,
        random_state=SEED,
        stratify=df["label"],
    )
    print(f"train: {len(train_df)}, eval: {len(eval_df)}")

    # 3. Tokenize
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

    def tokenize(batch):
        return tokenizer(
            batch["proposition"],
            truncation=True,
            max_length=128,
            padding=False,
        )

    train_ds = Dataset.from_pandas(train_df[["proposition", "label"]], preserve_index=False).map(tokenize, batched=True)
    eval_ds = Dataset.from_pandas(eval_df[["proposition", "label"]], preserve_index=False).map(tokenize, batched=True)

    # 4. Model
    device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
    print(f"device: {device}")

    model = AutoModelForSequenceClassification.from_pretrained(
        BASE_MODEL,
        num_labels=2,
        id2label={0: "not_high_controversy", 1: "high_controversy"},
        label2id={"not_high_controversy": 0, "high_controversy": 1},
    )

    # 5. Class-weighted loss to mitigate imbalance.
    # Give MORE weight to the minority class (non-high, label=0).
    pos_count = int(train_df["label"].sum())
    neg_count = len(train_df) - pos_count
    total = len(train_df)
    weight_0 = total / (2 * max(neg_count, 1))  # minority gets higher weight
    weight_1 = total / (2 * max(pos_count, 1))
    print(f"class weights: non-high={weight_0:.3f}, high={weight_1:.3f}")

    class WeightedTrainer(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
            weight = torch.tensor([weight_0, weight_1], device=logits.device, dtype=logits.dtype)
            loss_fct = torch.nn.CrossEntropyLoss(weight=weight)
            loss = loss_fct(logits, labels)
            return (loss, outputs) if return_outputs else loss

    # 6. Training arguments
    args = TrainingArguments(
        output_dir=str(OUT_DIR / "_checkpoints"),
        num_train_epochs=5,
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        learning_rate=3e-5,
        weight_decay=0.01,
        warmup_ratio=0.1,
        logging_steps=10,
        eval_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=1,
        load_best_model_at_end=True,
        metric_for_best_model="f1",
        greater_is_better=True,
        report_to="none",
        seed=SEED,
    )

    def compute_metrics(eval_pred):
        preds = np.argmax(eval_pred.predictions, axis=1)
        labels = eval_pred.label_ids
        return {
            "accuracy": accuracy_score(labels, preds),
            "f1": f1_score(labels, preds, average="macro"),
            "f1_high": f1_score(labels, preds, pos_label=1, average="binary"),
        }

    collator = DataCollatorWithPadding(tokenizer=tokenizer)
    trainer = WeightedTrainer(
        model=model,
        args=args,
        train_dataset=train_ds,
        eval_dataset=eval_ds,
        processing_class=tokenizer,
        data_collator=collator,
        compute_metrics=compute_metrics,
    )

    trainer.train()

    # 7. Final eval
    eval_metrics = trainer.evaluate()
    preds = trainer.predict(eval_ds)
    pred_labels = np.argmax(preds.predictions, axis=1)
    cm = confusion_matrix(eval_df["label"].values, pred_labels).tolist()
    cls_report = classification_report(eval_df["label"].values, pred_labels, output_dict=True, zero_division=0)

    # 8. Baselines
    baseline_majority = int((eval_df["label"] == 1).sum()) / len(eval_df)
    baseline_random = 0.5

    summary = {
        "base_model": BASE_MODEL,
        "n_train": len(train_df),
        "n_eval": len(eval_df),
        "eval_accuracy": float(eval_metrics["eval_accuracy"]),
        "eval_f1_macro": float(eval_metrics["eval_f1"]),
        "eval_f1_high": float(eval_metrics["eval_f1_high"]),
        "confusion_matrix": cm,
        "baseline_always_predict_high": baseline_majority,
        "baseline_random": baseline_random,
        "classification_report": cls_report,
        "training_args": {
            "epochs": 5,
            "batch_size": 8,
            "learning_rate": 3e-5,
            "weight_decay": 0.01,
            "warmup_ratio": 0.1,
            "seed": SEED,
        },
    }
    METRICS_PATH.write_text(json.dumps(summary, indent=2))
    print(f"\n=== final eval ===")
    print(json.dumps(summary, indent=2))

    # 9. Save model
    trainer.save_model(str(OUT_DIR))
    tokenizer.save_pretrained(str(OUT_DIR))
    print(f"\nmodel saved to {OUT_DIR}")


if __name__ == "__main__":
    main()