#!/usr/bin/env python3
"""Put a housing exported in *another* frame into the sensor frame, once, reproducibly.

The finger sensor's design export (``orisys-sensor-design/build/export/finger/r01/compute``)
holds ``mesh.npz`` and ``housing.npz`` (the mounting frame, ``part_names = ['frame']``) in a
**camera/world frame**, while ``compile_scene.py`` compiles from the SDK's asset directory,
whose ``mesh.npz`` is already in the sensor frame (millimetres, the unit stated by
``--meters-per-source-unit``).

The two meshes are the same surface, and their ``faces`` arrays are identical — so vertex *i*
corresponds to vertex *i* and the frame relationship is a similarity transform that can be
**fitted** rather than guessed. This tool:

1. loads both ``mesh.npz`` files,
2. fits the similarity (rotation + uniform scale + translation) from design to sensor vertices,
3. refuses to write anything if the fit does not hold to ``--tolerance`` of the mesh diagonal,
4. applies that transform to the housing vertices and writes a sensor-frame ``housing.npz``
   (millimetres, same as the SDK's mesh) ready for ``compile_scene.py --housing``.

Why fit instead of hand-transcribing a matrix: a wrong matrix still produces a plausible
picture, and this one is checked against 1032 known correspondences every time it runs.

    python tools/scene_assets/align_housing.py \
        --design D:/orisys_ws/orisys-sensor-design/build/export/finger/r01/compute \
        --sensor D:/orisys_ws/orisys-sdk/sensors/finger/assets/30fps-0703 \
        --output tools/scene_assets/housing/finger_r01_housing.npz
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np

#: Residual budget, as a fraction of the sensor-frame mesh diagonal. The fit is exact in
#: practice (float noise); anything above this means the two meshes are not the same shape and
#: the housing must not be transformed as if they were.
DEFAULT_TOLERANCE = 1e-4


def load_vertices(path: Path) -> np.ndarray:
    with np.load(path, allow_pickle=False) as data:
        if "vertices" not in data.files:
            raise SystemExit(f"{path} has no 'vertices' array (found {list(data.files)})")
        return np.asarray(data["vertices"], dtype=np.float64)


def load_faces(path: Path) -> np.ndarray:
    with np.load(path, allow_pickle=False) as data:
        return np.asarray(data["faces"])


def fit_similarity(source: np.ndarray, target: np.ndarray) -> tuple[np.ndarray, np.ndarray, float]:
    """Rotation, translation and uniform scale mapping ``source`` onto ``target``.

    ``target = scale * (source @ rotation) + translation``. Uses the closed-form Kabsch
    solution with a scale factor; a reflection is never the answer here.
    """
    source_centroid = source.mean(axis=0)
    target_centroid = target.mean(axis=0)
    source_centered = source - source_centroid
    target_centered = target - target_centroid
    u, singular, vt = np.linalg.svd(source_centered.T @ target_centered)
    rotation = u @ vt
    if np.linalg.det(rotation) < 0.0:
        u = u.copy()
        u[:, -1] *= -1.0
        rotation = u @ vt
    scale = float(np.sum(singular) / np.sum(source_centered**2))
    translation = target_centroid - scale * (source_centroid @ rotation)
    return rotation, translation, scale


def apply_similarity(
    points: np.ndarray, rotation: np.ndarray, translation: np.ndarray, scale: float
) -> np.ndarray:
    return scale * (points @ rotation) + translation


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--design", type=Path, required=True, help="export dir with mesh+housing")
    parser.add_argument("--sensor", type=Path, required=True, help="export dir whose mesh is sensor-frame")
    parser.add_argument("--output", type=Path, required=True, help="sensor-frame housing npz to write")
    parser.add_argument("--tolerance", type=float, default=DEFAULT_TOLERANCE)
    parser.add_argument(
        "--matrix",
        type=Path,
        default=None,
        help="also write the fitted 4x4 (source -> sensor) as a .npy, for the record",
    )
    args = parser.parse_args(argv)

    design_mesh = args.design / "mesh.npz"
    sensor_mesh = args.sensor / "mesh.npz"
    housing_path = args.design / "housing.npz"
    for path in (design_mesh, sensor_mesh, housing_path):
        if not path.is_file():
            raise SystemExit(f"missing {path}")

    source = load_vertices(design_mesh)
    target = load_vertices(sensor_mesh)
    if source.shape != target.shape:
        raise SystemExit(
            f"the meshes disagree: {source.shape} vs {target.shape} — the correspondence the "
            "fit relies on does not exist"
        )
    if not np.array_equal(load_faces(design_mesh), load_faces(sensor_mesh)):
        raise SystemExit(
            "the faces arrays differ, so vertex i is not vertex i; refusing to fit a transform"
        )

    rotation, translation, scale = fit_similarity(source, target)
    residual = apply_similarity(source, rotation, translation, scale) - target
    rms = float(np.sqrt((residual**2).sum(axis=1).mean()))
    diagonal = float(np.linalg.norm(target.max(axis=0) - target.min(axis=0)))
    relative = rms / max(diagonal, 1e-12)
    print(f"fit over {len(source)} vertices: scale={scale:.8g} rms={rms:.3g} "
          f"({relative:.3g} of the {diagonal:.3f}-unit diagonal)")
    if relative > float(args.tolerance):
        raise SystemExit(
            f"the fit is off by {relative:.3g} of the mesh diagonal (budget {args.tolerance:g}): "
            "these two meshes are not the same shape, so transforming the housing with this "
            "matrix would put it in the wrong place"
        )

    with np.load(housing_path, allow_pickle=False) as data:
        housing_vertices = np.asarray(data["vertices"], dtype=np.float64)
        housing_faces = np.asarray(data["faces"])
        part_names = data["part_names"] if "part_names" in data.files else None
    aligned = apply_similarity(housing_vertices, rotation, translation, scale)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    np.savez(
        args.output,
        vertices=aligned.astype(np.float32),
        faces=housing_faces.astype(np.int32),
        # Provenance, so a later reader can see this was fitted rather than exported.
        source_frame=np.array(str(args.design)),
        sensor_frame=np.array(str(args.sensor)),
        fitted_scale=np.array(scale, dtype=np.float64),
        fitted_rms=np.array(rms, dtype=np.float64),
        part_names=(part_names if part_names is not None else np.array(["housing"])),
    )
    print(f"wrote {args.output}")
    print(f"  housing in the sensor frame: low={np.round(aligned.min(axis=0), 3)} "
          f"high={np.round(aligned.max(axis=0), 3)}"
          f"  ({len(housing_faces)} triangles)")

    if args.matrix is not None:
        matrix = np.eye(4)
        matrix[:3, :3] = scale * rotation
        matrix[:3, 3] = translation
        args.matrix.parent.mkdir(parents=True, exist_ok=True)
        np.save(args.matrix, matrix)
        print(f"wrote {args.matrix}")
    return 0


if __name__ == "__main__":
    sys.exit(main())