"""Import a real Stringman room from a naavox LeRobot dataset into the twin: geometry, named places, the floor image.

    hf download naavox/playroom-sep29-2 --repo-type dataset --local-dir data/hf/playroom-sep29-2
    python scripts/import_room.py --dataset data/hf/playroom-sep29-2 --conf conf_playroom --calibrations <fw>/src/nf_robot/ml/calibrations --name playroom
    python scripts/import_room.py --dataset data/hf/bedroom-sep6 --conf conf_bedroom_move_clutter_2 --name bedroom
    uv run python scripts/import_room.py --building commercial_real_estate --office r5 --name cre_r5

A BUILDING (the room layer, milestone R, 2026-10-13): one office cut out of a scanned building
(~/robot-lab/buildings/<name>/, shared by every robot; see import_building's docstring). Writes configs/rooms/<name>.yaml
and data/rooms/<name>/ (visual meshes, the fixed map's points; outside git, rebuilt by this command).

Writes configs/rooms/<name>.yaml and assets/rooms/<name>/floor.png. Everything comes from the recording itself, in its
own room frame (fw ml/calibrations/readme.md: a config is only valid for the sessions it was recorded with):
  - anchors: the `anchor_poses` feature (datasets from August 2026 on carry it per frame; constant within a session)
  - named places: observation.state `<place>_bearing` / `_distance`, which are in the GRIPPER frame:
        place = gripper_pos_xy + distance * (sin(a), cos(a)),  a = bearing + spin
    (found 2026-10-08 by trying the conventions: this one puts every frame of playroom-sep29-2 on the same point, std
    0.00 m; the others scatter by 0.6-1.8 m). A place whose solution is the origin is unset.
  - floor: the median of the episodes' first overhead frames. The overhead camera is the firmware's orthographic floor
    view (host/floor_view.py: room origin at the image centre, +y up, 5 m across), so it drops onto the twin's floor
    at its true scale. Static clutter stays in it: it is the room's look, not something to pick.
  - eyelets (the indirect lines' far corners) are NOT in the recording: taken from the room's calibration file
    (fw ml/calibrations) through the 2D similarity that maps that file's two anchors onto the recording's, z = the
    recording's anchor height. APPROXIMATE (that file is an older calibration of the same room).
"""
from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

import sys

import numpy as np
import yaml

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
# the firmware's calibration files (cranebot3-firmware src/nf_robot/ml/calibrations): --calibrations or $NF_CALIBRATIONS
FW_CAL = Path(os.environ.get("NF_CALIBRATIONS", "nf_robot/ml/calibrations"))
PLACES = ("hamper", "toybox", "trashcan", "gamepad", "parking_location")
EXTENT_M = 5.0          # fw host/floor_view.EXTENT_M


def similarity_2d(src: np.ndarray, dst: np.ndarray):
    """Scale, rotation, translation mapping 2 src points onto 2 dst points (complex-number form)."""
    s0, s1 = complex(*src[0]), complex(*src[1])
    d0, d1 = complex(*dst[0]), complex(*dst[1])
    a = (d1 - d0) / (s1 - s0)
    b = d0 - a * s0
    return lambda p: np.array([(a * complex(*p) + b).real, (a * complex(*p) + b).imag]), abs(a)


