import argparse
from pathlib import Path

import numpy as np
import onnxruntime as ort
from PIL import Image


IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"}


def parse_args():
    parser = argparse.ArgumentParser(description="Run DepthPolyp ONNX inference on images.")
    parser.add_argument("--onnx", default="checkpoints/DepthPolyp_Kvasir.onnx")
    parser.add_argument("--input", default="samples/kvasir/images")
    parser.add_argument("--output", default="samples/kvasir/outputs")
    parser.add_argument("--image-size", type=int, default=224)
    parser.add_argument("--threshold", type=float, default=0.3)
    return parser.parse_args()


def list_images(input_path: Path):
    if input_path.is_file():
        return [input_path]
    return sorted(path for path in input_path.rglob("*") if path.suffix.lower() in IMAGE_EXTENSIONS)


def preprocess(image_path: Path, image_size: int):
    image = Image.open(image_path).convert("RGB")
    original_size = image.size
    resized = image.resize((image_size, image_size), Image.BILINEAR)
    array = np.asarray(resized).astype(np.float32) / 255.0
    tensor = np.transpose(array, (2, 0, 1))[None, ...]
    return image, original_size, tensor


def to_grayscale(probability: np.ndarray, size):
    probability = np.clip(probability, 0.0, 1.0)
    image = Image.fromarray((probability * 255).astype(np.uint8), mode="L")
    return image.resize(size, Image.BILINEAR)


def colorize_purple_yellow(probability: np.ndarray, size):
    probability = np.clip(probability, 0.0, 1.0)
    stops = np.array(
        [
            [38, 5, 84],
            [86, 33, 132],
            [141, 48, 140],
            [203, 71, 119],
            [245, 135, 48],
            [252, 231, 37],
        ],
        dtype=np.float32,
    )
    scaled = probability * (len(stops) - 1)
    lower = np.floor(scaled).astype(np.int32)
    upper = np.clip(lower + 1, 0, len(stops) - 1)
    alpha = (scaled - lower)[..., None]
    colored = stops[lower] * (1.0 - alpha) + stops[upper] * alpha
    image = Image.fromarray(colored.astype(np.uint8), mode="RGB")
    return image.resize(size, Image.BILINEAR)


def make_overlay(image: Image.Image, mask: Image.Image):
    base = image.convert("RGBA")
    mask_array = np.asarray(mask).astype(np.float32) / 255.0
    color = np.zeros((mask_array.shape[0], mask_array.shape[1], 4), dtype=np.uint8)
    color[..., 0] = 252
    color[..., 1] = 231
    color[..., 2] = 37
    color[..., 3] = (mask_array * 155).astype(np.uint8)
    return Image.alpha_composite(base, Image.fromarray(color, mode="RGBA")).convert("RGB")


def main():
    args = parse_args()
    input_path = Path(args.input)
    output_root = Path(args.output)
    mask_dir = output_root / "masks"
    depth_dir = output_root / "depth"
    overlay_dir = output_root / "overlay"
    for directory in (mask_dir, depth_dir, overlay_dir):
        directory.mkdir(parents=True, exist_ok=True)

    session = ort.InferenceSession(args.onnx, providers=["CPUExecutionProvider"])
    input_name = session.get_inputs()[0].name
    images = list_images(input_path)
    if not images:
        raise FileNotFoundError(f"No images found under {input_path}")

    for image_path in images:
        image, original_size, tensor = preprocess(image_path, args.image_size)
        segmentation, depth = session.run(None, {input_name: tensor})
        seg_prob = segmentation[0, 0]
        depth_prob = depth[0, 0]

        seg_image = to_grayscale(seg_prob, original_size)
        depth_image = colorize_purple_yellow(depth_prob, original_size)
        binary_mask = seg_image.point(lambda value: 255 if value >= int(args.threshold * 255) else 0)
        overlay = make_overlay(image, seg_image)

        stem = image_path.stem
        binary_mask.save(mask_dir / f"{stem}.png")
        depth_image.save(depth_dir / f"{stem}.png")
        overlay.save(overlay_dir / f"{stem}.jpg", quality=95)

    print(f"Processed {len(images)} image(s). Outputs saved to {output_root}")


if __name__ == "__main__":
    main()