"""
Fine-tune DistilBERT for Samsung Health Chart Intent Classification
===================================================================
Optimized for Apple Silicon (M5/M4/M3/M2/M1) with MPS acceleration.

Usage:
    python finetune.py

Requirements:
    pip install torch transformers datasets scikit-learn pandas
"""

import os
import json
import pandas as pd
import numpy as np
from pathlib import Path

import torch
from torch.utils.data import Dataset, DataLoader
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    get_linear_schedule_with_warmup,
)
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix

# ─────────────────────────────────────────────────────────
# STEP 0: Detect device
# ─────────────────────────────────────────────────────────

def get_device():
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print(f"Using CUDA GPU: {torch.cuda.get_device_name(0)}")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
        print("Using Apple Silicon MPS (Metal Performance Shaders)")
    else:
        device = torch.device("cpu")
        print("Using CPU")
    return device


# ─────────────────────────────────────────────────────────
# STEP 1: Config — all hyperparameters in one place
# ─────────────────────────────────────────────────────────

CONFIG = {
    # Model
    # distilbert-base-multilingual-cased: 134M params, handles English well,
    # and will generalize better if you later add Chinese/mixed queries.
    # If your queries will always be English-only, you can swap to:
    # "distilbert-base-uncased" (66M params, slightly faster, English only)
    "model_name": "distilbert-base-multilingual-cased",

    # Data
    "data_path": "samsung_health_intent.csv",
    "max_length": 64,          # User queries are short; 64 is plenty and ~8x faster than 512
    "test_size": 0.15,         # 15% test set  (~15 samples from 100)
    "val_size": 0.15,          # 15% val set   (~15 samples from 100)

    # Training — tuned for M5 MPS
    "batch_size": 16,          # Safe for 16GB+ unified memory; lower to 8 if you get OOM
    "epochs": 8,               # More epochs needed for small dataset (100 samples)
    "learning_rate": 3e-5,
    "weight_decay": 0.01,
    "warmup_ratio": 0.1,

    # MPS-specific
    "num_workers": 0,          # MUST be 0 for MPS (multiprocessing + MPS = deadlock)
    "fp16": False,             # MPS does NOT support fp16 training

    # Output
    "output_dir": "./chart_intent_model",
    "seed": 42,
}


# ─────────────────────────────────────────────────────────
# STEP 2: Dataset class
# ─────────────────────────────────────────────────────────

class HealthIntentDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length):
        self.encodings = tokenizer(
            texts,
            truncation=True,
            padding="max_length",
            max_length=max_length,
            return_tensors="pt",
        )
        self.labels = torch.tensor(labels, dtype=torch.long)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return {
            "input_ids":      self.encodings["input_ids"][idx],
            "attention_mask": self.encodings["attention_mask"][idx],
            "labels":         self.labels[idx],
        }


# ─────────────────────────────────────────────────────────
# STEP 3: Metrics
# ─────────────────────────────────────────────────────────

def compute_metrics(preds, labels):
    preds = np.array(preds)
    labels = np.array(labels)
    acc = (preds == labels).mean()
    report = classification_report(
        labels, preds,
        target_names=["no_chart", "chart"],
        output_dict=True,
    )
    return {
        "accuracy":  acc,
        "f1_chart":  report["chart"]["f1-score"],
        "precision": report["chart"]["precision"],
        "recall":    report["chart"]["recall"],
    }


# ─────────────────────────────────────────────────────────
# STEP 4: Training loop
# ─────────────────────────────────────────────────────────

def train_epoch(model, loader, optimizer, scheduler, device):
    model.train()
    total_loss = 0
    all_preds, all_labels = [], []

    for batch in loader:
        input_ids      = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels         = batch["labels"].to(device)

        optimizer.zero_grad()
        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        loss.backward()

        # Gradient clipping — stabilizes training on small datasets
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        optimizer.step()
        scheduler.step()

        total_loss += loss.item()
        preds = outputs.logits.argmax(dim=-1).cpu().numpy()
        all_preds.extend(preds)
        all_labels.extend(labels.cpu().numpy())

    metrics = compute_metrics(all_preds, all_labels)
    metrics["loss"] = total_loss / len(loader)
    return metrics


