#!/usr/bin/env python3
"""
vl_compare.py — Run PaddleOCR-VL on one newspaper page and compare its
layout / reading-order output against the existing PaddleOCR LayoutDetection
regions (page_XXX.regions.json) produced by extractor.py.

Goal: see whether VL's reading-order blocks line up with article boundaries
better than the current "box soup" before deciding to swap engines.

USAGE
    # auto-pick the most recent output/<job>/pages/page_001.png and its regions.json
    python vl_compare.py

    # or point at a specific page + regions file
    python vl_compare.py output/20260529_004651/pages/page_001.png \
                         output/20260529_004651/pages/page_001.regions.json

OUTPUT (written next to the page image)
    page_XXX.vl_layout.json        normalized VL blocks (id, type, bbox, reading order, text)
    page_XXX.vl_debug.png          page with VL boxes + reading-order numbers drawn
    page_XXX.vl_vs_paddle.json     side-by-side summary (counts by type, totals)
    page_XXX.vl_raw.json           raw VL result (best-effort, for debugging)

NOTES
  * PaddleOCR-VL is a ~0.9B vision-language model — expect a GPU and a model
    download on first run. It is SEPARATE from the lightweight LayoutDetection
    used in extractor.py.
  * The VL entry point has moved across releases. This script tries several
    known import paths and prints a clear hint if none are found so you can
    adapt the one line to your installed version.
  * Run inside your project venv:  ./venv/bin/python vl_compare.py
"""

import os
import sys
import glob
import json
import traceback


# --------------------------------------------------------------------------
# 1. Resolve input page image + matching regions.json
# --------------------------------------------------------------------------
def resolve_inputs(argv):
    if len(argv) >= 2:
        page_png = argv[1]
        regions_json = argv[2] if len(argv) >= 3 else os.path.splitext(page_png)[0] + ".regions.json"
        return page_png, regions_json

    # auto-discover: most recently modified page_001.png under output/
    candidates = sorted(
        glob.glob("output/*/pages/page_001.png"),
        key=os.path.getmtime,
        reverse=True,
    )
    if not candidates:
        sys.exit("No page_001.png found under output/*/pages/. "
                 "Pass a page image path explicitly:  python vl_compare.py <page.png> [regions.json]")
    page_png = candidates[0]
    regions_json = os.path.splitext(page_png)[0] + ".regions.json"
    return page_png, regions_json


# --------------------------------------------------------------------------
# 2. Load the VL model (defensive across releases)
# --------------------------------------------------------------------------
def load_vl_model():
    import paddleocr
    print(f"paddleocr version: {getattr(paddleocr, '__version__', 'unknown')}")

    # Try the documented entry points, newest-first. Adjust if your release differs.
    attempts = [
        ("paddleocr.PaddleOCRVL", lambda: paddleocr.PaddleOCRVL()),
        ("paddleocr.PPDocVL",     lambda: paddleocr.PPDocVL()),
        ("paddleocr.DocVLM",      lambda: paddleocr.DocVLM()),
    ]
    for name, ctor in attempts:
        try:
            model = ctor()
            print(f"Loaded VL model via: {name}")
            return model
        except Exception as e:  # noqa: BLE001
            print(f"  - {name} unavailable: {e}")

    # Last resort: PaddleX pipeline name
    try:
        from paddlex import create_pipeline
        model = create_pipeline(pipeline="PaddleOCR-VL")
        print("Loaded VL model via: paddlex.create_pipeline('PaddleOCR-VL')")
        return model
    except Exception as e:  # noqa: BLE001
        print(f"  - paddlex pipeline unavailable: {e}")

    sys.exit(
        "\nCould not load PaddleOCR-VL with any known entry point.\n"
        "Check the exact class name for your installed version, e.g.:\n"
        "    python -c \"import paddleocr; print([n for n in dir(paddleocr) if 'VL' in n or 'Doc' in n])\"\n"
        "then edit the `attempts` list near the top of load_vl_model()."
    )


# --------------------------------------------------------------------------
# 3. Run VL + normalize its output to the same shape as regions.json
# --------------------------------------------------------------------------
def to_bbox(coord):
    """Accept [x1,y1,x2,y2] or polygon [x1,y1,...] -> [x1,y1,x2,y2] ints."""
    if coord is None:
        return [0, 0, 0, 0]
    if hasattr(coord, "tolist"):
        coord = coord.tolist()
    coord = [float(v) for v in coord]
    if len(coord) == 4:
        return [int(v) for v in coord]
    xs = coord[0::2]
    ys = coord[1::2]
    return [int(min(xs)), int(min(ys)), int(max(xs)), int(max(ys))]


