"""WaveSeg (A1) Gradio demo - Hugging Face Space app.

WaveSeg = SegFormer-B0 + Frequency-Boundary Adapter (FBA) gating only, with
NO boundary-frequency loss - the boundary loss was tested at
lambda_boundary in {0.1, 0.25, 0.5, 1.0} and did not beat gating alone after
seed-averaging on either dataset, so it is excluded here (see
MODEL_CARD.md's "Honest negative ablation" section).

Runs on a Hugging Face ZeroGPU Space: `spaces.GPU` is imported optionally
so this also runs locally on CPU-only machines (e.g. for development), and
no CUDA call happens at import time - only inside the GPU-decorated
`predict()`, which resolves its own device each call.
"""
from __future__ import annotations

import os
from typing import Optional, Tuple

import numpy as np
import torch
import gradio as gr
from PIL import Image

try:
    import spaces
except ImportError:  # running outside a ZeroGPU Space (e.g. local dev/testing)

    class _SpacesShim:
        """No-op replacement for the `spaces` package's GPU decorator."""

        @staticmethod
        def GPU(*args, **kwargs):
            def decorator(fn):
                return fn

            # Support both bare `@spaces.GPU` and parametrized `@spaces.GPU(duration=60)`.
            if len(args) == 1 and callable(args[0]) and not kwargs:
                return args[0]
            return decorator

    spaces = _SpacesShim()

from waveseg_model import build_waveseg_a1

HF_REPO_ID = "Sarvarbek13/WaveSeg"
IMG_SIZE = 256
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)

DOMAIN_TO_KEY = {"Brain MRI (LGG)": "lgg", "Polyp (Kvasir-SEG)": "kvasir"}
OVERLAY_COLOR = (52, 152, 219)  # blue

_MODELS: dict = {}


def _load_state_dict_from_path(path: str) -> dict:
    """Load a state dict from either a .safetensors file or a training-repo
    .pth checkpoint (which wraps the state dict under a "model_state" key).
    """
    if path.endswith(".safetensors"):
        from safetensors.torch import load_file

        return load_file(path)
    ckpt = torch.load(path, map_location="cpu", weights_only=False)
    return ckpt["model_state"] if isinstance(ckpt, dict) and "model_state" in ckpt else ckpt


def _resolve_checkpoint_path(domain_key: str) -> str:
    """Resolve the A1 checkpoint path for a domain ("lgg" or "kvasir").

    Checks WAVESEG_LOCAL_CHECKPOINTS_ROOT first (local dev/testing - expects
    ``<root>/a1_<domain>/best.pth``, matching the training repo's own
    checkpoints/ layout, so no network call happens); otherwise downloads
    from the HF Hub model repo (the production path once deployed).
    """
    local_root = os.environ.get("WAVESEG_LOCAL_CHECKPOINTS_ROOT")
    if local_root:
        local_path = os.path.join(local_root, f"a1_{domain_key}", "best.pth")
        if os.path.exists(local_path):
            return local_path

    from huggingface_hub import hf_hub_download

    return hf_hub_download(repo_id=HF_REPO_ID, filename=f"a1_{domain_key}.safetensors")


def _get_model(domain_key: str) -> torch.nn.Module:
    """Load (once) and cache the A1 model for a domain, on CPU."""
    if domain_key not in _MODELS:
        model = build_waveseg_a1()
        state_dict = _load_state_dict_from_path(_resolve_checkpoint_path(domain_key))
        model.load_state_dict(state_dict)
        model.eval()
        _MODELS[domain_key] = model
    return _MODELS[domain_key]


def _preprocess(image: Image.Image) -> torch.Tensor:
    """Resize + ImageNet-normalize an image into a (1, 3, 256, 256) tensor."""
    resized = image.convert("RGB").resize((IMG_SIZE, IMG_SIZE), Image.BILINEAR)
    arr = np.asarray(resized, dtype=np.float32) / 255.0
    arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
    return torch.from_numpy(arr.transpose(2, 0, 1)).unsqueeze(0).float()


def _overlay_mask(image: Image.Image, mask: np.ndarray, color=OVERLAY_COLOR, alpha: float = 0.45) -> Image.Image:
    """Alpha-blend ``color`` onto a resized copy of ``image`` wherever ``mask`` is truthy."""
    base = np.asarray(image.convert("RGB").resize((IMG_SIZE, IMG_SIZE), Image.BILINEAR), dtype=np.float32)
    out = base.copy()
    m = mask.astype(bool)
    out[m] = out[m] * (1 - alpha) + np.array(color, dtype=np.float32) * alpha
    return Image.fromarray(out.astype(np.uint8))


def _dice(pred: np.ndarray, gt: np.ndarray, smooth: float = 1e-6) -> float:
    """Dice coefficient between two binary masks."""
    pred_b, gt_b = pred.astype(bool), gt.astype(bool)
    intersection = np.logical_and(pred_b, gt_b).sum()
    total = pred_b.sum() + gt_b.sum()
    return float((2.0 * intersection + smooth) / (total + smooth))


