#!/usr/bin/env python3
# Nombres en inglés; comentarios y prints en español.

from __future__ import annotations
from pathlib import Path
import sys
import math
import numpy as np
import pyzed.sl as sl
import laspy

# ─────────────────────────────────────────────────────────────────────────────
# CONFIG
# ─────────────────────────────────────────────────────────────────────────────
SVO_PATH = Path("Input/idioteria_dia4_cam2_35689352_Patinete_3.svo2")
FRAME_LIST = [1312, 2515]              # frames a exportar
PITCH_DEG = -20.0                      # rotación inicial (pitch negativo)
OUTPUT_PREFIX = "camframe_initpitch_m20"  # prefijo salida .laz

# Cámara: bloqueo de auto-exposición/ganancia/WB para consistencia (opcional)
EXPOSURE_0_100 = 60
GAIN_0_100 = 60
WB_TEMP_K = 4600

# ─────────────────────────────────────────────────────────────────────────────
# HELPERS
# ─────────────────────────────────────────────────────────────────────────────
def rx_matrix(deg: float) -> np.ndarray:
    """Matriz Rx (eje X) en grados, para depuración/prints (no aplicada a la nube)."""
    a = math.radians(deg)
    return np.array([
        [1.0, 0.0, 0.0],
        [0.0, math.cos(a), -math.sin(a)],
        [0.0, math.sin(a),  math.cos(a)],
    ], dtype=np.float32)

def rotation_to_euler_xyz(R: np.ndarray) -> tuple[float, float, float]:
    """Convierte matriz R a Euler XYZ (grados), útil para depuración humana."""
    sy = math.hypot(R[0, 0], R[1, 0])
    singular = sy < 1e-6
    if not singular:
        x = math.atan2(R[2, 1], R[2, 2])
        y = math.atan2(-R[2, 0], sy)
        z = math.atan2(R[1, 0], R[0, 0])
    else:
        x = math.atan2(-R[1, 2], R[1, 1])
        y = math.atan2(-R[2, 0], sy)
        z = 0.0
    return tuple(math.degrees(v) for v in (x, y, z))

def remap_camera_to_las(xyz_cam: np.ndarray) -> np.ndarray:
    """Frame CAMERA (X der, Y abajo, Z adelante) → LAS (Z arriba) para Potree."""
    X_out = xyz_cam[:, 0]
    Y_out = xyz_cam[:, 2]
    Z_out = -xyz_cam[:, 1]
    return np.stack([X_out, Y_out, Z_out], axis=1)

def rgba_to_rgb_u16(rgba_flat: np.ndarray) -> np.ndarray:
    """Convierte RGBA uint8 [0..255] a RGB uint16 [0..65535]."""
    return rgba_flat[:, :3].astype(np.uint16) * 256

def write_laz(points_xyz: np.ndarray, colors_rgb_u16: np.ndarray, out_path: Path) -> None:
    """Exporta a LAZ si hay backend; si no, a LAS con aviso."""
    header = laspy.LasHeader(point_format=7, version="1.4")
    las = laspy.LasData(header)
    mins = points_xyz.min(axis=0)
    las.header.offsets = mins
    las.header.scales = np.array([0.001, 0.001, 0.001])  # 1 mm
    las.x, las.y, las.z = points_xyz[:, 0], points_xyz[:, 1], points_xyz[:, 2]
    las.red, las.green, las.blue = colors_rgb_u16[:, 0], colors_rgb_u16[:, 1], colors_rgb_u16[:, 2]
    try:
        las.write(out_path.as_posix())
        print(f"[OK] Nube guardada en: {out_path}")
    except Exception as e:
        print(f"[AVISO] LAZ no disponible ({e}), escribiendo .las", file=sys.stderr)
        out_las = out_path.with_suffix(".las")
        las.write(out_las.as_posix())
        print(f"[OK] Nube guardada en: {out_las}")