def floor_median(ds: Path, n_max: int = 60) -> np.ndarray:
    import av
    import pandas as pd
    ep = pd.read_parquet(ds / "meta/episodes/chunk-000/file-000.parquet")
    col = "videos/observation.images.overhead_camera/from_timestamp"
    fidx = "videos/observation.images.overhead_camera/file_index"
    ts = sorted(ep.loc[ep[fidx] == 0, col].values[:n_max] + 0.5)
    c = av.open(str(ds / "videos/observation.images.overhead_camera/chunk-000/file-000.mp4"))
    s = c.streams.video[0]
    s.thread_type = "AUTO"
    out, i = [], 0
    for f in c.decode(s):
        t = float(f.pts * s.time_base)
        while i < len(ts) and t >= ts[i]:
            out.append(f.to_ndarray(format="rgb24"))
            i += 1
        if i >= len(ts):
            break
    return np.median(np.stack(out), 0).astype(np.uint8), len(out)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dataset")
    ap.add_argument("--conf", help="fw ml/calibrations/<conf>.json for the eyelets")
    ap.add_argument("--calibrations", help="the firmware's ml/calibrations folder (default $NF_CALIBRATIONS)")
    ap.add_argument("--building", help="a scanned building (~/robot-lab/buildings/<name>): import one office of it")
    ap.add_argument("--office", help="the office's id in building.yaml (e.g. r5)")
    ap.add_argument("--anchor-z", type=float, default=2.65, help="mount height: ceiling 2.85 less the winch block")
    ap.add_argument("--name", required=True)
    args = ap.parse_args()
    if args.building:
        return import_building(args)
    import pandas as pd
    ds = Path(args.dataset)
    info = json.loads((ds / "meta/info.json").read_text())
    df = pd.read_parquet(ds / "data/chunk-000/file-000.parquet")
    names = info["features"]["observation.state"]["names"]
    S = np.stack(df["observation.state"].values)
    col = lambda k: S[:, names.index(k)]  # noqa: E731
    ap_ = np.stack(df["anchor_poses"].values)
    assert len(np.unique(ap_.round(4), axis=0)) == 1, "anchors moved within the recording"
    a = ap_[0].reshape(2, 6)                                     # rx, ry, rz, x, y, z per anchor
    anchors = a[:, 3:6]

    places = {}
    g = np.c_[col("gripper_pos_x"), col("gripper_pos_y")]
    for p in PLACES:
        ang = col(f"{p}_bearing") + col("spin")
        P = g + col(f"{p}_distance")[:, None] * np.c_[np.sin(ang), np.cos(ang)]
        med, sd = np.median(P, 0), P.std(0).max()
        if sd < 0.02 and np.linalg.norm(med) > 0.05:
            places[p] = [round(float(med[0]), 3), round(float(med[1]), 3)]

    conf = json.loads((Path(args.calibrations or FW_CAL) / f"{args.conf}.json").read_text())
    ca = np.array([[x["pose"]["position"][k] if "pose" in x else x["position"][k] for k in "xy"] for x in conf["anchors"]])
    ce = np.array([[x["indirectLine"]["eyeletPos"][k] for k in "xy"] for x in conf["anchors"]])
    f, scale = similarity_2d(ca, anchors[:, :2])
    eyelets = np.array([f(e) for e in ce])
    z = anchors[:, 2]
    # line order (fw-sim): 0 anchor0 direct, 1 anchor0 indirect via its eyelet, 2 anchor1 direct, 3 anchor1 indirect
    mounts = [anchors[0, :2], eyelets[0], anchors[1, :2], eyelets[1]]
    mz = [z[0], z[0], z[1], z[1]]
    runs = [0.0, float(np.linalg.norm(eyelets[0] - anchors[0, :2])), 0.0, float(np.linalg.norm(eyelets[1] - anchors[1, :2]))]

    img, n = floor_median(ds)
    out_dir = ROOT / "assets/rooms" / args.name
    out_dir.mkdir(parents=True, exist_ok=True)
    from PIL import Image
    Image.fromarray(img).save(out_dir / "floor.png")
    gan = np.c_[col("gantry_position_x"), col("gantry_position_y"), col("gantry_position_z")]
    room = {
        "source": {"dataset": f"naavox/{ds.name}", "frames": int(len(df)), "conf_for_eyelets": args.conf,
                   "eyelet_similarity_scale": round(float(scale), 3), "floor_median_of": int(n)},
        "mounts_xyz": [[round(float(m[0]), 3), round(float(m[1]), 3), round(float(zz), 3)] for m, zz in zip(mounts, mz)],
        "indirect_run_m": [round(r, 3) for r in runs],
        "ceiling_m": round(float(z.max()) + 0.13, 3),           # docs: anchors ~13 cm below the ceiling
        "named_places": places,
        "floor_image": {"path": f"assets/rooms/{args.name}/floor.png", "extent_m": EXTENT_M, "centre": [0.0, 0.0]},
        "recorded_gantry_box": [[round(float(gan[:, i].min()), 2), round(float(gan[:, i].max()), 2)] for i in range(3)],
        "anchor_yaw_rad": [round(float(a[0, 2]), 3), round(float(a[1, 2]), 3)],
    }
    path = ROOT / "configs/rooms" / f"{args.name}.yaml"
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(f"# Imported by scripts/import_room.py from naavox/{ds.name} (see its docstring for each field's source).\n"
                    + yaml.safe_dump(room, sort_keys=False))
    print(path.read_text())