@spaces.GPU(duration=60)
def predict(
    image: Optional[Image.Image], domain: str, gt_mask: Optional[Image.Image]
) -> Tuple[Optional[Image.Image], str]:
    """Run WaveSeg (A1) inference on ``image`` for the chosen ``domain``.

    Returns the mask-overlay image and a status string (includes Dice if a
    ground-truth mask was supplied). Device is resolved inside this
    function (ZeroGPU pattern) - never at module import time.
    """
    if image is None:
        raise gr.Error("Please upload an image first.")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    domain_key = DOMAIN_TO_KEY[domain]
    model = _get_model(domain_key).to(device)

    x = _preprocess(image).to(device)
    with torch.no_grad():
        output = model(x)
        logits = output[0] if isinstance(output, tuple) else output
        prob = torch.sigmoid(logits)[0, 0].detach().cpu().numpy()
    model.to("cpu")  # ZeroGPU: release GPU memory once inference is done

    pred_mask = prob > 0.5
    overlay = _overlay_mask(image, pred_mask)
    status = f"Predicted foreground: {pred_mask.mean() * 100:.1f}% of image."

    if gt_mask is not None:
        gt_arr = np.asarray(gt_mask.convert("L").resize((IMG_SIZE, IMG_SIZE), Image.NEAREST))
        dice = _dice(pred_mask, gt_arr > 127)
        status += f" | Dice vs uploaded GT mask: {dice:.4f}"

    return overlay, status


ABOUT_MD = """
**WaveSeg** adds a tiny (<0.2M-parameter) **Frequency-Boundary Adapter (FBA)** between a
SegFormer-B0 decoder and its segmentation head. FBA applies a **Haar Discrete Wavelet
Transform** to the decoder features, exposing their high-frequency sub-bands (LH/HL/HH) -
where object boundaries live - gates them through a 1x1 conv into a **boundary attention
map**, and uses that map to re-weight the features before the final prediction:
`out = feat * (1 + alpha * attn)`, with `alpha` a learnable scalar.

A boundary-frequency auxiliary loss (explicitly supervising the attention map toward
ground-truth edges) was also tested, at several weights - it did **not** improve results
after seed-averaging on either dataset, so the shipped model (A1) uses FBA gating alone,
trained end-to-end via the segmentation loss only. This negative result is reported
honestly in [MODEL_CARD.md](https://huggingface.co/Sarvarbek13/WaveSeg) rather than
dropped.
"""

with gr.Blocks(title="WaveSeg - Frequency-Boundary Segmentation") as demo:
    gr.Markdown("# WaveSeg")
    gr.Markdown("*Segmentation that listens to the high frequencies.*")
    gr.Markdown(
        "Upload a brain MRI (FLAIR) slice or a colonoscopy image, pick the matching domain, "
        "and WaveSeg (SegFormer-B0 + Frequency-Boundary Adapter gating) will segment it. "
        "Optionally upload a ground-truth mask to see the Dice score."
    )

    with gr.Row():
        with gr.Column():
            image_in = gr.Image(type="pil", label="Input image")
            domain_in = gr.Radio(list(DOMAIN_TO_KEY.keys()), value="Brain MRI (LGG)", label="Domain")
            gt_in = gr.Image(type="pil", label="Ground-truth mask (optional)")
            run_btn = gr.Button("Run WaveSeg", variant="primary")
        with gr.Column():
            image_out = gr.Image(type="pil", label="Predicted mask overlay")
            status_out = gr.Textbox(label="Result", interactive=False)

    run_btn.click(predict, inputs=[image_in, domain_in, gt_in], outputs=[image_out, status_out])

    gr.Examples(
        examples=[
            [
                os.path.join(os.path.dirname(__file__), "examples", "lgg_1.png"),
                "Brain MRI (LGG)",
                os.path.join(os.path.dirname(__file__), "examples", "lgg_1_mask.png"),
            ],
            [
                os.path.join(os.path.dirname(__file__), "examples", "lgg_2.png"),
                "Brain MRI (LGG)",
                os.path.join(os.path.dirname(__file__), "examples", "lgg_2_mask.png"),
            ],
            [
                os.path.join(os.path.dirname(__file__), "examples", "kvasir_1.jpg"),
                "Polyp (Kvasir-SEG)",
                os.path.join(os.path.dirname(__file__), "examples", "kvasir_1_mask.png"),
            ],
            [
                os.path.join(os.path.dirname(__file__), "examples", "kvasir_2.jpg"),
                "Polyp (Kvasir-SEG)",
                os.path.join(os.path.dirname(__file__), "examples", "kvasir_2_mask.png"),
            ],
        ],
        inputs=[image_in, domain_in, gt_in],
        label="Examples (real test-set images)",
    )

    with gr.Accordion("About FBA (Frequency-Boundary Adapter)", open=False):
        gr.Markdown(ABOUT_MD)

# Eagerly load both domain models at startup, on CPU (ZeroGPU rule: never
# call .cuda()/.to("cuda") at import time - only inside the @spaces.GPU
# function, which resolves its own device per call).
for _domain_key in DOMAIN_TO_KEY.values():
    _get_model(_domain_key)

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