#!/usr/bin/env python3
"""
3b_corner_marker_poses.py
=========================
Produktiver Pipeline-Schritt: leitet aus den 4 ArUco-Ecken jedes Markers eine
volle Marker-Pose ab (Position + gemessene Normale), statt nur den Center zu
triangulieren.

Validiert in benchmark/stage0_corner_normals.py: die aus triangulierten Ecken
abgeleitete Normale ist ~1 deg genau (Median), auch fuer Finger-Marker.

Input:
  --evalDir   Ordner mit render_*_aruco_detection.json + _camera_pose.json
  --robot     robot.json (fuer marker_id -> link Zuordnung)
Output:
  <evalDir>/aruco_marker_poses.json   (pro Marker: position, gemessene normal,
                                        4 triangulierte Ecken, #Kameras, Kantenlaenge)

Das Format ist kompatibel mit robot_viewer.html (marker_id, position_m/mm, normal)
und mit 9_evaluateMarker.py (position_m), erweitert um die gemessene Orientierung.
"""
from __future__ import annotations

import argparse
import glob
import json
import os
import re
import time
from typing import Dict, List, Tuple

import numpy as np
import cv2


# ------------------------------------------------------------------
# Loading
# ------------------------------------------------------------------

def load_cameras(eval_dir: str) -> Dict[str, dict]:
    cams: Dict[str, dict] = {}
    for det_path in glob.glob(os.path.join(eval_dir, "*_aruco_detection.json")):
        base = os.path.basename(det_path)
        m = re.match(r"render_([A-Za-z0-9]+)_aruco_detection\.json", base)
        if not m:
            continue
        cam_id = m.group(1)
        pose_path = os.path.join(eval_dir, f"render_{cam_id}_camera_pose.json")
        if not os.path.exists(pose_path):
            print(f"[WARN] no pose for camera {cam_id}, skipping")
            continue
        det = json.load(open(det_path, "r", encoding="utf-8"))
        pose = json.load(open(pose_path, "r", encoding="utf-8"))
        K = np.array(det["camera"]["camera_matrix"], dtype=float).reshape(3, 3)
        D = np.array(det["camera"]["distortion_coefficients"], dtype=float).reshape(-1, 1)
        w2c = pose["camera_pose"]["world_to_camera"]
        R = np.array(w2c["rotation_matrix"], dtype=float).reshape(3, 3)
        t = np.array(w2c["translation_m"], dtype=float).reshape(3)
        markers: Dict[int, np.ndarray] = {}
        for d in det.get("detections", []):
            pts = d.get("image_points_px")
            if pts is not None:
                markers[int(d["marker_id"])] = np.array(pts, dtype=float).reshape(4, 2)
        cams[cam_id] = dict(K=K, D=D, R=R, t=t, markers=markers)
    return cams


def load_marker_links(robot_path: str) -> Dict[int, str]:
    robot = json.load(open(robot_path, "r", encoding="utf-8"))
    out: Dict[int, str] = {}
    for link_name, link in (robot.get("links", {}) or {}).items():
        for mk in link.get("markers", []) or []:
            mid = int(mk.get("id", -1))
            if mid >= 0:
                out[mid] = link_name
    return out


# ------------------------------------------------------------------
# Geometry  (validated in stage0)
# ------------------------------------------------------------------

def triangulate_multiview(observations) -> np.ndarray:
    A = []
    for K, D, R, t, uv in observations:
        und = cv2.undistortPoints(np.array([[uv]], dtype=np.float32), K, D).reshape(2)
        x, y = float(und[0]), float(und[1])
        P = np.hstack([R, t.reshape(3, 1)])
        A.append(x * P[2] - P[0])
        A.append(y * P[2] - P[1])
    _, _, Vt = np.linalg.svd(np.asarray(A, dtype=float))
    X = Vt[-1]
    return np.array([np.nan] * 3) if abs(X[3]) < 1e-12 else X[:3] / X[3]


