"""Simple Gradio interface for testing the trained Gabor-UNet hyphae segmentation model.

Reproduces the exact preprocessing/model architecture from notebook-4
(vastai_gabor_unet) so the pretrained checkpoint loads and runs unmodified.
"""

import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
import segmentation_models_pytorch as smp

# On HF Spaces free tier the app runs on ZeroGPU, which requires the entry point
# to be declared with @spaces.GPU — without it the Space exits with
# "No @spaces.GPU function detected during startup". Locally `spaces` isn't
# installed, so fall back to plain CPU execution and a no-op decorator.
try:
    import spaces

    ON_ZERO_GPU = True
except ImportError:
    spaces = None
    ON_ZERO_GPU = False


def gpu_task(duration=120):
    """Mark the inference entry point for ZeroGPU; no-op when running locally."""
    def wrap(fn):
        return spaces.GPU(duration=duration)(fn) if ON_ZERO_GPU else fn
    return wrap


CKPT_PATH = "gabor_unet_best.pth"
TILE_SIZE = 512
GABOR_ORIENTATIONS = 6
GABOR_SCALES = 3
GABOR_KSIZE = 31
GABOR_STORE_SIZE = 64
DEFAULT_THRESHOLD = 0.75  # calibrated value, see threshold_calibration_results.csv

# ZeroGPU reports no CUDA device in the main process at import time, so this can't
# be probed with torch.cuda.is_available() — the `spaces` runtime attaches the GPU
# when the decorated function is called.
DEVICE = torch.device("cuda") if ON_ZERO_GPU else torch.device("cpu")


# ── Gabor filter bank (matches notebook-4 exactly) ──────────────────────────
def build_gabor_bank(orientations=6, scales=3, ksize=31):
    kernels = []
    for wl in [4, 8, 16][:scales]:
        for i in range(orientations):
            theta = (i / orientations) * np.pi
            kernel = cv2.getGaborKernel(
                ksize=(ksize, ksize), sigma=wl * 0.56,
                theta=theta, lambd=float(wl),
                gamma=0.5, psi=0, ktype=cv2.CV_32F,
            )
            kernel /= (kernel.sum() + 1e-8)
            kernels.append(kernel)
    return kernels


def compute_gabor_small(img_bgr, kernels, store_size=GABOR_STORE_SIZE):
    small = cv2.resize(img_bgr, (128, 128), interpolation=cv2.INTER_AREA)
    gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY).astype(np.float32) / 255.0
    responses = np.stack([np.abs(cv2.filter2D(gray, cv2.CV_32F, k)) for k in kernels], axis=0)
    max_r = responses.max(axis=0)
    ori = responses.argmax(axis=0).astype(np.float32) / len(kernels)
    feat = np.stack([max_r, ori], axis=-1)
    if store_size != 128:
        feat = cv2.resize(feat, (store_size, store_size), interpolation=cv2.INTER_LINEAR)
    return feat.astype(np.float32)


GABOR_KERNELS = build_gabor_bank(GABOR_ORIENTATIONS, GABOR_SCALES, GABOR_KSIZE)


# ── Model (matches notebook-4 exactly) ──────────────────────────────────────
class GaborAttentionGate(nn.Module):
    def __init__(self, feat_channels):
        super().__init__()
        self.gate = nn.Sequential(
            nn.Conv2d(feat_channels + 1, feat_channels, kernel_size=1, bias=False),
            nn.BatchNorm2d(feat_channels),
            nn.Sigmoid(),
        )

    def forward(self, feat, gabor_max):
        g = F.interpolate(gabor_max, size=feat.shape[2:], mode="bilinear", align_corners=False)
        return feat * self.gate(torch.cat([feat, g], dim=1))


class GaborUNet(nn.Module):
    def __init__(self, in_channels=5, pretrained=False):
        super().__init__()
        self.unet = smp.Unet(
            encoder_name="efficientnet-b4",
            encoder_weights="imagenet" if pretrained else None,
            in_channels=in_channels,
            classes=1,
            activation=None,
        )
        self.attn_gate = GaborAttentionGate(feat_channels=16)

    def forward(self, x):
        gabor_max = x[:, 3:4, :, :]
        features = self.unet.encoder(x)
        decoder_output = self.unet.decoder(features)
        decoder_output = self.attn_gate(decoder_output, gabor_max)
        return self.unet.segmentation_head(decoder_output)


print("Loading model...")
model = GaborUNet(in_channels=5, pretrained=False)
ckpt = torch.load(CKPT_PATH, map_location="cpu", weights_only=False)
model.load_state_dict(ckpt["model_state"])
model.eval().to(DEVICE)
print(f"Loaded checkpoint — epoch {ckpt.get('epoch')}  Val IoU: {ckpt.get('val_iou'):.4f}")


def load_image_any_format(file_path):
    """Load virtually any image file (jpg, png, tiff, webp, jp2/JPEG2000, bmp, ...) as RGB uint8."""
    try:
        from PIL import Image
        im = Image.open(file_path)
        im = im.convert("RGB")
        return np.array(im)
    except Exception:
        pass

    img = cv2.imread(file_path, cv2.IMREAD_COLOR)
    if img is not None:
        return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    raise gr.Error(f"Could not read this file as an image: {file_path}")


def pad_to_multiple(img, tile_size):
    """Reflect-pad so H and W are exact multiples of tile_size (matches notebook-1 tiling)."""
    h, w = img.shape[:2]
    pad_h = (tile_size - h % tile_size) % tile_size
    pad_w = (tile_size - w % tile_size) % tile_size
    if img.ndim == 3:
        return np.pad(img, ((0, pad_h), (0, pad_w), (0, 0)), mode="reflect")
    return np.pad(img, ((0, pad_h), (0, pad_w)), mode="reflect")


