"""
MAUDE NLP Classifier v2 — Bio_ClinicalBERT Inference Demo
Hugging Face Space: mukundisb/maude-classifier-bert

Tabs:
  1. Single Inference — classify a free-text MAUDE narrative
  2. Model Card      — metrics comparison Phase 1 (TF-IDF) vs Phase 2 (ClinicalBERT)
"""

import json
import os

import numpy as np
import streamlit as st
import torch
import torch.nn as nn
import matplotlib
import matplotlib.pyplot as plt
from transformers import AutoModel, AutoTokenizer
from huggingface_hub import hf_hub_download

matplotlib.use("Agg")

HUB_REPO    = "mukundisb/maude-clinicalbert"
LABEL_ORDER = ["D", "I", "M", "O"]
LABEL_NAMES = {"D": "Death", "I": "Injury", "M": "Malfunction", "O": "Other"}
SEVERITY_COLORS = {
    "D": "#d62728",   # red
    "I": "#ff7f0e",   # orange
    "M": "#1f77b4",   # blue
    "O": "#7f7f7f",   # grey
}
POOLING    = "cls_mean_concat"
MAX_LENGTH = 512


# ── Model class (must match Kaggle training architecture exactly) ─────────────

class BertSeverityClassifier(nn.Module):
    def __init__(self, bert, num_labels: int, dropout: float, pooling: str):
        super().__init__()
        self.bert       = bert
        H               = bert.config.hidden_size
        self.dropout    = nn.Dropout(dropout)
        self.pooling    = pooling
        head_in         = H * 2 if pooling == "cls_mean_concat" else H
        self.classifier = nn.Linear(head_in, num_labels)

    def forward(self, input_ids, attention_mask, token_type_ids=None):
        out  = self.bert(input_ids=input_ids, attention_mask=attention_mask,
                         token_type_ids=token_type_ids)
        last = out.last_hidden_state
        cls  = last[:, 0, :]
        if self.pooling == "cls":
            pooled = cls
        elif self.pooling == "mean":
            mask   = attention_mask.unsqueeze(-1).float()
            pooled = (last * mask).sum(1) / mask.sum(1).clamp(min=1)
        else:   # cls_mean_concat
            mask   = attention_mask.unsqueeze(-1).float()
            mean   = (last * mask).sum(1) / mask.sum(1).clamp(min=1)
            pooled = torch.cat([cls, mean], dim=-1)
        return self.classifier(self.dropout(pooled))


# ── Model loader ──────────────────────────────────────────────────────────────

@st.cache_resource(show_spinner="Loading Bio_ClinicalBERT from Hub (first run ~30 s)…")
def load_model_and_tokenizer():
    device    = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    tokenizer = AutoTokenizer.from_pretrained(HUB_REPO, do_lower_case=True)
    bert      = AutoModel.from_pretrained(HUB_REPO)
    head_path = hf_hub_download(repo_id=HUB_REPO, filename="classifier_head.pt")
    head_data = torch.load(head_path, map_location="cpu", weights_only=True)

    model = BertSeverityClassifier(
        bert=bert,
        num_labels=head_data["num_labels"],
        dropout=head_data["dropout_p"],
        pooling=head_data["pooling"],
    )
    model.classifier.load_state_dict(head_data["classifier_state_dict"])
    model = model.to(device)
    model.eval()
    return model, tokenizer, device


def predict(model, tokenizer, device, text: str) -> dict:
    enc = tokenizer(
        text,
        truncation=True,
        padding="max_length",
        max_length=MAX_LENGTH,
        return_tensors="pt",
    )
    input_ids      = enc["input_ids"].to(device)
    attention_mask = enc["attention_mask"].to(device)
    token_type_ids = enc.get("token_type_ids")
    if token_type_ids is not None:
        token_type_ids = token_type_ids.to(device)

    with torch.no_grad():
        logits = model(input_ids, attention_mask, token_type_ids)
        probs  = torch.softmax(logits, dim=-1).squeeze().cpu().numpy()

    label_order = LABEL_ORDER
    predicted   = label_order[int(np.argmax(probs))]
    return {
        "predicted_label": predicted,
        "probabilities":   {lbl: float(probs[i]) for i, lbl in enumerate(label_order)},
    }


# ── Page config ───────────────────────────────────────────────────────────────

st.set_page_config(
    page_title="MAUDE ClinicalBERT Classifier",
    page_icon="🏥",
    layout="wide",
    initial_sidebar_state="collapsed",
)

