"""
Head-to-head Gradio baseline for the VoxPixel paper (Section V).

This file deliberately implements the SRMA-Mamba liver segmentation model
in three idiomatic Gradio configurations to expose, on the same hardware
and with the same checkpoints used by VoxPixel, the three failure modes
called out by reviewers:

    V1: gr.Image(...)            -> failure mode 1: NIfTI ingestion not
                                    supported. The component rejects
                                    .nii / .nii.gz at upload time.

    V2: gr.File + nibabel +      -> failure mode 3: CUDA OOM. A single
        single-shot full-volume     forward pass on the full 3D volume
        forward pass.               exceeds GPU memory on standard scans
                                    because no sliding-window aggregator
                                    is available out of the box.

    V3: gr.File + nibabel +      -> failure mode 2: spatial metadata
        manually-wired MONAI        loss. Inference succeeds, but the
        SlidingWindowInferer.       output mask is written with
                                    affine = identity because Gradio's
                                    I/O contract does not preserve the
                                    NIfTI affine. Reported liver volume
                                    in ml is therefore biased relative
                                    to the spatially-correct VoxPixel
                                    output.

This is NOT a production app. See app.py for VoxPixel itself. This file
exists so reviewers can reproduce the head-to-head table in the paper.

Usage (interactive):
    python app_gradio_baseline.py --variant v1
    python app_gradio_baseline.py --variant v2
    python app_gradio_baseline.py --variant v3

Usage (programmatic, called by bench_gradio_baseline.py):
    from app_gradio_baseline import predict_v2, predict_v3
"""

import os
import sys
import time
import argparse
from pathlib import Path

import numpy as np
import torch
import nibabel as nib


# ---------------------------------------------------------------------------
# Compatibility shims (mirror app.py).
# Must run BEFORE `import gradio as gr` for the HfFolder shim, and AFTER for
# the schema bug fix.
# ---------------------------------------------------------------------------
def _patch_hf_folder():
    """Gradio 4.44.x imports HfFolder from huggingface_hub, which was removed
    in huggingface_hub 1.0+. Inject a minimal shim."""
    import huggingface_hub as hh
    if getattr(hh, "HfFolder", None) is not None:
        return
    try:
        from huggingface_hub import get_token
    except ImportError:
        get_token = lambda: None  # noqa: E731

    class HfFolder:
        @staticmethod
        def get_token():
            return get_token()

    hh.HfFolder = HfFolder


_patch_hf_folder()

import gradio as gr


def _fix_gradio_schema_bug():
    """Patch gradio_client.utils.get_type to handle boolean schemas
    (Gradio 4.44.x crashes on additionalProperties: True)."""
    try:
        import gradio_client.utils as gu
        if not hasattr(gu, "get_type"):
            return
        original = gu.get_type

        def patched(schema):
            if isinstance(schema, bool):
                return "Any"
            if isinstance(schema, dict):
                if schema.get("additionalProperties") is True:
                    schema["additionalProperties"] = {}
                elif schema.get("additionalProperties") is False:
                    schema.pop("additionalProperties", None)
            return original(schema)

        gu.get_type = patched
    except Exception:
        pass


_fix_gradio_schema_bug()


_model_cache: dict = {}


def _load_bare_model(modality: str = "T1") -> torch.nn.Module:
    """Load architecture + checkpoint with no sliding-window inferer,
    no warm-up, no torch.compile, no channels-last layout. This is what
    a researcher writing a naive Gradio wrapper would produce."""
    # Imported lazily so V1 (which never calls this) can run on a laptop
    # without mamba_ssm / selective_scan_cuda_oflex installed.
    from config import (
        build_SRMAMamba,
        BUILD_SRMAMAMBA_AVAILABLE,
        SRMA_MAMBA_DIR,
    )

    if not BUILD_SRMAMAMBA_AVAILABLE or build_SRMAMamba is None:
        raise ImportError(
            "SRMA-Mamba builder not available. Run setup.sh / install "
            "mamba_ssm and selective_scan_cuda_oflex first."
        )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if SRMA_MAMBA_DIR:
        original_cwd = os.getcwd()
        try:
            os.chdir(SRMA_MAMBA_DIR)
            model = build_SRMAMamba()
        finally:
            os.chdir(original_cwd)
    else:
        model = build_SRMAMamba()

    ckpt_path = f"checkpoint_{modality}.pth"
    if not os.path.exists(ckpt_path):
        local = os.path.join(os.path.dirname(__file__), ckpt_path)
        if os.path.exists(local):
            ckpt_path = local
        else:
            from huggingface_hub import hf_hub_download
            repo_id = os.environ.get(
                "HF_MODEL_REPO", "HarshithReddy01/srmamamba-liver-segmentation"
            )
            ckpt_path = hf_hub_download(
                repo_id=repo_id, filename=f"checkpoint_{modality}.pth", cache_dir="."
            )

    ckpt = torch.load(ckpt_path, map_location=device)
    state = (
        ckpt["state_dict"]
        if isinstance(ckpt, dict) and "state_dict" in ckpt
        else ckpt
    )
    model.load_state_dict(state)
    model.eval().to(device)
    return model


