#!/usr/bin/env python3
"""Run the optional fast CPU 2D baseline on accessibility samples.

The script is intentionally non-destructive: it only reads dataset samples and
writes a mirrored result tree under --output-dir.
"""

from __future__ import annotations

import argparse
import json
import random
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

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

try:
    from tools.accessibility_dataset_layout import iter_sample_dirs, resolve_sample_dir
except ModuleNotFoundError:  # Direct execution from tools/
    from accessibility_dataset_layout import iter_sample_dirs, resolve_sample_dir


PROJECT_ROOT = Path(__file__).resolve().parents[1]
IMAGE_CANDIDATES = ("image.jpg", "image.jpeg", "image.png")


@dataclass(frozen=True)
class Sample:
    sample_id: str
    sample_dir: Path
    image_path: Path
    metadata: dict[str, Any]
    target_visible_path: Path | None = None
    target_amodal_path: Path | None = None
    hidden_path: Path | None = None
    obstacle_path: Path | None = None


def resolve_path(path: str | Path) -> Path:
    candidate = Path(path)
    if candidate.is_absolute():
        return candidate
    return PROJECT_ROOT / candidate


def read_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def read_rgb(path: Path) -> np.ndarray:
    return np.array(ImageOps.exif_transpose(Image.open(path)).convert("RGB"))


def read_mask(path: Path | None, shape: tuple[int, int]) -> np.ndarray | None:
    if path is None or not path.is_file():
        return None
    mask = np.array(ImageOps.exif_transpose(Image.open(path)).convert("L")) > 127
    h, w = shape
    if mask.shape != (h, w):
        raise ValueError(
            f"Mask/RGB raster mismatch for {path}: mask={mask.shape}, rgb={(h, w)}. "
            "Refusing to resize because this can hide EXIF-orientation misalignment."
        )
    return mask.astype(bool)


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


def save_rgb(path: Path, image: np.ndarray) -> None:
    Image.fromarray(np.clip(image, 0, 255).astype(np.uint8), mode="RGB").save(path)


def maybe_path(sample_dir: Path, name: str) -> Path | None:
    path = sample_dir / name
    return path if path.is_file() else None


def find_image_path(sample_dir: Path, metadata: dict[str, Any], dataset_root: Path) -> Path:
    image_file = metadata.get("image_file")
    if isinstance(image_file, str):
        candidates = [
            sample_dir / Path(image_file).name,
            dataset_root / image_file,
            PROJECT_ROOT / image_file,
        ]
        for candidate in candidates:
            if candidate.is_file():
                return candidate
    for name in IMAGE_CANDIDATES:
        candidate = sample_dir / name
        if candidate.is_file():
            return candidate
    raise FileNotFoundError(f"No image file found in {sample_dir}")


def sample_from_dir(sample_dir: Path, dataset_root: Path) -> Sample:
    metadata_path = sample_dir / "metadata.json"
    metadata = read_json(metadata_path) if metadata_path.is_file() else {}
    sample_id = str(metadata.get("sample_id") or sample_dir.name)
    return Sample(
        sample_id=sample_id,
        sample_dir=sample_dir,
        image_path=find_image_path(sample_dir, metadata, dataset_root),
        metadata=metadata,
        target_visible_path=maybe_path(sample_dir, "target_visible.png"),
        target_amodal_path=maybe_path(sample_dir, "target_amodal.png"),
        hidden_path=maybe_path(sample_dir, "hidden.png"),
        obstacle_path=maybe_path(sample_dir, "obstacle.png"),
    )