st.title("🏥 MAUDE Adverse Event Classifier — Bio_ClinicalBERT")
st.markdown(
    "Classify FDA MAUDE adverse event narratives by severity using a fine-tuned "
    "[Bio_ClinicalBERT](https://huggingface.co/emilyalsentzer/Bio_ClinicalBERT) model. "
    "Classes: **D**eath · **I**njury · **M**alfunction · **O**ther"
)

# Load model once — cached across all users
try:
    model, tokenizer, device = load_model_and_tokenizer()
    model_ready = True
    hw = "GPU" if device.type == "cuda" else "CPU"
    st.caption(f"Model loaded · {hw} · fold 3 (Death F1 = 0.842, Weighted F1 = 0.869)")
except Exception as exc:
    model_ready = False
    st.error(f"Could not load model from {HUB_REPO}: {exc}")

# ── Tabs ──────────────────────────────────────────────────────────────────────

tab_infer, tab_card = st.tabs(["🔍 Single Inference", "📋 Model Card"])


# ══════════════════════════════════════════════════════════════════════════════
# TAB 1 — Single Inference
# ══════════════════════════════════════════════════════════════════════════════

with tab_infer:
    st.subheader("Classify an Adverse Event Narrative")

    SAMPLES = {
        "Select a sample…": "",
        "Pump malfunction (no injury)": (
            "The infusion pump alarmed and stopped delivering medication. "
            "No patient injury was reported. The device was returned for analysis."
        ),
        "Patient death — pacemaker": (
            "The patient experienced cardiac arrest following implantation of the pacemaker. "
            "The patient passed away 48 hours post-procedure. The event was considered device-related."
        ),
        "Serious injury — burn": (
            "The patient sustained third-degree burns on the left forearm during an "
            "electrosurgical procedure due to a grounding pad malfunction."
        ),
        "Catheter — minor injury": (
            "During catheter insertion the patient reported significant pain and minor bleeding. "
            "The procedure was discontinued and the patient was monitored overnight."
        ),
    }

    sample_choice   = st.selectbox("Load a sample narrative", list(SAMPLES.keys()))
    narrative_input = st.text_area(
        "Narrative text",
        value=SAMPLES[sample_choice],
        height=160,
        placeholder="Paste an MDR adverse event narrative here…",
    )

    classify_btn = st.button(
        "🔎 Classify",
        use_container_width=True,
        disabled=not model_ready,
    )

    if classify_btn:
        text = narrative_input.strip()
        if not text:
            st.warning("Please enter or select a narrative text.")
        else:
            with st.spinner("Running inference… (~3–8 s on CPU)"):
                result    = predict(model, tokenizer, device, text)
            predicted     = result["predicted_label"]
            label_name    = LABEL_NAMES.get(predicted, predicted)
            color         = SEVERITY_COLORS.get(predicted, "#333333")
            top_prob      = result["probabilities"][predicted]

            st.markdown(
                f"<div style='background:{color};padding:18px;border-radius:10px;"
                f"color:white;font-size:24px;font-weight:bold;text-align:center'>"
                f"Predicted: {label_name} &nbsp;({predicted}) &nbsp;·&nbsp; "
                f"confidence {top_prob:.1%}</div>",
                unsafe_allow_html=True,
            )

            st.markdown("**Class Probabilities**")
            proba_items = sorted(
                result["probabilities"].items(), key=lambda x: -x[1]
            )
            fig, ax = plt.subplots(figsize=(6, 3))
            labels  = [LABEL_NAMES[c] for c, _ in proba_items]
            values  = [v for _, v in proba_items]
            colors  = [SEVERITY_COLORS.get(c, "#aaa") for c, _ in proba_items]
            bars    = ax.barh(labels, values, color=colors)
            ax.bar_label(bars, fmt="%.3f", padding=4)
            ax.set_xlim(0, 1.05)
            ax.set_xlabel("Probability")
            ax.invert_yaxis()
            plt.tight_layout()
            st.pyplot(fig)
            plt.close()

    st.divider()
    st.caption(
        "Inference is CPU-only on free Spaces — first classify may take up to 10 s while "
        "PyTorch warms up. Subsequent calls are faster."
    )


# ══════════════════════════════════════════════════════════════════════════════
# TAB 2 — Model Card
# ══════════════════════════════════════════════════════════════════════════════

