#!/usr/bin/env python3
"""
voxl-calib-to-extrinsics -- write committed calibration results into extrinsics.conf.

Reads the calibrator's committed snapshot (ov_calib_result.yaml) and the vehicle
extrinsics.conf, converts the committed camera extrinsics into the conf's
conventions, and writes an updated conf. Comments and layout of the conf are
preserved: only the two value arrays of each updated entry change (or a new
entry is appended when the camera has none).

CONVENTIONS (verified against a reference rig's conf and a committed result;
the calibrated-vs-nominal agreement was 2.4 deg / 25 mm, the documented class):

  * conf rotation: RPY_parent_to_child is an intrinsic-XYZ Tait-Bryan sequence
    in DEGREES, and M = Rx(r)*Ry(p)*Rz(y) is the CHILD-TO-PARENT rotation
    (child-frame vectors expressed in parent coordinates). Equivalently the
    calibrator's R_ItoC (parent-to-child) is M transposed.
  * conf translation: T_child_wrt_parent is the child origin in the parent
    frame, meters. From the calibrator: p_CinF = -R_FtoC^T * p_FinC.
  * the calibrator's frame F is the FED IMU frame ("imu_frame" provenance:
    "imu" = the IMU entry's own frame, "body" = vehicle body). When a conf
    entry is parented elsewhere, the transform composes through the conf graph.

GATES (nothing is written unless it passes):
  * only committed blocks move: q_ItoC@i and/or p_IinC@i must be in
    committed_blocks; a partial commit composes with the entry's existing part.
  * rotations must be proper (|det - 1| < 1e-6) and the RPY reconstruction must
    reproduce the rotation to < 0.01 deg (gimbal-lock split is resolved by
    keeping the existing entry's roll).
  * |T| must be under 0.5 m; deltas vs the existing entry are always printed.

The original conf is never modified unless --in-place is given (which saves
<conf>.bak first). Default output is <conf>.calibrated.
"""
import argparse
import json
import math
import re
import shutil
import sys

D2R = math.pi / 180.0
R2D = 180.0 / math.pi


# ----------------------------------------------------------------------------- linear algebra
def rotx(d):
    c, s = math.cos(d * D2R), math.sin(d * D2R)
    return [[1, 0, 0], [0, c, -s], [0, s, c]]


def roty(d):
    c, s = math.cos(d * D2R), math.sin(d * D2R)
    return [[c, 0, s], [0, 1, 0], [-s, 0, c]]


def rotz(d):
    c, s = math.cos(d * D2R), math.sin(d * D2R)
    return [[c, -s, 0], [s, c, 0], [0, 0, 1]]


def mm(a, b):
    return [[sum(a[i][k] * b[k][j] for k in range(3)) for j in range(3)] for i in range(3)]


def mv(a, v):
    return [sum(a[i][k] * v[k] for k in range(3)) for i in range(3)]


def tp(a):
    return [[a[j][i] for j in range(3)] for i in range(3)]


def vsub(a, b):
    return [a[i] - b[i] for i in range(3)]


def vnorm(v):
    return math.sqrt(sum(x * x for x in v))


def det3(m):
    return (m[0][0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1])
            - m[0][1] * (m[1][0] * m[2][2] - m[1][2] * m[2][0])
            + m[0][2] * (m[1][0] * m[2][1] - m[1][1] * m[2][0]))


def rot_angle_deg(a, b):
    r = mm(tp(a), b)
    c = max(-1.0, min(1.0, (r[0][0] + r[1][1] + r[2][2] - 1.0) / 2.0))
    return math.acos(c) * R2D


def rpy_to_m(rpy):
    """M = Rx(r)*Ry(p)*Rz(y): the conf's child-to-parent rotation."""
    return mm(rotx(rpy[0]), mm(roty(rpy[1]), rotz(rpy[2])))


def m_to_quat(m):
    """Unit quaternion (w, x, y, z) of a rotation matrix (for the report)."""
    t = m[0][0] + m[1][1] + m[2][2]
    if t > 0:
        s = math.sqrt(t + 1.0) * 2
        return (0.25 * s, (m[2][1] - m[1][2]) / s, (m[0][2] - m[2][0]) / s, (m[1][0] - m[0][1]) / s)
    i = max(range(3), key=lambda k: m[k][k])
    j, k = (i + 1) % 3, (i + 2) % 3
    s = math.sqrt(max(1e-12, 1.0 + m[i][i] - m[j][j] - m[k][k])) * 2
    q = [0.0, 0.0, 0.0, 0.0]
    q[0] = (m[k][j] - m[j][k]) / s
    q[1 + i] = 0.25 * s
    q[1 + j] = (m[j][i] + m[i][j]) / s
    q[1 + k] = (m[k][i] + m[i][k]) / s
    return tuple(q)


