#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
单张图片快速测试：三角形四角标记识别 + 单应性落点 + PnP 估距

用法（在板子上）：
  python3 test/test_triangle_one_image.py --image /root/phot/xxx.jpg --out /root/phot/tri_out.jpg

调参对比（不改代码，临时覆盖 config.TRIANGLE_*）：
  python3 test/test_triangle_one_image.py --image /root/phot/xxx.jpg --preset shadow
  python3 test/test_triangle_one_image.py --image /root/phot/xxx.jpg --max-interior-gray 160 --min-dark-ratio 0.20
"""

import argparse
import json
import os
import time
from typing import Any, Dict, Tuple

import cv2
import numpy as np

import config
import triangle_target as tri_mod
from triangle_target import (
    detect_triangle_markers,
    load_camera_from_xml,
    load_triangle_positions,
    try_triangle_scoring,
)


def _apply_overrides(args) -> None:
    # 预设：阴影/低对比度场景更宽松（尽量保持速度：不启 CLAHE）
    if args.preset == "shadow":
        setattr(config, "TRIANGLE_ENABLE_CLAHE_FALLBACK", False)
        setattr(config, "TRIANGLE_MIN_CONTRAST_DIFF", 0)
        setattr(config, "TRIANGLE_MAX_INTERIOR_GRAY", 160)
        setattr(config, "TRIANGLE_DARK_PIXEL_GRAY", 160)
        setattr(config, "TRIANGLE_MIN_DARK_RATIO", 0.20)
        # adaptive 只在 Otsu 失败时尝试，保持尝试次数很少
        setattr(config, "TRIANGLE_ADAPTIVE_BLOCK_SIZES", (21,))

    # 手动覆盖（优先级高于 preset）
    if args.max_interior_gray is not None:
        setattr(config, "TRIANGLE_MAX_INTERIOR_GRAY", int(args.max_interior_gray))
    if args.dark_pixel_gray is not None:
        setattr(config, "TRIANGLE_DARK_PIXEL_GRAY", int(args.dark_pixel_gray))
    if args.min_dark_ratio is not None:
        setattr(config, "TRIANGLE_MIN_DARK_RATIO", float(args.min_dark_ratio))
    if args.min_contrast_diff is not None:
        setattr(config, "TRIANGLE_MIN_CONTRAST_DIFF", int(args.min_contrast_diff))
    if args.detect_scale is not None:
        setattr(config, "TRIANGLE_DETECT_SCALE", float(args.detect_scale))
    if args.adaptive_blocks is not None:
        bs = tuple(int(x) for x in args.adaptive_blocks.split(",") if x.strip())
        setattr(config, "TRIANGLE_ADAPTIVE_BLOCK_SIZES", bs)


def _dump_config() -> Dict[str, Any]:
    keys = [
        "TRIANGLE_DETECT_SCALE",
        "TRIANGLE_SIZE_RANGE",
        "TRIANGLE_MAX_INTERIOR_GRAY",
        "TRIANGLE_DARK_PIXEL_GRAY",
        "TRIANGLE_MIN_DARK_RATIO",
        "TRIANGLE_MIN_CONTRAST_DIFF",
        "TRIANGLE_ADAPTIVE_BLOCK_SIZES",
        "TRIANGLE_MAX_FILTERED_FOR_COMBO",
        "TRIANGLE_EARLY_EXIT_CANDIDATES",
        "TRIANGLE_ENABLE_CLAHE_FALLBACK",
    ]
    out = {}
    for k in keys:
        out[k] = getattr(config, k, None)
    return out


def _draw_tri_debug(img_bgr: np.ndarray, tri: Dict[str, Any]) -> np.ndarray:
    out = img_bgr.copy()
    markers = tri.get("markers") or []

    # 画三角形轮廓 + center + id
    for m in markers:
        corners = np.array(m.get("corners", []), dtype=np.int32)
        if corners.size == 0:
            continue
        cv2.polylines(out, [corners], True, (0, 255, 0), 2)
        c = m.get("center") or (corners[:, 0].mean(), corners[:, 1].mean())
        cx, cy = int(c[0]), int(c[1])
        cv2.circle(out, (cx, cy), 4, (0, 0, 255), -1)
        mid = m.get("id", "?")
        cv2.putText(out, f"T{mid}", (cx - 18, cy - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 255, 0), 1)

    # 若有 homography，画靶心（把 (0,0) 反投影到图像）
    H = tri.get("homography")
    if H is not None:
        try:
            H = np.array(H, dtype=np.float64)
            H_inv = np.linalg.inv(H)
            c_img = cv2.perspectiveTransform(np.array([[[0.0, 0.0]]], dtype=np.float32), H_inv)[0][0]
            ocx, ocy = int(c_img[0]), int(c_img[1])
            cv2.circle(out, (ocx, ocy), 5, (0, 0, 255), -1)
            cv2.circle(out, (ocx, ocy), 10, (0, 0, 255), 1)
        except Exception:
            pass

    # 叠加结果信息
    lines = []
    if tri.get("ok"):
        lines.append("tri_ok=True")
        if tri.get("dx_cm") is not None and tri.get("dy_cm") is not None:
            lines.append(f"dx,dy=({tri['dx_cm']:.2f},{tri['dy_cm']:.2f})cm")
        if tri.get("distance_m") is not None:
            lines.append(f"dist={float(tri['distance_m']):.2f}m")
    else:
        lines.append("tri_ok=False")

    y0 = 22
    for i, t in enumerate(lines):
        cv2.putText(out, t, (10, y0 + i * 18), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 255, 0), 1)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--image", required=True, help="输入图片路径（jpg/png）")
    ap.add_argument("--out", default="", help="输出标注图片路径（可选）")
    ap.add_argument("--laser-x", type=int, default=-1, help="激光点 x（像素），默认用图像中心")
    ap.add_argument("--laser-y", type=int, default=-1, help="激光点 y（像素），默认用图像中心")
    ap.add_argument("--preset", choices=["", "shadow"], default="", help="调参预设（shadow=阴影更鲁棒，不启 CLAHE）")
    ap.add_argument("--max-interior-gray", type=int, default=None)
    ap.add_argument("--dark-pixel-gray", type=int, default=None)
    ap.add_argument("--min-dark-ratio", type=float, default=None)
    ap.add_argument("--min-contrast-diff", type=int, default=None)
    ap.add_argument("--detect-scale", type=float, default=None)
    ap.add_argument("--adaptive-blocks", default=None, help="例如: 11,21 ；为空表示不改")
    ap.add_argument("--verbose", action="store_true", help="输出更多检测阶段信息")
    args = ap.parse_args()

    _apply_overrides(args)
    # triangle_target.py 的日志默认写到 logger_manager；在离线脚本里 logger 可能未初始化。
    # verbose 模式下把 _log 重定向为 print，方便直接看到诊断信息。
    if args.verbose:
        try:
            tri_mod._log = lambda msg: print(str(msg))
        except Exception:
            pass

    img_bgr = cv2.imread(args.image, cv2.IMREAD_COLOR)
    if img_bgr is None:
        raise SystemExit(f"读图失败：{args.image}")
    # triangle_target.try_triangle_scoring 约定输入为 RGB；OpenCV imread 返回 BGR
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

    h, w = img_bgr.shape[:2]
    if args.laser_x >= 0 and args.laser_y >= 0:
        laser_point = (int(args.laser_x), int(args.laser_y))
    else:
        laser_point = (w // 2, h // 2)

    K, dist = load_camera_from_xml(getattr(config, "CAMERA_CALIB_XML", ""))
    pos = load_triangle_positions(getattr(config, "TRIANGLE_POSITIONS_JSON", ""))

    print("[tri-test] image:", args.image, "shape:", (h, w))
    print("[tri-test] laser_point:", laser_point)
    print("[tri-test] calib_ok:", bool(K is not None and dist is not None), "pos_ok:", bool(pos))
    print("[tri-test] config:", json.dumps(_dump_config(), ensure_ascii=False))

    # 先单独跑一次三角形候选检测，便于区分“没找到候选” vs “找到候选但评分/单应性失败”
    scale = float(getattr(config, "TRIANGLE_DETECT_SCALE", 0.5) or 0.5)
    if not (0.05 <= scale <= 1.0):
        scale = 0.5
    long_side = max(h, w)
    max_dim = max(64, int(long_side * scale))
    if long_side > max_dim:
        det_scale = max_dim / long_side
        det_w = int(w * det_scale)
        det_h = int(h * det_scale)
        img_det = cv2.resize(img_bgr, (det_w, det_h), interpolation=cv2.INTER_LINEAR)
        inv_scale = 1.0 / det_scale
        size_range_det = (
            max(4, int(getattr(config, "TRIANGLE_SIZE_RANGE", (8, 500))[0] * det_scale)),
            max(8, int(getattr(config, "TRIANGLE_SIZE_RANGE", (8, 500))[1] * det_scale)),
        )
    else:
        img_det = img_bgr
        inv_scale = 1.0
        size_range_det = getattr(config, "TRIANGLE_SIZE_RANGE", (8, 500))

    gray = cv2.cvtColor(img_det, cv2.COLOR_BGR2GRAY)
    markers_det = detect_triangle_markers(
        gray,
        orig_gray=gray,
        size_range=size_range_det,
        verbose=bool(args.verbose),
    )
    if inv_scale != 1.0 and markers_det:
        for m in markers_det:
            m["center"] = [m["center"][0] * inv_scale, m["center"][1] * inv_scale]
            m["corners"] = [[c[0] * inv_scale, c[1] * inv_scale] for c in m["corners"]]

    print("[tri-test] markers_found:", len(markers_det), "ids:", [m.get("id") for m in markers_det])

    t0 = time.time()
    tri = try_triangle_scoring(
        img_rgb,  # try_triangle_scoring 期望 RGB
        laser_point,
        pos,
        K,
        dist,
        size_range=getattr(config, "TRIANGLE_SIZE_RANGE", (8, 500)),
    )
    dt_ms = int(round((time.time() - t0) * 1000))

    print("[tri-test] elapsed_ms:", dt_ms)
    print(json.dumps(tri, ensure_ascii=False, indent=2))

    if args.out:
        out_path = args.out
        # 允许传目录（如 ./），自动生成文件名；未带扩展名时默认 .jpg
        if out_path.endswith("/") or out_path.endswith("\\") or os.path.isdir(out_path):
            out_path = os.path.join(out_path, "tri_out.jpg")
        root, ext = os.path.splitext(out_path)
        if not ext:
            out_path = root + ".jpg"

        # 若 try_triangle_scoring 失败且没带回 markers，至少把候选 markers 画出来，方便肉眼判断
        tri_for_draw = tri if isinstance(tri, dict) else {"ok": False}
        if not tri_for_draw.get("markers") and markers_det:
            tri_for_draw = dict(tri_for_draw)
            tri_for_draw["markers"] = markers_det
        out_img = _draw_tri_debug(img_bgr, tri_for_draw)
        ok = cv2.imwrite(out_path, out_img)
        if not ok:
            raise SystemExit(f"写图失败（可能是不支持的扩展名）：{out_path}")
        print("[tri-test] wrote:", out_path)


if __name__ == "__main__":
    main()

