"""Gradio web app for the HMoE dopaminergic-subtype classifier.

Drag-and-drop an .h5ad (or Seurat .rds) of raw sc/snRNA-seq counts (dopaminergic
neurons) and download a CSV assigning each cell to a subtype. Deployable as a
Hugging Face Space (Gradio SDK); also runs locally with `python app.py`.
"""
from __future__ import annotations

import tempfile
from pathlib import Path

import gradio as gr

from hmoe_da import load_model
from hmoe_da.io import read_adata

MODEL = load_model()  # bundled weights, loads once at startup

DESCRIPTION = """
# Dopaminergic-neuron subtype classifier (HMoE)

Assign each cell in your **raw** sc/snRNA-seq dataset to one of 18 dopaminergic
subtypes across three families (**Sox6**, **Calb1**, **Gad2**).

**Input:** an `.h5ad` (raw integer counts in `.X`) or a Seurat `.rds`, with gene
symbols in `var_names` (human UPPER or mouse Title case both work).
**Output:** a CSV of per-cell subtype, family, confidence, and margin.

_Awatramani Lab, Northwestern University (Schonfeld et al., manuscript in preparation)._
"""


def _counts_markdown(labels) -> str:
    vc = labels.value_counts()
    total = int(vc.sum())
    lines = ["| subtype | cells | % |", "|---|---:|---:|"]
    for name, n in vc.items():
        lines.append(f"| {name} | {int(n)} | {100*n/total:.1f}% |")
    lines.append(f"| **total** | **{total}** | |")
    return "\n".join(lines)


def run(data_file, exclude_gad2_family, count_source):
    if data_file is None:
        raise gr.Error("Please upload an .h5ad or .rds file.")
    try:
        ad = read_adata(data_file.name)
    except Exception as e:
        raise gr.Error(f"Could not read file: {e}")

    is_h5ad = str(data_file.name).lower().endswith(".h5ad")
    if is_h5ad and count_source == "adata.raw":
        if ad.raw is None:
            raise gr.Error("adata.raw is None; choose 'adata.X' or upload raw counts.")
        ad = ad.raw.to_adata()
    elif is_h5ad and count_source.startswith("layer:"):
        layer = count_source.split("layer:", 1)[1]
        if layer not in ad.layers:
            raise gr.Error(f"Layer {layer!r} not found. Available: {list(ad.layers)}")
        ad.X = ad.layers[layer]

    try:
        ad = MODEL.predict(ad, include_gad2=not exclude_gad2_family, verbose=False)
    except ValueError as e:
        raise gr.Error(str(e))

    cols = [c for c in ad.obs.columns if c.startswith("hmoe")]
    df = ad.obs[cols].copy()
    df.index.name = "cell"
    out = Path(tempfile.mkdtemp()) / "hmoe_predictions.csv"
    df.to_csv(out)

    cov = ad.uns["hmoe_gene_coverage"]
    header = (
        f"**Predicted {ad.n_obs} cells.** "
        f"Gene coverage {cov['n_matched']}/{cov['n_features']} "
        f"({cov['pct_matched']:.1f}%). "
        f"Gad2 family {'excluded' if exclude_gad2_family else 'included'}.\n\n"
    )
    return header + _counts_markdown(ad.obs["hmoe_subtype"]), str(out)


def build():
    with gr.Blocks(title="HMoE DA subtype classifier") as demo:
        gr.Markdown(DESCRIPTION)
        with gr.Row():
            with gr.Column():
                f = gr.File(label="Upload .h5ad or .rds (raw counts)",
                            file_types=[".h5ad", ".rds"])
                gr.Markdown(
                    "<sub>Seurat `.rds` is read where R + SeuratObject are "
                    "installed (local runs). On the hosted app, please upload "
                    "`.h5ad` — convert in R with `sceasy` or `SeuratDisk`.</sub>"
                )
                src = gr.Dropdown(
                    ["adata.X", "adata.raw", "layer:counts"],
                    value="adata.X",
                    label="Where are the raw counts? (.h5ad only)",
                )
                excl = gr.Checkbox(
                    label="Exclude Gad2 family (reassign Gad2 calls to best non-Gad2 subtype)",
                    value=False,
                )
                btn = gr.Button("Predict subtypes", variant="primary")
            with gr.Column():
                msg = gr.Markdown()
                dl = gr.File(label="Download per-cell CSV")
        btn.click(run, inputs=[f, excl, src], outputs=[msg, dl])
    return demo


if __name__ == "__main__":
    build().launch()