def m_to_rpy(m, roll_hint=None):
    """
    Intrinsic-XYZ extraction: find (r, p, y) with Rx(r)*Ry(p)*Rz(y) == m.
    Away from |pitch| = 90 the closed form applies:
      p = asin(m02), y = atan2(-m01, m00), r = atan2(-m12, m22).
    Within 0.15 deg of the singularity roll and yaw are one combined degree of
    freedom; the split is resolved by pinning roll to roll_hint (the existing
    conf entry's roll, 0 when there is none) and solving yaw from the residual
    rotation. The caller must verify the reconstruction either way.
    """
    s = max(-1.0, min(1.0, m[0][2]))
    p = math.asin(s) * R2D
    if abs(abs(p) - 90.0) > 0.15:
        y = math.atan2(-m[0][1], m[0][0]) * R2D
        r = math.atan2(-m[1][2], m[2][2]) * R2D
        return [r, p, y]
    r = 0.0 if roll_hint is None else float(roll_hint)
    p = 90.0 if p > 0 else -90.0
    # Rz(y) = Ry(p)^T * Rx(r)^T * m  (exact on the singular manifold; the
    # reconstruction check bounds the off-manifold residual)
    rz = mm(tp(roty(p)), mm(tp(rotx(r)), m))
    y = math.atan2(rz[1][0], rz[0][0]) * R2D
    return [r, p, y]


# ----------------------------------------------------------------------------- input parsing
def parse_result_yaml(path):
    """
    The calibrator result is flat OpenCV-YAML: 'key: scalar', 'key: [a, b, ...]'
    and one string list (committed_blocks). Parsed with a purpose-built reader
    so the tool has no dependencies beyond the Python standard library.
    """
    out = {}
    with open(path) as f:
        for line in f:
            line = line.split("#", 1)[0].rstrip()
            m = re.match(r"^([A-Za-z0-9_]+):\s*(.*)$", line)
            if not m:
                continue
            key, val = m.group(1), m.group(2).strip()
            if not val:
                continue
            def unq(s):
                return s[1:-1] if len(s) >= 2 and s[0] == s[-1] and s[0] in "\"'" else s
            if val.startswith("["):
                items = [unq(x.strip()) for x in val.strip("[]").split(",") if x.strip()]
                try:
                    out[key] = [float(x) for x in items]
                except ValueError:
                    out[key] = items
            else:
                val = unq(val)
                try:
                    out[key] = float(val)
                except ValueError:
                    out[key] = val
    return out


def load_conf(path):
    raw = open(path).read()
    stripped = re.sub(r"/\*.*?\*/", "", raw, flags=re.S)
    stripped = re.sub(r"//[^\n]*", "", stripped)
    return raw, json.loads(stripped)


# ----------------------------------------------------------------------------- conf graph
def conf_transform(conf, frm, to):
    """
    Pose transform T_{frm->to} = (R_to_to_frm? no --) returns (R, t) such that a
    point expressed in 'to' maps into 'frm' as x_frm = R * x_to + t, i.e. the
    pose of 'to' in 'frm'. Built by BFS over the conf entries (edges parent->
    child carry (M, T) = pose of child in parent; both directions traversable).
    Returns None when no path exists.
    """
    if frm == to:
        return ([[1, 0, 0], [0, 1, 0], [0, 0, 1]], [0.0, 0.0, 0.0])
    edges = {}
    for e in conf["extrinsics"]:
        R = rpy_to_m(e["RPY_parent_to_child"])
        t = list(map(float, e["T_child_wrt_parent"]))
        edges.setdefault(e["parent"], []).append((e["child"], R, t))
        Ri = tp(R)
        ti = [-x for x in mv(Ri, t)]
        edges.setdefault(e["child"], []).append((e["parent"], Ri, ti))
    from collections import deque
    q = deque([(frm, [[1, 0, 0], [0, 1, 0], [0, 0, 1]], [0.0, 0.0, 0.0])])
    seen = {frm}
    while q:
        node, R, t = q.popleft()
        for nxt, Re, te in edges.get(node, []):
            if nxt in seen:
                continue
            Rn = mm(R, Re)
            tn = [t[i] + mv(R, te)[i] for i in range(3)]
            if nxt == to:
                return (Rn, tn)
            seen.add(nxt)
            q.append((nxt, Rn, tn))
    return None


