from __future__ import annotations

from pathlib import Path

import gradio as gr
from inference_core import CLASS_NAMES, run_demo_inference

EXAMPLES_DIR = Path(__file__).resolve().parent / "examples"
EXAMPLES = {
    "2999": {
        "image": EXAMPLES_DIR / "2999_LR_crop.png",
        "geojson": EXAMPLES_DIR / "2999_contours_lr.geojson",
    },
    "3096": {
        "image": EXAMPLES_DIR / "3096_LR_crop.png",
        "geojson": EXAMPLES_DIR / "3096_contours_lr.geojson",
    },
}


def get_example_preview(example_id: str):
    return EXAMPLES[example_id]["image"]


def get_input_preview(example_id: str, uploaded_image):
    if uploaded_image is not None:
        return uploaded_image
    return get_example_preview(example_id)


def predict(
    example_id: str,
    uploaded_image,
    uploaded_geojson,
    use_roi: bool,
    seed: int,
    progress: gr.Progress = gr.Progress(track_tqdm=True),
):
    progress(0.05, desc="Preparing inputs")
    if uploaded_image is not None:
        image_path = Path(uploaded_image)
        geojson_path = Path(uploaded_geojson) if uploaded_geojson else None
        source = "uploaded image"
    else:
        example = EXAMPLES[example_id]
        image_path = example["image"]
        geojson_path = example["geojson"] if use_roi else None
        source = f"bundled example {example_id}"

    if not use_roi:
        geojson_path = None

    progress(0.15, desc=f"Running density inference on {source}")
    result = run_demo_inference(
        image_path=image_path,
        geojson_path=geojson_path,
        use_roi=use_roi,
        seed=int(seed),
    )
    progress(0.95, desc="Preparing visual outputs")

    counts = [
        [row["class"], row["density_sum"], row["sampled_count"]]
        for row in result["counts"]
    ]
    density_maps = result["density_maps"]
    sampled_maps = result["sampled_maps"]

    return (
        result["original"],
        result["combined_density"],
        density_maps[0],
        density_maps[1],
        density_maps[2],
        result["combined_points"],
        sampled_maps[0],
        sampled_maps[1],
        sampled_maps[2],
        counts,
    )


with gr.Blocks(title="CALHippo Demo") as demo:
    gr.Markdown(
        "# CALHippo Demo\n"
        "Low-resolution hippocampus WSI density estimation for pyramidal cells, "
        "interneurons, and astrocytes. Use a bundled example or upload a LR PNG "
        "with an optional ROI GeoJSON."
    )

    with gr.Row():
        with gr.Column(scale=2):
            original = gr.Image(
                value=get_example_preview("3096"),
                label="Input used for inference",
                type="numpy",
            )
            with gr.Tabs():
                with gr.Tab("All Classes Density"):
                    combined_density = gr.Image(
                        label="All classes density overlay",
                        type="numpy",
                    )
                density_outputs = []
                for class_name in CLASS_NAMES:
                    with gr.Tab(f"{class_name} Density"):
                        density_outputs.append(
                            gr.Image(label=f"{class_name} density", type="numpy")
                        )

            gr.Markdown("## Sampled Points")
            with gr.Tabs():
                with gr.Tab("All Classes"):
                    combined_points = gr.Image(
                        label="All classes sampled points",
                        type="numpy",
                    )
                sampled_outputs = []
                for class_name in CLASS_NAMES:
                    with gr.Tab(class_name):
                        sampled_outputs.append(
                            gr.Image(
                                label=f"{class_name} sampled points",
                                type="numpy",
                            )
                        )

            counts = gr.Dataframe(
                headers=["Class", "Density sum", "Sampled count"],
                datatype=["str", "number", "number"],
                label="Predicted counts",
            )

        with gr.Column(scale=1, min_width=320):
            gr.Markdown("## Inputs")
            example_id = gr.Dropdown(
                choices=list(EXAMPLES),
                value="3096",
                label="Bundled example",
            )
            gr.Examples(
                examples=[["3096"], ["2999"]],
                inputs=[example_id],
                label="Ready-made examples",
                cache_examples=False,
            )
            use_roi = gr.Checkbox(value=True, label="Use ROI GeoJSON when available")
            seed = gr.Number(value=42, precision=0, label="Sampling seed")
            uploaded_image = gr.File(
                label="Optional LR crop PNG upload",
                file_types=[".png"],
                type="filepath",
            )
            uploaded_geojson = gr.File(
                label="Optional ROI GeoJSON upload",
                file_types=[".geojson", ".json"],
                type="filepath",
            )
            run_button = gr.Button("Run Inference", variant="primary")

    example_id.change(
        get_input_preview,
        inputs=[example_id, uploaded_image],
        outputs=original,
    )
    uploaded_image.change(
        get_input_preview,
        inputs=[example_id, uploaded_image],
        outputs=original,
    )

    run_button.click(
        predict,
        inputs=[example_id, uploaded_image, uploaded_geojson, use_roi, seed],
        outputs=[
            original,
            combined_density,
            *density_outputs,
            combined_points,
            *sampled_outputs,
            counts,
        ],
    )


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