app/visualization.py

from __future__ import annotations

from typing import Dict, Tuple

import numpy as np
from PIL import Image, ImageDraw

from .depth_pipeline import visualize_depth

GRAD_ALPHA = 0.35
FLAT_ALPHA = 0.25


def _outline_mask(mask: np.ndarray | None) -> np.ndarray:
    """Compute a 1px outline around a boolean mask without external deps."""
    if mask is None:
        return np.zeros((1, 1), dtype=bool)
    mask_bool = mask.astype(bool)
    h, w = mask_bool.shape
    padded = np.pad(mask_bool, 1, mode="edge")
    dilated = np.zeros_like(mask_bool, dtype=bool)
    for dy in (-1, 0, 1):
        for dx in (-1, 0, 1):
            if dx == 0 and dy == 0:
                continue
            dilated |= padded[1 + dy : 1 + dy + h, 1 + dx : 1 + dx + w]
    border = dilated & (~mask_bool)
    return border

def make_safety_heatmap(
    rgb: Image.Image,
    safe_mask: np.ndarray,
    hazard_mask: np.ndarray,
    risk_map: np.ndarray,
    risk_threshold: float = 0.35,
):
    safe = np.clip(safe_mask.astype(np.float32), 0.0, 1.0)
    hazard = hazard_mask.astype(bool)
    risk = np.clip(risk_map.astype(np.float32), 0.0, 1.0)

    h, w = safe.shape
    safe_overlay = np.zeros((h, w, 4), dtype=np.uint8)
    safe_mask_bool = safe > 0.0
    # Semi-transparent fill plus thicker outline for visibility (#00BF00) at ~80% alpha
    safe_overlay[safe_mask_bool, :] = (0, 191, 0, 204)
    safe_border = _outline_mask(safe_mask_bool)
    safe_border |= _outline_mask(safe_border)  # thicken outline to ~2px
    safe_overlay[safe_border, :] = (0, 191, 0, 255)

    risk_focus = np.zeros_like(risk)
    risk_focus[risk > risk_threshold] = risk[risk > risk_threshold]
    hazard_intensity = np.where(hazard, np.maximum(risk_focus, 1.0), risk_focus)
    hazard_alpha = (np.clip(hazard_intensity, 0.0, 1.0) * 255).astype(np.uint8)
    hazard_overlay = np.zeros((h, w, 4), dtype=np.uint8)
    hazard_overlay[..., 0] = 255  # pure red for depth-based hazards
    hazard_overlay[..., 3] = hazard_alpha
    hazard_mask_bool = hazard
    hazard_border = _outline_mask(hazard_mask_bool)
    hazard_overlay[hazard_border, :] = (255, 0, 0, 255)

    safe_img = Image.fromarray(safe_overlay, mode="RGBA").resize(rgb.size, resample=Image.NEAREST)
    hazard_img = Image.fromarray(hazard_overlay, mode="RGBA").resize(rgb.size, resample=Image.NEAREST)
    score_gray = Image.fromarray((safe * 255).astype(np.uint8)).resize(rgb.size, resample=Image.NEAREST)
    return safe_img, hazard_img, score_gray


def make_flatness_heatmap(std_map_vis: np.ndarray, target_size: tuple[int, int]) -> Image.Image:
    # Normalize and map to a simple turbo-like palette
    std_norm = std_map_vis
    if std_norm.max() > std_norm.min():
        std_norm = (std_norm - std_norm.min()) / (np.ptp(std_norm) + 1e-6)
    cmap = np.array(
        [
            [48, 18, 59],
            [65, 68, 135],
            [42, 120, 142],
            [34, 168, 132],
            [122, 209, 81],
            [253, 231, 36],
        ],
        dtype=np.float32,
    )
    idx = np.clip((std_norm * (len(cmap) - 1)).astype(np.int32), 0, len(cmap) - 1)
    heat_rgb = cmap[idx]
    heat_overlay = np.zeros((std_norm.shape[0], std_norm.shape[1], 4), dtype=np.uint8)
    heat_overlay[..., :3] = heat_rgb.astype(np.uint8)
    heat_overlay[..., 3] = (np.clip(std_norm, 0.0, 1.0) * 160).astype(np.uint8)
    return Image.fromarray(heat_overlay, mode="RGBA").resize(target_size, resample=Image.BILINEAR)


