#!/usr/bin/env python3
"""
5_camera_z_refine.py
-----------------------
Refines camera z-positions using the FK-predicted z of the Ellbow link.

Background
----------
Camera poses from step 2 are estimated from board markers alone.  Board
markers are roughly coplanar (z ≈ 0), so the camera's z-height above the
board is poorly constrained and can be 20–40 % off.

Once we have the elbow angle (step 3b, from 4b_revolute_angle.py),
FK predicts where the Ellbow markers *should* be in world space.  The
difference between FK-predicted z and triangulated z gives a global
z-offset for the cameras.

Algorithm
---------
1. Load elbow-angle result (v8_ellbow_angle.json) → joint state dict
2. Run FK with that state → predicted world positions for Ellbow markers (mm)
3. Load triangulated Ellbow marker positions from aruco_positions_optimized.json
4. delta_z_mm = median(fk_z - triangulated_z)  over matched markers
5. For each camera pose file:
   a. camera center world  C = -R_wc.T @ t_wc
   b. C_new = C + [0, 0, delta_z_mm / 1000]
   c. t_wc_new = -R_wc @ C_new
   d. Save as *_camera_pose_v8.json
6. Save v8_z_correction.json summary

Usage
-----
python 5_camera_z_refine.py \\
    --angle  path/to/v8_ellbow_angle.json \\
    --robot  path/to/robot.json \\
    --aruco  path/to/aruco_positions_optimized.json \\
    -pose    path/to/render_a_camera_pose.json \\
    -pose    path/to/render_b_camera_pose.json \\
    --outDir path/to/output_dir \\
    [--elbowLink Ellbow]
"""

from __future__ import annotations

import argparse
import copy
import json
import os
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional

import numpy as np

# robot_fk lives in the same directory as this script
sys.path.insert(0, str(Path(__file__).parent))
from robot_fk import RobotFK


# ──────────────────────────────────────────────────────────────
# I/O helpers
# ──────────────────────────────────────────────────────────────

def load_json(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def save_json(path: str, data: dict) -> None:
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)


def load_triangulated_mm(aruco_json: dict) -> Dict[int, np.ndarray]:
    """Read marker world positions (mm) from aruco_positions_*.json."""
    result: Dict[int, np.ndarray] = {}
    for m in aruco_json.get("markers", []):
        mid = int(m.get("marker_id", m.get("id", -1)))
        if mid < 0:
            continue
        if "position_mm" in m:
            result[mid] = np.array(m["position_mm"], dtype=float)
        elif "position_m" in m:
            result[mid] = np.array(m["position_m"], dtype=float) * 1000.0
    return result


