#!/usr/bin/env python3
"""Compare full-frame vs TILED inference for the fine-tuned footage model.

Tiled: run the detector on each tile of a grid (weeds appear bigger → in the
model's resolvable range), map boxes back, NMS. Report per-frame detection
rate + recall vs the tiled GT (fraction of GT weeds with an overlapping pred).
"""
import argparse
import glob
from pathlib import Path
from collections import defaultdict

from ultralytics import YOLO
from PIL import Image

GT = Path("/opt/weeds/footage/gt_tiled")


def iou(a, b):
    ix1, iy1 = max(a[0], b[0]), max(a[1], b[1])
    ix2, iy2 = min(a[2], b[2]), min(a[3], b[3])
    iw, ih = max(0, ix2-ix1), max(0, iy2-iy1)
    inter = iw*ih
    ua = (a[2]-a[0])*(a[3]-a[1]) + (b[2]-b[0])*(b[3]-b[1]) - inter
    return inter/ua if ua > 0 else 0


def nms(boxes, thr=0.5):
    boxes = sorted(boxes, key=lambda b: -b[4])
    keep = []
    for b in boxes:
        if all(iou(b[:4], k[:4]) < thr for k in keep):
            keep.append(b)
    return keep


def full_preds(model, img, conf, imgsz=1280):
    r = model.predict(img, imgsz=imgsz, conf=conf, verbose=False)[0]
    return [[*map(float, b.xyxy[0].tolist()), float(b.conf[0])] for b in r.boxes]


def tiled_preds(model, img_path, conf, grid=(3, 2), imgsz=640, overlap=0.12):
    im = Image.open(img_path); W, H = im.size
    gx, gy = (grid[1], grid[0]) if H > W else grid
    tw, th = W//gx, H//gy
    ox, oy = int(tw*overlap), int(th*overlap)
    boxes = []
    for j in range(gy):
        for i in range(gx):
            x0 = max(0, i*tw-ox); y0 = max(0, j*th-oy)
            x1 = min(W, (i+1)*tw+ox); y1 = min(H, (j+1)*th+oy)
            tile = im.crop((x0, y0, x1, y1))
            tmp = Path("/tmp")/f"ti_{Path(img_path).stem}_{i}_{j}.jpg"
            tile.convert("RGB").save(tmp, quality=90)
            r = model.predict(str(tmp), imgsz=imgsz, conf=conf, verbose=False)[0]
            tmp.unlink(missing_ok=True)
            for b in r.boxes:
                bx = b.xyxy[0].tolist()
                boxes.append([x0+bx[0], y0+bx[1], x0+bx[2], y0+bx[3], float(b.conf[0])])
    return nms(boxes, 0.5)


def gt_boxes(stem):
    W = H = None
    ip = GT/"images"/f"{stem}.jpg"
    with Image.open(ip) as im:
        W, H = im.size
    out = []
    for ln in (GT/"labels"/f"{stem}.txt").read_text().splitlines():
        if not ln.strip():
            continue
        _, cx, cy, bw, bh = map(float, ln.split())
        out.append([(cx-bw/2)*W, (cy-bh/2)*H, (cx+bw/2)*W, (cy+bh/2)*H])
    return out


def recall(preds, gts, thr=0.3):
    if not gts:
        return None
    hit = 0
    for g in gts:
        if any(iou(p[:4], g) >= thr for p in preds):
            hit += 1
    return hit/len(gts)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--imgsz-full", type=int, default=1280)
    ap.add_argument("--imgsz-tile", type=int, default=640)
    ap.add_argument("--clip", default="", help="restrict to one clip prefix; blank = all frames")
    ap.add_argument("--grid", default="3x2")
    args = ap.parse_args()
    gx, gy = (int(v) for v in args.grid.lower().split("x"))
    model = YOLO(args.model)
    pat = f"{args.clip}*.txt" if args.clip else "*.txt"
    stems = [Path(p).stem for p in glob.glob(str(GT/"labels"/pat))]
    agg = defaultdict(lambda: [0.0, 0])
    frames_hit = defaultdict(int)
    for stem in stems:
        gts = gt_boxes(stem)
        ip = str(GT/'images'/f'{stem}.jpg')
        for mode, preds in (("full", full_preds(model, ip, 0.25, args.imgsz_full)),
                            ("tiled", tiled_preds(model, ip, 0.25, (gx, gy), args.imgsz_tile))):
            r = recall(preds, gts)
            if r is not None:
                agg[mode][0] += r; agg[mode][1] += 1
            if preds:
                frames_hit[mode] += 1
    n = len(stems)
    for mode in ("full", "tiled"):
        mr = agg[mode][0]/agg[mode][1] if agg[mode][1] else 0
        print(f"{mode:6s} inference: mean per-frame weed-recall={mr:.3f} "
              f"frame-detection-rate={frames_hit[mode]/n:.3f}  (n={n})")
    print("TILED_EVAL_DONE")


if __name__ == "__main__":
    main()