"""
Samsung Health Chart Intent — Inference + OOD Detection
========================================================
Loads the fine-tuned DistilBERT model and an embedding-based OOD detector.
Run this directly to test queries in the terminal, or import the
`full_predict` function into app.py.

Usage:
    python test.py                        # interactive mode
    python test.py --query "plot my HR"   # single query
    python test.py --batch queries.txt    # one query per line in a text file

Requirements:
    pip install torch transformers sentence-transformers pandas
"""

import argparse
import pandas as pd
import numpy as np
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from sentence_transformers import SentenceTransformer

# ─────────────────────────────────────────────────────────
# Config
# ─────────────────────────────────────────────────────────

MODEL_DIR       = "./chart_intent_model"
TRAINING_DATA   = "./samsung_health_intent.csv"
MAX_LENGTH      = 64
OOD_PERCENTILE  = 95    # 95th percentile of training distances as threshold
CONF_THRESHOLD  = 0.70  # below this → uncertain even if in-domain

# ─────────────────────────────────────────────────────────
# Device
# ─────────────────────────────────────────────────────────

def get_device():
    if torch.cuda.is_available():
        return torch.device("cuda"), "CUDA"
    elif torch.backends.mps.is_available():
        return torch.device("mps"), "Apple MPS"
    return torch.device("cpu"), "CPU"

# ─────────────────────────────────────────────────────────
# OOD Detector
# ─────────────────────────────────────────────────────────

class OODDetector:
    """
    Embedding-distance OOD detector.
    Fits on training data. At inference, rejects inputs whose
    cosine distance from the training centroid exceeds the threshold.
    """
    def __init__(self, percentile: int = 95):
        self.percentile = percentile
        print("  Loading sentence encoder...")
        self.encoder = SentenceTransformer(
            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
        )
        self.centroid  = None
        self.threshold = None

    def fit(self, texts: list):
        print(f"  Fitting OOD detector on {len(texts)} training examples...")
        embs = self.encoder.encode(texts, normalize_embeddings=True, show_progress_bar=False)

        # Centroid of in-domain distribution
        self.centroid = embs.mean(axis=0)
        self.centroid /= np.linalg.norm(self.centroid)

        # Distance of each training example from centroid
        dists = 1 - (embs @ self.centroid)          # cosine distance, 0=identical

        # Threshold = Nth percentile → 95% of training samples are within it
        self.threshold = float(np.percentile(dists, self.percentile))
        print(f"  OOD threshold ({self.percentile}th pct): {self.threshold:.4f}")
        print(f"  Training distance range: [{dists.min():.4f}, {dists.max():.4f}]")

    def check(self, text: str) -> tuple:
        """Returns (is_in_domain: bool, distance: float)."""
        emb  = self.encoder.encode([text], normalize_embeddings=True)[0]
        dist = float(1 - (emb @ self.centroid))
        return dist <= self.threshold, dist

# ─────────────────────────────────────────────────────────
# Classifier
# ─────────────────────────────────────────────────────────

class ChartIntentClassifier:
    def __init__(self, model_dir: str, device):
        self.device    = device
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
        self.model     = AutoModelForSequenceClassification.from_pretrained(model_dir)
        self.model.to(device)
        self.model.eval()

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

        inputs.pop("token_type_ids", None)   # DistilBERT has no token_type_ids
        logits = self.model(**inputs).logits
        probs  = torch.softmax(logits, dim=-1).cpu().numpy()[0]
        label  = int(probs.argmax())

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

# ─────────────────────────────────────────────────────────
# Full pipeline
# ─────────────────────────────────────────────────────────