@torch.no_grad()
def evaluate(model, loader, device):
    model.eval()
    total_loss = 0
    all_preds, all_labels = [], []

    for batch in loader:
        input_ids      = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels         = batch["labels"].to(device)

        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)

        total_loss += outputs.loss.item()
        preds = outputs.logits.argmax(dim=-1).cpu().numpy()
        all_preds.extend(preds)
        all_labels.extend(labels.cpu().numpy())

    metrics = compute_metrics(all_preds, all_labels)
    metrics["loss"] = total_loss / len(loader)
    return metrics, all_preds, all_labels


# ─────────────────────────────────────────────────────────
# STEP 5: Inference helper (run after training)
# ─────────────────────────────────────────────────────────

class ChartIntentClassifier:
    """
    Load the fine-tuned model and classify new queries.
    Use this in your Samsung Health chatbot.
    """
    def __init__(self, model_dir: str, device=None):
        self.device = device or get_device()
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_dir)
        self.model.to(self.device)
        self.model.eval()

    @torch.no_grad()
    def predict(self, text: str) -> dict:
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            max_length=CONFIG["max_length"],
            padding=True,
        ).to(self.device)

        logits = self.model(**inputs).logits
        probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]
        label = int(probs.argmax())

        return {
            "text":       text,
            "label":      label,                         # 0 or 1
            "intent":     "chart" if label == 1 else "no_chart",
            "confidence": float(probs[label]),
            "prob_chart": float(probs[1]),
        }

    def predict_batch(self, texts: list) -> list:
        return [self.predict(t) for t in texts]


# ─────────────────────────────────────────────────────────
# MAIN
# ─────────────────────────────────────────────────────────

