from collections import namedtuple, deque
from enum import Enum

import numpy as np

from mm2d.core.spatial import Transform, Twist, State
from mm2d.core.body import Body
from mm2d.core.robot.state import RobotState
from mm2d.core.robot.description import parse_frame_info


class InputType(Enum):
    VELOCITY = 0
    TORQUE = 1


JointInfo = namedtuple(
    "JointInfo", ["link1", "link2", "r_j_l1_l1", "r_j_l2_l2", "actuated"]
)


# TODO it would be more consistent to parse a "BodyInfo" tuple
def parse_bodies(description):
    bodies = {}

    def add_body(link, pose):
        assert (
            link.name not in bodies
        ), f"Link with name {link.name} already attached to this robot."

        body = Body(pose, link.fixtures)
        bodies[link.name] = body

    # breadth-first search through the links (and joints)
    queue = deque([(Transform.eye(), description.base)])
    while len(queue) > 0:
        link, T_w_j1 = queue.popleft()

        T_j1_l1 = link.pose
        T_w_l1 = T_w_j1 @ T_j1_l1
        add_body(bodies, link, T_w_l1)

        for joint in link.child_joints:
            T_l1_j2 = joint.pose
            T_w_j2 = T_w_l1 @ T_l1_j2
            queue.append((T_w_j2, joint.child_link))


def parse_joint_info(description):
    joint_info = {}

    for link in description.links():
        for joint in link.child_joints:
            assert (
                joint.name not in joint_info
            ), f"Joint with name {joint.name} already attached to this robot."
            assert (
                joint.child_link is not None
            ), f"Joint {joint.name} has no child link."

            joint_info[joint.name] = JointInfo(
                link.name,
                joint.child_link.name,
                joint.pose.position,
                joint.child_link.pose.inverse().position,
                joint.actuated,
            )

    return joint_info


class Robot:
    """Simulated robot.

    Once added to a simulation environment, that environment adds a backend to
    the robot that allows it to interface with that particular simulator.
    """

    def __init__(self, description, input_type=InputType.VELOCITY):
        self.input_type = input_type
        self.base_link = description.base

        self.num_joint_dofs = np.sum([joint.actuated for joint in description.joints()])
        self.num_inputs = np.sum(self.base.dofs) + self.num_joint_dofs

        self.bodies = parse_bodies(description)
        self.frame_info = parse_frame_info(description)

        # TODO: we could conceivable get rid of this, just using it for the
        # backend
        self.joint_info = parse_joint_info(description)

    def command_base(self, base_cmd: np.ndarray) -> None:
        """Send control commands to the base.

        This applies velocities or torques (depending on input_type) directly
        to the body of the base.
        """
        assert (
            len(base_cmd) == self.base.num_dofs
        ), "Base command vector must be same length as number of base DOFs."
        v = np.zeros(3)
        v[self.base_link.nz] = base_cmd
        self.backend.command_base(v)

    def command_joints(self, joint_cmd: np.ndarray) -> None:
        """Send control commands to the robot's joints.

        The length of the joint_cmds array must be the same as the number of
        actuated joints. The order corresponds to the order of this robot's
        joint_names attribute.
        """
        assert (
            len(joint_cmd) == self.num_joint_dofs
        ), "Joint command vector must be same length as number of actuated joints."
        self.backend.command_joints(joint_cmd)

    def state(self) -> RobotState:
        """Get the current state of the robot."""
        base_state = self.bodies[self.base_link.name].state()
        joint_state = self.backend.joint_state()
        return RobotState(base_state, joint_state)

    def frame(self, name: str) -> State:
        """Get the current state of a frame."""
        link_name, T_lf = self.frame_info[name]
        S_wl = self.bodies[link_name].state()

        # if the frame is at the link origin, we can return directly
        if name == link_name:
            return S_wl

        # otherwise, we need to calculate
        T_wf = S_wl.pose @ T_lf
        V_wf = Twist.rigid(S_wl.twist, S_wl.pose, T_lf)
        return State(T_wf, V_wf)