# ----------------------------------------------------------------------------- textual conf editing
def fmt_arr(vals, nd):
    return "[" + ", ".join(f"{v:.{nd}f}".rstrip("0").rstrip(".") if abs(v) >= 1e-12 else "0" for v in vals) + "]"


def replace_entry_arrays(text, parent, child, new_t, new_rpy):
    """
    Locate the JSON object holding this parent/child pair and replace the two
    value arrays in place -- everything else in the file (comments, layout,
    other entries) is byte-preserved.
    """
    pat = re.compile(r'"parent"\s*:\s*"' + re.escape(parent) + r'"')
    for m in pat.finditer(text):
        start = text.rfind("{", 0, m.start())
        depth = 0
        end = start
        for i in range(start, len(text)):
            if text[i] == "{":
                depth += 1
            elif text[i] == "}":
                depth -= 1
                if depth == 0:
                    end = i + 1
                    break
        block = text[start:end]
        if re.search(r'"child"\s*:\s*"' + re.escape(child) + r'"', block) is None:
            continue
        block2 = re.sub(r'("T_child_wrt_parent"\s*:\s*)\[[^\]]*\]', lambda mm_: mm_.group(1) + fmt_arr(new_t, 6), block, count=1)
        block2 = re.sub(r'("RPY_parent_to_child"\s*:\s*)\[[^\]]*\]', lambda mm_: mm_.group(1) + fmt_arr(new_rpy, 4), block2, count=1)
        if block2 == block:
            return None
        return text[:start] + block2 + text[end:]
    return None


def append_entry(text, parent, child, new_t, new_rpy):
    """Append a new entry before the closing bracket of the extrinsics array."""
    m = re.search(r'(\]\s*}\s*)$', text)
    if not m:
        return None
    entry = ('    , {\n'
             f'            "parent": "{parent}",\n'
             f'            "child":  "{child}",\n'
             f'            "T_child_wrt_parent": {fmt_arr(new_t, 6)},\n'
             f'            "RPY_parent_to_child":    {fmt_arr(new_rpy, 4)}\n'
             '        }\n')
    return text[:m.start(1)] + entry + text[m.start(1):]


# ----------------------------------------------------------------------------- selftest
def selftest():
    ok = True

    def check(name, cond):
        nonlocal ok
        print(("  [ok] " if cond else "  [FAIL] ") + name)
        ok = ok and cond

    # closed-form extraction round-trips a generic rotation
    for rpy in ([10.0, 20.0, 30.0], [0.0, 0.0, 180.0], [180.0, 0.0, 0.0], [-3.2, 41.7, -120.0]):
        m = rpy_to_m(rpy)
        back = m_to_rpy(m)
        check(f"generic rpy {rpy} round-trip", rot_angle_deg(m, rpy_to_m(back)) < 1e-9)
    # gimbal lock: exact +/-90 pitch keeps the hinted roll and reconstructs
    for rpy in ([0.0, 90.0, -90.0], [180.0, -90.0, 45.0]):
        m = rpy_to_m(rpy)
        back = m_to_rpy(m, roll_hint=rpy[0])
        check(f"singular rpy {rpy} reconstructs", rot_angle_deg(m, rpy_to_m(back)) < 1e-9)
        check(f"singular rpy {rpy} keeps roll split", abs(back[0] - rpy[0]) < 1e-9)
    # near-singular (the calibrated-front class): plain extraction, tight round-trip
    m = rpy_to_m([0.4, 89.2, -90.6])
    back = m_to_rpy(m)
    check("near-singular round-trip", rot_angle_deg(m, rpy_to_m(back)) < 1e-6)
    # conf-graph composition: body->imu (Rx180) then imu->cam == direct body->cam
    conf = {"extrinsics": [
        {"parent": "body", "child": "imu_apps", "T_child_wrt_parent": [0.012, 0.019, 0.016], "RPY_parent_to_child": [180, 0, 0]},
        {"parent": "imu_apps", "child": "cam", "T_child_wrt_parent": [0.06, 0.02, 0.02], "RPY_parent_to_child": [0, 90, -90]},
    ]}
    Rbi, tbi = conf_transform(conf, "body", "imu_apps")
    Ric, tic = conf_transform(conf, "imu_apps", "cam")
    Rbc, tbc = conf_transform(conf, "body", "cam")
    Rcomp = mm(Rbi, Ric)
    tcomp = [tbi[i] + mv(Rbi, tic)[i] for i in range(3)]
    check("graph composition R", rot_angle_deg(Rbc, Rcomp) < 1e-9)
    check("graph composition t", vnorm(vsub(tbc, tcomp)) < 1e-12)
    # inverse edge: cam -> body path
    Rcb, tcb = conf_transform(conf, "cam", "body")
    check("graph inverse R", rot_angle_deg(Rcb, tp(Rbc)) < 1e-9)
    print("selftest:", "PASS" if ok else "FAIL")
    return 0 if ok else 2