with tab_card:
    st.subheader("Model Card — Bio_ClinicalBERT fine-tuned on MAUDE")

    st.markdown(
        """
### Background

The FDA's [MAUDE database](https://open.fda.gov/apis/device/event/) contains millions of
Medical Device Reports (MDRs). Each report includes a free-text narrative describing the
adverse event. This classifier assigns a severity label to each narrative.

**Phase 1** used a TF-IDF + Logistic Regression baseline (v1 Space).
**Phase 2** (this Space) fine-tuned [Bio_ClinicalBERT](https://huggingface.co/emilyalsentzer/Bio_ClinicalBERT)
on the same 154 k-record dataset using 5-fold stratified cross-validation on a Kaggle T4×2 GPU.
"""
    )

    st.markdown("### 5-Fold CV Results")

    import pandas as pd

    df_metrics = pd.DataFrame({
        "Metric":           ["Weighted F1", "Macro F1", "Death F1", "Injury F1",
                             "Malfunction F1", "Other F1"],
        "Phase 1 (TF-IDF)": ["0.848 ± 0.001", "0.738 ± 0.003", "0.770 ± 0.011",
                              "0.847 ± 0.002", "0.893 ± 0.002", "0.441 ± 0.005"],
        "Phase 2 (BERT)":   ["0.869 ± 0.001", "0.773 ± 0.002", "0.832 ± 0.008",
                              "0.873 ± 0.002", "0.907 ± 0.002", "0.479 ± 0.007"],
        "Δ":                ["+0.021", "+0.035", "+0.062", "+0.025", "+0.014", "+0.038"],
    })
    st.dataframe(df_metrics, hide_index=True, use_container_width=True)

    st.info(
        "**Death-class F1 +6.2 pp** is the headline result. "
        "The lift of +0.021 on weighted F1 against a combined σ of ~0.0016 is ~13σ — "
        "a statistically robust, reproducible improvement."
    )

    st.markdown("### Per-Fold Breakdown (Phase 2)")

    df_folds = pd.DataFrame({
        "Fold": [1, 2, 3, 4, 5],
        "Weighted F1": [0.8671, 0.8700, 0.8693, 0.8693, 0.8698],
        "Macro F1":    [0.7709, 0.7748, 0.7747, 0.7698, 0.7733],
        "Death F1":    [0.8346, 0.8364, 0.8418, 0.8249, 0.8212],
        "Best Epoch":  [3, 3, 3, 3, 3],
        "Deployed":    ["", "", "✓", "", ""],
    })
    st.dataframe(df_folds, hide_index=True, use_container_width=True)

    st.markdown("### Training Configuration")

    cfg = {
        "Base model":          "emilyalsentzer/Bio_ClinicalBERT",
        "Pooling":             "cls_mean_concat ([CLS] + mean → 1536-dim head)",
        "Max token length":    512,
        "Batch size":          "16 (×2 GPUs = effective 32 with grad_accum=2)",
        "Learning rate":       "2e-5 (AdamW + linear warmup 10%)",
        "Epochs":              3,
        "Class weights":       "Inverse-frequency (CrossEntropyLoss)",
        "Training hardware":   "Kaggle T4×2, DataParallel",
        "Training dataset":    "154 776 MAUDE records (post-cleaning, UNKNOWN dropped)",
        "CV strategy":         "5-fold StratifiedKFold, locked splits (SHA1 fingerprinted)",
        "Deployed fold":       "3 (highest Death F1 = 0.8418)",
    }
    for k, v in cfg.items():
        st.markdown(f"- **{k}:** {v}")

    st.markdown("### Honest Caveats")
    st.markdown(
        """
- All folds peaked at epoch 3 (the training cap) — the model was still improving.
  Extending to 5 epochs would likely add another +0.003 to +0.005 weighted F1.
  Current numbers are **conservative**.
- **Other-class F1 ≈ 0.479** remains the weakest class. "Other" is a small, semantically
  heterogeneous bucket; this is a known limitation, not an optimisation target.
- This model is a **research prototype**. Do not use for clinical decision-making.
"""
    )

    st.markdown("### Links")
    st.markdown(
        "- [GitHub repo](https://github.com/mukundp1/maude-nlp-classifier) — "
        "source code, CV splits, runbooks\n"
        "- [Model weights](https://huggingface.co/mukundisb/maude-clinicalbert) — "
        "HF Hub model card\n"
        "- [Phase 1 Space](https://huggingface.co/spaces/mukundisb/maude-classifier) — "
        "TF-IDF baseline demo"
    )