#!/usr/bin/env python3
"""Rebuild the UV/unroll map for a design export, using the SDK's own builder.

Why this exists: ``unroll_map_uv_s0.4.npz`` — ``map_x``, ``map_y``, ``valid_mask``, ``world_xyz`` —
is what decides which part of the sensing mesh the application treats as **sensing**; everything else
renders grey. Three facts make it worth rebuilding rather than trusting:

* the Blender export writes **no mask** (``blender/scripts/export.py`` writes geometry, optics,
  intrinsics and pose only);
* the **SDK builds it** — ``orisys.unroll.precompute.build_uv_domain`` rasterises the mesh's authored
  UV faces through the camera with a depth buffer, and ``build_or_load_domain`` builds-then-caches and
  already knows how to spot a stale cache;
* the map **records its own parameters** (``out_width``/``out_height``/``flip_v``/``source_scale``/
  ``edge_band``/``edge_expand``/``projection_K``), so a rebuild can be tied to the same inputs.

So this is a thin wrapper: resolve the paths, take the parameters from an existing map or the
command line (never guess them), call the SDK, and report what came out.

    python tools/build_uv_map.py \
        --export-dir build/export/finger/r01/compute \
        --params-from D:/orisys_ws/orisys-sdk/sensors/finger/assets/30fps-0703/unroll_map_uv_s0.4.npz \
        --output build/export/finger/r01/compute

Then compare it with what ships:

    python D:/orisys_ws/orisys-app/tools/scene_assets/compare_uv_map.py \
        --map build/export/finger/r01/compute/unroll_map_uv_s0.4.npz \
        --against D:/orisys_ws/orisys-sdk/sensors/finger/assets/30fps-0703/unroll_map_uv_s0.4.npz
"""

from __future__ import annotations

import argparse
import hashlib
import sys
from pathlib import Path

import numpy as np

REQUIRED_EXPORT_FILES = ("mesh.npz", "camera_intrinsics.npz", "world_to_cam.npy")
#: Parameters the map must state for a faithful rebuild. ``source_width``/``source_height`` are
#: deliberately *not* here: they are derived from the export's own calibration (`K`'s image size
#: times ``source_scale``), so a map built from a 1920x1080 camera is not rebuilt as if it were
#: 1280x720. Getting that wrong shrinks the island (measured: 3.5 % coverage instead of 40.8 %),
#: which is the whole reason this tool exists.
PARAM_KEYS = ("out_width", "out_height", "flip_v", "source_scale", "edge_band", "edge_expand")


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def params_from_map(path: Path) -> dict:
    """Read the parameters a built map recorded about itself."""
    with np.load(path, allow_pickle=False) as data:
        found: dict = {}
        for key in PARAM_KEYS:
            if key in data.files:
                value = np.asarray(data[key]).reshape(-1)[0].item()
                found[key] = value
    missing = [key for key in PARAM_KEYS if key not in found]
    if missing:
        raise SystemExit(
            f"{path.name} does not record {missing}; it predates the self-describing metadata, so "
            "pass the parameters explicitly instead"
        )
    found["edge_expand"] = float(found["edge_expand"])
    found["edge_band"] = float(found["edge_band"])
    found["source_scale"] = float(found["source_scale"])
    found["flip_v"] = bool(found["flip_v"])
    for key in ("out_width", "out_height"):
        found[key] = int(found[key])
    return found


def _warn_if_upsampled(source_width: int, source_height: int,
                       out_width: int, out_height: int) -> None:
    """Say it out loud when the output grid is finer than the source it samples."""
    if source_width < out_width or source_height < out_height:
        print(
            f"WARNING: the source frame {source_width} x {source_height} is smaller than the "
            f"{out_width} x {out_height} output grid, so the unrolled picture is upsampled "
            "(several output pixels read one source pixel). Raise source_scale to at least "
            f"{derive_source_scale(source_width, source_height, out_width, out_height) * 1.0:.3f} "
            "of the *native* frame, or accept the softer picture."
        )


