"""
差速履带运动学/动力学模型。
"""

from __future__ import annotations

import dataclasses
import math
import random
from typing import Tuple

GRAVITY = 9.80665  # m/s^2


def _clamp(value: float, min_value: float, max_value: float) -> float:
    return max(min_value, min(max_value, value))


def _wrap_angle(rad: float) -> float:
    """Wrap angle to [-pi, pi)."""
    pi2 = math.tau  # 2*pi
    while rad >= math.pi:
        rad -= pi2
    while rad < -math.pi:
        rad += pi2
    return rad


@dataclasses.dataclass
class DifferentialDriveState:
    east: float = 0.0
    north: float = 0.0
    up: float = 0.0
    heading: float = 0.0  # rad, 数学坐标系 (东=0, CCW>0)
    linear_velocity: float = 0.0  # m/s
    angular_velocity: float = 0.0  # rad/s
    pitch_deg: float = 0.0
    roll_deg: float = 0.0
    east_velocity: float = 0.0
    north_velocity: float = 0.0
    up_velocity: float = 0.0
    body_accel_g: Tuple[float, float, float] = (0.0, 0.0, -1.0)
    gyro_deg_s: Tuple[float, float, float] = (0.0, 0.0, 0.0)
    temperature_c: float = 30.0

    def copy(self) -> "DifferentialDriveState":
        return dataclasses.replace(self)


class DifferentialDriveModel:
    """准静态差速履带模型，可模拟原地转向。"""

    def __init__(
        self,
        track_width: float = 0.8,
        max_linear_speed: float = 2.0,
        max_linear_accel: float = 1.5,
        max_angular_speed: float = math.radians(120.0),
        max_angular_accel: float = math.radians(180.0),
    ):
        self.track_width = track_width
        self.max_linear_speed = max_linear_speed
        self.max_linear_accel = max_linear_accel
        self.max_angular_speed = max_angular_speed
        self.max_angular_accel = max_angular_accel
        self.state = DifferentialDriveState()
        self._rng = random.Random(42)

    def reset(self, east: float = 0.0, north: float = 0.0, up: float = 0.0, heading_deg: float = 0.0):
        self.state = DifferentialDriveState(east=east, north=north, up=up, heading=math.radians(heading_deg))

    def step(self, target_linear: float, target_angular: float, dt: float) -> DifferentialDriveState:
        if dt <= 0.0:
            return self.state

        target_linear = _clamp(target_linear, -self.max_linear_speed, self.max_linear_speed)
        target_angular = _clamp(target_angular, -self.max_angular_speed, self.max_angular_speed)

        dv = target_linear - self.state.linear_velocity
        max_dv = self.max_linear_accel * dt
        dv = _clamp(dv, -max_dv, max_dv)
        linear_acc = dv / dt
        self.state.linear_velocity += dv

        # 符号定义：
        # - 罗盘方向：负的控制信号 → 逆时针旋转，正的控制信号 → 顺时针旋转
        # - 数学坐标系：东为0°，逆时针为正（CCW>0）
        # - 所以：负的turn_rate（逆时针）→ angular_velocity为正，正的turn_rate（顺时针）→ angular_velocity为负
        # - 需要取反：angular_velocity = -target_angular
        target_angular_math = -target_angular  # 取反以匹配数学坐标系
        
        dw = target_angular_math - self.state.angular_velocity
        max_dw = self.max_angular_accel * dt
        dw = _clamp(dw, -max_dw, max_dw)
        angular_acc = dw / dt
        self.state.angular_velocity += dw

        # 更新姿态与位置
        self.state.heading = _wrap_angle(self.state.heading + self.state.angular_velocity * dt)
        cos_h = math.cos(self.state.heading)
        sin_h = math.sin(self.state.heading)
        self.state.east += self.state.linear_velocity * cos_h * dt
        self.state.north += self.state.linear_velocity * sin_h * dt
        self.state.up += 0.0  # 可扩展地形模型

        # 速度 (ENU)
        self.state.east_velocity = self.state.linear_velocity * cos_h
        self.state.north_velocity = self.state.linear_velocity * sin_h
        self.state.up_velocity = 0.0

        # 体轴加速度：前向加速度 + 向心加速度
        lateral_acc = self.state.linear_velocity * self.state.angular_velocity
        ax_g = linear_acc / GRAVITY
        ay_g = lateral_acc / GRAVITY
        az_g = -1.0  # 假设水平地面
        self.state.body_accel_g = (
            ax_g + self._rng.gauss(0.0, 0.002),
            ay_g + self._rng.gauss(0.0, 0.002),
            az_g + self._rng.gauss(0.0, 0.0005),
        )

        # 姿态近似：由加速度分量映射成小角度
        self.state.pitch_deg = math.degrees(math.atan2(ax_g, 1.0))
        self.state.roll_deg = math.degrees(math.atan2(-ay_g, 1.0))

        # 惯导（IMU）符号定义：顺时针正，逆时针负（从上往下看Z轴，右手螺旋定则）
        # 在数学坐标系中：顺时针 → angular_velocity为负，逆时针 → angular_velocity为正
        # 所以：GyroZ = -angular_velocity（取反以匹配惯导定义）
        # 例如：angular_velocity = -0.8 rad/s（顺时针）→ GyroZ = 0.8 deg/s（正值，顺时针）
        self.state.gyro_deg_s = (
            self._rng.gauss(0.0, 0.01),
            self._rng.gauss(0.0, 0.01),
            -math.degrees(self.state.angular_velocity) + self._rng.gauss(0.0, 0.1),  # 取反以匹配惯导定义
        )

        temp_drift = self._rng.gauss(0.0, 0.0005)
        self.state.temperature_c += temp_drift

        return self.state