"""Data models for the Windsuite SDK."""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import ClassVar

# ! -------------
# ! USER FACING
# ! -------------


@dataclass
class Vec3:
    """3D vector with x, y, z components."""

    x: float
    y: float
    z: float


@dataclass
class Quat:
    """Quaternion with w, x, y, z components."""

    w: float
    x: float
    y: float
    z: float


@dataclass
class WindProbeData:
    """
    Represents all data from a wind probe.

    This includes:
    - Raw data
    - Data corrected for windshaper orientation
    - Data corrected for windshaper and probe movement / orientation

    Notes:
    Depending on the tracking quality, some data might contain NaN values.

    If you want the best possible, you either need to check the tracking data yourself
    or use the BestWindProbeData callback that will give you the best possible Wind Data and Tracking Data directly.

    All wind velocity data is in meters per second (mps).

    """

    timestamp_s: float
    wind_velocity_mps_probe_ref: Vec3  # ! Raw data
    wind_velocity_mps_windshaper_ref: Vec3  # ! Corrected in windshaper reference frame
    wind_velocity_mps_windshaper_ref_corrected: Vec3  # ! Corrected for probe movement / orientation in windshaper reference frame

    temperature_celcius: float
    atmospheric_pressure_hpascal: float
    static_pressure_pascal: float

    conversion_status: int
    data_is_valid: bool


@dataclass
class TrackingData:
    """
    Represents tracking data for a single object.

    - Position in meters in the world reference frame,
    - Orientation as a quaternion in the world reference frame,

    ! Special case for the windshaper itself that will have the same as it's his reference frame
    - Position in meters in the windshaper reference frame,
    - Orientation as a quaternion in the windshaper reference frame,

    - Velocity in meters per second,
    - Angular velocity in degrees per second.
    """

    timestamp: float

    position_meters_world_ref: Vec3
    rotation_world_ref: Quat

    position_meters_windshaper_ref: Vec3
    rotation_windshaper_ref: Quat

    velocity_mps_world_ref: Vec3
    angular_velocity_degps_world_ref: Vec3

    velocity_mps_windshaper_ref: Vec3
    angular_velocity_degps_windshaper_ref: Vec3

    is_tracked: bool = False


@dataclass
class TrackedProbeData:
    """
    Represents the best Wind data associated with the current position of the probe.

    Notes:
    The best data is chosen based on the tracking quality to provide the most accurate wind data possible.
    If there is no tracking, the raw wind data is used.
    If there is unstable tracking data, the orientation-corrected wind data is used.
    If there is stable tracking data, the movement / orientation corrected wind data is used.

    The position is given in the windshaper reference frame (if possible) in millimeters.

    """

    timestamp_s: float

    wind_velocity_mps: Vec3
    position_data_mm: Vec3

    temperature_celcius: float
    atmospheric_pressure_hpascal: float
    static_pressure_pascal: float


@dataclass
class ModuleInfo:
    """
    Represents the current state of a Windshaper Module.

    Notes:
    It's possible to receive rpms and pwms in different formats depending on the Module type.

    For simplicity of use, the following conventions are used:
        Every received information is a list of lists, where each inner list represents a layer and each of these list contains the values for each fan in that layer.

        e.g.
            ! Module type 2420
            pwms = [[25.0]] -> Module with 1 layer and 1 fan with 25% pwm
            rpms = [[3445]] -> Module with 1 layer and 1 fan with 3445 rpm

            ! Module type 0816 / 0812
            pwms = [
                        [10, 10, 20, 20, 10, 10, 20, 20, 30], # ! UPSTREAM LAYER
                        [40, 40, 50, 50, 40, 40, 50, 50, 30], # ! DOWNSTREAM LAYER
                    ] -> Module with 2 layers and 9 fans per layer
            rpms = [
                        [1200, 1250, 1300, 1280, 1190, 1230, 1310, 1290, 1150], # ! UPSTREAM LAYER
                        [2200, 2250, 2300, 2280, 2190, 2230, 2310, 2290, 2150], # ! DOWNSTREAM LAYER
                    ] -> Module with 2 layers and 9 fans per layer

    """

    INDEX_DOWNSTREAM: ClassVar[int] = 0
    INDEX_UPSTREAM: ClassVar[int] = 1

    row: int
    col: int
    mac: str
    ip: str
    type: str

    lifepoints: int

    target_pwm: list[list[float]]
    current_pwm: list[list[float]]

    current_rpm: list[list[float]]

    target_psu_state: bool
    current_psu_state: bool

    is_connected: bool = False

    esc_rpm: int | None = None
    applied_pwm_us: int | None = None
    esc_voltage_mv: int | None = None
    esc_bus_current_ma: int | None = None
    esc_motor_current_ma: int | None = None
    esc_index: int | None = None
    esc_temp_mos_celcius: int | None = None
    esc_temp_capacitor_celcius: int | None = None
    esc_temp_mcu_celcius: int | None = None
    esc_temp_motor_celcius: int | None = None
    esc_running_error: int | None = None
    esc_uptime: int | None = None
    applied_psu: int | None = None
    psu_v_in: float | None = None
    psu_i_in: float | None = None
    psu_p_in: float | None = None
    psu_v_out: float | None = None
    psu_i_out: float | None = None
    psu_p_out: float | None = None
    psu_temp1: float | None = None
    psu_temp2: float | None = None
    psu_temp3: float | None = None


# ! -------------
# ! INTERNAL
# ! -------------


@dataclass
class FullTrackedWindProbeData:
    """
    Represents windprobe data associated with the current position of the probe.
    """

    probe_tracking_data: TrackingData
    windprobe_data: WindProbeData
    curated_tracked_probe_data: BestFullTrackedWindProbeData | None = None


@dataclass
class BestWindProbeData:
    """
    Represents curated data from a wind probe.
    The velocity is the most accurate depending on the tracking.
    """

    timestamp_s: float
    best_wind_velocity_mps: Vec3
    temperature_celcius: float
    atmospheric_pressure_hpascal: float
    static_pressure_pascal: float
    conversion_status: int


@dataclass
class BestFullTrackedWindProbeData:
    """
    Represents curated windprobe data associated with the current position of the probe.
    """

    probe_position_mm_windshaper_ref: Vec3
    best_windprobe_data: BestWindProbeData


class ModuleType(str, Enum):
    """
    Enum representing the type of a Windshaper Module.
    """

    MODULE_0812 = "0812"
    MODULE_0816 = "0816"
    MODULE_2420 = "2420"
