#!/usr/bin/env python3
"""
3_board_anchor_patch.py
-----------------------
DROP-IN PATCH for  3_multiview_bundle_adjustment_v4.py

What it adds
------------
Board markers have KNOWN world positions (they are fixed to the physical board,
which defines the coordinate frame).  Letting the optimizer move them freely
makes them drift, which degrades all arm-marker estimates.

This patch adds a  PositionAnchorConstraint  that pins each observed Board
marker to its robot.json world position during bundle adjustment.

How to integrate
----------------
Copy the three sections labelled  [PATCH – copy into v4]  into the v4 file:

  1. The  PositionAnchorConstraint  dataclass  →  after the JointAxisConstraint dataclass
  2. Update the  Constraint  type alias         →  replace the existing line
  3. The  load_board_anchors()  function        →  anywhere before main()
  4. The change to  bundle_adjustment_residuals →  add one branch inside the for-loop
  5. The three lines in  main()                 →  just before "STEP 4"

Quick-test without editing v4
------------------------------
You can also run this file standalone — it imports from v4 and patches it at
runtime.  Then call  main_anchored()  instead of main() for the patched run.

    python 3_board_anchor_patch.py -det ... -pose ... -robot ... -outDir ...
"""

from __future__ import annotations

# ══════════════════════════════════════════════════════════════
# [PATCH section 1 – copy into v4 after JointAxisConstraint]
# ══════════════════════════════════════════════════════════════
from dataclasses import dataclass
import numpy as np
from typing import Dict, Any, List, Optional, Tuple


@dataclass
class PositionAnchorConstraint:
    """
    Pins a single marker to a known world-space position.
    Used to anchor Board markers whose world positions are read from robot.json.
    """
    marker_id:       int
    target_world_m:  np.ndarray   # shape (3,), in metres
    weight:          float = 100.0
    source:          str   = "board_anchor"


# ══════════════════════════════════════════════════════════════
# [PATCH section 2 – replace the Constraint type alias in v4]
# ══════════════════════════════════════════════════════════════
#   OLD:  Constraint = MarkerDistanceConstraint | JointAxisConstraint
#   NEW:
#         from 3_board_anchor_patch import PositionAnchorConstraint
#         Constraint = MarkerDistanceConstraint | JointAxisConstraint | PositionAnchorConstraint


# ══════════════════════════════════════════════════════════════
# [PATCH section 3 – new function, copy anywhere before main()]
# ══════════════════════════════════════════════════════════════

def load_board_anchors(
    robot_data: Dict[str, Any],
    length_scale: float,
    weight: float = 100.0,
) -> List[PositionAnchorConstraint]:
    """
    Build PositionAnchorConstraints for every marker on the 'Board' link.

    Board is the root link; its marker positions are in the Board (world) frame.
    length_scale = 1/1000 when robot.json uses mm (the standard).

    Parameters
    ----------
    robot_data   : loaded robot.json
    length_scale : float  (1/1000 to convert mm → m)
    weight       : constraint weight (100 ≈ 1 mm anchor tolerance at λ=100)
    """
    links  = robot_data.get("links", {}) or {}
    board  = links.get("Board", {}) or {}
    anchors: List[PositionAnchorConstraint] = []

    for m in board.get("markers", []):
        mid = int(m.get("id", -1))
        pos = m.get("position", None)
        if mid < 0 or pos is None or len(pos) != 3:
            continue
        world_m = np.array(pos, dtype=np.float64) * float(length_scale)
        anchors.append(
            PositionAnchorConstraint(
                marker_id=mid,
                target_world_m=world_m,
                weight=weight,
                source="board_anchor",
            )
        )

    return anchors


# ══════════════════════════════════════════════════════════════
# [PATCH section 4 – add this branch inside bundle_adjustment_residuals,
#  in the  "for constraint in constraints:"  loop, after the JointAxisConstraint branch]
# ══════════════════════════════════════════════════════════════

def _anchor_residuals(constraint: PositionAnchorConstraint,
                      marker_dict: Dict[int, np.ndarray],
                      lambda_constraint: float) -> List[float]:
    """
    Returns 3 residual components  w·λ·(X_observed - X_target)  for each axis.

    Paste the body of this function into bundle_adjustment_residuals like this:

        elif isinstance(constraint, PositionAnchorConstraint):
            if constraint.marker_id in marker_dict:
                diff = marker_dict[constraint.marker_id] - constraint.target_world_m
                for d in diff:
                    residuals.append(float(d) * constraint.weight * lambda_constraint)
    """
    residuals: List[float] = []
    if constraint.marker_id in marker_dict:
        diff = marker_dict[constraint.marker_id] - constraint.target_world_m
        for d in diff:
            residuals.append(float(d) * constraint.weight * lambda_constraint)
    return residuals


# ══════════════════════════════════════════════════════════════
# [PATCH section 5 – add these lines in main(), just BEFORE
#  "print('\n[STEP 4] Bundle adjustment with constraints...')"]
# ══════════════════════════════════════════════════════════════
#
#   board_anchors = load_board_anchors(robot_data, length_scale, weight=100.0)
#   constraints  += board_anchors
#   print(f"[INFO] Board anchors added: {len(board_anchors)}")
#
# ══════════════════════════════════════════════════════════════