# ─────────────────────────────────────────────────────────────────────────────
# MAIN
# ─────────────────────────────────────────────────────────────────────────────
def main() -> None:
    if not SVO_PATH.exists():
        print(f"[ERROR] No existe el fichero SVO/SVO2: {SVO_PATH}", file=sys.stderr)
        sys.exit(1)

    # 1) Apertura “congelada” + rango 0.5–10 m
    print("[INFO] Abriendo SVO con profundidad CONGELADA (sin estabilización, sin auto-calib, sin realce)...")
    init = sl.InitParameters()
    init.set_from_svo_file(SVO_PATH.as_posix())
    init.coordinate_units = sl.UNIT.METER
    init.depth_mode = sl.DEPTH_MODE.NEURAL_PLUS
    init.depth_minimum_distance = 0.5
    init.depth_maximum_distance = 10.0
    init.depth_stabilization = 0
    init.enable_image_enhancement = False
    init.camera_disable_self_calib = True
    init.svo_real_time_mode = False
    init.sdk_verbose = True

    zed = sl.Camera()
    err = zed.open(init)
    if err != sl.ERROR_CODE.SUCCESS:
        print(f"[ERROR] No se pudo abrir la SVO: {err}", file=sys.stderr)
        sys.exit(2)

    # (Opcional) Bloqueo de imagen para consistencia
    print("[INFO] Bloqueando exposición/ganancia/WB (sin auto) para consistencia...")
    zed.set_camera_settings(sl.VIDEO_SETTINGS.AEC_AGC, 0)
    zed.set_camera_settings(sl.VIDEO_SETTINGS.EXPOSURE, EXPOSURE_0_100)
    zed.set_camera_settings(sl.VIDEO_SETTINGS.GAIN, GAIN_0_100)
    zed.set_camera_settings(sl.VIDEO_SETTINGS.WHITEBALANCE_AUTO, 0)
    zed.set_camera_settings(sl.VIDEO_SETTINGS.WHITEBALANCE_TEMPERATURE, WB_TEMP_K)

    # 2) Positional Tracking: rotación inicial y SIN gravity-as-origin
    print("[INFO] Activando Positional Tracking con rotación inicial Rx(-20°) y SIN gravity-as-origin...")
    track = sl.PositionalTrackingParameters()
    # Muy importante: no alinear automáticamente con gravedad
    try:
        track.set_gravity_as_origin = False
    except Exception:
        pass
    # (Opcional) si quieres evitar cualquier mezcla de IMU:
    # track.enable_imu_fusion = False

    # Transform inicial Rx(-20°)
    rot = sl.Rotation()
    # set_rotation_vector usa Rodrigues: vector = (angulo_rad, 0, 0) para rotar en X
    rot.set_rotation_vector(math.radians(PITCH_DEG), 0.0, 0.0)
    trans = sl.Translation()
    trans.init_vector(0.0, 0.0, 0.0)
    tf0 = sl.Transform()
    tf0.init_rotation_translation(rot, trans)
    track.set_initial_world_transform(tf0)

    err = zed.enable_positional_tracking(track)
    if err != sl.ERROR_CODE.SUCCESS:
        print(f"[ERROR] No se pudo habilitar tracking: {err}", file=sys.stderr)
        zed.close()
        sys.exit(3)

    # 3) Runtime: MEDIDAS 3D EN CAMERA (NO WORLD)
    runtime = sl.RuntimeParameters(
        enable_depth=True,
        enable_fill_mode=False,
        confidence_threshold=50,
        texture_confidence_threshold=100,
        measure3D_reference_frame=sl.REFERENCE_FRAME.CAMERA,  # ← Nube en CAMERA
        remove_saturated_areas=True,
    )

    pc_mat = sl.Mat()
    img_mat = sl.Mat()

    # Depuración: imprime la Rx(-20°) declarada
    R_init = rx_matrix(PITCH_DEG)
    ex, ey, ez = rotation_to_euler_xyz(R_init)
    print("\n[INFO] Rotación inicial declarada Rx({:.3f}°):".format(PITCH_DEG))
    print(R_init)
    print(f"[INFO] Euler equivalente (X,Y,Z) = ({ex:.3f}, {ey:.3f}, {ez:.3f})°")

    # Warmup breve para que el tracking reporte algo estable (solo inspección de poses)
    print("[INFO] Warmup del tracking con 3 grabs (inspección de poses)...")
    for i in range(3):
        if zed.grab(runtime) != sl.ERROR_CODE.SUCCESS:
            continue
        pose_w = sl.Pose()
        pose_c = sl.Pose()
        zed.get_position(pose_w, sl.REFERENCE_FRAME.WORLD)
        zed.get_position(pose_c, sl.REFERENCE_FRAME.CAMERA)
        ew = pose_w.get_euler_angles(); ec = pose_c.get_euler_angles()
        print(f"[DEBUG] Warmup {i+1}: Euler WORLD (deg)=({math.degrees(ew[0]):.3f},{math.degrees(ew[1]):.3f},{math.degrees(ew[2]):.3f}) | "
              f"Euler CAMERA (deg)=({math.degrees(ec[0]):.3f},{math.degrees(ec[1]):.3f},{math.degrees(ec[2]):.3f})")

    # 4) Export por frame: recuperar nube en CAMERA (sin rotaciones nuestras)
    for fr in FRAME_LIST:
        print(f"\n[INFO] Procesando frame {fr} → NUBE EN CAMERA → remap a LAS → export...")
        zed.set_svo_position(fr)

        if zed.grab(runtime) != sl.ERROR_CODE.SUCCESS:
            print(f"[ERROR] grab() falló en frame {fr}.", file=sys.stderr)
            continue

        # (Solo inspección) Pose en WORLD y CAMERA que reporta el tracking en ese instante
        pose_w = sl.Pose()
        pose_c = sl.Pose()
        zed.get_position(pose_w, sl.REFERENCE_FRAME.WORLD)
        zed.get_position(pose_c, sl.REFERENCE_FRAME.CAMERA)
        ew = pose_w.get_euler_angles(); ec = pose_c.get_euler_angles()
        tw = pose_w.get_translation().get(); tc = pose_c.get_translation().get()
        print(f"[POSE] Frame {fr} WORLD euler (deg)=({math.degrees(ew[0]):.3f},{math.degrees(ew[1]):.3f},{math.degrees(ew[2]):.3f})  "
              f"trans (m)=({tw[0]:.3f},{tw[1]:.3f},{tw[2]:.3f})")
        print(f"[POSE] Frame {fr} CAMERA euler (deg)=({math.degrees(ec[0]):.3f},{math.degrees(ec[1]):.3f},{math.degrees(ec[2]):.3f}) "
              f"trans (m)=({tc[0]:.3f},{tc[1]:.3f},{tc[2]:.3f})")

        # Recuperar nube en CAMERA e imagen
        zed.retrieve_measure(pc_mat, sl.MEASURE.XYZRGBA)   # ← XYZ en CAMERA por runtime
        zed.retrieve_image(img_mat, sl.VIEW.LEFT)

        pc_cam = pc_mat.get_data()[..., :3].reshape(-1, 3)
        rgba   = img_mat.get_data().reshape(-1, 4)

        # Filtrar NaN/Inf
        valid = np.isfinite(pc_cam).all(axis=1)
        pc_cam = pc_cam[valid]
        rgba   = rgba[valid]
        print(f"[DEBUG] Puntos válidos (CAMERA) tras filtrado: {pc_cam.shape[0]}")

        # Remap CAMERA → LAS (Z arriba) para Potree (sin rotaciones nuestras)
        pc_las = remap_camera_to_las(pc_cam)
        zmin, zmax, zmean = float(np.min(pc_las[:,2])), float(np.max(pc_las[:,2])), float(np.mean(pc_las[:,2]))
        print(f"[DEBUG] Stats Z (CAMERA→LAS) → minZ={zmin:.3f}, maxZ={zmax:.3f}, meanZ={zmean:.3f}")

        # Export
        rgb16 = rgba_to_rgb_u16(rgba)
        out = Path(f"{OUTPUT_PREFIX}_frame{fr}.laz")
        write_laz(pc_las.astype(np.float64), rgb16, out)

    # Cierre
    zed.disable_positional_tracking()
    zed.close()
    print("\n[LISTO] Export completado: nube en CAMERA (sin rotaciones propias), con tracking inicial Rx(-20°) y sin gravity-as-origin.")
# ─────────────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n[INFO] Cancelado por el usuario.")