def _get_model(modality: str = "T1") -> torch.nn.Module:
    if modality not in _model_cache:
        _model_cache[modality] = _load_bare_model(modality)
    return _model_cache[modality]


def _normalize(vol: np.ndarray) -> np.ndarray:
    """Z-score normalization over nonzero voxels. Mirrors what
    NormalizeIntensityd(nonzero=True) does in the VoxPixel pipeline,
    inlined here so the comparison is about deployment, not preprocessing."""
    vol = vol.astype(np.float32)
    nz = vol[vol > 1e-6]
    if nz.size == 0:
        return vol
    mu = float(nz.mean())
    sigma = float(nz.std()) + 1e-8
    out = np.where(vol > 1e-6, (vol - mu) / sigma, 0.0)
    return out.astype(np.float32)


def _threshold_for(modality: str) -> float:
    if modality.upper() == "T1":
        return float(os.environ.get("T1_THRESHOLD", "0.65"))
    return float(os.environ.get("SEGMENTATION_THRESHOLD", "0.5"))


def _summary(img: nib.Nifti1Image, mask: np.ndarray, t0: float) -> dict:
    """Build the result record. naive_volume_ml uses the identity affine
    that the V2/V3 predict() functions write to disk; true_volume_ml uses
    the original affine the researcher should have preserved."""
    runtime = time.time() - t0
    if torch.cuda.is_available():
        peak = torch.cuda.max_memory_allocated() / (1024 ** 3)
    else:
        peak = 0.0
    voxels = int(mask.sum())
    naive_ml = voxels / 1000.0
    true_voxel_mm3 = abs(float(np.linalg.det(img.affine[:3, :3])))
    true_ml = voxels * true_voxel_mm3 / 1000.0
    return {
        "status": "OK",
        "runtime_s": runtime,
        "peak_vram_gb": peak,
        "voxel_count": voxels,
        "naive_volume_ml": naive_ml,
        "true_volume_ml": true_ml,
    }


def predict_v2(nifti_path: str, modality: str = "T1") -> dict:
    """V2: gr.File + nibabel + single-shot full-volume forward pass.

    Demonstrates failure mode 3 (CUDA OOM on full volumes) and, when the
    forward pass does fit, failure mode 2 (output mask written with
    identity affine, losing voxel spacing)."""
    t0 = time.time()
    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats()
        torch.cuda.empty_cache()

    img = nib.load(nifti_path)
    vol = np.asarray(img.get_fdata(dtype=np.float32))
    norm = _normalize(vol)
    x = torch.from_numpy(norm)[None, None].contiguous()

    model = _get_model(modality)
    device = next(model.parameters()).device
    x = x.to(device)

    with torch.no_grad():
        if device.type == "cuda":
            from torch.amp import autocast
            with autocast(device_type="cuda"):
                y_tuple = model(x)
        else:
            y_tuple = model(x)

    y1 = y_tuple[0] if isinstance(y_tuple, (list, tuple)) else y_tuple
    pred = torch.sigmoid(y1)[0, 0]
    mask = (pred > _threshold_for(modality)).to(torch.uint8).cpu().numpy()

    out_path = os.path.join(
        os.path.dirname(nifti_path) or ".", "v2_pred.nii.gz"
    )
    nib.Nifti1Image(mask, affine=np.eye(4)).to_filename(out_path)

    res = _summary(img, mask, t0)
    res["out_path"] = out_path
    return res


def predict_v3(nifti_path: str, modality: str = "T1") -> dict:
    """V3: gr.File + nibabel + manually-wired MONAI sliding window.

    Inference succeeds even on full volumes, but the output is still
    written with affine = identity, demonstrating failure mode 2."""
    from monai.inferers import SlidingWindowInferer

    t0 = time.time()
    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats()
        torch.cuda.empty_cache()

    img = nib.load(nifti_path)
    vol = np.asarray(img.get_fdata(dtype=np.float32))
    norm = _normalize(vol)
    x = torch.from_numpy(norm)[None, None].contiguous()

    model = _get_model(modality)
    device = next(model.parameters()).device
    x = x.to(device)

    inferer = SlidingWindowInferer(
        roi_size=[224, 224, 64],
        sw_batch_size=1,
        overlap=0.1,
    )

    def predictor(inp: torch.Tensor) -> torch.Tensor:
        out = model(inp)
        return out[0] if isinstance(out, (list, tuple)) else out

    with torch.no_grad():
        if device.type == "cuda":
            from torch.amp import autocast
            with autocast(device_type="cuda"):
                y = inferer(x, predictor)
        else:
            y = inferer(x, predictor)

    pred = torch.sigmoid(y)[0, 0]
    mask = (pred > _threshold_for(modality)).to(torch.uint8).cpu().numpy()

    out_path = os.path.join(
        os.path.dirname(nifti_path) or ".", "v3_pred.nii.gz"
    )
    nib.Nifti1Image(mask, affine=np.eye(4)).to_filename(out_path)

    res = _summary(img, mask, t0)
    res["out_path"] = out_path
    return res