# ----------------------------------------------------------------------------------------------- a scanned building
def floor_field(P, box, fz, cell=0.25):
    """The scanned floor's height as a smooth field f(x, y): the scan is not flat. r5's floor points wander from -6.7 to
    +1.0 cm (0.6 m cell medians, 2026-10-13; a RANSAC plane came out 1.05 deg tilted and left the open floor at
    +1.7 cm median): Polycam drift, not a sloping office. Per 25 cm cell: the median of the points within 4 cm of
    the cell's 10th percentile (the floor, not a chair leg), cells with < 20 such points filled from the nearest
    measured ones (under the table), then smoothed (sigma 2 cells). Returns f (callable on (n, 2)), its grid."""
    from scipy import ndimage
    from scipy.interpolate import NearestNDInterpolator, RegularGridInterpolator
    x0, x1, y0, y1 = box
    Q = P[(P[:, 0] > x0) & (P[:, 0] < x1) & (P[:, 1] > y0) & (P[:, 1] < y1) & (np.abs(P[:, 2] - fz) < 0.12)]
    nx, ny = int(np.ceil((x1 - x0) / cell)), int(np.ceil((y1 - y0) / cell))
    ij = np.minimum(((Q[:, :2] - [x0, y0]) / cell).astype(int), [nx - 1, ny - 1])
    G = np.full((nx, ny), np.nan)
    key = ij[:, 0] * ny + ij[:, 1]
    order = np.argsort(key)
    ks, starts = np.unique(key[order], return_index=True)
    for k, a, b in zip(ks, starts, np.r_[starts[1:], len(order)]):
        z = Q[order[a:b], 2]
        lo = np.percentile(z, 10)
        z = z[z < lo + 0.04]
        if len(z) >= 20:
            G[k // ny, k % ny] = np.median(z)
    ok = ~np.isnan(G)
    ii, jj = np.nonzero(ok)
    fill = NearestNDInterpolator(np.c_[ii, jj], G[ok])
    ai, aj = np.nonzero(~ok)
    if len(ai):
        G[~ok] = fill(np.c_[ai, aj])
    G = ndimage.gaussian_filter(G, 2.0, mode="nearest")
    xs, ys = x0 + (np.arange(nx) + 0.5) * cell, y0 + (np.arange(ny) + 0.5) * cell
    f = RegularGridInterpolator((xs, ys), G, bounds_error=False, fill_value=None)
    return lambda xy: f(np.clip(xy, [xs[0], ys[0]], [xs[-1], ys[-1]])), G, float(ok.mean())


def import_building(args):
    """One office of a scanned building -> a room for the cable robot.

      - mounts: the office's four corners at anchor height, found from the scan, not typed in: the free space in the
        band 1.0 m .. anchor_z + 0.15 m (wall points occupy a 4 cm grid; the anchor band alone leaked out of r5
        through two scan holes in its walls, 2026-10-13), the connected part around the office centre,
        eroded 5 cm (fw definitions: eyelets 5 cm off the walls); the corner = the free cell furthest along
        (+-1, +-1). Order as scripts/make_room3d.py (CCW from +x+y; lines 0, 2 direct, 1, 3 via eyelets along a wall).
      - the room frame: origin on the floor under the mounts' centre, axes = the building's (pose in building.yaml)
      - the visual mesh: the OBJ's faces in the office box (+30 cm: its walls), the ceiling dropped (the overhead
        camera hangs under it), one mesh per texture. VISUAL ONLY in the plant (contype 0).
      - the fixed map's points: the PLY in the same crop (data/rooms/<name>/fixed_points.npy, xyz + rgb).
      - supports: the floor and each large horizontal surface (upward mesh faces 0.3-1.5 m, clustered by height and
        place, >= 0.25 m^2): a table top's height and box. The plant gets a thin collision box under each so waste
        can lie on it (the mesh itself has no contacts).
      - the bin (named place `hamper`): free floor near a wall, >= 0.6 m inside the mount polygon, nothing of the
        scan 0.1-1.5 m above within 0.45 m."""
    from scipy import ndimage
    from cablepilot.roommap.scan import load_building, load_mtl, load_obj, load_ply, obj_to_frame, write_obj
    bdir, B = load_building(args.building)
    off = B["offices"][args.office]
    x0, x1, y0, y1 = off["box"]
    fz = float(B["floor_z"])
    P, C = load_ply(bdir / B["files"]["points"])
    pad = 0.3
    near = (P[:, 0] > x0 - pad) & (P[:, 0] < x1 + pad) & (P[:, 1] > y0 - pad) & (P[:, 1] < y1 + pad)
    # the office's own floor, FOUND (armpilot scan_furniture: RANSAC, "found, not given"): the building's floor_z is
    # the whole floor's histogram peak, and r5's floor lies 2.2 cm under it (median, 2026-10-13); a level-floor
    # assumption left a speckle of "obstacle" voxels 2-4 cm up all over the room
    ffield, G, measured = floor_field(P, (x0 - pad, x1 + pad, y0 - pad, y1 + pad), fz)
    P = P.copy()
    P[near, 2] -= ffield(P[near, :2])                        # flatten (the same correction for the mesh below)
    fz_mean = float(np.mean(G))

    # mounts: free space at anchor height
    res = 0.04
    lo = np.array([x0 - pad, y0 - pad])
    shape = (np.array([x1 - x0, y1 - y0]) + 2 * pad) / res
    shape = shape.astype(int) + 1
    band = near & (P[:, 2] > 1.0) & (P[:, 2] < args.anchor_z + 0.15)
    occ = np.zeros(shape, bool)
    ij = ((P[band, :2] - lo) / res).astype(int)
    occ[ij[:, 0], ij[:, 1]] = True
    occ = ndimage.binary_closing(occ, iterations=2)          # the scan's 2.2 cm points leave pinholes in a wall
    lab, _ = ndimage.label(~occ)
    c_ij = ((np.array([(x0 + x1) / 2, (y0 + y1) / 2]) - lo) / res).astype(int)
    free = ndimage.binary_erosion(lab == lab[tuple(c_ij)], iterations=2)   # 2 cells = 8 cm from a wall point's cell
    fi, fj = np.nonzero(free)
    cells = lo + (np.c_[fi, fj] + 0.5) * res
    M = []
    for d in ((1, 1), (-1, 1), (-1, -1), (1, -1)):
        M.append(cells[np.argmax(cells @ np.array(d, float))])
    M = np.array(M)
    O = M.mean(0).round(3)                                   # the room frame's origin (floor under the mounts' centre)
    mounts = np.c_[M - O, np.full(4, args.anchor_z)]
    runs = [0.0, float(np.linalg.norm(M[1] - M[0])), 0.0, float(np.linalg.norm(M[3] - M[2]))]

    out = ROOT / "data/rooms" / args.name
    out.mkdir(parents=True, exist_ok=True)
    keep = near & (P[:, 2] > -0.05) & (P[:, 2] < off["ceiling_m"] - 0.1)
    pts = np.c_[P[keep] - [*O, 0], C[keep]].astype(np.float32)
    np.save(out / "fixed_points.npy", pts)

    V, T, Fc, Mt = load_obj(bdir / B["files"]["mesh"])
    V = obj_to_frame(V)                                      # the mesh shares the PLY's frame (same capture)
    V[:, 2] -= ffield(V[:, :2])
    V = V - [*O, 0]
    A = V[Fc[:, :, 0]]
    nrm = np.cross(A[:, 1] - A[:, 0], A[:, 2] - A[:, 0])
    area = np.linalg.norm(nrm, axis=1) / 2
    up = nrm[:, 2] / (2 * area + 1e-12)
    cen = A.mean(1)
    bx0, bx1, by0, by1 = x0 - O[0] - pad, x1 - O[0] + pad, y0 - O[1] - pad, y1 - O[1] + pad
    inbox = (cen[:, 0] > bx0) & (cen[:, 0] < bx1) & (cen[:, 1] > by0) & (cen[:, 1] < by1)
    keepF = inbox & ~((cen[:, 2] > off["ceiling_m"] - 0.25) & (up < -0.5))
    mtl = load_mtl(bdir / B["files"]["mesh"].replace(" ", "_").replace(".obj", ".mtl"))
    meshes = []
    for k, mat in enumerate(sorted(set(Mt[keepF]))):
        sel = keepF & (Mt == mat)
        n = write_obj(out / f"mesh_{k}.obj", V, T, Fc[sel])
        from PIL import Image                                 # MuJoCo reads PNG textures, not the scan's JPEGs
        Image.open(bdir / mtl[mat]).convert("RGB").save(out / f"mesh_{k}.png")
        meshes.append({"obj": f"mesh_{k}.obj", "texture": f"mesh_{k}.png", "faces": int(n)})

    # supports: large horizontal surfaces above the floor
    hz = keepF & (up > 0.95) & (cen[:, 2] > 0.3) & (cen[:, 2] < 1.5)
    supports = [{"name": "floor", "z": 0.0}]
    g = 0.05
    zb = np.round(cen[hz, 2] / 0.02).astype(int)
    for zk in np.unique(zb):
        s = np.nonzero(hz)[0][np.abs(zb - zk) <= 1]
        if area[s].sum() < 0.25:
            continue
        q = ((cen[s, :2] - [bx0, by0]) / g).astype(int)
        grid = np.zeros((int((bx1 - bx0) / g) + 1, int((by1 - by0) / g) + 1), bool)
        grid[q[:, 0], q[:, 1]] = True
        grid = ndimage.binary_closing(grid, iterations=2)
        lab, nl = ndimage.label(grid)
        for L in range(1, nl + 1):
            mem = lab[q[:, 0], q[:, 1]] == L
            if area[s[mem]].sum() < 0.25 or any(abs(sp["z"] - np.average(cen[s[mem], 2], weights=area[s[mem]])) < 0.03
                                                 and sp.get("_l") == (zk, L) for sp in supports):
                continue
            xy = cen[s[mem], :2]
            zz = float(np.average(cen[s[mem], 2], weights=area[s[mem]]))
            if any(sp["name"] != "floor" and abs(sp["z"] - zz) < 0.03 and np.all(np.abs(np.mean(xy, 0) - np.mean(sp["box"], 1)) < 0.3)
                   for sp in supports):
                continue                                      # the same table seen in the neighbouring height bin
            bxy = np.percentile(xy, [3, 97], axis=0)
            supports.append({"name": f"top{len(supports)}", "z": round(zz, 3),
                             "box": [[round(float(bxy[0, 0]), 3), round(float(bxy[1, 0]), 3)],
                                     [round(float(bxy[0, 1]), 3), round(float(bxy[1, 1]), 3)]],
                             "area_m2": round(float(area[s[mem]].sum()), 2)})

    # one table, one support: the height bins overlap (r5's table came out three times, 0.710 / 0.717 / 0.718 m,
    # boxes inside each other, 2026-10-13): merge same-height (3 cm) supports whose boxes overlap into their union
    tops = [sp for sp in supports if "box" in sp]
    merged = True
    while merged:
        merged = False
        for i in range(len(tops)):
            for j in range(i + 1, len(tops)):
                a, b = tops[i], tops[j]
                (ax0, ax1), (ay0, ay1) = a["box"]
                (bx0_, bx1_), (by0_, by1_) = b["box"]
                if abs(a["z"] - b["z"]) < 0.03 and ax0 < bx1_ and bx0_ < ax1 and ay0 < by1_ and by0_ < ay1:
                    w = a["area_m2"] + b["area_m2"]
                    a["z"] = round((a["z"] * a["area_m2"] + b["z"] * b["area_m2"]) / w, 3)
                    a["box"] = [[min(ax0, bx0_), max(ax1, bx1_)], [min(ay0, by0_), max(ay1, by1_)]]
                    a["area_m2"] = round(max(a["area_m2"], b["area_m2"]), 2)       # nested bins: the larger, not the sum
                    tops.pop(j)
                    merged = True
                    break
            if merged:
                break
    supports = [supports[0]] + [{**t, "name": f"table{k}"} for k, t in enumerate(tops)]

    # the bin: free floor near a wall
    from cablepilot.sim.kin import ccw_polygon
    poly = ccw_polygon(mounts[:, :2])
    e = np.roll(poly, -1, 0) - poly
    nr = np.c_[-e[:, 1], e[:, 0]] / np.linalg.norm(e, axis=1)[:, None]
    tall = pts[(pts[:, 2] > 0.1) & (pts[:, 2] < 1.5), :2]
    from scipy.spatial import cKDTree
    tree = cKDTree(tall)
    best = None
    for x in np.arange(poly[:, 0].min(), poly[:, 0].max(), 0.05):
        for y in np.arange(poly[:, 1].min(), poly[:, 1].max(), 0.05):
            q = np.array([x, y])
            ins = ((q - poly) * nr).sum(1).min()
            if ins < 0.6 or tree.query(q)[0] < 0.45:
                continue
            if best is None or ins < best[0]:                 # nearest a wall (yet >= 0.6 m in)
                best = (ins, q)
    hamper = best[1].round(3).tolist()

    room = {"source": {"building": str(bdir).replace(str(Path.home()), "~"), "office": args.office, "imported_by": "scripts/import_room.py --building",
                       "scan_points_kept": int(keep.sum())},
            "floor_flattened": {"note": "scan z - f(x, y), f = the scanned floor's smooth height (floor_field)",
                                "f_range_m": [round(float(G.min()), 4), round(float(G.max()), 4)],
                                "cells_measured": round(measured, 2)},
            "building_pose": {"xyz": [float(O[0]), float(O[1]), round(fz_mean, 4)], "yaw_rad": 0.0,
                              "note": "room frame = building frame shifted by xyz (floor at z 0); same axes"},
            "size_m": [round(float(np.ptp(M[:, 0])), 3), round(float(np.ptp(M[:, 1])), 3)],
            "mounts_xyz": mounts.round(3).tolist(), "indirect_run_m": [round(r, 3) for r in runs],
            "ceiling_m": float(off["ceiling_m"]),
            "named_places": {"hamper": hamper},
            "scan": {"dir": f"data/rooms/{args.name}", "meshes": meshes, "points": "fixed_points.npy",
                     "box": [[round(float(bx0), 3), round(float(bx1), 3)], [round(float(by0), 3), round(float(by1), 3)]]},
            "supports": supports}
    path = ROOT / "configs/rooms" / f"{args.name}.yaml"
    path.write_text(f"# One office of a scanned building ({B['name']}, office {args.office}), imported by "
                    f"scripts/import_room.py --building (see import_building's docstring for each field's source).\n"
                    + yaml.safe_dump(room, sort_keys=False))
    print(path.read_text())


if __name__ == "__main__":
    main()