"""Contact sheet of detections with predicted distance against ground truth.

The numeric evaluation says the median error is a few metres. This says what
that looks like on a frame, which is the only way to notice that the misses are
all one scene, or all one class, or all boxes clipped at the image edge.

Each detection is drawn with its class, the predicted distance, and -- where a
ground-truth box overlaps it -- the true distance and the error. Boxes are
coloured by how far off the estimate is, so a frame can be read at a glance.

    python -m src.depth.view_detections --weights trained/yolo11n/weights/best.pt
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import yaml
from PIL import Image, ImageDraw, ImageFont

from src.common import calib, paths, schema
from src.depth import estimate as depth_estimate
from src.depth.evaluate import iou

# Green when the estimate is close, amber when it is loose, red when it is wrong.
# Relative, not absolute, because 3 m at 50 m is fine and 3 m at 5 m is not.
BANDS = [(0.10, (60, 220, 90)), (0.25, (245, 190, 60)), (1e9, (240, 70, 70))]
NO_TRUTH = (150, 150, 160)


def colour_for(relative_error: float | None):
    if relative_error is None:
        return NO_TRUTH
    for limit, colour in BANDS:
        if relative_error <= limit:
            return colour
    return BANDS[-1][1]


def load_font(size: int):
    for candidate in ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
                      "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"):
        if Path(candidate).exists():
            return ImageFont.truetype(candidate, size)
    return ImageFont.load_default()


def render(image_path: str, detections: list[dict], root: Path,
           cell_width: int) -> Image.Image | None:
    full = paths.resolve_image(root, image_path)
    try:
        image = Image.open(full).convert("RGB")
    except OSError:
        return None

    scale = cell_width / image.width
    image = image.resize((cell_width, max(1, round(image.height * scale))))
    draw = ImageDraw.Draw(image)
    font = load_font(max(11, cell_width // 48))

    for detection in detections:
        colour = colour_for(detection["relative_error"])
        box = [detection["x1"] * scale, detection["y1"] * scale,
               detection["x2"] * scale, detection["y2"] * scale]
        draw.rectangle(box, outline=colour, width=2)

        label = f"{detection['class']} {detection['predicted_m']:.1f}m"
        if detection["truth_m"] is not None:
            label += f" (gt {detection['truth_m']:.1f}, {detection['error_m']:+.1f})"
        position = (box[0] + 2, max(0.0, box[1] - font.size - 3))
        draw.rectangle(draw.textbbox(position, label, font=font), fill=colour)
        draw.text(position, label, fill=(0, 0, 0), font=font)

    matched = [d for d in detections if d["truth_m"] is not None]
    caption = (f"{Path(image_path).name[:38]}   {len(detections)} det, "
               f"{len(matched)} matched")
    draw.rectangle(draw.textbbox((4, 4), caption, font=font), fill=(0, 0, 0))
    draw.text((4, 4), caption, fill=(255, 255, 255), font=font)
    return image


def tile(cells, columns: int, padding: int = 6) -> Image.Image:
    width = max(c.width for c in cells)
    height = max(c.height for c in cells)
    rows = (len(cells) + columns - 1) // columns
    sheet = Image.new("RGB", (columns * width + (columns + 1) * padding,
                              rows * height + (rows + 1) * padding), (22, 22, 26))
    for index, cell in enumerate(cells):
        column, row = index % columns, index // columns
        sheet.paste(cell, (padding + column * (width + padding),
                           padding + row * (height + padding)))
    return sheet


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--weights", required=True)
    parser.add_argument("--unified-root", type=Path, default=None)
    parser.add_argument("--dataset", default="nuscenes")
    parser.add_argument("--split", default="val")
    parser.add_argument("--frames", type=int, default=6)
    parser.add_argument("--conf", type=float, default=0.25)
    parser.add_argument("--min-iou", type=float, default=0.5)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--cell-width", type=int, default=760)
    parser.add_argument("--cols", type=int, default=2)
    parser.add_argument("--out", type=Path, default=Path("qa/detections_depth.png"))
    parser.add_argument("--only-matched", action="store_true",
                        help="only draw detections that overlap a ground-truth box")
    args = parser.parse_args()

    from ultralytics import YOLO

    root = paths.unified_root(args.unified_root)
    manifest = schema.objects_only(schema.read_manifest(root / "manifest.parquet"))
    manifest = manifest[(manifest["source"] == args.dataset)
                        & (manifest["split"] == args.split)
                        & manifest["gt_distance_m"].notna()]
    priors = yaml.safe_load((root / "class_priors.yaml").read_text())
    calibrations = calib.load_all(root)
    model = YOLO(args.weights)

    # One frame per scene, richest first within each, then scenes in turn. A
    # sheet drawn from whichever frames have the most objects comes entirely
    # from one location, which is exactly the wrong sample when the error is
    # known to vary by scene.
    counts = manifest.groupby(["scene_id", "image_path"]).size().reset_index(name="n")
    counts = counts.sort_values("n", ascending=False)
    order = counts.drop_duplicates("scene_id")["image_path"].tolist()
    order += [p for p in counts["image_path"] if p not in set(order)]

    cells = []
    for image_path in order:
        if len(cells) >= args.frames:
            break
        truth = manifest[manifest["image_path"] == image_path]
        calibration = calibrations.get(truth.iloc[0]["sensor_id"])
        if calibration is None:
            continue

        result = model.predict(str(paths.resolve_image(root, image_path)),
                               conf=args.conf, device=args.device, verbose=False)[0]
        if not len(result.boxes):
            continue
        boxes = result.boxes.xyxy.cpu().numpy()
        classes = result.boxes.cls.cpu().numpy().astype(int)
        truth_boxes = truth[["x1", "y1", "x2", "y2"]].to_numpy()

        detections = []
        for box, class_index in zip(boxes, classes):
            name = model.names[int(class_index)]
            estimate = depth_estimate.estimate(tuple(box), name, calibration, priors)
            scores = iou(box, truth_boxes) if len(truth_boxes) else np.zeros(0)
            best = int(np.argmax(scores)) if len(scores) else -1
            matched = best >= 0 and scores[best] >= args.min_iou
            if args.only_matched and not matched:
                continue
            truth_m = float(truth.iloc[best]["gt_distance_m"]) if matched else None
            error = (estimate.distance_m - truth_m) if matched else None
            detections.append({
                "class": name, "x1": box[0], "y1": box[1], "x2": box[2], "y2": box[3],
                "predicted_m": estimate.distance_m, "truth_m": truth_m,
                "error_m": error,
                "relative_error": abs(error) / truth_m if matched and truth_m else None,
            })

        cell = render(image_path, detections, root, args.cell_width)
        if cell is not None:
            cells.append(cell)

    if not cells:
        raise SystemExit("no frames with detections")
    args.out.parent.mkdir(parents=True, exist_ok=True)
    tile(cells, args.cols).save(args.out)
    print(f"wrote {args.out}  ({len(cells)} frames)")
    print("  green <=10% error, amber <=25%, red above, grey no ground truth")


if __name__ == "__main__":
    main()