from abc import ABC, abstractmethod
import numpy as np

from ...core.types import GlobalMotionPlan, RobotCmd, RobotState, JointStates
from .agent_base import AgentBase


class MLAgentBase(AgentBase, ABC):
    """An abstract base class machine-learned controller agents.

    See docstrings for each method to be implemented.
    """

    @abstractmethod
    def initialise_robot_cmd(self, joint_states: JointStates):
        """Override this method in child class if custom initilisation is required (e.g. joint name order).
        By default, this method sets the joint name and position values to be equal to the input joint states
        object, with kp and kd set to be the default (in constructor).
        """

    @abstractmethod
    def update_input_to_model(
        self,
        robot_state: RobotState,
        global_plan: GlobalMotionPlan,
        t: float,
        dt: float,
    ) -> np.ndarray:
        """Should update (`self._input_ndarray`) using appropriate values (input to model). This method
        has access to `self._latest_ctrl_cmd` (type `RobotCmd`) as well if needed.
        (NOTE: `self._latest_ctrl_cmd` is set to be the initial robot joint positions (zero velocities
        and efforts commands) with `default_kp` and `default_kd` at start).
        """

    @abstractmethod
    def update_cmd_from_model_output(
        self,
        model_output: np.ndarray,
        robot_state: RobotState,
        global_plan: GlobalMotionPlan,
        t: float,
        dt: float,
    ) -> None:
        """Should update `self._latest_robot_cmd` (type `RobotCmd`) using the output from the NN model.

        Args:
            model_output (np.ndarray): the numpy array created from the output tensor from the model
                after the inference query was done. This is the output of the neural network. This
                method should use this object to update `self._latest_robot_cmd` to be sent to the
                robot.
        """

    @abstractmethod
    def get_action(
        self,
        robot_state: RobotState,
        global_plan: GlobalMotionPlan,
        t: float = None,
        dt: float = None,
    ) -> RobotCmd:
        """Should return the control command given the current robot state and global plan.

        Args:
            robot_state (RobotState): The current state information from the robot.
            global_plan (GlobalMotionPlan): the latest global plan generated by the global planner
                used in the loop.
            t (float, optional): the current time signature of the control loop. Defaults to None
                (controllers may or may not need this).
            dt (float, optional): the time since the last control loop. Defaults to None
                (controllers may or may not need this).

        Returns:
            RobotCmd: The output control command to be sent to the robot.
        """