def _format_log(res: dict) -> str:
    return (
        f"voxels:                              {res['voxel_count']:,}\n"
        f"naive volume (ml, identity affine):  {res['naive_volume_ml']:.2f}\n"
        f"true  volume (ml, original affine):  {res['true_volume_ml']:.2f}\n"
        f"peak VRAM (GB):                      {res['peak_vram_gb']:.2f}\n"
        f"runtime (s):                         {res['runtime_s']:.2f}"
    )


def _wrap_v2(file_obj):
    if file_obj is None:
        return None, "Please upload a NIfTI file."
    res = predict_v2(file_obj.name)
    return res["out_path"], _format_log(res)


def _wrap_v3(file_obj):
    if file_obj is None:
        return None, "Please upload a NIfTI file."
    res = predict_v3(file_obj.name)
    return res["out_path"], _format_log(res)


def build_v1_demo() -> gr.Interface:
    """V1: gr.Image. The component does not accept .nii / .nii.gz; the
    upload is rejected before predict() is ever invoked. This is the
    exact failure observed when a researcher reaches for the default
    image component for a volumetric workflow."""

    def _identity(image):
        return image

    return gr.Interface(
        fn=_identity,
        inputs=gr.Image(label="Upload an image (will not accept .nii.gz)"),
        outputs=gr.Image(label="Output"),
        title="V1: Idiomatic Gradio (gr.Image)",
        description=(
            "gr.Image only accepts standard 2D image formats (PNG, JPG, "
            "TIFF). Uploading a NIfTI volume (.nii / .nii.gz) is rejected "
            "at the upload step. This is failure mode 1 in the paper: "
            "native NIfTI ingestion is not supported by Gradio's image "
            "components."
        ),
    )


def build_v2_demo() -> gr.Interface:
    return gr.Interface(
        fn=_wrap_v2,
        inputs=gr.File(label="Upload .nii / .nii.gz"),
        outputs=[
            gr.File(label="Predicted mask (.nii.gz)"),
            gr.Textbox(label="Log", lines=6),
        ],
        title="V2: gr.File + naive single-shot full-volume inference",
        description=(
            "Loads the NIfTI manually with nibabel and runs a single "
            "forward pass on the full 3D volume. Expect CUDA OOM on "
            "standard scans on most GPUs (failure mode 3). When inference "
            "does succeed, the mask is written with affine=identity, so "
            "the reported liver volume in ml has identity-affine bias "
            "(failure mode 2)."
        ),
    )


def build_v3_demo() -> gr.Interface:
    return gr.Interface(
        fn=_wrap_v3,
        inputs=gr.File(label="Upload .nii / .nii.gz"),
        outputs=[
            gr.File(label="Predicted mask (.nii.gz)"),
            gr.Textbox(label="Log", lines=6),
        ],
        title="V3: gr.File + manually-wired MONAI sliding window",
        description=(
            "Adds MONAI SlidingWindowInferer manually to avoid OOM. "
            "Inference now succeeds on full volumes, but the mask is "
            "still written with affine=identity because Gradio's I/O "
            "contract does not preserve spatial metadata (failure mode 2)."
        ),
    )


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Head-to-head Gradio baseline (paper Section V)."
    )
    parser.add_argument(
        "--variant", choices=["v1", "v2", "v3"], required=True,
        help="Which baseline variant to launch.",
    )
    parser.add_argument("--port", type=int, default=7861)
    parser.add_argument("--share", action="store_true")
    args = parser.parse_args()

    if args.variant == "v1":
        demo = build_v1_demo()
    elif args.variant == "v2":
        demo = build_v2_demo()
    else:
        demo = build_v3_demo()

    server_name = os.environ.get(
        "GRADIO_SERVER_NAME",
        "127.0.0.1" if os.name == "nt" else "0.0.0.0",
    )
    demo.launch(
        server_name=server_name,
        server_port=args.port,
        share=args.share,
    )


if __name__ == "__main__":
    main()