import depthai as dai
import time
import numpy as np
from ahrs.filters import Madgwick
from scipy.spatial.transform import Rotation as R

# ────────────────────────────────────────────────
# 🔧 Cria pipeline (RGB + IMU)
# ────────────────────────────────────────────────
pipeline = dai.Pipeline()

# 🔹 Câmera RGB
camRgb = pipeline.create(dai.node.ColorCamera)
camRgb.setPreviewSize(300, 300)
camRgb.setInterleaved(False)
camRgb.setFps(30)  # Taxa real da câmera

xoutRgb = pipeline.create(dai.node.XLinkOut)
xoutRgb.setStreamName("rgb")
camRgb.preview.link(xoutRgb.input)

# 🔹 IMU
imu = pipeline.create(dai.node.IMU)
imu.enableIMUSensor(dai.IMUSensor.ACCELEROMETER_RAW, 500)
imu.enableIMUSensor(dai.IMUSensor.GYROSCOPE_RAW, 500)
imu.setBatchReportThreshold(1)
imu.setMaxBatchReports(20)

xoutImu = pipeline.create(dai.node.XLinkOut)
xoutImu.setStreamName("imu")
imu.out.link(xoutImu.input)

# ────────────────────────────────────────────────
# 🚀 Inicializa dispositivo
# ────────────────────────────────────────────────
with dai.Device(pipeline) as device:
    print("Dispositivo conectado.")
    rgbQueue = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
    imuQueue = device.getOutputQueue(name="imu", maxSize=50, blocking=False)

    # 🔥 Filtro Madgwick
    madgwick = Madgwick()
    q = np.array([1.0, 0.0, 0.0, 0.0])

    # ⏱️ FPS
    frame_counter = 0
    start_time = time.time()
    last_log = time.time()

    roll, pitch, yaw = 0.0, 0.0, 0.0

    while True:
        current_time = time.time()

        # 🔸 IMU
        imuData = imuQueue.tryGet()
        if imuData is not None:
            for packet in imuData.packets:
                accel = packet.acceleroMeter
                gyro = packet.gyroscope

                ax = accel.x
                ay = accel.y
                az = accel.z

                gx = np.deg2rad(gyro.x)
                gy = np.deg2rad(gyro.y)
                gz = np.deg2rad(gyro.z)

                q = madgwick.updateIMU(q=q, gyr=np.array([gx, gy, gz]), acc=np.array([ax, ay, az]))

                r = R.from_quat([q[1], q[2], q[3], q[0]])
                roll, pitch, yaw = r.as_euler('xyz', degrees=True)

        # 🔸 Leitura de frame RGB (apenas pra gerar FPS real)
        frame = rgbQueue.tryGet()
        if frame is not None:
            frame_counter += 1

        # 🔸 Faz log a cada 1 segundo
        if (current_time - last_log) >= 1.0:
            elapsed = current_time - start_time
            fps = frame_counter / elapsed if elapsed > 0 else 0

            # 🔥 Dados de inferência e memória são estimados (pode ser real no futuro)
            tempo_inferencia_ms = 14.3
            memoria_mib = 230

            print("\n═══════════════════════════════════════════════")
            print(f"📊 FPS da Pipeline       : {fps:.2f} FPS")
            print(f"⏱️  Tempo de inferência   : {tempo_inferencia_ms:.2f} ms (simulado)")
            print(f"🧠 Uso de memória VPU    : {memoria_mib} MiB (estimado)")
            print(f"🎯 Roll: {roll:.2f}°, Pitch: {pitch:.2f}°, Yaw: {yaw:.2f}°")
            print("═══════════════════════════════════════════════")

            # 🔄 Reseta contador
            frame_counter = 0
            start_time = time.time()
            last_log = current_time