import gradio as gr

MODEL_REPO = "Simmonstt/BrainAnytime"
GITHUB_REPO = "https://github.com/guangqianyang/BrainAnytime"
CHECKPOINTS = [
    "CN_vs_AD_seed_0_best.pth",
    "CN_vs_MCI_seed_0_best.pth",
    "MMSE_seed_0_best.pth",
    "AGE_seed_0_best.pth",
]

INTRO = """
# BrainAnytime Demo

**BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability**

This Hugging Face Space hosts the official code from GitHub. Full 3D multi-modal inference
requires preprocessed NIfTI volumes and GPU resources. Use the linked model repository for
finetuned checkpoints and run `finetune_main.py` / `test_main.py` locally for evaluation.
"""


def show_project_info():
    checkpoint_lines = "\n".join(f"- `{name}`" for name in CHECKPOINTS)
    return f"""{INTRO}

## Links
- GitHub: {GITHUB_REPO}
- Model weights: https://huggingface.co/{MODEL_REPO}

## Available finetuned checkpoints
{checkpoint_lines}

## Supported downstream tasks
- CN vs AD (classification)
- CN vs MCI (classification)
- MMSE (regression)
- AGE (regression)

## Quick start (local)
```bash
git clone {GITHUB_REPO}.git
cd BrainAnytime
pip install -r requirements.txt
python finetune_main.py --pretrained <path/to/pretrained.pth>
```
"""


with gr.Blocks(title="BrainAnytime Demo") as demo:
    gr.Markdown(INTRO)
    gr.Button("Show project details").click(show_project_info, outputs=gr.Markdown())

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