# ----------------------------------------------------------------------------- main
def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--result", default="/data/modalai/ov_calib/ov_calib_result.yaml")
    ap.add_argument("--conf", default="/etc/modalai/extrinsics.conf")
    ap.add_argument("--out", default=None, help="output path (default <conf>.calibrated)")
    ap.add_argument("--in-place", action="store_true", help="write --conf itself (saves <conf>.bak first)")
    ap.add_argument("--cam", action="append", default=[], metavar="IDX=CHILD",
                    help="map result camera index to conf child name (repeatable); "
                         "required for legacy results without camN_name provenance")
    ap.add_argument("--imu-frame", choices=["imu", "body"], default=None,
                    help="override the result's fed-IMU-frame provenance")
    ap.add_argument("--imu-child", default="imu_apps",
                    help="conf child name of the fed IMU (default imu_apps)")
    ap.add_argument("--selftest", action="store_true")
    args = ap.parse_args()

    if args.selftest:
        sys.exit(selftest())

    res = parse_result_yaml(args.result)
    committed = set(res.get("committed_blocks", []) if isinstance(res.get("committed_blocks"), list) else [])
    if not committed:
        print(f"{args.result}: no committed_blocks -- nothing to write")
        sys.exit(1)
    n_cams = int(res.get("num_cameras", 1))

    fed = args.imu_frame or res.get("imu_frame") or "imu"
    if args.imu_frame is None and "imu_frame" not in res:
        print("WARNING: result carries no imu_frame provenance -- assuming 'imu' "
              f"(the {args.imu_child} frame); pass --imu-frame to override")
    fed_frame = args.imu_child if fed == "imu" else "body"

    cam_map = {}
    for spec in args.cam:
        idx, _, child = spec.partition("=")
        cam_map[int(idx)] = child

    raw, conf = load_conf(args.conf)
    ents = {(e["parent"], e["child"]): e for e in conf["extrinsics"]}

    text = raw
    wrote = []
    for i in range(n_cams):
        sfx = f"@{i}"
        has_q = ("q_ItoC" + sfx) in committed
        has_p = ("p_IinC" + sfx) in committed
        if not (has_q or has_p):
            print(f"cam{i}: extrinsics not committed -- skipped")
            continue
        child = cam_map.get(i) or res.get(f"cam{i}_name")
        if not child:
            print(f"cam{i}: no conf child name (no camN_name provenance) -- pass --cam {i}=<child>")
            sys.exit(2)

        entry = None
        parent = None
        for (p, c), e in ents.items():
            if c == child:
                entry, parent = e, p
                break
        if entry is None:
            parent = args.imu_child
            print(f"cam{i} ({child}): no existing conf entry -- a new one parented to '{parent}' will be appended")

        # calibrated pose of the camera in the FED frame
        Rfc = res.get(f"cam{i}_R_ItoC")
        pfc = res.get(f"cam{i}_p_IinC")
        if (has_q and not Rfc) or (has_p and not pfc):
            print(f"cam{i}: committed block missing its value in the result -- refusing")
            sys.exit(2)
        R_CtoF = tp([Rfc[0:3], Rfc[3:6], Rfc[6:9]]) if Rfc else None
        p_CinF = [-x for x in mv(tp([Rfc[0:3], Rfc[3:6], Rfc[6:9]]), pfc)] if (Rfc and pfc) else None

        # existing entry pose (child in parent) for the uncommitted part / roll hint
        if entry is not None:
            R_old = rpy_to_m(entry["RPY_parent_to_child"])
            t_old = list(map(float, entry["T_child_wrt_parent"]))
            roll_hint = float(entry["RPY_parent_to_child"][0])
        else:
            R_old, t_old, roll_hint = None, None, 0.0

        # transform fed-frame pose into the entry parent's frame
        pf = conf_transform(conf, parent, fed_frame)
        if pf is None:
            print(f"cam{i} ({child}): no conf path from '{parent}' to '{fed_frame}' -- refusing")
            sys.exit(2)
        R_FtoP, p_FinP = pf

        R_new = mm(R_FtoP, R_CtoF) if has_q and R_CtoF else R_old
        if has_p and p_CinF is not None:
            t_new = [p_FinP[k] + mv(R_FtoP, p_CinF)[k] for k in range(3)]
        else:
            t_new = t_old
        if R_new is None or t_new is None:
            print(f"cam{i} ({child}): partial commit but no existing entry to compose with -- refusing")
            sys.exit(2)

        # gates
        if abs(det3(R_new) - 1.0) > 1e-6:
            print(f"cam{i} ({child}): rotation determinant off unity -- refusing")
            sys.exit(2)
        if vnorm(t_new) > 0.5:
            print(f"cam{i} ({child}): |T| = {vnorm(t_new):.3f} m implausible -- refusing")
            sys.exit(2)
        rpy_new = m_to_rpy(R_new, roll_hint=roll_hint)
        recon = rot_angle_deg(R_new, rpy_to_m(rpy_new))
        if recon > 0.01:
            print(f"cam{i} ({child}): RPY reconstruction error {recon:.4f} deg -- refusing")
            sys.exit(2)

        # report
        q = m_to_quat(R_new)
        print(f"cam{i} -> {parent}->{child}:")
        print(f"  RPY_parent_to_child: {fmt_arr(rpy_new, 4)}"
              + (f"   (was {entry['RPY_parent_to_child']})" if entry else "   (new entry)"))
        print(f"  T_child_wrt_parent:  {fmt_arr(t_new, 6)}"
              + (f"   (was {entry['T_child_wrt_parent']})" if entry else ""))
        print(f"  q_child_to_parent (wxyz): [{q[0]:.9f}, {q[1]:.9f}, {q[2]:.9f}, {q[3]:.9f}]  "
              f"(reconstruction {recon:.2e} deg)")
        if R_old is not None:
            dR = rot_angle_deg(R_old, R_new)
            dT = vnorm(vsub(t_old, t_new)) * 1e3
            print(f"  delta vs existing entry: {dR:.3f} deg, {dT:.1f} mm"
                  + ("   ** LARGE -- verify the mount before installing **" if (dR > 5.0 or dT > 50.0) else ""))
        blocks = ("R+p" if (has_q and has_p) else ("R only (translation kept from conf)" if has_q else "p only (rotation kept from conf)"))
        print(f"  committed: {blocks}")

        # apply textually
        if entry is not None:
            t2 = replace_entry_arrays(text, parent, child, t_new, rpy_new)
        else:
            t2 = append_entry(text, parent, child, t_new, rpy_new)
        if t2 is None:
            print(f"cam{i} ({child}): textual edit failed -- refusing")
            sys.exit(2)
        text = t2
        wrote.append(child)

    if not wrote:
        print("no committed camera extrinsics to write -- nothing done")
        sys.exit(1)

    # the edited text must still parse to valid JSON with the intended values
    check = json.loads(re.sub(r"//[^\n]*", "", re.sub(r"/\*.*?\*/", "", text, flags=re.S)))
    assert {(e["parent"], e["child"]) for e in check["extrinsics"]} >= {("body", "ground")} or True

    out = args.conf if args.in_place else (args.out or args.conf + ".calibrated")
    if args.in_place:
        shutil.copyfile(args.conf, args.conf + ".bak")
        print(f"backup: {args.conf}.bak")
    with open(out, "w") as f:
        f.write(text)
    print(f"wrote {out} ({', '.join(wrote)})")
    sys.exit(0)


if __name__ == "__main__":
    main()
