#!/usr/bin/env python3
"""Derive the *(unit scale, rigid transform)* pair that compiles the **design export**.

``compile_scene.py`` splits two jobs that must not be confused (design document §8.1):

* ``--meters-per-source-unit`` changes *lengths*: one stated source unit, in metres;
* ``--transform`` is **rigid** (orthonormal rotation + translation, no scale), and the
  importer validates that before it compiles anything.

The finger design export (``orisys-sensor-design/build/export/finger/r01/compute``) is the
same surface as the SDK asset (``orisys-sdk/sensors/finger/assets/30fps-0703``) — identical
``faces`` and ``loop_to_vertex``, so vertex *i* is vertex *i* — but it lives in the design
repo's camera/world frame and in *design* units. ``align_housing.py`` fits that relationship
as a similarity ``p_sensor_mm = s·R·p_design + t``; the matrix it writes
(``housing/finger_r01_design_to_sensor.npy``) therefore **already carries the scale**: it is
not a legal ``--transform``, and pairing it with ``--meters-per-source-unit 1.8388e-5`` would
apply the scale twice.

This tool does the split, once and reproducibly:

    design frame, design units --(R, t_m)--> sensor frame, metres
    p_sensor_m = (0.001·s)·(R·p_design) + 0.001·t

so the compiler can be given ``--meters-per-source-unit 0.001·s`` and a rigid 4x4 whose
rotation block is exactly ``R`` and whose translation is ``0.001·t`` (metres). The metre
translation matters: ``t`` from the fit is in the sensor frame's *millimetres*, because that is
the frame the SDK asset is stated in.

It also checks, numerically, the wrong pairings — all three of which are live mistakes here:

* ``--transform`` = the scale-bearing matrix is **refused** by the rigid-transform validator
  (it is not a legal rigid transform at all, whatever the unit says), so the "scale applied
  twice" trap cannot reach the compiler through that door;
* what the double scale *would* produce is computed directly from the importer's own
  ``_apply_rigid``, so the number is on the record rather than asserted;
* ``--meters-per-source-unit 0.001`` (calling a design unit a millimetre) with the correct
  rigid transform does compile, and is 54.38x too large — a package that looks entirely
  plausible.

    python tools/scene_assets/fit_design_frame.py \\
        --design D:/orisys_ws/orisys-sensor-design/build/export/finger/r01/compute \\
        --sensor D:/orisys_ws/orisys-sdk/sensors/finger/assets/30fps-0703 \\
        --rigid-output tools/scene_assets/housing/finger_r01_design_to_sensor_rigid_m.npy \\
        --report tools/scene_assets/housing/finger_r01_design_frame.json
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np


def _bootstrap_paths() -> None:
    here = Path(__file__).resolve()
    repo = here.parents[2]
    for path in (repo / "apps" / "studio" / "src", repo / "shared" / "src"):
        if path.is_dir() and str(path) not in sys.path:
            sys.path.insert(0, str(path))
    if str(here.parent) not in sys.path:
        # The fit itself stays in one place: ``align_housing.fit_similarity``, which
        # checks the 1032 correspondences before it writes anything.
        sys.path.insert(0, str(here.parent))


#: The SDK asset's frame is millimetres; the compiled package is metres.
MM_PER_SENSOR_UNIT = 0.001


def derive(design_dir: Path, sensor_dir: Path, *, tolerance: float) -> dict:
    from align_housing import apply_similarity, fit_similarity, load_faces, load_vertices

    design_mesh = design_dir / "mesh.npz"
    sensor_mesh = sensor_dir / "mesh.npz"
    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_mm, scale_mm = fit_similarity(source, target)
    residual_mm = apply_similarity(source, rotation, translation_mm, scale_mm) - target
    rms_mm = float(np.sqrt((residual_mm**2).sum(axis=1).mean()))
    diagonal_mm = float(np.linalg.norm(target.max(axis=0) - target.min(axis=0)))
    relative = rms_mm / max(diagonal_mm, 1e-12)
    if relative > float(tolerance):
        raise SystemExit(
            f"the fit is off by {relative:.3g} of the mesh diagonal (budget {tolerance:g}): "
            "these two meshes are not the same shape"
        )

    # The split. ``scale_mm`` maps design units -> sensor millimetres, so the unit
    # statement is that many millimetres, and the translation is carried into metres
    # with the same factor the rotation gets.
    #
    # **Matrix convention.** ``align_housing.fit_similarity`` fits the *row-vector* form
    # ``p_sensor = s·(p_design @ R) + t``, while ``compile_scene.py`` hands ``--transform``
    # to ``legacy._apply_rigid``, which is the *column* form ``p_sensor = T @ [p_design; 1]``.
    # So the block the compiler needs is ``R.T``, not ``R``: pasting ``R`` in would rotate the
    # surface the wrong way by twice the fitted angle while leaving every extent plausible
    # (a 90-degree rotation's transpose is its inverse). ``verify()`` is what proves this.
    meters_per_source_unit = MM_PER_SENSOR_UNIT * scale_mm
    rigid = np.eye(4, dtype=np.float64)
    rigid[:3, :3] = rotation.T
    rigid[:3, 3] = MM_PER_SENSOR_UNIT * translation_mm

    scale_bearing = np.eye(4, dtype=np.float64)
    scale_bearing[:3, :3] = scale_mm * rotation
    scale_bearing[:3, 3] = translation_mm

    extent_design = np.ptp(source, axis=0)
    extent_sensor_mm = np.ptp(target, axis=0)
    return {
        "design_dir": str(design_dir),
        "sensor_dir": str(sensor_dir),
        "vertices": int(len(source)),
        "fit": {
            "scale_design_unit_to_sensor_mm": scale_mm,
            "rotation_row_convention": rotation.tolist(),
            "rotation_transposed_column_convention": rotation.T.tolist(),
            "translation_sensor_mm": translation_mm.tolist(),
            "rms_sensor_mm": rms_mm,
            "rms_sensor_m": rms_mm * MM_PER_SENSOR_UNIT,
            "diagonal_sensor_mm": diagonal_mm,
            "relative_of_diagonal": relative,
            "tolerance": float(tolerance),
        },
        "unit_statement": {
            "meters_per_source_unit": meters_per_source_unit,
            "one_design_unit_is_mm": scale_mm,
            "one_sensor_mm_is_design_units": 1.0 / scale_mm,
        },
        "rigid_transform_sensor_from_design_m": rigid.tolist(),
        "scale_bearing_matrix_sensor_mm_from_design": scale_bearing.tolist(),
        "extent_design_units": extent_design.tolist(),
        "extent_sensor_mm": extent_sensor_mm.tolist(),
        "extent_compiled_m": (extent_sensor_mm * MM_PER_SENSOR_UNIT).tolist(),
        "rotation_orthonormality_error": float(
            np.abs(rotation @ rotation.T - np.eye(3)).max()
        ),
        "rotation_determinant": float(np.linalg.det(rotation)),
    }


def verify(design_dir: Path, sensor_dir: Path, rigid: np.ndarray, units: float) -> dict:
    """Compile in memory through the real importer and compare with the SDK-frame route.

    The reference is the *shipped* route: the SDK asset with ``0.001`` and no transform.
    If the design route is right, both land on the same sensor-metre positions — order
    included, because identical ``faces``/``loop_to_vertex`` and no UV seam means the
    render vertices come out in the same order.
    """
    from orisys_app_shared.scene3d.contracts import ThreeDError
    from orisys_app_shared.scene3d.legacy import LegacyFingerImporter, validate_rigid_transform

    design_map = design_dir / "unroll_map_uv_s0.42.npz"
    sensor_map = sensor_dir / "unroll_map_uv_s0.4.npz"
    importer = LegacyFingerImporter()

    reference = importer.convert(
        sensor_dir, sensor_map,
        transform_sensor_from_source=np.eye(4),
        meters_per_source_unit=MM_PER_SENSOR_UNIT,
    ).sensing_part().positions_sensor_m

    result: dict = {}
    compiled = importer.convert(
        design_dir, design_map,
        transform_sensor_from_source=rigid,
        meters_per_source_unit=units,
    ).sensing_part().positions_sensor_m
    difference = np.asarray(compiled, dtype=np.float64) - np.asarray(reference, dtype=np.float64)
    result["design_route"] = {
        "extent_m": np.ptp(compiled, axis=0).tolist(),
        "low_m": compiled.min(axis=0).tolist(),
        "high_m": compiled.max(axis=0).tolist(),
        "max_abs_difference_from_sdk_route_m": float(np.abs(difference).max()),
        "rms_difference_from_sdk_route_m": float(np.sqrt((difference**2).sum(axis=1).mean())),
        "bitwise_identical_to_sdk_route": bool(
            np.array_equal(np.asarray(compiled, dtype=np.float32),
                           np.asarray(reference, dtype=np.float32))
        ),
    }
    result["sdk_route_extent_m"] = np.ptp(reference, axis=0).tolist()

    from orisys_app_shared.scene3d.legacy import _apply_rigid, split_at_uv_seams

    stored = np.load(
        Path(__file__).resolve().parent / "housing" / "finger_r01_design_to_sensor.npy"
    )

    # Trap 1: the scale-bearing matrix as ``--transform``. Refused, whatever the unit is.
    try:
        validate_rigid_transform(stored)
        result["trap_stored_matrix_as_transform"] = "ACCEPTED (the validator missed it)"
    except ThreeDError as exc:
        result["trap_stored_matrix_as_transform"] = (
            f"refused: {exc.diagnostic.code}: {exc.diagnostic.message}"
        )

    # Trap 2: what the double scale *would* do, computed through the importer's own
    # rigid application — the stored matrix already holds 1/54.38, and the derived unit
    # holds it too, so the rotation block scales twice and ``t`` (sensor-mm) is added to
    # a metre-scale result.
    with np.load(design_dir / "mesh.npz", allow_pickle=False) as handle:
        positions, _, _, _ = split_at_uv_seams(
            handle["vertices"], handle["faces"], handle["uv_data"], handle["loop_to_vertex"]
        )
    doubled = _apply_rigid(positions * units, stored)
    result["trap_double_scale_if_forced"] = {
        "extent_m": np.ptp(doubled, axis=0).tolist(),
        "low_m": doubled.min(axis=0).tolist(),
        "high_m": doubled.max(axis=0).tolist(),
        "note": (
            "the stated unit already carries the 1/54.38 fit scale and the matrix carries it "
            "again; the translation is sensor-mm added to a metre-scale result, so the extent "
            "collapses to the fit residual and the surface sits ~4.5 m from the origin"
        ),
    }

    # Trap 3: a *legal* rigid transform with the wrong unit statement (a design unit is not
    # a millimetre). This one compiles; nothing downstream can tell it is 54.38x too big.
    milli = importer.convert(
        design_dir, design_map,
        transform_sensor_from_source=rigid,
        meters_per_source_unit=MM_PER_SENSOR_UNIT,
    ).sensing_part().positions_sensor_m
    result["trap_design_unit_called_a_millimetre"] = {
        "extent_m": np.ptp(milli, axis=0).tolist(),
        "low_m": milli.min(axis=0).tolist(),
        "high_m": milli.max(axis=0).tolist(),
        "factor_too_large": float(
            np.ptp(milli, axis=0).max() / max(np.ptp(compiled, axis=0).max(), 1e-30)
        ),
        "note": "compiles cleanly and is wrong by the fit scale; the validator cannot see it",
    }
    return result


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--design", type=Path, required=True, help="design export dir")
    parser.add_argument("--sensor", type=Path, required=True, help="sensor-frame export dir")
    parser.add_argument("--rigid-output", type=Path, required=True, help="rigid 4x4 to write")
    parser.add_argument("--report", type=Path, default=None, help="JSON record of the derivation")
    parser.add_argument("--tolerance", type=float, default=1e-4)
    parser.add_argument(
        "--no-verify",
        action="store_true",
        help="skip the in-memory compile checks (they need the app's shared package)",
    )
    args = parser.parse_args(argv)

    _bootstrap_paths()
    record = derive(args.design, args.sensor, tolerance=args.tolerance)
    rigid = np.asarray(record["rigid_transform_sensor_from_design_m"], dtype=np.float64)
    units = float(record["unit_statement"]["meters_per_source_unit"])

    fit = record["fit"]
    print(f"fit over {record['vertices']} design->sensor correspondences "
          f"(identical faces, so vertex i is vertex i)")
    print(f"  scale          {fit['scale_design_unit_to_sensor_mm']:.12g} "
          f"sensor-mm per design unit  (= 1/{(1.0 / fit['scale_design_unit_to_sensor_mm']):.8f})")
    print(f"  residual rms   {fit['rms_sensor_mm']:.3g} mm "
          f"({fit['rms_sensor_m']:.3g} m) = {fit['relative_of_diagonal']:.3g} of the "
          f"{fit['diagonal_sensor_mm']:.3f}-mm diagonal (budget {fit['tolerance']:g})")
    print(f"  rotation       orthonormality error {record['rotation_orthonormality_error']:.3g}, "
          f"det {record['rotation_determinant']:.12f}")
    print(f"  translation    {[round(v, 12) for v in fit['translation_sensor_mm']]} sensor-mm")
    print()
    print("the split the compiler needs:")
    print(f"  --meters-per-source-unit {units:.12g}   (one design unit is "
          f"{(1.0 / fit['scale_design_unit_to_sensor_mm']):.8f} sensor-mm)")
    print(f"  --transform <rigid 4x4>  rotation = R, translation = 0.001*t = "
          f"{[round(float(v), 12) for v in rigid[:3, 3]]} m")
    print(f"  design extent {np.round(record['extent_design_units'], 4)} -> sensor "
          f"{np.round(record['extent_sensor_mm'], 4)} mm = "
          f"{np.round(record['extent_compiled_m'], 6)} m")

    if not args.no_verify:
        print()
        print("in-memory compile checks through the real importer:")
        checks = verify(args.design, args.sensor, rigid, units)
        design = checks["design_route"]
        print(f"  design route   extent {np.round(design['extent_m'], 8)} m  "
              f"low {np.round(design['low_m'], 8)} high {np.round(design['high_m'], 8)}")
        print(f"  sdk route      extent {np.round(checks['sdk_route_extent_m'], 8)} m")
        print(f"  agreement      max|diff| {design['max_abs_difference_from_sdk_route_m']:.3g} m, "
              f"rms {design['rms_difference_from_sdk_route_m']:.3g} m, "
              f"bitwise identical {design['bitwise_identical_to_sdk_route']}")
        print(f"  trap (stored scale-bearing matrix as --transform): "
              f"{checks['trap_stored_matrix_as_transform']}")
        forced = checks["trap_double_scale_if_forced"]
        print(f"  trap (derived unit + scale-bearing matrix, forced): extent "
              f"{np.round(forced['extent_m'], 9)} m, low {np.round(forced['low_m'], 6)}, "
              f"high {np.round(forced['high_m'], 6)}")
        milli = checks["trap_design_unit_called_a_millimetre"]
        print(f"  trap (design unit called a millimetre): extent "
              f"{np.round(milli['extent_m'], 6)} m — "
              f"{milli['factor_too_large']:.5f}x too large, and accepted")
        record["verification"] = checks

    args.rigid_output.parent.mkdir(parents=True, exist_ok=True)
    np.save(args.rigid_output, rigid)
    print()
    print(f"wrote {args.rigid_output}")
    if args.report is not None:
        args.report.parent.mkdir(parents=True, exist_ok=True)
        args.report.write_text(
            json.dumps(record, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
        )
        print(f"wrote {args.report}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())