def corner_plane_normal(corners3d: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    center = corners3d.mean(axis=0)
    _, _, Vt = np.linalg.svd(corners3d - center)
    n = Vt[-1]
    # ArUco corners clockwise from the front: outward (camera-facing) normal,
    # matching the Blender/robot.json convention, points opposite cross(e01,e02).
    cross = np.cross(corners3d[1] - corners3d[0], corners3d[2] - corners3d[0])
    if np.dot(n, cross) > 0:
        n = -n
    nn = np.linalg.norm(n)
    return (n / nn if nn > 1e-12 else n), center


# ------------------------------------------------------------------
# Main
# ------------------------------------------------------------------

def main() -> None:
    ap = argparse.ArgumentParser(description="Derive marker poses (position + measured normal) from ArUco corners")
    ap.add_argument("--evalDir", required=True, help="folder with detection + camera_pose JSONs")
    ap.add_argument("--robot", required=True, help="robot.json (for marker->link)")
    ap.add_argument("--minCams", type=int, default=2, help="min cameras to triangulate a marker")
    ap.add_argument("--out", default=None, help="output path (default <evalDir>/aruco_marker_poses.json)")
    args = ap.parse_args()

    cams = load_cameras(args.evalDir)
    if len(cams) < 2:
        print("[ERROR] need >=2 cameras")
        return
    links = load_marker_links(args.robot)
    print(f"[INFO] Cameras: {sorted(cams.keys())} | marker-link entries: {len(links)}")

    marker_cams: Dict[int, List[str]] = {}
    for cid, cam in cams.items():
        for mid in cam["markers"]:
            marker_cams.setdefault(mid, []).append(cid)

    markers_out = []
    for mid, cam_ids in sorted(marker_cams.items()):
        if len(cam_ids) < args.minCams:
            continue
        corners3d, ok = [], True
        for ci in range(4):
            obs = [(cams[c]["K"], cams[c]["D"], cams[c]["R"], cams[c]["t"], cams[c]["markers"][mid][ci])
                   for c in cam_ids]
            X = triangulate_multiview(obs)
            if not np.all(np.isfinite(X)):
                ok = False
                break
            corners3d.append(X)
        if not ok:
            continue
        corners3d = np.array(corners3d)
        normal, center = corner_plane_normal(corners3d)
        edge_mm = float(np.mean([np.linalg.norm(corners3d[(i + 1) % 4] - corners3d[i]) for i in range(4)]) * 1000.0)

        markers_out.append({
            "marker_id": int(mid),
            "link": links.get(mid, "unknown"),
            "position_m": [float(v) for v in center],
            "position_mm": [float(v * 1000.0) for v in center],
            "normal": [float(v) for v in normal],
            "corners_m": [[float(v) for v in c] for c in corners3d],
            "num_cameras": len(cam_ids),
            "edge_length_mm": edge_mm,
        })

    # camera poses in world (for viewer frusta): centre C = -R^T t, view axis = R[2]
    cameras_out = []
    for cid in sorted(cams.keys()):
        cam = cams[cid]
        C = -cam["R"].T @ cam["t"]
        cameras_out.append({
            "camera_id": cid,
            "position_m": [float(v) for v in C],
            "position_mm": [float(v * 1000.0) for v in C],
            "direction": [float(v) for v in cam["R"][2]],
        })

    out_path = args.out or os.path.join(args.evalDir, "aruco_marker_poses.json")
    output = {
        "schema_version": "1.1",
        "stage": "corner_marker_poses",
        "created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "summary": {"num_cameras": len(cams), "num_markers": len(markers_out)},
        "cameras": cameras_out,
        "markers": markers_out,
    }
    json.dump(output, open(out_path, "w", encoding="utf-8"), indent=2)
    print(f"[INFO] {len(markers_out)} marker poses -> {out_path}")


if __name__ == "__main__":
    main()