def build_result_layers(
    image: Image.Image,
    depth_raw: np.ndarray,
    std_map_vis: np.ndarray,
    grad_norm: np.ndarray,
    grad_thresh: float,
    safe_mask: np.ndarray,
    risk_map: np.ndarray,
    footprint_img_px: int,
    center_img: Tuple[int, int],
    water_mask: np.ndarray | None,
    road_mask: np.ndarray | None,
    roof_mask: np.ndarray | None,
    tree_mask: np.ndarray | None,
    hazard_mask: np.ndarray,
) -> Dict[str, Image.Image]:
    depth_vis = Image.fromarray(visualize_depth(depth_raw, cmap="Spectral")).resize(
        image.size, resample=Image.BILINEAR
    )
    flatness_img = Image.fromarray((std_map_vis / (std_map_vis.max() + 1e-6) * 255).astype(np.uint8)).resize(
        image.size, resample=Image.NEAREST
    )
    grad_img = Image.fromarray((grad_norm * 255).astype(np.uint8)).resize(image.size, resample=Image.BILINEAR)
    grad_mask_img = Image.fromarray(((grad_norm < grad_thresh).astype(np.uint8) * 255)).resize(
        image.size, resample=Image.NEAREST
    )

    def _mask_to_image(mask: np.ndarray | None) -> Image.Image:
        if mask is None:
            return Image.new("L", image.size, 0)
        return Image.fromarray((mask.astype(np.uint8) * 255)).resize(image.size, resample=Image.NEAREST)

    water_mask_img = _mask_to_image(water_mask)
    road_mask_img = _mask_to_image(road_mask)
    roof_mask_img = _mask_to_image(roof_mask)
    tree_mask_img = _mask_to_image(tree_mask)

    def _hatched_overlay(
        mask: np.ndarray | None,
        color: tuple[int, int, int],
        alpha: int = 180,
        hatch_step: int = 8,
        hatch_thickness: int = 3,
    ) -> Image.Image:
        if mask is None:
            return Image.new("RGBA", image.size, (0, 0, 0, 0))
        m = np.array(
            Image.fromarray((mask.astype(np.uint8) * 255)).resize(image.size, resample=Image.NEAREST)
        ).astype(bool)
        overlay = np.zeros((image.height, image.width, 4), dtype=np.uint8)
        overlay[..., :3] = np.array(color, dtype=np.uint8)
        overlay[..., 3] = np.where(m, alpha, 0).astype(np.uint8)
        if hatch_step > 0 and hatch_thickness > 0:
            xs = np.arange(image.width, dtype=np.int32)
            ys = np.arange(image.height, dtype=np.int32)
            grid_x, grid_y = np.meshgrid(xs, ys)
            stripe = ((grid_x + grid_y) % hatch_step) < hatch_thickness
            stripe_alpha = (alpha // 3)
            overlay[..., 3] = np.where(m & stripe, stripe_alpha, overlay[..., 3])
        return Image.fromarray(overlay, mode="RGBA")

    # Segmentation hazards: black with hatch to separate from depth-risk red.
    hazard_color = (0, 0, 0)
    overlays = []
    for mask, alpha in (
        (water_mask, 130),
        (road_mask, 130),
        (tree_mask, 130),
        (roof_mask, 150),
    ):
        ov = _hatched_overlay(mask, hazard_color, alpha=alpha, hatch_step=24, hatch_thickness=2)
        overlays.append(ov)
    water_hazard_overlay, road_hazard_overlay, tree_hazard_overlay, roof_hazard_overlay = overlays
    for overlay in overlays:
        mask = np.array(overlay.getchannel("A")) > 0
        if not mask.any():
            continue
        border = _outline_mask(mask)
        if not border.any():
            continue
        arr = np.array(overlay)
        arr[border, :] = (0, 0, 0, 255)
        arr[np.logical_and(border, arr[..., 3] > 0), 3] = 255
        overlay.paste(Image.fromarray(arr, mode="RGBA"))

    safe_overlay, hazard_overlay, heat_gray = make_safety_heatmap(image, safe_mask, hazard_mask, risk_map)
    flat_heat_overlay = make_flatness_heatmap(std_map_vis, image.size)

    spot_overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(spot_overlay)
    cx_img, cy_img = center_img
    side_img = max(3, footprint_img_px | 1)
    half_img = side_img // 2
    bx0 = cx_img - half_img
    by0 = cy_img - half_img
    bx1 = bx0 + side_img - 1
    by1 = by0 + side_img - 1
    clipped_x = False
    clipped_y = False
    if bx0 < 0:
        shift = -bx0
        bx0 = 0
        bx1 += shift
        clipped_x = True
    if bx1 >= image.width:
        shift = bx1 - (image.width - 1)
        bx1 = image.width - 1
        bx0 = max(0, bx0 - shift)
        clipped_x = True
    if by0 < 0:
        shift = -by0
        by0 = 0
        by1 += shift
        clipped_y = True
    if by1 >= image.height:
        shift = by1 - (image.height - 1)
        by1 = image.height - 1
        by0 = max(0, by0 - shift)
        clipped_y = True
    if clipped_x:
        cx_draw = int(round((bx0 + bx1) / 2.0))
    else:
        cx_draw = int(round(min(max(cx_img, bx0), bx1)))
    if clipped_y:
        cy_draw = int(round((by0 + by1) / 2.0))
    else:
        cy_draw = int(round(min(max(cy_img, by0), by1)))
    overlay_box = Image.new("RGBA", image.size, (0, 0, 0, 0))
    box_draw = ImageDraw.Draw(overlay_box)
    fill = (255, 140, 0, 90)  # translucent orange
    outline = (255, 140, 0, 255)

    # Crosshair sized 3x the landing box for clearer focus, thicker lines.
    cross_half = int(round(side_img * 1.5))
    hx0 = max(0, cx_draw - cross_half)
    hx1 = min(image.width - 1, cx_draw + cross_half)
    hy0 = max(0, cy_draw - cross_half)
    hy1 = min(image.height - 1, cy_draw + cross_half)
    cross_width = 6
    draw.line((hx0, cy_draw, hx1, cy_draw), fill=outline, width=cross_width)
    draw.line((cx_draw, hy0, cx_draw, hy1), fill=outline, width=cross_width)

    box_draw.rectangle((bx0, by0, bx1, by1), fill=fill, outline=outline, width=6)
    box_draw.line((cx_draw, by0, cx_draw, by1), fill=outline, width=3)
    box_draw.line((bx0, cy_draw, bx1, cy_draw), fill=outline, width=3)
    radius = 10
    box_draw.ellipse((cx_draw - radius, cy_draw - radius, cx_draw + radius, cy_draw + radius), fill=outline)

    return {
        "RGB": image,
        "Depth": depth_vis,
        "Flatness map (std)": flatness_img,
        "Depth gradient": grad_img,
        "Gradient mask": grad_mask_img,
        "Water mask": water_mask_img,
        "Road mask": road_mask_img,
        "Roof mask": roof_mask_img,
        "Tree mask": tree_mask_img,
        "Safety heatmap overlay": safe_overlay,
        "Hazard overlay": hazard_overlay,
        "Water hazard overlay": water_hazard_overlay,
        "Road hazard overlay": road_hazard_overlay,
        "Tree hazard overlay": tree_hazard_overlay,
        "Roof hazard overlay": roof_hazard_overlay,
        "Flatness heatmap overlay": flat_heat_overlay,
        "Safety score": heat_gray,
        "Landing spot overlay": Image.alpha_composite(spot_overlay, overlay_box),
    }


def compose_view(
    images_dict: dict,
    base_view: str,
    heat_on: bool,
    heat_alpha: float,
    risk_on: bool,
    risk_alpha: float,
    hazards_on: bool,
    grad_on: bool,
    flat_on: bool,
    flat_heat_on: bool,
    spot_on: bool,
) -> Image.Image:
    import gradio as gr

    if not images_dict:
        raise gr.Error("Run inference first, then select a view.")
    if base_view not in images_dict:
        raise gr.Error(f"Unknown view: {base_view}")

    base = images_dict.get(base_view)
    if base is None:
        raise gr.Error(f"No image for view: {base_view}")
    out = base.convert("RGBA")

    if heat_on and "Safety heatmap overlay" in images_dict:
        safe_overlay = images_dict["Safety heatmap overlay"]
        if safe_overlay is not None:
            safe_rgba = safe_overlay.convert("RGBA")
            alpha_factor = max(0.0, min(1.0, heat_alpha))
            alpha_channel = np.array(safe_rgba.getchannel("A"), dtype=np.uint8)
            alpha_channel = (alpha_channel.astype(np.float32) * alpha_factor).astype(np.uint8)
            safe_rgba.putalpha(Image.fromarray(alpha_channel, mode="L"))
            out = Image.alpha_composite(out, safe_rgba)

    if risk_on and "Hazard overlay" in images_dict:
        hazard = images_dict.get("Hazard overlay")
        if hazard is not None:
            hazard_rgba = hazard.convert("RGBA")
            alpha_factor = max(0.0, min(1.0, risk_alpha))
            alpha_channel = np.array(hazard_rgba.getchannel("A"), dtype=np.uint8)
            alpha_channel = (alpha_channel.astype(np.float32) * alpha_factor).astype(np.uint8)
            hazard_rgba.putalpha(Image.fromarray(alpha_channel, mode="L"))
            out = Image.alpha_composite(out, hazard_rgba)

    if hazards_on:
        if "Water hazard overlay" in images_dict:
            water = images_dict.get("Water hazard overlay")
            if water is not None:
                out = Image.alpha_composite(out, water.convert("RGBA"))
        if "Road hazard overlay" in images_dict:
            road = images_dict.get("Road hazard overlay")
            if road is not None:
                out = Image.alpha_composite(out, road.convert("RGBA"))
        if "Tree hazard overlay" in images_dict:
            tree = images_dict.get("Tree hazard overlay")
            if tree is not None:
                out = Image.alpha_composite(out, tree.convert("RGBA"))
        if "Roof hazard overlay" in images_dict:
            roof = images_dict.get("Roof hazard overlay")
            if roof is not None:
                out = Image.alpha_composite(out, roof.convert("RGBA"))

    if grad_on and "Depth gradient" in images_dict:
        grad_img = images_dict["Depth gradient"]
        if grad_img is not None:
            grad_rgba = grad_img.convert("RGBA")
            grad_rgba.putalpha(int(GRAD_ALPHA * 255))
            out = Image.alpha_composite(out, grad_rgba)

    if flat_on and "Flatness map (std)" in images_dict:
        flat_img = images_dict["Flatness map (std)"]
        if flat_img is not None:
            flat_rgba = flat_img.convert("RGBA")
            flat_rgba.putalpha(int(FLAT_ALPHA * 255))
            out = Image.alpha_composite(out, flat_rgba)

    if flat_heat_on and "Flatness heatmap overlay" in images_dict:
        flat_heat = images_dict["Flatness heatmap overlay"]
        if flat_heat is not None:
            flat_heat_rgba = flat_heat.convert("RGBA")
            out = Image.alpha_composite(out, flat_heat_rgba)

    if spot_on and "Landing spot overlay" in images_dict:
        spot = images_dict["Landing spot overlay"]
        if spot is not None:
            out = Image.alpha_composite(out, spot.convert("RGBA"))

    return out.convert("RGB")


__all__ = ["build_result_layers", "compose_view", "make_safety_heatmap"]