def main():
    torch.manual_seed(CONFIG["seed"])
    np.random.seed(CONFIG["seed"])
    device = get_device()

    # ── Load data ──────────────────────────────────────────
    print("\n── Loading data ──────────────────────────────────")
    df = pd.read_csv(CONFIG["data_path"])
    print(f"Total samples : {len(df)}")
    print(f"Label dist    : {df['label'].value_counts().to_dict()}")

    texts  = df["text"].tolist()
    labels = df["label"].tolist()

    # Train / val / test split (stratified)
    X_train, X_temp, y_train, y_temp = train_test_split(
        texts, labels,
        test_size=CONFIG["test_size"] + CONFIG["val_size"],
        stratify=labels,
        random_state=CONFIG["seed"],
    )
    val_ratio_of_temp = CONFIG["val_size"] / (CONFIG["test_size"] + CONFIG["val_size"])
    X_val, X_test, y_val, y_test = train_test_split(
        X_temp, y_temp,
        test_size=1 - val_ratio_of_temp,
        stratify=y_temp,
        random_state=CONFIG["seed"],
    )
    print(f"Train: {len(X_train)}, Val: {len(X_val)}, Test: {len(X_test)}")

    # ── Tokenizer & model ──────────────────────────────────
    print(f"\n── Loading model: {CONFIG['model_name']} ────────────")
    print("(Downloading ~500MB on first run, cached after that)")
    tokenizer = AutoTokenizer.from_pretrained(CONFIG["model_name"])
    model = AutoModelForSequenceClassification.from_pretrained(
        CONFIG["model_name"],
        num_labels=2,
        id2label={0: "no_chart", 1: "chart"},
        label2id={"no_chart": 0, "chart": 1},
    )
    model.to(device)

    total_params = sum(p.numel() for p in model.parameters())
    trainable    = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Parameters: {total_params:,} total, {trainable:,} trainable")

    # ── Datasets & loaders ─────────────────────────────────
    print("\n── Tokenizing ────────────────────────────────────")
    train_ds = HealthIntentDataset(X_train, y_train, tokenizer, CONFIG["max_length"])
    val_ds   = HealthIntentDataset(X_val,   y_val,   tokenizer, CONFIG["max_length"])
    test_ds  = HealthIntentDataset(X_test,  y_test,  tokenizer, CONFIG["max_length"])

    train_loader = DataLoader(train_ds, batch_size=CONFIG["batch_size"],
                              shuffle=True,  num_workers=CONFIG["num_workers"])
    val_loader   = DataLoader(val_ds,   batch_size=CONFIG["batch_size"],
                              shuffle=False, num_workers=CONFIG["num_workers"])
    test_loader  = DataLoader(test_ds,  batch_size=CONFIG["batch_size"],
                              shuffle=False, num_workers=CONFIG["num_workers"])

    # ── Optimizer & scheduler ──────────────────────────────
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=CONFIG["learning_rate"],
        weight_decay=CONFIG["weight_decay"],
    )
    total_steps   = len(train_loader) * CONFIG["epochs"]
    warmup_steps  = int(total_steps * CONFIG["warmup_ratio"])
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=warmup_steps,
        num_training_steps=total_steps,
    )

    # ── Training loop ──────────────────────────────────────
    print(f"\n── Training for {CONFIG['epochs']} epochs ────────────────────")
    best_val_f1 = 0
    best_epoch  = 0
    history     = []

    for epoch in range(1, CONFIG["epochs"] + 1):
        train_m = train_epoch(model, train_loader, optimizer, scheduler, device)
        val_m, _, _ = evaluate(model, val_loader, device)

        history.append({"epoch": epoch, "train": train_m, "val": val_m})

        print(
            f"Epoch {epoch:2d}/{CONFIG['epochs']} │ "
            f"Train loss={train_m['loss']:.3f} acc={train_m['accuracy']:.3f} │ "
            f"Val   loss={val_m['loss']:.3f} acc={val_m['accuracy']:.3f} f1={val_m['f1_chart']:.3f}"
        )

        # Save best checkpoint
        if val_m["f1_chart"] >= best_val_f1:
            best_val_f1 = val_m["f1_chart"]
            best_epoch  = epoch
            Path(CONFIG["output_dir"]).mkdir(parents=True, exist_ok=True)
            model.save_pretrained(CONFIG["output_dir"])
            tokenizer.save_pretrained(CONFIG["output_dir"])

    print(f"\nBest model: epoch {best_epoch}, val F1={best_val_f1:.4f}")

    # ── Test set evaluation (load best model) ──────────────
    print("\n── Test set evaluation ───────────────────────────")
    best_model = AutoModelForSequenceClassification.from_pretrained(CONFIG["output_dir"])
    best_model.to(device)
    test_m, test_preds, test_labels = evaluate(best_model, test_loader, device)

    print(f"Test accuracy : {test_m['accuracy']:.4f}")
    print(f"Test F1 (chart): {test_m['f1_chart']:.4f}")
    print()
    print(classification_report(test_labels, test_preds,
                                target_names=["no_chart", "chart"]))
    print("Confusion matrix:")
    print(confusion_matrix(test_labels, test_preds))

    # ── Save training history ──────────────────────────────
    with open(f"{CONFIG['output_dir']}/history.json", "w") as f:
        json.dump(history, f, indent=2)

    # ── Quick inference demo ───────────────────────────────
    print("\n── Inference demo ────────────────────────────────")
    clf = ChartIntentClassifier(CONFIG["output_dir"], device)
    demo_queries = [
        "plot my heart rate over the last 7 days",
        "what is my resting heart rate",
        "show me a chart of my weekly step count",
        "how many calories did I burn today",
        "visualize my sleep stages breakdown",
        "did I hit my step goal yesterday",
    ]
    for q in demo_queries:
        result = clf.predict(q)
        bar = "█" * int(result["prob_chart"] * 20)
        print(f"  [{result['intent']:10s}] ({result['confidence']:.2f}) {bar:<20} {q}")

    print(f"\nModel saved to: {CONFIG['output_dir']}/")
    print("Done.")


if __name__ == "__main__":
    main()