# ──────────────────────────────────────────────────────────────
# STANDALONE RUNTIME PATCH  (alternative to editing v4)
# Run this file directly and it patches v4 in memory.
# ──────────────────────────────────────────────────────────────

def _monkey_patch_v4():
    """
    Import the v4 module under an alias and patch it so Board markers are
    anchored without touching the original file.
    """
    import importlib.util, sys, types, pathlib

    v4_path = pathlib.Path(__file__).parent / "3_multiview_bundle_adjustment_v4.py"
    if not v4_path.exists():
        raise FileNotFoundError(
            f"Expected v4 at {v4_path}.\n"
            "Place this patch file alongside 3_multiview_bundle_adjustment_v4.py."
        )

    spec   = importlib.util.spec_from_file_location("ba_v4", v4_path)
    mod    = importlib.util.module_from_spec(spec)
    sys.modules["ba_v4"] = mod          # must be registered before exec_module for @dataclass
    spec.loader.exec_module(mod)

    # Inject PositionAnchorConstraint into the module namespace
    mod.PositionAnchorConstraint = PositionAnchorConstraint

    # Patch bundle_adjustment_residuals ─────────────────────────
    _orig_residuals = mod.bundle_adjustment_residuals

    def _patched_residuals(marker_positions_flat, marker_ids, marker_observations,
                           cameras, constraints, obs_weights, lambda_constraint=100.0):
        r = _orig_residuals(marker_positions_flat, marker_ids, marker_observations,
                            cameras, constraints, obs_weights, lambda_constraint)
        # Rebuild marker_dict (same as inside the original function)
        marker_dict = {
            mid: marker_positions_flat[i*3:(i+1)*3]
            for i, mid in enumerate(marker_ids)
        }
        extra = []
        for c in constraints:
            if isinstance(c, PositionAnchorConstraint):
                extra.extend(_anchor_residuals(c, marker_dict, lambda_constraint))
        if extra:
            r = np.concatenate([r, np.array(extra, dtype=np.float64)])
        return r

    mod.bundle_adjustment_residuals = _patched_residuals

    # Patch main to inject board anchors ────────────────────────
    _orig_main = mod.main

    def _patched_main():
        import sys as _sys
        # We hijack the module-level function calls by patching compile_constraints
        _orig_compile  = mod.compile_constraints
        _orig_initial  = mod.initial_triangulation

        # robot_data / length_scale are captured here so the initial-triangulation
        # wrapper below can pin Board markers to their true world positions.
        # In v4.main(), compile_constraints() runs BEFORE initial_triangulation(),
        # so this dict is populated by the time the triangulation wrapper fires.
        captured: Dict[str, Any] = {}

        def _compile_with_anchors(robot_data, length_scale, cfg):
            captured["robot_data"]   = robot_data
            captured["length_scale"] = length_scale

            marker_to_link, link_markers, constraints, issues, marker_meta = \
                _orig_compile(robot_data, length_scale, cfg)

            anchors = load_board_anchors(robot_data, length_scale, weight=100.0)
            print(f"[PATCH] Board anchors added: {len(anchors)}")
            constraints = constraints + anchors
            return marker_to_link, link_markers, constraints, issues, marker_meta

        def _initial_with_board(marker_observations, cameras):
            """
            Run the normal multi-view triangulation, then overwrite every
            OBSERVED Board marker with its exact robot.json world position.
            Board markers define the world frame, so they must not be left at
            their (z-noisy) triangulated values — not even in the initial JSON.
            """
            tri = _orig_initial(marker_observations, cameras)

            robot_data   = captured.get("robot_data")
            length_scale = captured.get("length_scale", 1.0)
            if robot_data is None:
                return tri

            board = (robot_data.get("links", {}) or {}).get("Board", {}) or {}
            pinned = 0
            for m in board.get("markers", []):
                mid = int(m.get("id", -1))
                pos = m.get("position", None)
                if mid < 0 or pos is None or len(pos) != 3:
                    continue
                # Only pin markers that were actually observed/triangulated,
                # so we don't introduce markers that no camera ever saw.
                if mid in tri:
                    tri[mid] = np.array(pos, dtype=np.float64) * float(length_scale)
                    pinned += 1
            print(f"[PATCH] Board markers pinned to robot.json in initial triangulation: {pinned}")
            return tri

        mod.compile_constraints   = _compile_with_anchors
        mod.initial_triangulation = _initial_with_board
        _orig_main()
        mod.compile_constraints   = _orig_compile   # restore
        mod.initial_triangulation = _orig_initial   # restore

    mod.main = _patched_main
    return mod


def main():
    """Run the patched v4 pipeline with Board marker anchoring."""
    import sys
    mod = _monkey_patch_v4()
    sys.argv[0] = "3_multiview_bundle_adjustment_v4.py"  # cosmetic
    mod.main()


if __name__ == "__main__":
    main()
