#!/usr/bin/env python3
"""
Wrapper ORB-SLAM3 — nœud Python minimal.
Lit les frames stéréo, appelle ORB-SLAM3 via bindings C++ (à compiler).
Ce fichier est un placeholder fonctionnel jusqu'à compilation ORB-SLAM3.
"""

import cv2
import time
import os
import sys
import yaml
import numpy as np

CONFIG_YAML = os.path.join(
    os.path.dirname(__file__), "..", "..", "config", "stereo_calib.yaml"
)

def load_calib(path):
    """Charge les paramètres de calibration stéréo."""
    if not os.path.isfile(path):
        print(f"[ERREUR] {path} introuvable. Lancer la calibration d'abord.")
        sys.exit(1)
    with open(path) as f:
        return yaml.safe_load(f)


def rectify_pair(frame_l, frame_r, calib):
    """Rectifie une paire stéréo."""
    sz = tuple(calib["image_size"])
    K1 = np.array(calib["K1"])
    D1 = np.array(calib["D1"])
    K2 = np.array(calib["K2"])
    D2 = np.array(calib["D2"])
    R1 = np.array(calib["R1"])
    R2 = np.array(calib["R2"])
    P1 = np.array(calib["P1"])
    P2 = np.array(calib["P2"])

    map1_l, map2_l = cv2.initUndistortRectifyMap(K1, D1, R1, P1, sz, cv2.CV_16SC2)
    map1_r, map2_r = cv2.initUndistortRectifyMap(K2, D2, R2, P2, sz, cv2.CV_16SC2)

    rect_l = cv2.remap(frame_l, map1_l, map2_l, cv2.INTER_LINEAR)
    rect_r = cv2.remap(frame_r, map1_r, map2_r, cv2.INTER_LINEAR)
    return rect_l, rect_r


def compute_disparity(rect_l, rect_r):
    """Calcule la carte de disparité (SGBM)."""
    gray_l = cv2.cvtColor(rect_l, cv2.COLOR_BGR2GRAY)
    gray_r = cv2.cvtColor(rect_r, cv2.COLOR_BGR2GRAY)

    stereo = cv2.StereoSGBM_create(
        minDisparity=0,
        numDisparities=64,       # multiple de 16
        blockSize=9,
        P1=8 * 3 * 9 ** 2,
        P2=32 * 3 * 9 ** 2,
        disp12MaxDiff=1,
        uniquenessRatio=10,
        speckleWindowSize=100,
        speckleRange=32,
        mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY
    )
    disparity = stereo.compute(gray_l, gray_r).astype(np.float32) / 16.0
    return disparity


def main():
    print("[SLAM] Démarrage du nœud SLAM stéréo...")
    calib = load_calib(CONFIG_YAML)

    cap_l = cv2.VideoCapture(0)
    cap_r = cv2.VideoCapture(1)

    if not cap_l.isOpened() or not cap_r.isOpened():
        print("[ERREUR] Caméras non disponibles")
        sys.exit(1)

    print("[SLAM] Caméras ouvertes. Ctrl+C pour arrêter.")

    try:
        while True:
            ret_l, frame_l = cap_l.read()
            ret_r, frame_r = cap_r.read()

            if not ret_l or not ret_r:
                time.sleep(0.05)
                continue

            t = time.time()

            # Rectification
            rect_l, rect_r = rectify_pair(frame_l, frame_r, calib)

            # Disparité (fallback si ORB-SLAM3 pas compilé)
            disparity = compute_disparity(rect_l, rect_r)

            # TODO: appeler ORB-SLAM3 C++ bindings quand disponible
            # pose = orbslam3.track_stereo(rect_l, rect_r, t)

            fps = 1.0 / (time.time() - t + 1e-9)
            print(f"\r[SLAM] fps={fps:.1f} disp_mean={disparity[disparity>0].mean():.1f}px",
                  end="", flush=True)

    except KeyboardInterrupt:
        print("\n[SLAM] Arrêt.")
    finally:
        cap_l.release()
        cap_r.release()


if __name__ == "__main__":
    main()