def full_predict(text: str, ood: OODDetector, clf: ChartIntentClassifier) -> dict:
    """
    Two-stage inference:
      1. OOD check  → reject if not a health query
      2. Classifier → chart vs no_chart
    """
    # Stage 1: domain gate
    in_domain, dist = ood.check(text)
    if not in_domain:
        return {
            "text":      text,
            "intent":    "out_of_domain",
            "label":     -1,
            "confident": False,
            "distance":  dist,
            "message":   "Not a health query — rejected by OOD detector",
        }

    # Stage 2: classify
    result = clf.predict(text)
    uncertain = result["confidence"] < CONF_THRESHOLD

    return {
        "text":      text,
        "intent":    result["intent"],
        "label":     result["label"],
        "confident": not uncertain,
        "confidence":    result["confidence"],
        "prob_chart":    result["prob_chart"],
        "prob_no_chart": result["prob_no_chart"],
        "distance":  dist,
        "message":   "uncertain — confidence below threshold" if uncertain else "ok",
    }

# ─────────────────────────────────────────────────────────
# Pretty print
# ─────────────────────────────────────────────────────────

def print_result(r: dict):
    intent = r["intent"]

    if intent == "out_of_domain":
        icon   = "🚫"
        label  = "OUT OF DOMAIN"
        detail = f"distance={r['distance']:.4f} (exceeds OOD threshold)"
    elif not r["confident"]:
        icon   = "⚠️ "
        label  = f"UNCERTAIN → {intent.upper()}"
        detail = f"confidence={r['confidence']:.1%}  prob_chart={r['prob_chart']:.3f}  dist={r['distance']:.4f}"
    elif intent == "chart":
        icon   = "📊"
        label  = "CHART"
        detail = f"confidence={r['confidence']:.1%}  prob_chart={r['prob_chart']:.3f}  dist={r['distance']:.4f}"
    else:
        icon   = "💬"
        label  = "NO CHART"
        detail = f"confidence={r['confidence']:.1%}  prob_no_chart={r['prob_no_chart']:.3f}  dist={r['distance']:.4f}"

    print(f"\n  {icon}  [{label}]")
    print(f"      Query : {r['text']}")
    print(f"      Detail: {detail}")

# ─────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────

def build_pipeline():
    """Load everything. Call once, reuse the returned objects."""
    device, device_name = get_device()
    print(f"\n── Device: {device_name}")

    # Load training data for OOD fitting
    print(f"\n── Loading training data: {TRAINING_DATA}")
    df  = pd.read_csv(TRAINING_DATA)
    texts = df["text"].tolist()

    # OOD detector
    print("\n── Building OOD detector")
    ood = OODDetector(percentile=OOD_PERCENTILE)
    ood.fit(texts)

    # Classifier
    print(f"\n── Loading classifier: {MODEL_DIR}")
    clf = ChartIntentClassifier(MODEL_DIR, device)
    print("  Model loaded")

    return ood, clf


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--query", type=str, default=None,
                        help="Single query to classify")
    parser.add_argument("--batch", type=str, default=None,
                        help="Path to a .txt file with one query per line")
    args = parser.parse_args()

    ood, clf = build_pipeline()

    # ── Single query from CLI ──────────────────────────────
    if args.query:
        r = full_predict(args.query, ood, clf)
        print_result(r)
        return

    # ── Batch from file ────────────────────────────────────
    if args.batch:
        with open(args.batch) as f:
            queries = [l.strip() for l in f if l.strip()]
        print(f"\n── Batch: {len(queries)} queries")
        print("─" * 60)
        for q in queries:
            print_result(full_predict(q, ood, clf))
        return

    # ── Interactive mode ───────────────────────────────────
    print("\n── Interactive mode  (type 'quit' to exit)")
    print(f"    OOD threshold percentile : {OOD_PERCENTILE}")
    print(f"    Confidence threshold     : {CONF_THRESHOLD:.0%}")
    print("─" * 60)

    while True:
        try:
            query = input("\nQuery: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nBye.")
            break
        if not query:
            continue
        if query.lower() in ("quit", "exit", "q"):
            print("Bye.")
            break
        print_result(full_predict(query, ood, clf))


if __name__ == "__main__":
    main()