DroneMind/voice_drone/core/command.py

from pydantic import BaseModel, Field
from typing import Dict, Any, Optional, Literal
from datetime import datetime
from voice_drone.core.configuration import (
    TAKEOFF_CONFIG,
    LAND_CONFIG,
    FOLLOW_CONFIG,
    FORWARD_CONFIG,
    BACKWARD_CONFIG,
    LEFT_CONFIG,
    RIGHT_CONFIG,
    UP_CONFIG,
    DOWN_CONFIG,
    HOVER_CONFIG,
    RETURN_HOME_CONFIG,
)
import warnings
warnings.filterwarnings("ignore")


class CommandParams(BaseModel):
    """
    命令参数
    """
    distance: Optional[float] = Field(
        None,
        description="飞行距离,单位：米(m),必须大于等于0(land/hover 可以为0)",
        ge=0
    )
    speed: Optional[float] = Field(
        None,
        description="飞行速度,单位：米每秒(m/s),必须大于等于0(land/hover 可以为0)",
        ge=0
    )
    duration: Optional[float] = Field(
        None,
        description="飞行持续时间,单位：秒(s),必须大于0",
        gt=0
    )

class Command(BaseModel):
    """
    无人机控制命令
    """
    command: Literal[
        "takeoff",
        "follow",
        "forward",
        "backward",
        "left",
        "right",
        "up",
        "down",
        "hover",
        "land",
        "return_home",
    ] = Field(
        ...,
        description="无人机控制动作: takeoff(起飞), follow(跟随), forward(向前), backward(向后), left(向左), right(向右), up(向上), down(向下), hover(悬停), land(降落), return_home(返航)",
    )
    params: CommandParams = Field(..., description="命令参数")
    timestamp: str = Field(..., description="命令生成时间戳,ISO 8601 格式(如:2024-01-01T12:00:00.000Z)")
    sequence_id: int = Field(..., description="命令序列号,用于保证命令顺序和去重")

    # 命令配置映射字典
    _CONFIG_MAP = {
        "takeoff": TAKEOFF_CONFIG,
        "follow": FOLLOW_CONFIG,
        "land": LAND_CONFIG,
        "forward": FORWARD_CONFIG,
        "backward": BACKWARD_CONFIG,
        "left": LEFT_CONFIG,
        "right": RIGHT_CONFIG,
        "up": UP_CONFIG,
        "down": DOWN_CONFIG,
        "hover": HOVER_CONFIG,
        "return_home": RETURN_HOME_CONFIG,
    }

    # 创建命令
    @classmethod
    def create(
        cls,
        command: str,
        sequence_id: int,
        distance: Optional[float] = None,
        speed: Optional[float] = None,
        duration: Optional[float] = None
    ) -> "Command":
        return cls(
            command=command,
            params=CommandParams(distance=distance, speed=speed, duration=duration),
            timestamp=datetime.utcnow().isoformat() + "Z",
            sequence_id=sequence_id,
        )
    
    def _get_default_config(self):
        """获取当前命令的默认配置"""
        return self._CONFIG_MAP.get(self.command)
    
    # 填充默认值
    def fill_defaults(self) -> None:
        """填充缺失的参数值"""
        # 如果所有参数都已提供，直接返回
        if (self.params.distance is not None and 
            self.params.speed is not None and 
            self.params.duration is not None):
            return
        
        # 如果有缺失的参数，调用智能填充方法
        self._fill_smart_params()

    def _fill_smart_params(self):
        """智能填充缺失的参数值"""
        default = self._get_default_config()
        if default is None:
            # 若命令未知，则直接返回不填充
            return

        d = self.params.distance
        s = self.params.speed
        t = self.params.duration

        # 统计 None 个数
        none_cnt = sum(x is None for x in [d, s, t])

        # 三个都为None，直接填默认值
        if none_cnt == 3:
            self.params.distance = default["distance"]
            self.params.speed = default["speed"]
            self.params.duration = default["duration"]
            return

        # 只有一个参数有值的情况
        if none_cnt == 2:
            if s is not None and d is None and t is None:
                # 仅速度：使用默认持续时间，计算距离
                self.params.duration = default["duration"]
                self.params.distance = s * self.params.duration
                return
            
            if t is not None and d is None and s is None:
                # 仅持续时间：使用默认速度，计算距离
                self.params.speed = default["speed"]
                self.params.distance = self.params.speed * t
                return
            
            if d is not None and s is None and t is None:
                # 仅距离：使用默认速度，计算持续时间
                self.params.speed = default["speed"]
                # 防止除以0
                if self.params.speed == 0:
                    self.params.duration = 0
                else:
                    self.params.duration = d / self.params.speed
                return

        # 两个参数有值，一个None，自动计算缺失的参数
        if none_cnt == 1:
            if d is None and s is not None and t is not None:
                # 缺失距离：distance = speed * duration
                self.params.distance = s * t
                return
            
            if s is None and d is not None and t is not None:
                # 缺失速度：speed = distance / duration
                if t == 0:
                    self.params.speed = 0
                else:
                    self.params.speed = d / t
                return
            
            if t is None and d is not None and s is not None:
                # 缺失持续时间：duration = distance / speed
                if s == 0:
                    self.params.duration = 0
                else:
                    self.params.duration = d / s
                return


    # 转换为字典
    def to_dict(self) -> dict:

        result = {
            "command": self.command,
            "params": {},
            "timestamp": self.timestamp,
            "sequence_id": self.sequence_id
        }
        
        if self.params.distance is None or self.params.speed is None or self.params.duration is None:
            self.fill_defaults()

        result["params"]["distance"] = self.params.distance
        result["params"]["speed"] = self.params.speed
        result["params"]["duration"] = self.params.duration

        return result



if __name__ == "__main__":
    command = Command.create("takeoff", 1, speed=2)
    print(command.to_dict())