"""Evaluate FFDNet-L with SAHI tiling on checkbox documents.

Usage:
    python datasets/checkbox/eval_ffdnet_sahi.py
    python datasets/checkbox/eval_ffdnet_sahi.py --tile-size 1024 --overlap 0.15 --confidence 0.2
    python datasets/checkbox/eval_ffdnet_sahi.py --pdf pdfs/01-practice.pdf --page 0
"""

import argparse
import json
import time
from pathlib import Path

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


def get_detector():
    from natural_pdf.analyzers.checkbox.onnx_engine import OnnxCheckboxDetector
    from natural_pdf.analyzers.checkbox.checkbox_options import OnnxCheckboxOptions

    opts = OnnxCheckboxOptions(
        model_repo="jbarrow/FFDNet-L-cpu",
        model_file="FFDNet-L.onnx",
        input_size=1216,
        class_names=["textbox", "choice_button", "signature"],
        checkbox_class_indices=[1],
        confidence=0.2,
        nms_threshold=0.45,
    )
    return OnnxCheckboxDetector(), opts


def nms(boxes: np.ndarray, scores: np.ndarray, iou_thresh: float = 0.45):
    """Non-maximum suppression on [x0, y0, x1, y1] boxes."""
    if len(boxes) == 0:
        return []
    order = np.argsort(-scores)
    keep = []
    while len(order) > 0:
        i = order[0]
        keep.append(i)
        if len(order) == 1:
            break
        rest = order[1:]
        xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
        yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
        xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
        yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
        inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1)
        area_i = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
        area_r = (boxes[rest, 2] - boxes[rest, 0]) * (boxes[rest, 3] - boxes[rest, 1])
        iou = inter / (area_i + area_r - inter + 1e-6)
        order = rest[iou < iou_thresh]
    return keep


def sahi_detect(detector, image, opts, tile_size=1024, overlap=0.15):
    """Sliding-window detection with global NMS merge.

    Returns (detections, num_tiles).
    """
    W, H = image.size
    stride = int(tile_size * (1 - overlap))

    # Build tile grid
    tiles = []
    for y in range(0, H, stride):
        for x in range(0, W, stride):
            x1, y1 = min(x + tile_size, W), min(y + tile_size, H)
            x0, y0 = max(0, x1 - tile_size), max(0, y1 - tile_size)
            tiles.append((x0, y0, x1, y1))

    all_dets = []
    for x0, y0, x1, y1 in tiles:
        crop = image.crop((x0, y0, x1, y1))
        for d in detector.detect(crop, opts):
            bx0, by0, bx1, by1 = d["bbox"]
            d["bbox"] = (bx0 + x0, by0 + y0, bx1 + x0, by1 + y0)
            all_dets.append(d)

    if not all_dets:
        return all_dets, len(tiles)

    boxes = np.array([d["bbox"] for d in all_dets])
    scores = np.array([d["confidence"] for d in all_dets])
    keep = nms(boxes, scores, opts.nms_threshold)
    return [all_dets[i] for i in keep], len(tiles)


def match_detections(gt_boxes, detections, dist_thresh=20):
    """Match GT boxes to detections by center distance.

    Returns (matched_gt_indices, matched_det_indices, unmatched_gt, unmatched_det).
    """
    matched_gt, matched_det = [], []
    used_dets = set()

    for gi, gt in enumerate(gt_boxes):
        gx = (gt["bbox"][0] + gt["bbox"][2]) / 2
        gy = (gt["bbox"][1] + gt["bbox"][3]) / 2
        best_di, best_dist = None, float("inf")
        for di, d in enumerate(detections):
            if di in used_dets:
                continue
            dx = (d["bbox"][0] + d["bbox"][2]) / 2
            dy = (d["bbox"][1] + d["bbox"][3]) / 2
            dist = ((gx - dx) ** 2 + (gy - dy) ** 2) ** 0.5
            if dist < best_dist:
                best_dist = dist
                best_di = di
        if best_di is not None and best_dist < dist_thresh:
            matched_gt.append(gi)
            matched_det.append(best_di)
            used_dets.add(best_di)

    unmatched_gt = [i for i in range(len(gt_boxes)) if i not in matched_gt]
    unmatched_det = [i for i in range(len(detections)) if i not in used_dets]
    return matched_gt, matched_det, unmatched_gt, unmatched_det