def derive_source_scale(native_width: int, native_height: int,
                        out_width: int, out_height: int) -> float:
    """The smallest scale whose source grid still covers the output grid 1:1.

    The remap *gathers* from the source frame into the ``out_width x out_height`` domain, so if
    the source is smaller than the output in either axis, several output pixels read the same
    source pixel and the unrolled picture is blocky. Requiring

        min(native_width * s, native_height * s) >= max(out_width, out_height)   (per axis)

    gives ``s = max(out_w / native_w, out_h / native_h)``: coverage without upsampling, and — the
    useful property — a source size that follows the *output* requirement, so two cameras of
    different resolution produce the same output quality.

    Above this the scale only buys interpolation smoothness and costs per-frame pixel work
    (the live frame is resized by ``s`` before undistort/remap), so the derived value is the
    sensible floor rather than the optimum: a profile may pin a larger one deliberately.
    """
    if native_width <= 0 or native_height <= 0:
        raise SystemExit("the export's camera_intrinsics.npz has a non-positive image size")
    scale = max(out_width / float(native_width), out_height / float(native_height))
    # Keep the cache name readable (it embeds the scale) and never exceed the native frame.
    return min(1.0, round(scale + 1e-9, 3))


def params_from_profile(path: Path) -> dict:
    """Read the ``uv_map:`` block from a profile yaml.

    Hand-parsed on purpose: the block is six scalars, and this must keep working with the
    interpreter that runs the design tools rather than depending on PyYAML being installed.
    Unknown keys and a missing block are reported rather than ignored.
    """
    if not path.is_file():
        raise SystemExit(f"no such profile: {path}")
    found: dict = {}
    inside = False
    for raw in path.read_text(encoding="utf-8").splitlines():
        line = raw.split("#", 1)[0].rstrip()
        if not line.strip():
            continue
        if not line.startswith((" ", "\t")):
            inside = line.strip() == "uv_map:"
            continue
        if not inside or ":" not in line:
            continue
        key, _, value = line.strip().partition(":")
        value = value.strip()
        if value.lower() in ("true", "false"):
            found[key] = value.lower() == "true"
        else:
            try:
                found[key] = float(value) if "." in value else int(value)
            except ValueError:
                found[key] = value
    if not found:
        raise SystemExit(f"{path} has no uv_map: block (add out_width/out_height/flip_v/"
                         "source_scale/edge_band/edge_expand)")
    unknown = sorted(set(found) - set(PARAM_KEYS))
    if unknown:
        raise SystemExit(f"{path}: uv_map has keys this tool does not understand: {unknown}")
    for key in ("out_width", "out_height"):
        found[key] = int(found[key])
    found["flip_v"] = bool(found["flip_v"])
    for key in ("source_scale", "edge_band", "edge_expand"):
        found[key] = float(found[key])
    return found


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--export-dir", type=Path, required=True,
                        help="directory holding mesh.npz, camera_intrinsics.npz and world_to_cam.npy")
    parser.add_argument("--profile", type=Path, default=None,
                        help="profile yaml carrying the uv_map: parameters (preferred)")
    parser.add_argument("--params-from", type=Path, default=None,
                        help="an existing map to take the parameters from (cross-check)")
    parser.add_argument("--output", type=Path, default=None,
                        help="directory to write into (default: the export dir)")
    parser.add_argument("--force", action="store_true",
                        help="delete an existing map first, so the SDK must rebuild rather than load")
    parser.add_argument("--edge-expand", type=float, default=None,
                        help="override; > 1 pins it, otherwise the SDK resolves it from distortion")
    for key in ("out-width", "out-height", "source-width", "source-height"):
        parser.add_argument(f"--{key}", type=int, default=None)
    parser.add_argument("--flip-v", dest="flip_v", action="store_true", default=None)
    parser.add_argument("--no-flip-v", dest="flip_v", action="store_false")
    parser.add_argument("--source-scale", type=float, default=None)
    parser.add_argument("--edge-band", type=float, default=None)
    args = parser.parse_args(argv)

    export_dir: Path = args.export_dir
    missing = [name for name in REQUIRED_EXPORT_FILES if not (export_dir / name).is_file()]
    if missing:
        raise SystemExit(
            f"{export_dir} is missing {missing}: this needs the payload "
            "blender/scripts/export.py writes, not the source .blend"
        )

    params: dict = {}
    if args.profile is not None:
        params.update(params_from_profile(args.profile))
    if args.params_from is not None:
        if not args.params_from.is_file():
            raise SystemExit(f"no such map to take parameters from: {args.params_from}")
        recorded = params_from_map(args.params_from)
        # A disagreement is the interesting case: the profile is the source of truth, but a map
        # that was built with different numbers has to be visible, not resolved silently.
        disagree = {
            key: (recorded.get(key), params.get(key))
            for key in params
            if key in recorded and recorded[key] != params[key]
        }
        if disagree:
            print(f"NOTE: {args.params_from.name} records different parameters than the profile "
                  f"(map, profile): {disagree}")
        params.update({key: value for key, value in recorded.items() if key not in params})
    overrides = {
        "out_width": args.out_width, "out_height": args.out_height,
        "flip_v": args.flip_v, "source_scale": args.source_scale, "edge_band": args.edge_band,
        "edge_expand": args.edge_expand,
    }
    params.update({key: value for key, value in overrides.items() if value is not None})
    absent = [key for key in PARAM_KEYS if key not in params]
    if absent:
        raise SystemExit(
            f"no value for {absent}: pass --params-from an existing map or give them on the command "
            "line. The importer never guesses what a source unit or an output size is."
        )

    try:
        import orisys
        from orisys.unroll.geometry import scale_camera_geometry, scaled_map_cache_path
        from orisys.unroll.precompute import build_or_load_domain
    except Exception as exc:  # noqa: BLE001
        raise SystemExit(
            f"the SDK is not importable ({type(exc).__name__}: {exc}); install the orisys checkout "
            "you intend to build with — the map must come from the same implementation that "
            "built the shipped one"
        )

    out_dir: Path = args.output or export_dir
    out_dir.mkdir(parents=True, exist_ok=True)

    with np.load(export_dir / "camera_intrinsics.npz", allow_pickle=False) as handle:
        K_native = np.asarray(handle["K"], dtype=np.float64)
        native_width = int(np.asarray(handle["width"]).reshape(-1)[0])
        native_height = int(np.asarray(handle["height"]).reshape(-1)[0])

    # The scale is a decision, so it is reported as one: the profile (or the CLI) may pin it, and
    # this says what the output grid would otherwise require.
    derived_scale = derive_source_scale(
        native_width, native_height, int(params["out_width"]), int(params["out_height"])
    )
    pinned_scale = "source_scale" in params
    if not pinned_scale:
        params["source_scale"] = derived_scale
    effective_scale = float(params["source_scale"])
    source_K, source_width, source_height = scale_camera_geometry(
        K_native, native_width, native_height, effective_scale
    )
    if args.source_width is not None:
        source_width = int(args.source_width)
    if args.source_height is not None:
        source_height = int(args.source_height)
    _warn_if_upsampled(source_width, source_height,
                       int(params["out_width"]), int(params["out_height"]))

    base = str(out_dir / "unroll_map_uv.npz")
    target = Path(scaled_map_cache_path(base, float(params["source_scale"])))
    if args.force and target.is_file():
        target.unlink()
        print(f"removed {target.name} (--force: rebuild rather than load)")

    print(f"export         {export_dir}")
    print(f"target         {target}")
    print(f"SDK            orisys {getattr(orisys, '__version__', '?')}")
    print(f"parameters     {params}")
    print(f"native camera  {native_width} x {native_height}, f={K_native[0, 0]:.1f}")
    print(f"source frame   {source_width} x {source_height} "
          f"(scale {effective_scale}), f={source_K[0, 0]:.1f}")
    print(f"scale          {'pinned' if pinned_scale else 'derived'} {effective_scale}"
          f" (output {params['out_width']} x {params['out_height']} needs >= {derived_scale})")

    domain = build_or_load_domain(
        str(export_dir / "mesh.npz"),
        str(export_dir / "camera_intrinsics.npz"),
        str(export_dir / "world_to_cam.npy"),
        str(target),
        source_K=source_K,
        source_width=int(source_width),
        source_height=int(source_height),
        out_width=int(params["out_width"]),
        out_height=int(params["out_height"]),
        flip_v=bool(params["flip_v"]),
        edge_band=float(params["edge_band"]),
        edge_expand=float(params["edge_expand"]),
        include_surface=True,
        verbose=True,
        source_scale=float(params["source_scale"]),
    )

    with np.load(target, allow_pickle=False) as written:
        mask = np.asarray(written["valid_mask"]).astype(bool)
        surface = "world_xyz" in written.files
    print(f"\nwrote          {target}")
    print(f"  shape        {mask.shape[1]} x {mask.shape[0]} (W x H)")
    print(f"  coverage     {mask.mean() * 100:.3f} % of the domain ({int(mask.sum())} pixels)")
    print(f"  surface map  {'present' if surface else 'MISSING (the 3D UV->surface mapping needs it)'}")
    print(f"  sha256       {_sha256(target)}")
    print("\ncompare it with the shipped map:")
    print("  python D:/orisys_ws/orisys-app/tools/scene_assets/compare_uv_map.py "
          f"--map \"{target}\" --against \"<shipped>.npz\"")
    return 0


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