def run_vl(model, page_png):
    if hasattr(model, "predict"):
        raw = model.predict(page_png)
    elif callable(model):
        raw = model(page_png)
    else:
        raw = model.predict(input=page_png)

    # raw is typically a list of result objects (dict-like). Iterate generically.
    results = list(raw) if not isinstance(raw, list) else raw

    blocks = []
    raw_dump = []

    for res in results:
        # try to get a plain dict for dumping/parsing
        d = None
        for attr in ("json", "res", "_to_dict"):
            try:
                v = getattr(res, attr, None)
                d = v() if callable(v) else v
                if d:
                    break
            except Exception:  # noqa: BLE001
                pass
        if d is None and isinstance(res, dict):
            d = res
        raw_dump.append(_jsonable(d if d is not None else str(res)))

        # VL layout blocks commonly live under one of these keys
        layout = None
        if isinstance(d, dict):
            for key in ("layout_det_res", "parsing_res_list", "layout", "blocks", "boxes"):
                if key in d and d[key]:
                    layout = d[key]
                    break
        if layout is None:
            continue

        # layout may itself be a dict containing 'boxes'
        if isinstance(layout, dict):
            layout = layout.get("boxes", layout.get("parsing_res_list", []))

        for i, b in enumerate(layout):
            if not isinstance(b, dict):
                continue
            coord = b.get("coordinate", b.get("bbox", b.get("block_bbox", b.get("layout_bbox"))))
            label = str(b.get("label", b.get("type", b.get("block_label", "unknown")))).lower()
            score = float(b.get("score", b.get("confidence", 0)) or 0)
            order = b.get("reading_order", b.get("order", b.get("index", i)))
            text = b.get("text", b.get("block_content", b.get("rec_text", "")))
            if isinstance(text, list):
                text = " ".join(map(str, text))
            blocks.append({
                "id": len(blocks) + 1,
                "type": label,
                "bbox": to_bbox(coord),
                "confidence": round(score, 4),
                "reading_order": order,
                "text_preview": (str(text)[:120] if text else ""),
            })

    # sort by reading order if available
    try:
        blocks.sort(key=lambda x: (x["reading_order"] if isinstance(x["reading_order"], (int, float)) else 0))
        for i, b in enumerate(blocks, 1):
            b["id"] = i
    except Exception:  # noqa: BLE001
        pass

    return blocks, raw_dump


def _jsonable(obj):
    try:
        json.dumps(obj)
        return obj
    except Exception:  # noqa: BLE001
        return str(obj)


# --------------------------------------------------------------------------
# 4. Compare against existing PaddleOCR regions
# --------------------------------------------------------------------------
def type_counts(regions):
    c = {}
    for r in regions:
        t = r.get("type", "unknown")
        c[t] = c.get(t, 0) + 1
    return dict(sorted(c.items(), key=lambda x: -x[1]))


def draw_debug(page_png, blocks, out_png):
    try:
        from PIL import Image, ImageDraw, ImageFont
    except Exception:  # noqa: BLE001
        print("Pillow not available — skipping debug image.")
        return
    img = Image.open(page_png).convert("RGB")
    draw = ImageDraw.Draw(img)
    try:
        font = ImageFont.truetype("DejaVuSans-Bold.ttf", 60)
    except Exception:  # noqa: BLE001
        font = ImageFont.load_default()
    for b in blocks:
        x1, y1, x2, y2 = b["bbox"]
        draw.rectangle([x1, y1, x2, y2], outline=(220, 30, 30), width=6)
        tag = f"{b['id']}:{b['type']}"
        draw.text((x1 + 8, y1 + 8), tag, fill=(220, 30, 30), font=font)
    img.save(out_png)
    print(f"  wrote {out_png}")


def main():
    page_png, regions_json = resolve_inputs(sys.argv)
    if not os.path.exists(page_png):
        sys.exit(f"Page image not found: {page_png}")
    print(f"Page image : {page_png}")
    print(f"Regions JSON: {regions_json}  (exists={os.path.exists(regions_json)})")

    stem = os.path.splitext(page_png)[0]

    # existing paddle regions
    paddle_regions = []
    if os.path.exists(regions_json):
        with open(regions_json) as f:
            paddle_regions = json.load(f).get("regions", [])

    # run VL
    print("\nLoading PaddleOCR-VL (first run downloads the model)...")
    model = load_vl_model()
    print("Running VL on the page...")
    blocks, raw_dump = run_vl(model, page_png)

    # write outputs
    with open(stem + ".vl_layout.json", "w") as f:
        json.dump({"page_image": os.path.basename(page_png), "blocks": blocks}, f, indent=2, ensure_ascii=False)
    with open(stem + ".vl_raw.json", "w") as f:
        json.dump(raw_dump, f, indent=2, ensure_ascii=False)

    summary = {
        "page_image": os.path.basename(page_png),
        "paddle_layoutdetection": {
            "total_regions": len(paddle_regions),
            "by_type": type_counts(paddle_regions),
        },
        "paddleocr_vl": {
            "total_blocks": len(blocks),
            "by_type": type_counts(blocks),
            "has_reading_order": any(isinstance(b.get("reading_order"), (int, float)) for b in blocks),
        },
    }
    with open(stem + ".vl_vs_paddle.json", "w") as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)

    draw_debug(page_png, blocks, stem + ".vl_debug.png")

    print("\n" + "=" * 60)
    print("COMPARISON SUMMARY")
    print("=" * 60)
    print(json.dumps(summary, indent=2, ensure_ascii=False))
    print("\nWrote:")
    for ext in (".vl_layout.json", ".vl_raw.json", ".vl_vs_paddle.json", ".vl_debug.png"):
        print("  " + stem + ext)
    print("\nNext: open page_XXX_regions_debug.png (existing) next to page_XXX.vl_debug.png")
    print("and eyeball whether VL's blocks track article boundaries better.")


if __name__ == "__main__":
    try:
        main()
    except SystemExit:
        raise
    except Exception:  # noqa: BLE001
        traceback.print_exc()
        sys.exit("\nFailed. If the error is about the VL class name, see load_vl_model().")