def draw_results(image, detections, gt_boxes=None):
    """Draw detections (red) and GT (green=unchecked, lime=checked) on image."""
    vis = image.copy()
    draw = ImageDraw.Draw(vis)

    if gt_boxes:
        for cb in gt_boxes:
            x0, y0, x1, y1 = cb["bbox"]
            color = "lime" if cb.get("state") == "checked" else "green"
            for i in range(2):
                draw.rectangle([x0 - i, y0 - i, x1 + i, y1 + i], outline=color)

    for d in detections:
        x0, y0, x1, y1 = d["bbox"]
        pad = 3 if gt_boxes else 0  # Offset so both are visible
        for i in range(2):
            draw.rectangle(
                [x0 - i - pad, y0 - i - pad, x1 + i + pad, y1 + i + pad],
                outline="red",
            )

    return vis


def eval_training_docs(args):
    """Evaluate on training documents with ground truth annotations."""
    from natural_pdf import PDF

    detector, opts = get_detector()
    opts.confidence = args.confidence

    ann_dir = Path("datasets/checkbox/raw/page_annotations")
    pdf_dir = Path("datasets/checkbox/raw/pdfs")
    out_dir = Path("temp/ffdnet_eval")
    out_dir.mkdir(parents=True, exist_ok=True)

    # Collect pages with annotations
    ann_files = sorted(ann_dir.glob("*.json"))
    ann_files = [f for f in ann_files if not f.name.endswith(".raw.json")]

    # Group by doc
    docs = {}
    for f in ann_files:
        parts = f.stem.rsplit("_p", 1)
        if len(parts) != 2:
            continue
        doc_id, page_str = parts
        page_idx = int(page_str)
        with open(f) as fh:
            data = json.load(fh)
        cbs = data.get("checkboxes", [])
        if not cbs:
            continue
        docs.setdefault(doc_id, []).append((page_idx, cbs))

    # Pick a diverse sample
    sample = []
    for doc_id, pages in sorted(docs.items()):
        for page_idx, gt_boxes in pages[:2]:  # Max 2 pages per doc
            sample.append((doc_id, page_idx, gt_boxes))
    sample = sample[:20]  # Cap at 20 pages

    print(f"Evaluating {len(sample)} pages from {len(docs)} documents\n")

    totals = {"gt": 0, "gt_checked": 0, "gt_unchecked": 0, "det": 0,
              "matched": 0, "matched_checked": 0, "matched_unchecked": 0,
              "false_pos": 0, "missed": 0, "missed_checked": 0, "missed_unchecked": 0}

    for doc_id, page_idx, gt_boxes in tqdm(sample, desc="Evaluating"):
        pdf_path = pdf_dir / f"{doc_id}.pdf"
        if not pdf_path.exists():
            continue

        pdf = PDF(str(pdf_path))
        if page_idx >= len(pdf.pages):
            pdf.close()
            continue

        img = pdf.pages[page_idx].render(resolution=300)

        # Scale GT from staging coords (~770x1000) to 300 DPI coords
        staging_path = Path(f"datasets/checkbox/images/staging/{doc_id}_p{page_idx:03d}.png")
        if staging_path.exists():
            staging_img = Image.open(staging_path)
            sx = img.width / staging_img.width
            sy = img.height / staging_img.height
            for cb in gt_boxes:
                cb["bbox"] = [
                    cb["bbox"][0] * sx, cb["bbox"][1] * sy,
                    cb["bbox"][2] * sx, cb["bbox"][3] * sy,
                ]

        t0 = time.time()
        dets, n_tiles = sahi_detect(
            detector, img, opts,
            tile_size=args.tile_size, overlap=args.overlap,
        )
        elapsed = time.time() - t0

        # Match
        m_gt, m_det, um_gt, um_det = match_detections(gt_boxes, dets, dist_thresh=30)

        gt_chk = sum(1 for cb in gt_boxes if cb.get("state") == "checked")
        gt_unchk = len(gt_boxes) - gt_chk
        missed_chk = sum(1 for i in um_gt if gt_boxes[i].get("state") == "checked")
        missed_unchk = len(um_gt) - missed_chk
        matched_chk = gt_chk - missed_chk
        matched_unchk = gt_unchk - missed_unchk

        totals["gt"] += len(gt_boxes)
        totals["gt_checked"] += gt_chk
        totals["gt_unchecked"] += gt_unchk
        totals["det"] += len(dets)
        totals["matched"] += len(m_gt)
        totals["matched_checked"] += matched_chk
        totals["matched_unchecked"] += matched_unchk
        totals["false_pos"] += len(um_det)
        totals["missed"] += len(um_gt)
        totals["missed_checked"] += missed_chk
        totals["missed_unchecked"] += missed_unchk

        # Save visualization
        vis = draw_results(img, dets, gt_boxes)
        vis.save(out_dir / f"{doc_id}_p{page_idx:03d}.png")

        tqdm.write(
            f"  {doc_id} p{page_idx}: "
            f"GT={len(gt_boxes)} ({gt_chk}chk/{gt_unchk}unchk) "
            f"Det={len(dets)} "
            f"Recall={len(m_gt)}/{len(gt_boxes)} "
            f"MissedChk={missed_chk} MissedUnchk={missed_unchk} "
            f"FP={len(um_det)} "
            f"({elapsed:.1f}s, {n_tiles} tiles)"
        )

        pdf.close()

    # Summary
    print("\n" + "=" * 60)
    print("SUMMARY")
    print("=" * 60)
    recall = totals["matched"] / totals["gt"] if totals["gt"] else 0
    recall_chk = totals["matched_checked"] / totals["gt_checked"] if totals["gt_checked"] else 0
    recall_unchk = totals["matched_unchecked"] / totals["gt_unchecked"] if totals["gt_unchecked"] else 0
    precision = totals["matched"] / totals["det"] if totals["det"] else 0

    print(f"  GT total:       {totals['gt']} ({totals['gt_checked']} checked, {totals['gt_unchecked']} unchecked)")
    print(f"  Detections:     {totals['det']}")
    print(f"  Matched:        {totals['matched']}")
    print(f"  Recall overall: {recall:.1%}")
    print(f"  Recall checked: {recall_chk:.1%}  ({totals['matched_checked']}/{totals['gt_checked']})")
    print(f"  Recall unchecked: {recall_unchk:.1%}  ({totals['matched_unchecked']}/{totals['gt_unchecked']})")
    print(f"  Precision:      {precision:.1%}")
    print(f"  False positives: {totals['false_pos']}")
    print(f"  Missed checked:  {totals['missed_checked']}")
    print(f"  Missed unchecked: {totals['missed_unchecked']}")
    print(f"\n  Visualizations saved to: {out_dir}/")


