import sys
import json
import time
import math
import board
import busio
from adafruit_bno08x.i2c import BNO08X_I2C
from adafruit_bno08x import BNO_REPORT_ROTATION_VECTOR

POLL_INTERVAL = 0.1        # seconds between output lines
REPORT_INTERVAL_US = 50000 # sensor report rate (20 Hz)
STUCK_TIMEOUT = 10.0       # re-init if quaternion is bit-identical this long
MAX_CONSECUTIVE_ERRORS = 10


def quaternion_to_euler(x, y, z, w):
    """Adafruit BNO08X quaternion (x, y, z, w) -> yaw, pitch, roll in degrees."""
    sinr_cosp = 2.0 * (w * x + y * z)
    cosr_cosp = 1.0 - 2.0 * (x * x + y * y)
    roll = math.atan2(sinr_cosp, cosr_cosp)

    sinp = max(-1.0, min(1.0, 2.0 * (w * y - z * x)))
    pitch = math.asin(sinp)

    siny_cosp = 2.0 * (w * z + x * y)
    cosy_cosp = 1.0 - 2.0 * (y * y + z * z)
    yaw = math.atan2(siny_cosp, cosy_cosp)

    return math.degrees(yaw), math.degrees(pitch), math.degrees(roll)


def valid_quaternion(quat):
    """Return a normalized (x, y, z, w) tuple, or None if the data is garbage."""
    if not quat or len(quat) != 4:
        return None
    if any(v is None or math.isnan(v) or math.isinf(v) for v in quat):
        return None
    norm = math.sqrt(sum(v * v for v in quat))
    # A healthy unit quaternion has norm ~1. Corrupt I2C reads usually don't.
    if norm < 0.9 or norm > 1.1:
        return None
    return tuple(v / norm for v in quat)


def init_sensor(i2c):
    bno = BNO08X_I2C(i2c, address=0x4B)
    try:
        bno.enable_feature(BNO_REPORT_ROTATION_VECTOR, REPORT_INTERVAL_US)
    except TypeError:
        # Older library versions don't accept a report interval
        bno.enable_feature(BNO_REPORT_ROTATION_VECTOR)
    time.sleep(0.5)
    return bno


def main():
    try:
        i2c = busio.I2C(board.SCL, board.SDA)
        bno = init_sensor(i2c)

        last_quat = None
        last_change = time.time()
        errors = 0

        while True:
            try:
                raw = bno.quaternion  # (x, y, z, w) from rotation vector
                quat = valid_quaternion(raw)
                if quat is None:
                    raise ValueError("invalid quaternion")

                errors = 0

                # Stuck-sensor detection
                if raw != last_quat:
                    last_quat = raw
                    last_change = time.time()
                elif time.time() - last_change > STUCK_TIMEOUT:
                    raise RuntimeError("sensor output frozen")

                yaw, pitch, roll = quaternion_to_euler(*quat)
                if yaw < 0:
                    yaw += 360.0

                print(json.dumps({
                    "yaw": round(yaw, 2),
                    "pitch": round(pitch, 2),
                    "roll": round(roll, 2),
                }))
                sys.stdout.flush()

            except (ValueError, TypeError, IndexError, KeyError, OSError,
                    RuntimeError) as e:
                errors += 1
                frozen = isinstance(e, RuntimeError)
                if frozen or errors >= MAX_CONSECUTIVE_ERRORS:
                    print(json.dumps({"warning": f"re-initializing sensor ({e})"}),
                          file=sys.stderr)
                    sys.stderr.flush()
                    try:
                        bno = init_sensor(i2c)
                    except Exception as init_err:
                        print(json.dumps({"warning": f"re-init failed: {init_err}"}),
                              file=sys.stderr)
                        time.sleep(1.0)
                    errors = 0
                    last_quat = None
                    last_change = time.time()

            time.sleep(POLL_INTERVAL)

    except Exception as e:
        print(json.dumps({"error": str(e)}), file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()