"""
ECG SmartRead AI Microservice — Pure Gradio for Hugging Face ZeroGPU.

Exposes two API endpoints via Gradio:
  - predict_manual: BioLinkBERT classification
  - generate_report: SciFive T5 report generation

Called from the VPS backend via the Gradio Python Client.
"""

import json
import os
from functools import lru_cache

import spaces
import gradio as gr
import torch
from transformers import AutoTokenizer, T5ForConditionalGeneration, AutoModelForSequenceClassification
from huggingface_hub import snapshot_download

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_REPO_ID = "adelfr2009/ecg-smartread-models"
MODELS_DIR = os.path.join(BASE_DIR, "models")
GENERATION_DIR = os.path.join(MODELS_DIR, "generation")
MANUAL_TEST_DIR = os.path.join(MODELS_DIR, "test")


# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------
def _ensure_models_downloaded():
    gen_ok = os.path.exists(os.path.join(GENERATION_DIR, "model.safetensors"))
    test_ok = os.path.exists(os.path.join(MANUAL_TEST_DIR, "model.safetensors"))
    if not gen_ok or not test_ok:
        print(f"Downloading models from {MODEL_REPO_ID} ...")
        snapshot_download(repo_id=MODEL_REPO_ID, local_dir=MODELS_DIR)
        print("Download complete.")


@lru_cache(maxsize=1)
def _load_generation():
    _ensure_models_downloaded()
    tok = AutoTokenizer.from_pretrained(GENERATION_DIR)
    mdl = T5ForConditionalGeneration.from_pretrained(GENERATION_DIR)
    mdl.eval()
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    mdl.to(dev)
    return tok, mdl, dev


@lru_cache(maxsize=1)
def _load_classification():
    _ensure_models_downloaded()
    tok = AutoTokenizer.from_pretrained(MANUAL_TEST_DIR)
    mdl = AutoModelForSequenceClassification.from_pretrained(MANUAL_TEST_DIR)
    mdl.eval()
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    mdl.to(dev)
    with open(os.path.join(MANUAL_TEST_DIR, "thresholds.json")) as f:
        th = json.load(f)
    return tok, mdl, dev, th["labels"], th["thresholds"], th.get("weak_classes", [])


# ---------------------------------------------------------------------------
# GPU-decorated inference functions (ZeroGPU requirement)
# ---------------------------------------------------------------------------
@spaces.GPU
def predict_manual(
    heart_rate, qrs_duration, qt_interval, qtc_interval,
    pr_interval, p_axis, qrs_axis, t_axis, age, sex
):
    """Classify ECG from manual global metrics via BioLinkBERT."""
    def fmt(v):
        return "unknown" if v is None or str(v).strip() == "" else v

    text = (
        f"Heart Rate: {fmt(heart_rate)} BPM. "
        f"QRS Duration: {fmt(qrs_duration)} ms. "
        f"QT Interval: {fmt(qt_interval)} ms. "
        f"QTc Interval: {fmt(qtc_interval)} ms. "
        f"PR Interval: {fmt(pr_interval)} ms. "
        f"P axis: {fmt(p_axis)} degrees. QRS axis: {fmt(qrs_axis)} degrees. T axis: {fmt(t_axis)} degrees. "
        f"Age: {fmt(age)}. Sex: {fmt(sex)}."
    )

    tok, mdl, dev, labels, thresholds, weak_classes = _load_classification()

    inputs = tok(text, truncation=True, padding="max_length",
                 max_length=64, return_tensors="pt").to(dev)

    with torch.no_grad():
        logits = mdl(**inputs).logits
    probs = torch.sigmoid(logits)[0].cpu().numpy()

    probabilities = {}
    predictions = {}
    for label, prob, thresh in zip(labels, probs, thresholds):
        probabilities[label] = round(float(prob) * 100, 1)
        predictions[label] = bool(prob > thresh)

    return json.dumps({
        "input_text": text,
        "probabilities": probabilities,
        "predictions": predictions,
        "weak_classes": weak_classes,
    })


@spaces.GPU
def generate_report(input_text, max_new_tokens=200):
    """Generate medical report from clinical summary via SciFive T5."""
    tok, mdl, dev = _load_generation()

    max_tok = int(max_new_tokens) if max_new_tokens else 200
    inputs = tok(input_text, return_tensors="pt", truncation=True,
                 max_length=512).to(dev)

    with torch.no_grad():
        ids = mdl.generate(**inputs, max_new_tokens=max_tok,
                           num_beams=4, early_stopping=True)

    generated = tok.decode(ids[0], skip_special_tokens=True)
    return json.dumps({
        "input_text": input_text,
        "generated_report": generated,
    })


# ---------------------------------------------------------------------------
# Gradio interface (required for ZeroGPU to detect @spaces.GPU functions)
# ---------------------------------------------------------------------------
with gr.Blocks(title="ECG SmartRead AI") as demo:
    gr.Markdown("# 🫀 ECG SmartRead AI Microservice")
    gr.Markdown("API-only service. Use the API tab below or call via `gradio_client`.")

    with gr.Tab("Classification"):
        with gr.Row():
            hr  = gr.Number(label="Heart Rate (BPM)")
            qrs = gr.Number(label="QRS Duration (ms)")
            qt  = gr.Number(label="QT Interval (ms)")
            qtc = gr.Number(label="QTc Interval (ms)")
        with gr.Row():
            pr  = gr.Number(label="PR Interval (ms)")
            pa  = gr.Number(label="P Axis (°)")
            qa  = gr.Number(label="QRS Axis (°)")
            ta  = gr.Number(label="T Axis (°)")
        with gr.Row():
            age = gr.Number(label="Age")
            sex = gr.Textbox(label="Sex", placeholder="male / female")
        classify_btn = gr.Button("Classify", variant="primary")
        classify_out = gr.Textbox(label="Result (JSON)", lines=10)
        classify_btn.click(
            fn=predict_manual,
            inputs=[hr, qrs, qt, qtc, pr, pa, qa, ta, age, sex],
            outputs=classify_out,
        )

    with gr.Tab("Report Generation"):
        gen_input = gr.Textbox(label="Clinical Summary Input", lines=5,
                               placeholder="Heart Rate: 75 BPM. QRS Duration: ...")
        gen_tokens = gr.Number(label="Max New Tokens", value=200)
        gen_btn = gr.Button("Generate Report", variant="primary")
        gen_out = gr.Textbox(label="Generated Report (JSON)", lines=10)
        gen_btn.click(
            fn=generate_report,
            inputs=[gen_input, gen_tokens],
            outputs=gen_out,
        )

demo.launch()