"""Inference with the hlt-geometric-vision release (ONNXRuntime + OpenCV, no PyTorch).

Path A: the vertices come out of the model (1 polygon, training domain = clean rectangles).
Path B: model edge_map + geometric extractor -> N rectangles on real photographs.

    python inference_example.py --model ../model.onnx --image photo.jpg [--nested]

The extractor is a copy of hlt/postproc/quads.py from the source repository (numpy + OpenCV,
no weights): see documentation.md section 6 for the specification, section 8 for the parameters.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort

from hlt_postproc_quads import QuadConfig, extract_quads, hybrid_edge_map


def preprocess(bgr, size=512):
    """Letterbox to a square, aspect preserved, with a REPLICATED border (a constant gray fill
    creates artificial lines in the edge map). Also returns the real-image-area mask."""
    h0, w0 = bgr.shape[:2]
    scale = float(size) / float(max(h0, w0))
    nw, nh = max(1, round(w0 * scale)), max(1, round(h0 * scale))
    small = cv2.resize(bgr, (nw, nh), interpolation=cv2.INTER_AREA)
    ox, oy = (size - nw) // 2, (size - nh) // 2
    disp = cv2.copyMakeBorder(small, oy, size - nh - oy, ox, size - nw - ox, cv2.BORDER_REPLICATE)
    valid = np.zeros((size, size), bool)
    valid[oy:oy + nh, ox:ox + nw] = True
    x = cv2.cvtColor(disp, cv2.COLOR_BGR2RGB).astype(np.float32).transpose(2, 0, 1) / 255.0
    return np.ascontiguousarray(x)[None], disp, valid, scale, ox, oy


def to_original(corners, scale, ox, oy):
    """512 frame -> pixels of the original image. The half-pixel correction is not cosmetic: on
    large photographs the mapping error reaches ~1.6 px."""
    return (corners + 0.5 - np.array([ox, oy], np.float64)) / scale - 0.5


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="model.onnx")
    ap.add_argument("--image", required=True)
    ap.add_argument("--nested", action="store_true", help="also keep nested rectangles")
    ap.add_argument("--only-net", action="store_true", help="path A only (model vertices)")
    ap.add_argument("--no-canny", action="store_true", help="network edge map only (no Canny)")
    ap.add_argument("--threads", type=int, default=0, help="ONNXRuntime threads (0 = default)")
    ap.add_argument("--max-quads", type=int, default=64,
                    help="cap on kept rectangles (QuadConfig default: 64; the release examples use 12)")
    args = ap.parse_args()

    so = ort.SessionOptions()
    if args.threads:
        so.intra_op_num_threads = args.threads
    sess = ort.InferenceSession(args.model, sess_options=so, providers=["CPUExecutionProvider"])
    in_name = sess.get_inputs()[0].name
    out_names = [o.name for o in sess.get_outputs()]

    img = cv2.imread(args.image, cv2.IMREAD_COLOR)
    if img is None:
        raise SystemExit("image not decoded: " + args.image)
    x, disp, valid, scale, ox, oy = preprocess(img)
    out = dict(zip(out_names, sess.run(out_names, {in_name: x})))

    # ---- path A: a single polygon, strong on the training domain
    mask = out["vertex_mask"][0, 0] > 0.5
    net_corners = to_original(out["vertices"][0, 0][mask], scale, ox, oy)

    result = {"image": args.image, "size_wh": [img.shape[1], img.shape[0]],
              "path_A_model_polygon": np.round(net_corners, 2).tolist(),
              "path_A_active_vertices": int(mask.sum()),
              "note": "the model confidence is not calibrated: do not use it as a threshold"}

    if not args.only_net:
        # ---- path B: edge_map + geometry (rectangles on real photographs)
        # The model edge map lives in the 512 letterbox frame: the helper Canny must be computed on
        # THAT frame, not on the original image (different areas and different scales).
        g512 = cv2.cvtColor(disp, cv2.COLOR_BGR2GRAY)
        er = cv2.erode(valid.astype(np.uint8), np.ones((5, 5), np.uint8)).astype(bool)
        if args.no_canny:
            edge = out["edge_map"][0, 0] * er
            cfg = QuadConfig(nested=args.nested, max_quads=args.max_quads)
        else:
            edge = hybrid_edge_map(out["edge_map"][0, 0], g512, valid=er)
            cfg = QuadConfig.hybrid(nested=args.nested, max_quads=args.max_quads)
        quads = extract_quads(edge, cfg, gray=g512.astype(np.float32))
        result["path_B_rectangles"] = [
            {"corners_xy": to_original(q["corners"], scale, ox, oy).round(2).tolist(),
             "score": round(float(q["score"]), 3),
             "side_edge_coverage": [round(float(v), 2) for v in q["edge_cov"]]}
            for q in quads]
        result["path_B_count"] = len(quads)

    print(json.dumps(result, ensure_ascii=False, indent=2))
    Path(args.image).with_suffix(".quads.json").write_text(
        json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")


if __name__ == "__main__":
    main()