def eval_single_pdf(args):
    """Evaluate on a single PDF."""
    from natural_pdf import PDF

    detector, opts = get_detector()
    opts.confidence = args.confidence

    out_dir = Path("temp/ffdnet_eval")
    out_dir.mkdir(parents=True, exist_ok=True)

    pdf = PDF(args.pdf)
    page = pdf.pages[args.page]
    img = page.render(resolution=300)
    print(f"Page {args.page}: {img.size}")

    t0 = time.time()
    dets, n_tiles = sahi_detect(
        detector, img, opts,
        tile_size=args.tile_size, overlap=args.overlap,
    )
    elapsed = time.time() - t0

    print(f"  {len(dets)} detections, {n_tiles} tiles, {elapsed:.1f}s")
    for d in sorted(dets, key=lambda x: (x["bbox"][1], x["bbox"][0])):
        x0, y0, x1, y1 = d["bbox"]
        print(f"  conf={d['confidence']:.3f}  [{x0:.0f},{y0:.0f},{x1:.0f},{y1:.0f}]")

    vis = draw_results(img, dets)
    out = out_dir / f"single_{Path(args.pdf).stem}_p{args.page}.png"
    vis.save(out)
    print(f"  Saved: {out}")

    pdf.close()


def main():
    parser = argparse.ArgumentParser(description="Evaluate FFDNet-L + SAHI on checkboxes")
    parser.add_argument("--tile-size", type=int, default=1024)
    parser.add_argument("--overlap", type=float, default=0.15)
    parser.add_argument("--confidence", type=float, default=0.2)
    parser.add_argument("--pdf", type=str, default=None, help="Single PDF to evaluate")
    parser.add_argument("--page", type=int, default=0, help="Page index for single PDF")
    args = parser.parse_args()

    if args.pdf:
        eval_single_pdf(args)
    else:
        eval_training_docs(args)


if __name__ == "__main__":
    main()
