from dataclasses import dataclass, field
from enum import Enum

from .tf_types import Twist, RelativePose
from .robot_io import EndEffectorStates, JointStates
from .motion_datatypes import PointMotion


class ControlMode(Enum):
    """Request type from local planner to controller.

    MODES:
        - IDLE: controller should ignore any other values in the local plan.
        - CONTROL: controller can use values from the LocalMotionPlan to generate commands.
    """

    IDLE = 0
    """controller should ignore any other values in the local plan."""

    CONTROL = 1
    """controller will/should read rest of the LocalMotionPlan object only if in this mode."""


class PlannerMode(Enum):
    """Request type from global planner to local planner.

    MODES:
        - IDLE: No local plan should be generated by local planner.
        - CUSTOM: Local planner is free to use other attributes of the GlobalMotionPlan and
            generate a custom LocalMotionPlan message for the controller.
        - HOLD_POSITION: pre-defined action signal to keep the robot in the current position;
            other attributes in GlobalMotionPlan should be ignored by local planner, and planner
            should generate LocalMotionPlan that will keep the joints of the robot locked.
    """

    IDLE = 0
    """No local plan should be generated. No command is to generated by controller in this state."""

    HOLD_POSITION = 1
    """pre-defined action signal to keep the robot in the current position, other attributes in
    GlobalMotionPlan should be ignored by local planner (this state has to be handled by local
    planner or control loop; see examples.03).
    """

    CUSTOM = 2
    """In this mode, the local planner is free to use other attributes of the GlobalMotionPlan and
    generate a custom LocalMotionPlan message for the controller.
    """


@dataclass
class GlobalMotionPlan:
    """Object representing a global motion plan provided by a generic global motion planner.

    In our case, parsed user input is translated to this format for local planner.

    In full autonomous robot framework, this should be generated by a global world-aware planner,
    based on task.


    Attributes (see doc for each attribute):
        relative_pose (RelativePose): Relative pose target in teleop frame.
        twist (Twist): target velocity command for the base in base frame.
        joint_references (JointStates): joint targets.
        end_effector_references (EndEffectorStates): references for end-effectors of the robot.
        planner_mode (PlannerMode): available modes: [PlannerMode.IDLE, PlannerMode.HOLD_POSITION,
            PlannerMode.CUSTOM]
    """

    relative_pose: RelativePose = None
    """Pose in teleop frame of the robot.

    Teleop frame is equivalent to robot (base) frame, but has the roll, pitch and height values
    fixed in some neutral pre-defined values in the world coordinates.

    We define the following as our teleop frame. This may be changed as desired.
    [x,y,z] = [x_BF, y_BF, 0.0] in world
    [r,p,y] = [0.0, 0.0, yaw_BF] in world
    """

    twist: Twist = field(default_factory=Twist)
    """Twist in base/teleop frame of the robot"""

    # joint reference is not ideal in a global plan message; only left for directly sending
    # reference using joystick/gui (for testing new robots, or for controlling manipulator)
    joint_references: JointStates = field(default_factory=JointStates)
    """Desired joint references. Typically a local planner would only fill joint position or
    velocity values, and not effort values."""
    # ee references is not ideal in a global plan message; only left for directly sending reference
    # using joystick/gui (for testing new robots, or for controlling manipulator)
    end_effector_references: EndEffectorStates = field(default_factory=EndEffectorStates)
    """Desired end-effector pose (in global frame), contact states, contact forces etc. can be
    defined here."""

    planner_mode: PlannerMode = PlannerMode.IDLE
    """Defined by a global planner to let the local planner know if the rest of the
    GlobalMotionPlan message is usable, or if there is a specific plan to be generated (such as
    HOLD_POSITION)"""


@dataclass
class LocalMotionPlan:
    """Object representing a local motion plan provided by a generic local motion planner.

    This object will be used as reference by controllers.

    Attributes (see doc for each attribute):
        relative_pose (RelativePose): Relative pose target in teleop frame.
        twist (Twist): target velocity command for the base in base frame.
        joint_references (JointStates): references for joints (position, velocity)
        end_effector_references (EndEffectorStates): References for end-effectors (pos, force,
            contact states)
        control_mode (ControlMode): control modes (available modes: [ControlMode.IDLE,
            ControlMode.CONTROL])

    WARNING
    -------
    NOTE: as of now, the base pose (`relative_pose`) is defined in the teleop frame, and
    the end-effector poses (in `end_effector_references.ee_poses`) is in the world frame.
    This may have to be changed later if there is a need.
    """

    # setting default value of identity is wrong (due to definition of teleop frame)
    relative_pose: RelativePose = None
    """Pose in teleop frame of the robot.

    Teleop frame is equivalent to robot (base) frame, but has the roll, pitch and height values
    fixed in some neutral pre-defined values in the world coordinates.

    We define the following as our teleop frame. This may be changed as desired.
    [x,y,z] = [x_BF, y_BF, 0.0] in world
    [r,p,y] = [0.0, 0.0, yaw_BF] in world

    WARNING
    -------
    NOTE: The frame used may have to be changed later. We may need to use this in the global frame.
    """
    twist: Twist = field(default_factory=Twist)
    """Base velocity commands in the base frame."""
    joint_references: JointStates = field(default_factory=JointStates)
    """Desired joint references. Typically a local planner would only fill joint position or
    velocity values, and not effort values."""
    end_effector_references: EndEffectorStates = field(default_factory=EndEffectorStates)
    """Desired end-effector pose (in global frame), contact states, contact forces etc. can be
    defined here.

    WARNING
    -------
    NOTE: The frame used may have to be changed later, depending on whether it is more useful in
    global/base/teleop frame.
    """
    com_reference: PointMotion = None
    """Reference position ([x,y,z] in global frame), velocity and acceleration for center of
    mass of the robot for whole-body controllers."""
    control_mode: ControlMode = ControlMode.IDLE
    """Defined by a local planner to let the controller(s) know if the rest of the LocalMotionPlan
    message is usable, or if there is a specific control action to take (such as IDLE)."""