def tile_to_tensor(tile_rgb):
    """Build the 5-channel (RGB + Gabor) model input for one native-scale 512x512 tile."""
    tile_bgr = cv2.cvtColor(tile_rgb, cv2.COLOR_RGB2BGR)
    img_t = torch.from_numpy(tile_rgb).permute(2, 0, 1).float() / 255.0

    feat = compute_gabor_small(tile_bgr, GABOR_KERNELS)
    feat = cv2.resize(feat, (TILE_SIZE, TILE_SIZE), interpolation=cv2.INTER_LINEAR)
    gabor_t = torch.from_numpy(feat.transpose(2, 0, 1))

    return torch.cat([img_t, gabor_t], dim=0).unsqueeze(0)


CANDIDATE_SCALES = [1.0, 1.5, 2.0, 2.5]
EARLY_STOP_CONFIDENCE = 0.9


def run_tiled_inference(img_rgb):
    """Tile at native resolution (matching the training pipeline) and return the
    full-resolution sigmoid probability map."""
    padded = pad_to_multiple(img_rgb, TILE_SIZE)
    H, W = padded.shape[:2]
    prob_map = np.zeros((H, W), dtype=np.float32)
    for y in range(0, H, TILE_SIZE):
        for x in range(0, W, TILE_SIZE):
            tile = padded[y:y + TILE_SIZE, x:x + TILE_SIZE]
            x_in = tile_to_tensor(tile).to(DEVICE)
            logits = model(x_in)
            prob_map[y:y + TILE_SIZE, x:x + TILE_SIZE] = (
                torch.sigmoid(logits).squeeze().float().cpu().numpy()
            )
    return prob_map


@gpu_task(duration=120)
@torch.no_grad()
def predict(file_path, threshold, auto_scale=True):
    if file_path is None:
        return None, None, ""

    img_raw = load_image_any_format(file_path)

    # A photo's microns-per-pixel almost never matches the calibrated slide scanner
    # used for training, and that scale can't be inferred from the file itself.
    # Measured fix: scale alone (no color/denoise changes) took a real off-distribution
    # photo from max confidence 0.039 (undetectable) to 0.9950 at 2x scale. So: try
    # native scale first, and only pay for larger, slower re-tries if confidence is
    # low — this keeps already-good native-resolution uploads (like your JP2s) just
    # as fast as before, since they hit the early-stop on the very first pass.
    scales_to_try = CANDIDATE_SCALES if auto_scale else [1.0]

    best_scale, best_conf, best_prob_map, best_img_scaled = scales_to_try[0], -1.0, None, img_raw

    for s in scales_to_try:
        if s != 1.0:
            h, w = img_raw.shape[:2]
            new_h = max(TILE_SIZE, int(round(h * s)))
            new_w = max(TILE_SIZE, int(round(w * s)))
            img_s = cv2.resize(img_raw, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
        else:
            img_s = img_raw

        h0, w0 = img_s.shape[:2]
        prob_map = run_tiled_inference(img_s)[:h0, :w0]
        conf = float(prob_map.max())

        if conf > best_conf:
            best_scale, best_conf, best_prob_map, best_img_scaled = s, conf, prob_map, img_s

        if conf >= EARLY_STOP_CONFIDENCE:
            break

    pred = (best_prob_map > threshold).astype(np.uint8)

    mask_rgb = np.stack([pred * 255] * 3, axis=-1).astype(np.uint8)

    overlay = best_img_scaled.copy()
    overlay[pred == 1] = (
        0.5 * overlay[pred == 1] + 0.5 * np.array([255, 0, 0])
    ).astype(np.uint8)

    info = f"Scale used: {best_scale}x  |  peak confidence: {best_conf:.3f}"
    return mask_rgb, overlay, info


def preview_image(file_path):
    if file_path is None:
        return None
    return load_image_any_format(file_path)


with gr.Blocks(title="Gabor U-Net — Hyphae Segmentation") as demo:
    gr.Markdown(
        "### Gabor U-Net — Hyphae Segmentation\n\n"
        "<span style='font-size:0.85em;'>Upload any image (any format/size). It's reflect-padded and "
        "split into native-resolution 512x512 tiles (matching the training pipeline — no downscaling), "
        "each tile is run through the Gabor-attention U-Net (EfficientNet-B4 encoder), and predictions "
        "are stitched back together. Large images may take a while.</span>"
    )
    with gr.Row():
        with gr.Column():
            file_input = gr.File(label="Input image (any format — jpg, png, tiff, webp, jp2, bmp, ...)", type="filepath", height=120)
            preview = gr.Image(label="Preview of uploaded image", interactive=False)
            threshold = gr.Slider(0.0, 1.0, value=DEFAULT_THRESHOLD, step=0.01, label="Threshold (calibrated default: 0.75)")
            auto_scale = gr.Checkbox(
                value=True,
                label="Auto-scale (recommended)",
                info="Tries native scale first; only retries at 1.5x/2x/2.5x if confidence is low — "
                     "fixes photos where hyphae render too thin/thick vs. training tiles, at no extra "
                     "cost to already-good uploads.",
            )
            submit_btn = gr.Button("Submit", variant="primary")
        with gr.Column():
            mask_out = gr.Image(label="Predicted mask")
            overlay_out = gr.Image(label="Overlay")
            scale_info = gr.Textbox(label="Detection info", interactive=False)

    file_input.change(fn=preview_image, inputs=file_input, outputs=preview)
    submit_btn.click(fn=predict, inputs=[file_input, threshold, auto_scale], outputs=[mask_out, overlay_out, scale_info])

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