#!/usr/bin/env python3
"""Predict review-only hidden/amodal masks for one accessibility image.

The tiny adapter consumes RGB, a visible-target proposal, an obstacle proposal,
and a coarse stairs/non-stairs category plane. Predictions are constrained to
the obstacle support and written as candidates, never as ground truth.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

import cv2
import numpy as np
import torch
from PIL import Image, ImageOps


TOOLS_DIR = Path(__file__).resolve().parent
if str(TOOLS_DIR) not in sys.path:
    sys.path.insert(0, str(TOOLS_DIR))

import train_accessibility_amodal_adapter as trainer  # noqa: E402


CATEGORIES = ("curb_cut", "ramp", "stairs", "tactile_paving", "walkway")


def load_rgb(path: Path) -> Image.Image:
    return ImageOps.exif_transpose(Image.open(path)).convert("RGB")


def load_mask(path: Path, size: tuple[int, int]) -> np.ndarray:
    image = ImageOps.exif_transpose(Image.open(path)).convert("L")
    if image.size != size:
        raise ValueError(
            f"Mask/RGB raster mismatch for {path}: mask={image.size}, rgb={size}. "
            "Refusing to resize because this can hide EXIF-orientation misalignment."
        )
    return np.asarray(image) > 127


def save_mask(path: Path, mask: np.ndarray) -> None:
    Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save(path)


def retain_hidden_components(obstacle: np.ndarray, hidden: np.ndarray) -> np.ndarray:
    if not hidden.any():
        return np.zeros_like(obstacle)
    _, labels = cv2.connectedComponents(obstacle.astype(np.uint8), connectivity=8)
    keep = np.unique(labels[obstacle & hidden])
    keep = keep[keep != 0]
    return np.isin(labels, keep)


def checkpoint_category_names(checkpoint: dict[str, Any], model_input_channels: int) -> tuple[str, ...]:
    """Recover the category-plane order used to train a checkpoint.

    The original adapter used only ``stairs`` and ``walkway`` category planes
    (7 total input channels).  The current adapter uses one plane for each
    canonical category (10 total).  Keeping this explicit makes old reviewed
    runs reproducible while preventing a silent channel-order mismatch.
    """
    category_count = model_input_channels - 5  # RGB + visible + obstacle
    if category_count < 0:
        raise ValueError(
            f"Checkpoint expects {model_input_channels} inputs; at least 5 are required."
        )
    recorded = checkpoint.get("input_channels", [])
    if isinstance(recorded, list):
        names = tuple(name for name in recorded if name in CATEGORIES)
        if len(names) == category_count:
            return names
    if category_count == len(CATEGORIES):
        return CATEGORIES
    raise ValueError(
        "Checkpoint category-plane metadata is incompatible with its model input shape: "
        f"expected {category_count} category planes, recorded={recorded!r}."
    )


def category_planes(category: str, names: tuple[str, ...], size: int) -> np.ndarray:
    return np.stack(
        [np.full((size, size), category == name, dtype=np.float32) for name in names],
        axis=0,
    )


def predict_probability(
    model: torch.nn.Module,
    device: torch.device,
    image: Image.Image,
    visible: np.ndarray,
    obstacle: np.ndarray,
    category: str,
    category_names: tuple[str, ...],
    image_size: int,
) -> np.ndarray:
    rgb = np.asarray(
        image.resize((image_size, image_size), Image.Resampling.BILINEAR),
        dtype=np.float32,
    ) / 255.0
    visible_small = np.asarray(
        Image.fromarray(visible.astype(np.uint8) * 255).resize(
            (image_size, image_size), Image.Resampling.NEAREST
        )
    ) > 127
    obstacle_small = np.asarray(
        Image.fromarray(obstacle.astype(np.uint8) * 255).resize(
            (image_size, image_size), Image.Resampling.NEAREST
        )
    ) > 127
    inputs = np.concatenate(
        [
            rgb.transpose(2, 0, 1),
            visible_small[None].astype(np.float32),
            obstacle_small[None].astype(np.float32),
            category_planes(category, category_names, image_size),
        ],
        axis=0,
    )
    with torch.inference_mode():
        logits = model(torch.from_numpy(inputs[None]).to(device)).sigmoid()[0, 0]
    probability = np.asarray(
        Image.fromarray(logits.detach().cpu().numpy().astype(np.float32), mode="F").resize(
            image.size, Image.Resampling.BILINEAR
        )
    ).copy()
    probability *= (obstacle & ~visible).astype(np.float32)
    return probability


def overlay(
    image: Image.Image,
    visible: np.ndarray,
    hidden: np.ndarray,
    obstacle: np.ndarray,
) -> Image.Image:
    result = np.asarray(image, dtype=np.float32).copy()
    for mask, color, alpha in (
        (visible, np.asarray((30, 210, 70), dtype=np.float32), 0.42),
        (hidden, np.asarray((40, 100, 245), dtype=np.float32), 0.62),
        (obstacle, np.asarray((230, 45, 45), dtype=np.float32), 0.52),
    ):
        result[mask] = result[mask] * (1.0 - alpha) + color * alpha
    return Image.fromarray(np.clip(result, 0, 255).astype(np.uint8), mode="RGB")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--image", type=Path, required=True)
    parser.add_argument("--target-visible-mask", type=Path, required=True)
    parser.add_argument("--obstacle-candidate-mask", type=Path, required=True)
    parser.add_argument("--category", choices=CATEGORIES, required=True)
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--image-size", type=int, default=256)
    parser.add_argument(
        "--threshold",
        type=float,
        default=None,
        help="Defaults to validation_threshold stored in the checkpoint.",
    )
    return parser


def main() -> int:
    args = build_parser().parse_args()
    image_path = args.image.expanduser().resolve()
    checkpoint_path = args.checkpoint.expanduser().resolve()
    output_dir = args.output_dir.expanduser().resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    image = load_rgb(image_path)
    visible = load_mask(args.target_visible_mask.expanduser().resolve(), image.size)
    obstacle_candidate = load_mask(
        args.obstacle_candidate_mask.expanduser().resolve(), image.size
    )
    visible &= ~obstacle_candidate

    device = torch.device(args.device)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA requested but unavailable; run through Slurm or use --device cpu")
    checkpoint: dict[str, Any] = torch.load(
        checkpoint_path, map_location=device, weights_only=False
    )
    model_state = checkpoint.get("model_state")
    if not isinstance(model_state, dict) or "enc1.block.0.weight" not in model_state:
        raise ValueError("Checkpoint does not contain a TinyAmodalUNet model_state.")
    model_input_channels = int(model_state["enc1.block.0.weight"].shape[1])
    category_names = checkpoint_category_names(checkpoint, model_input_channels)
    threshold = float(
        args.threshold
        if args.threshold is not None
        else checkpoint.get("validation_threshold", 0.6)
    )
    model = trainer.TinyAmodalUNet(in_channels=model_input_channels).to(device)
    model.load_state_dict(model_state)
    model.eval()

    probability = predict_probability(
        model,
        device,
        image,
        visible,
        obstacle_candidate,
        args.category,
        category_names,
        args.image_size,
    )
    hidden = probability >= threshold
    obstacle = retain_hidden_components(obstacle_candidate & ~visible, hidden)
    hidden &= obstacle
    amodal = visible | hidden

    if np.any(visible & obstacle):
        raise AssertionError("target_visible overlaps obstacle")
    if np.any(hidden & ~obstacle):
        raise AssertionError("hidden lies outside obstacle")
    if np.any(hidden != (amodal & ~visible)):
        raise AssertionError("hidden formula failed")

    save_mask(output_dir / "target_visible.png", visible)
    save_mask(output_dir / "obstacle_all_detected.png", obstacle_candidate)
    save_mask(output_dir / "obstacle.png", obstacle)
    save_mask(output_dir / "hidden.png", hidden)
    save_mask(output_dir / "target_amodal.png", amodal)
    Image.fromarray(
        np.clip(probability * 255.0, 0, 255).astype(np.uint8), mode="L"
    ).save(output_dir / "hidden_probability.png")
    overlay(image, visible, hidden, obstacle).save(output_dir / "mask_overlay.png")

    metadata = {
        "image": str(image_path),
        "raster_orientation_policy": "RGB and every input mask are decoded with PIL ImageOps.exif_transpose and must match exactly.",
        "display_raster_size": {"width": image.width, "height": image.height},
        "category": args.category,
        "checkpoint": str(checkpoint_path),
        "model_input_channels": model_input_channels,
        "category_plane_names": list(category_names),
        "threshold": threshold,
        "target_visible_pixels": int(visible.sum()),
        "obstacle_candidate_pixels": int(obstacle_candidate.sum()),
        "obstacle_pixels": int(obstacle.sum()),
        "hidden_pixels": int(hidden.sum()),
        "target_amodal_pixels": int(amodal.sum()),
        "automatic_masks_are_ground_truth": False,
        "review_status": "single_image_candidate_requires_human_review",
        "mask_invariants_passed": True,
    }
    (output_dir / "metadata.json").write_text(
        json.dumps(metadata, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(metadata, ensure_ascii=False, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())