# ──────────────────────────────────────────────────────────────
# Main
# ──────────────────────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(
        description="5v8: refine camera z-positions using FK-predicted Ellbow marker z"
    )
    parser.add_argument("--angle",     required=True,
                        help="v8_ellbow_angle.json from 4b_revolute_angle.py")
    parser.add_argument("--robot",     required=True, help="robot.json")
    parser.add_argument("--aruco",     required=True,
                        help="aruco_positions_optimized.json (first-pass triangulation)")
    parser.add_argument("-pose", "--poses", action="append", required=True,
                        help="*_camera_pose.json files (repeat for each camera)")
    parser.add_argument("--outDir",    required=True, help="Output directory")
    parser.add_argument("--elbowLink", default="Ellbow",
                        help="Link name used for z estimation (default: Ellbow)")
    args = parser.parse_args()

    os.makedirs(args.outDir, exist_ok=True)

    # ── 1. Load angle / state ────────────────────────────────
    angle_data = load_json(args.angle)
    if angle_data.get("status") != "ok":
        print(f"[ERROR] Angle estimation was not ok: {angle_data.get('reason', '?')}")
        sys.exit(1)

    accumulated_state = angle_data.get("accumulated_state", {})
    elbow_link        = angle_data.get("link", args.elbowLink)

    print(f"[INFO] Joint state from 4b:  {accumulated_state}")
    print(f"[INFO] Link used for z-ref:  {elbow_link}")

    # ── 2. FK prediction ─────────────────────────────────────
    fk = RobotFK.from_file(args.robot)
    T  = fk.compute(accumulated_state)
    all_fk = fk.all_markers_world(T)      # marker_id -> {world_mm, link, …}

    elbow_fk = {mid: v for mid, v in all_fk.items() if v["link"] == elbow_link}

    if not elbow_fk:
        print(f"[ERROR] No FK markers found for link '{elbow_link}'")
        sys.exit(1)

    print(f"[INFO] FK markers for '{elbow_link}': {sorted(elbow_fk.keys())}")

    # ── 3. Load triangulated positions ───────────────────────
    aruco_data    = load_json(args.aruco)
    triangulated  = load_triangulated_mm(aruco_data)

    # ── 4. Compute z-deltas ──────────────────────────────────
    z_deltas: List[float] = []
    print("\n[INFO] Per-marker z comparison (FK vs triangulated):")
    for mid in sorted(elbow_fk.keys()):
        fk_z = float(elbow_fk[mid]["world_mm"][2])
        if mid not in triangulated:
            print(f"  Marker {mid:4d}: FK z={fk_z:8.1f} mm   [NOT triangulated – skip]")
            continue
        obs_z = float(triangulated[mid][2])
        delta = fk_z - obs_z
        z_deltas.append(delta)
        print(f"  Marker {mid:4d}: FK z={fk_z:8.1f} mm   obs z={obs_z:8.1f} mm   Δz={delta:+7.1f} mm")

    if not z_deltas:
        print(f"\n[ERROR] No matched markers for '{elbow_link}' — cannot compute z correction.")
        sys.exit(1)

    delta_z_mm = float(np.median(z_deltas))
    print(f"\n[INFO] z-correction (median): {delta_z_mm:+.1f} mm  "
          f"(from {len(z_deltas)} markers)")
    if abs(delta_z_mm) < 1.0:
        print("[INFO] Correction < 1 mm — cameras already well-calibrated in z.")

    # ── 5. Apply correction to each camera pose ──────────────
    corrected_files: List[str] = []

    for pose_file in args.poses:
        pose_data = load_json(pose_file)

        pose_section = pose_data.get("camera_pose", {}) or {}
        w2c          = pose_section.get("world_to_camera", {}) or {}

        R_wc = np.array(w2c.get("rotation_matrix", []), dtype=float).reshape(3, 3)
        t_wc = np.array(w2c.get("translation_m",   []), dtype=float).reshape(3)

        # Camera centre in world (metres)
        C_world     = -R_wc.T @ t_wc
        C_world_new =  C_world + np.array([0.0, 0.0, delta_z_mm / 1000.0])

        t_wc_new = -R_wc @ C_world_new

        pose_new = copy.deepcopy(pose_data)
        pose_new["camera_pose"]["world_to_camera"]["translation_m"] = \
            [float(v) for v in t_wc_new.tolist()]
        pose_new["camera_pose"]["camera_in_world"]["position_m"] = \
            [float(v) for v in C_world_new.tolist()]
        pose_new["camera_pose"]["camera_in_world"]["position_mm"] = \
            [float(v * 1000.0) for v in C_world_new.tolist()]
        pose_new["v8_z_correction_mm"] = delta_z_mm
        pose_new["created_utc"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())

        basename = os.path.basename(pose_file)
        out_name = basename.replace("_camera_pose.json", "_camera_pose_v8.json")
        if out_name == basename:
            out_name = os.path.splitext(basename)[0] + "_v8.json"

        out_path = os.path.join(args.outDir, out_name)
        save_json(out_path, pose_new)
        print(f"[INFO] Saved: {out_path}")
        corrected_files.append(out_path)

    # ── 6. Save correction summary ───────────────────────────
    summary = {
        "schema_version":     "1.0",
        "created_utc":        time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "elbow_link":         elbow_link,
        "z_correction_mm":    delta_z_mm,
        "n_markers_used":     len(z_deltas),
        "z_deltas_mm":        z_deltas,
        "corrected_pose_files": corrected_files,
    }
    summary_path = os.path.join(args.outDir, "v8_z_correction.json")
    save_json(summary_path, summary)
    print(f"\n[INFO] Correction summary → {summary_path}")


if __name__ == "__main__":
    main()