def discover_samples(
    dataset_root: Path,
    sample_ids: list[str],
    categories: set[str],
    splits: set[str],
) -> list[Sample]:
    sample_root = dataset_root / "samples" if (dataset_root / "samples").is_dir() else dataset_root
    if sample_ids:
        dirs = [resolve_sample_dir(sample_root, sample_id) for sample_id in sample_ids]
    else:
        dirs = list(iter_sample_dirs(sample_root))

    samples: list[Sample] = []
    for sample_dir in dirs:
        if not sample_dir.is_dir():
            raise FileNotFoundError(f"Missing sample directory: {sample_dir}")
        sample = sample_from_dir(sample_dir, dataset_root)
        category = str(sample.metadata.get("category") or sample.metadata.get("taxonomy_category") or "")
        split = str(sample.metadata.get("split") or sample.metadata.get("strict_gt_split") or "")
        if categories and category not in categories:
            continue
        if splits and split not in splits:
            continue
        samples.append(sample)
    return samples


def hidden_pixel_count(sample: Sample) -> int:
    value = sample.metadata.get("hidden_pixels")
    if isinstance(value, int):
        return value
    if sample.hidden_path and sample.hidden_path.is_file():
        return int((np.array(Image.open(sample.hidden_path).convert("L")) > 127).sum())
    return 0


def choose_samples(samples: list[Sample], args: argparse.Namespace) -> list[Sample]:
    eligible = [sample for sample in samples if hidden_pixel_count(sample) >= args.min_hidden_pixels]
    if args.sample_policy == "largest-hidden":
        eligible.sort(key=lambda sample: (-hidden_pixel_count(sample), sample.sample_id))
    elif args.sample_policy == "random":
        rng = random.Random(args.seed)
        rng.shuffle(eligible)
    else:
        eligible.sort(key=lambda sample: sample.sample_id)

    if args.all:
        return eligible
    return eligible[: args.limit]


def ellipse_kernel(radius: int) -> np.ndarray | None:
    if radius <= 0:
        return None
    size = radius * 2 + 1
    return cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (size, size))


def dilate_mask(mask: np.ndarray, radius: int) -> np.ndarray:
    kernel = ellipse_kernel(radius)
    if kernel is None or not mask.any():
        return mask.astype(bool)
    return cv2.dilate(mask.astype(np.uint8), kernel, iterations=1).astype(bool)


def close_mask(mask: np.ndarray, radius: int) -> np.ndarray:
    kernel = ellipse_kernel(radius)
    if kernel is None or not mask.any():
        return mask.astype(bool)
    return cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, kernel).astype(bool)


def derive_hidden_mask(
    hidden: np.ndarray | None,
    target_visible: np.ndarray | None,
    target_amodal: np.ndarray | None,
) -> tuple[np.ndarray, str]:
    if hidden is not None and hidden.any():
        derived = hidden.astype(bool).copy()
        source = "hidden mask"
        if target_visible is not None and target_amodal is not None:
            derived |= target_amodal & ~target_visible
            source += " union target_amodal-minus-target_visible"
        return derived, source
    if target_visible is not None and target_amodal is not None:
        return (target_amodal & ~target_visible).astype(bool), "target_amodal-minus-target_visible"
    if target_amodal is not None:
        return target_amodal.astype(bool), "target_amodal fallback"
    raise ValueError("Need hidden.png or target_amodal.png + target_visible.png to derive completion mask")


def build_inpaint_mask(
    hidden: np.ndarray,
    target_visible: np.ndarray | None,
    target_amodal: np.ndarray | None,
    obstacle: np.ndarray | None,
    mode: str,
    target_band_dilate: int,
    final_dilate: int,
    close_radius: int,
) -> tuple[np.ndarray, dict[str, Any]]:
    mask = hidden.astype(bool).copy()
    metadata: dict[str, Any] = {
        "mode": mode,
        "hidden_pixels": int(hidden.sum()),
        "added_obstacle_pixels": 0,
    }

    if mode == "hidden":
        pass
    elif mode == "hidden_dilated":
        mask = dilate_mask(mask, target_band_dilate)
    elif mode == "target_occluder":
        if obstacle is not None:
            obstacle_on_target = obstacle & dilate_mask(hidden, target_band_dilate)
            metadata["added_obstacle_pixels"] = int((obstacle_on_target & ~mask).sum())
            mask |= obstacle_on_target
    elif mode == "obstacle":
        if obstacle is not None:
            metadata["added_obstacle_pixels"] = int((obstacle & ~mask).sum())
            mask |= obstacle
    else:
        raise ValueError(f"Unsupported mask mode: {mode}")

    mask = close_mask(mask, close_radius)
    mask = dilate_mask(mask, final_dilate)
    metadata["final_pixels"] = int(mask.sum())
    return mask.astype(bool), metadata


