"""Crops of individual detections, grouped by confidence, with their distances.

`view_detections.py` draws whole frames, which shows context. This shows the
objects themselves, side by side, which is what makes a pattern visible: a wall
of confident detections that are all close and unoccluded next to a wall of
uncertain ones that are all small, clipped or half behind a car.

Reads the CSV that `accuracy.py` already wrote, so it needs no model and no
inference -- it is re-reading a scored run, not producing a new one.

    python -m src.depth.accuracy --weights best.pt --tag best
    python -m src.depth.view_gallery --tag best

Each crop is bordered by how far off the estimate was, relative to the true
distance, and captioned with confidence, the estimate, the truth and the error.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import pandas as pd
from PIL import Image, ImageDraw, ImageFont

from src.common import paths

BANDS = [(0.05, 0.25), (0.25, 0.50), (0.50, 0.75), (0.75, 1.01)]
# Relative, not absolute: 3 m out at 50 m is fine, 3 m out at 5 m is not.
COLOURS = [(0.10, (70, 200, 100)), (0.25, (240, 185, 60)), (1e9, (235, 75, 75))]
CELL = 190
PAD = 6


def colour_for(percent_off: float):
    for limit, colour in COLOURS:
        if percent_off / 100 <= limit:
            return colour
    return COLOURS[-1][1]


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


def crop(row, root: Path, context: float) -> Image.Image | None:
    """One detection, cropped with surrounding context, letterboxed to CELL."""
    try:
        image = Image.open(paths.resolve_image(root, row.image_path)).convert("RGB")
    except OSError:
        return None

    width, height = row.x2 - row.x1, row.y2 - row.y1
    # Context scales with the box, so a distant cone is not a single pixel and a
    # near barrier does not fill the frame with road.
    margin = max(width, height) * context
    box = (max(0, row.x1 - margin), max(0, row.y1 - margin),
           min(image.width, row.x2 + margin), min(image.height, row.y2 + margin))
    patch = image.crop(tuple(int(v) for v in box))
    if patch.width < 2 or patch.height < 2:
        return None

    scale = min(CELL / patch.width, CELL / patch.height)
    patch = patch.resize((max(1, int(patch.width * scale)),
                          max(1, int(patch.height * scale))))

    cell = Image.new("RGB", (CELL, CELL + 30), (28, 28, 32))
    cell.paste(patch, ((CELL - patch.width) // 2, (CELL - patch.height) // 2))
    draw = ImageDraw.Draw(cell)

    # The detection box, in crop coordinates.
    draw.rectangle([(row.x1 - box[0]) * scale + (CELL - patch.width) // 2,
                    (row.y1 - box[1]) * scale + (CELL - patch.height) // 2,
                    (row.x2 - box[0]) * scale + (CELL - patch.width) // 2,
                    (row.y2 - box[1]) * scale + (CELL - patch.height) // 2],
                   outline=colour_for(row.percent_off), width=2)

    small = font(12)
    draw.text((5, CELL + 2), f"{row.predicted_class}  conf {row.confidence:.2f}",
              fill=(225, 225, 230), font=small)
    draw.text((5, CELL + 15),
              f"ours {row.predicted_m:.1f}  true {row.truth_m:.1f}  "
              f"{row.error_m:+.1f}m", fill=colour_for(row.percent_off), font=small)
    return cell


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--tag", required=True)
    parser.add_argument("--results-dir", type=Path,
                        default=Path("results/depth_accuracy"))
    parser.add_argument("--unified-root", type=Path, default=None)
    parser.add_argument("--per-band", type=int, default=10)
    parser.add_argument("--context", type=float, default=0.6,
                        help="padding around each box, as a fraction of its size")
    parser.add_argument("--out", type=Path, default=None)
    parser.add_argument("--worst", action="store_true",
                        help="show the largest errors instead of a confidence "
                             "sweep -- the view that says what is going wrong")
    parser.add_argument("--unique", action="store_true",
                        help="one crop per ground-truth object. Several "
                             "detections of the same cone otherwise fill a row "
                             "with the same error repeated.")
    args = parser.parse_args()

    root = paths.unified_root(args.unified_root)
    csv = args.results_dir / f"{args.tag}_detections.csv"
    frame = pd.read_csv(csv)
    frame = frame[frame["correct"] & frame["truth_m"].notna()]
    print(f"{len(frame)} correct detections in {csv}")

    if args.unique:
        # A ground-truth object is identified by its frame and its true
        # distance; repeated detections of it share both.
        before = len(frame)
        frame = (frame.sort_values("confidence", ascending=False)
                      .drop_duplicates(["image_path", "truth_m"]))
        print(f"  {before} detections cover {len(frame)} distinct objects")

    if args.worst:
        worst = frame.sort_values("percent_off", ascending=False)
        per_row = args.per_band
        rows, labels = [], []
        for start in range(0, min(len(worst), per_row * 3), per_row):
            chunk = worst.iloc[start:start + per_row]
            cells = [c for c in (crop(r, root, args.context)
                                 for r in chunk.itertuples()) if c is not None]
            if cells:
                rows.append(cells)
                labels.append(f"worst {start + 1}-{start + len(chunk)} by relative "
                              f"error   {chunk.percent_off.min():.0f}-"
                              f"{chunk.percent_off.max():.0f}% off")
        draw_sheet(rows, labels, args.out or Path(f"qa/gallery_{args.tag}_worst.png"))
        return

    rows, labels = [], []
    for low, high in BANDS:
        band = frame[(frame.confidence >= low) & (frame.confidence < high)]
        band = band.sort_values("confidence", ascending=False)
        if band.empty:
            continue
        # Evenly spaced through the band, so the row spans it rather than
        # showing ten near-identical detections from its top edge.
        step = max(1, len(band) // args.per_band)
        picked = band.iloc[::step][:args.per_band]
        cells = [c for c in (crop(r, root, args.context)
                             for r in picked.itertuples()) if c is not None]
        if cells:
            rows.append(cells)
            labels.append(f"confidence {low:.2f}-{min(high, 1.0):.2f}   "
                          f"n={len(band)}   median error "
                          f"{band.error_m.abs().median():.1f} m")

    draw_sheet(rows, labels, args.out or Path(f"qa/gallery_{args.tag}.png"))


def draw_sheet(rows, labels, out: Path) -> None:
    if not rows:
        raise SystemExit("nothing to draw")

    columns = max(len(r) for r in rows)
    header = 26
    sheet = Image.new("RGB",
                      (PAD + columns * (CELL + PAD),
                       sum(header + CELL + 30 + PAD for _ in rows) + PAD),
                      (18, 18, 22))
    draw = ImageDraw.Draw(sheet)
    title = font(14)

    y = PAD
    for cells, label in zip(rows, labels):
        draw.text((PAD, y + 5), label, fill=(235, 235, 240), font=title)
        y += header
        for index, cell in enumerate(cells):
            sheet.paste(cell, (PAD + index * (CELL + PAD), y))
        y += CELL + 30 + PAD

    out.parent.mkdir(parents=True, exist_ok=True)
    sheet.save(out)
    print(f"wrote {out}  ({sum(len(r) for r in rows)} crops)")
    print("  border: green within 10% of truth, amber within 25%, red beyond")


if __name__ == "__main__":
    main()