def inpaint_opencv(rgb: np.ndarray, mask: np.ndarray, radius: float, method: str) -> np.ndarray:
    if not mask.any():
        return rgb.copy()
    flag = cv2.INPAINT_TELEA if method == "telea" else cv2.INPAINT_NS
    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    completed = cv2.inpaint(bgr, mask.astype(np.uint8) * 255, radius, flag)
    return cv2.cvtColor(completed, cv2.COLOR_BGR2RGB)


def inpaint_opencv_pyramid(
    rgb: np.ndarray,
    mask: np.ndarray,
    radius: float,
    method: str,
    levels: int,
    seam_radius: int,
) -> np.ndarray:
    if not mask.any() or levels <= 1:
        return inpaint_opencv(rgb, mask, radius, method)

    h, w = rgb.shape[:2]
    scale = 1.0 / float(2 ** (levels - 1))
    low_w = max(32, int(round(w * scale)))
    low_h = max(32, int(round(h * scale)))
    low_rgb = cv2.resize(rgb, (low_w, low_h), interpolation=cv2.INTER_AREA)
    low_mask = cv2.resize(mask.astype(np.uint8), (low_w, low_h), interpolation=cv2.INTER_NEAREST) > 0
    low_completed = inpaint_opencv(low_rgb, low_mask, radius, method)
    up_completed = cv2.resize(low_completed, (w, h), interpolation=cv2.INTER_CUBIC)

    composite = rgb.copy()
    composite[mask] = up_completed[mask]

    seam = dilate_mask(mask, seam_radius)
    inner = cv2.erode(
        mask.astype(np.uint8),
        ellipse_kernel(max(1, seam_radius // 2)),
        iterations=1,
    ).astype(bool)
    seam = seam & ~inner
    if seam.any():
        composite = inpaint_opencv(composite, seam, max(1.0, radius * 0.5), method)
    return composite


def run_inpaint(rgb: np.ndarray, mask: np.ndarray, args: argparse.Namespace) -> np.ndarray:
    if args.opencv_mode == "pyramid":
        return inpaint_opencv_pyramid(
            rgb,
            mask,
            args.inpaint_radius,
            args.method,
            args.pyramid_levels,
            args.seam_radius,
        )
    return inpaint_opencv(rgb, mask, args.inpaint_radius, args.method)


def blend_masks(rgb: np.ndarray, masks: list[tuple[np.ndarray, tuple[int, int, int], float]]) -> np.ndarray:
    out = rgb.astype(np.float32).copy()
    for mask, color, alpha in masks:
        if mask is not None and mask.any():
            out[mask] = out[mask] * (1.0 - alpha) + np.array(color, dtype=np.float32) * alpha
    return np.clip(out, 0, 255).astype(np.uint8)


def write_target_rgba(path: Path, completed: np.ndarray, target_amodal: np.ndarray | None, inpaint_mask: np.ndarray) -> None:
    alpha = target_amodal if target_amodal is not None and target_amodal.any() else inpaint_mask
    rgba = np.dstack([completed, alpha.astype(np.uint8) * 255])
    Image.fromarray(rgba, mode="RGBA").save(path)


def checkerboard(size: tuple[int, int], cell: int = 16) -> Image.Image:
    width, height = size
    yy, xx = np.indices((height, width))
    pattern = ((xx // cell + yy // cell) % 2).astype(np.uint8)
    values = np.where(pattern[..., None] == 0, 228, 188).astype(np.uint8)
    image = np.repeat(values, 3, axis=2)
    return Image.fromarray(image, mode="RGB")


def rgba_on_checker(rgba_path: Path, size: tuple[int, int]) -> Image.Image:
    image = Image.open(rgba_path).convert("RGBA")
    canvas = checkerboard(image.size)
    canvas.paste(image, (0, 0), image)
    return ImageOps.contain(canvas, size)


def label_panel(image: Image.Image, label: str, size: tuple[int, int]) -> Image.Image:
    body = ImageOps.contain(image.convert("RGB"), size)
    panel = Image.new("RGB", (size[0], size[1] + 28), "white")
    draw = ImageDraw.Draw(panel)
    draw.text((8, 8), label, fill=(0, 0, 0))
    panel.paste(body, ((size[0] - body.width) // 2, 28 + (size[1] - body.height) // 2))
    return panel


def write_contact_sheet(
    path: Path,
    original: np.ndarray,
    overlay: np.ndarray,
    completed: np.ndarray,
    rgba_path: Path,
    panel_width: int,
) -> None:
    h, w = original.shape[:2]
    panel_height = max(160, int(panel_width * h / max(w, 1)))
    size = (panel_width, panel_height)
    panels = [
        label_panel(Image.fromarray(original), "source image", size),
        label_panel(Image.fromarray(overlay), "mask guide", size),
        label_panel(Image.fromarray(completed), "completed RGB", size),
        label_panel(rgba_on_checker(rgba_path, size), "amodal target RGBA", size),
    ]
    sheet = Image.new("RGB", (sum(panel.width for panel in panels), max(panel.height for panel in panels)), "white")
    x = 0
    for panel in panels:
        sheet.paste(panel, (x, 0))
        x += panel.width
    sheet.save(path, quality=92)


def path_for_manifest(path: Path) -> str:
    try:
        return str(path.relative_to(PROJECT_ROOT))
    except ValueError:
        return str(path)


def process_sample(sample: Sample, output_dir: Path, args: argparse.Namespace) -> dict[str, Any]:
    sample_out = output_dir / sample.sample_id
    if sample_out.exists() and not args.overwrite:
        return {
            "sample_id": sample.sample_id,
            "status": "skipped_existing_output",
            "output_dir": path_for_manifest(sample_out),
        }
    sample_out.mkdir(parents=True, exist_ok=True)

    rgb = read_rgb(sample.image_path)
    shape = rgb.shape[:2]
    target_visible = read_mask(sample.target_visible_path, shape)
    target_amodal = read_mask(sample.target_amodal_path, shape)
    hidden_input = read_mask(sample.hidden_path, shape)
    obstacle = read_mask(sample.obstacle_path, shape)
    hidden, hidden_source = derive_hidden_mask(hidden_input, target_visible, target_amodal)
    inpaint_mask, mask_meta = build_inpaint_mask(
        hidden=hidden,
        target_visible=target_visible,
        target_amodal=target_amodal,
        obstacle=obstacle,
        mode=args.mask_mode,
        target_band_dilate=args.target_band_dilate,
        final_dilate=args.mask_dilate,
        close_radius=args.mask_close,
    )

    mask_ratio = float(inpaint_mask.sum()) / float(inpaint_mask.size)
    row: dict[str, Any] = {
        "sample_id": sample.sample_id,
        "category": sample.metadata.get("category") or sample.metadata.get("taxonomy_category"),
        "split": sample.metadata.get("split") or sample.metadata.get("strict_gt_split"),
        "source_image": path_for_manifest(sample.image_path),
        "source_sample_dir": path_for_manifest(sample.sample_dir),
        "hidden_source": hidden_source,
        "mask": {
            **mask_meta,
            "mask_ratio": mask_ratio,
            "max_mask_area_ratio": args.max_mask_area_ratio,
        },
        "backend": {
            "name": "opencv",
            "method": args.method,
            "opencv_mode": args.opencv_mode,
            "inpaint_radius": args.inpaint_radius,
            "pyramid_levels": args.pyramid_levels,
            "seam_radius": args.seam_radius,
        },
    }
    if not inpaint_mask.any():
        row["status"] = "skipped_empty_mask"
        (sample_out / "sample_manifest.json").write_text(json.dumps(row, indent=2, ensure_ascii=False), encoding="utf-8")
        return row
    if mask_ratio > args.max_mask_area_ratio and not args.allow_large_mask:
        save_mask(sample_out / "inpaint_mask.png", inpaint_mask)
        row["status"] = "skipped_large_mask"
        (sample_out / "sample_manifest.json").write_text(json.dumps(row, indent=2, ensure_ascii=False), encoding="utf-8")
        return row

    completed = run_inpaint(rgb, inpaint_mask, args)
    mask_overlay = blend_masks(
        rgb,
        [
            (target_amodal if target_amodal is not None else np.zeros(shape, dtype=bool), (0, 150, 255), 0.28),
            (target_visible if target_visible is not None else np.zeros(shape, dtype=bool), (0, 220, 80), 0.42),
            (inpaint_mask, (255, 48, 48), 0.65),
        ],
    )
    completion_delta = np.abs(completed.astype(np.int16) - rgb.astype(np.int16)).max(axis=2) > 8

    outputs = {
        "completed_rgb": sample_out / "completed_rgb.png",
        "amodal_target_rgba": sample_out / "amodal_target_rgba.png",
        "inpaint_mask": sample_out / "inpaint_mask.png",
        "hidden_mask": sample_out / "hidden_mask.png",
        "mask_overlay": sample_out / "mask_overlay.jpg",
        "completion_delta": sample_out / "completion_delta.png",
        "contact_sheet": sample_out / "contact_sheet.jpg",
        "manifest": sample_out / "sample_manifest.json",
    }
    save_rgb(outputs["completed_rgb"], completed)
    write_target_rgba(outputs["amodal_target_rgba"], completed, target_amodal, inpaint_mask)
    save_mask(outputs["inpaint_mask"], inpaint_mask)
    save_mask(outputs["hidden_mask"], hidden)
    save_rgb(outputs["mask_overlay"], mask_overlay)
    save_mask(outputs["completion_delta"], completion_delta)
    write_contact_sheet(
        outputs["contact_sheet"],
        rgb,
        mask_overlay,
        completed,
        outputs["amodal_target_rgba"],
        args.panel_width,
    )

    row["status"] = "completed"
    row["outputs"] = {key: path_for_manifest(value) for key, value in outputs.items() if key != "manifest"}
    outputs["manifest"].write_text(json.dumps(row, indent=2, ensure_ascii=False), encoding="utf-8")
    return row


def sample_from_single_image(args: argparse.Namespace) -> Sample:
    image_path = resolve_path(args.image)
    if not image_path.is_file():
        raise FileNotFoundError(image_path)
    sample_id = args.single_sample_id or image_path.stem
    return Sample(
        sample_id=sample_id,
        sample_dir=image_path.parent,
        image_path=image_path,
        metadata={"sample_id": sample_id, "source": "single_image_cli"},
        target_visible_path=resolve_path(args.target_visible_mask) if args.target_visible_mask else None,
        target_amodal_path=resolve_path(args.target_amodal_mask) if args.target_amodal_mask else None,
        hidden_path=resolve_path(args.hidden_mask) if args.hidden_mask else None,
        obstacle_path=resolve_path(args.obstacle_mask) if args.obstacle_mask else None,
    )


def write_run_manifest(output_dir: Path, rows: list[dict[str, Any]], args: argparse.Namespace) -> None:
    manifest = {
        "created_at_utc": datetime.now(timezone.utc).isoformat(),
        "tool": "tools/accessibility_fast_2d_baseline.py",
        "non_destructive_policy": "Read source samples only; write all derived RGB/RGBA outputs under output_dir.",
        "output_dir": path_for_manifest(output_dir),
        "args": {
            "dataset_root": args.dataset_root,
            "sample_id": args.sample_id,
            "image": args.image,
            "limit": args.limit,
            "all": args.all,
            "sample_policy": args.sample_policy,
            "mask_mode": args.mask_mode,
            "method": args.method,
        },
        "references": {
            "saraao_amodal": "https://github.com/saraao/amodal",
            "pix2gestalt": "https://github.com/cvlab-columbia/pix2gestalt",
            "amodal_completion_in_the_wild": "https://github.com/Championchess/Amodal-Completion-in-the-Wild",
            "local_amodal": "external backend; path supplied by user",
            "local_pix2gestalt": "external backend; path supplied by user",
            "local_amodal_wild": "external backend; path supplied by user",
        },
        "samples": rows,
    }
    (output_dir / "run_manifest.json").write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset-root", default="output/Accessibility", help="Dataset root containing samples/<sample_id>/ directories.")
    parser.add_argument("--sample-id", action="append", default=[], help="Specific sample id to process. Repeat for multiple samples.")
    parser.add_argument("--category", action="append", default=[], help="Optional category filter, e.g. stairs or ramp.")
    parser.add_argument("--split", action="append", default=[], help="Optional split filter.")
    parser.add_argument("--limit", type=int, default=8, help="Maximum samples when --all is not set.")
    parser.add_argument("--all", action="store_true", help="Process all eligible samples.")
    parser.add_argument("--sample-policy", choices=["largest-hidden", "random", "sorted"], default="largest-hidden")
    parser.add_argument("--seed", type=int, default=13)
    parser.add_argument("--min-hidden-pixels", type=int, default=32)

    parser.add_argument("--image", default=None, help="Single-image mode input image.")
    parser.add_argument("--single-sample-id", default=None)
    parser.add_argument("--target-visible-mask", default=None)
    parser.add_argument("--target-amodal-mask", default=None)
    parser.add_argument("--hidden-mask", default=None)
    parser.add_argument("--obstacle-mask", default=None)

    parser.add_argument("--output-dir", default="output/amodal2d_color_completion")
    parser.add_argument("--overwrite", action="store_true", help="Overwrite only files inside --output-dir.")
    parser.add_argument("--mask-mode", choices=["hidden", "hidden_dilated", "target_occluder", "obstacle"], default="hidden")
    parser.add_argument("--target-band-dilate", type=int, default=16, help="Pixels to dilate hidden completion support before intersecting obstacle mask.")
    parser.add_argument("--mask-dilate", type=int, default=3, help="Final dilation radius for the inpaint mask.")
    parser.add_argument("--mask-close", type=int, default=3, help="Closing radius for small holes in the inpaint mask.")
    parser.add_argument("--max-mask-area-ratio", type=float, default=0.18, help="Skip unexpectedly huge masks unless --allow-large-mask is set.")
    parser.add_argument("--allow-large-mask", action="store_true")
    parser.add_argument("--opencv-mode", choices=["single", "pyramid"], default="single")
    parser.add_argument("--pyramid-levels", type=int, default=3, help="Coarse-to-full resolution levels used with --opencv-mode pyramid.")
    parser.add_argument("--seam-radius", type=int, default=8, help="Boundary refinement radius used with --opencv-mode pyramid.")
    parser.add_argument("--method", choices=["telea", "ns"], default="telea", help="OpenCV inpainting method.")
    parser.add_argument("--inpaint-radius", type=float, default=5.0)
    parser.add_argument("--panel-width", type=int, default=420)
    return parser


def main() -> int:
    args = build_parser().parse_args()
    output_dir = resolve_path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    if args.image:
        samples = [sample_from_single_image(args)]
    else:
        dataset_root = resolve_path(args.dataset_root)
        samples = discover_samples(dataset_root, args.sample_id, set(args.category), set(args.split))
        samples = choose_samples(samples, args)

    rows = [process_sample(sample, output_dir, args) for sample in samples]
    write_run_manifest(output_dir, rows, args)
    completed = sum(1 for row in rows if row.get("status") == "completed")
    print(f"Wrote {completed}/{len(rows)} completed samples to {output_dir}")
    return 0


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