"""Base class for robotics environments."""

#!/usr/bin/python
#
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import importlib
import inspect
import os
from typing import Dict, Optional

import numpy as np

from d4rl_alt.kitchen.adept_envs import mujoco_env
from d4rl_alt.kitchen.adept_envs.base_robot import BaseRobot
from d4rl_alt.kitchen.adept_envs.utils.configurable import import_class_from_path
from d4rl_alt.kitchen.adept_envs.utils.constants import MODELS_PATH


class RobotEnv(mujoco_env.MujocoEnv):
    """Base environment for all adept robots."""

    # Mapping of robot name to fully qualified class path.
    # e.g. 'robot': 'adept_envs.dclaw.robot.Robot'
    # Subclasses should override this to specify the Robot classes they support.
    ROBOTS = {}

    # Mapping of device path to the calibration file to use. If the device path
    # is not found, the 'default' key is used.
    # This can be overriden by subclasses.
    CALIBRATION_PATHS = {}

    def __init__(
        self,
        model_path: str,
        robot: BaseRobot,
        frame_skip: int,
        camera_settings: Optional[Dict] = None,
    ):
        """Initializes a robotics environment.

        Args:
            model_path: The path to the model to run. Relative paths will be
              interpreted as relative to the 'adept_models' folder.
            robot: The Robot object to use.
            frame_skip: The number of simulation steps per environment step. On
              hardware this influences the duration of each environment step.
            camera_settings: Settings to initialize the simulation camera. This
              can contain the keys `distance`, `azimuth`, and `elevation`.
        """
        self._robot = robot

        # Initial pose for first step.
        self.desired_pose = np.zeros(self.n_jnt)

        if not model_path.startswith("/"):
            model_path = os.path.abspath(os.path.join(MODELS_PATH, model_path))

        self.remote_viz = None

        try:
            from adept_envs.utils.remote_viz import RemoteViz

            self.remote_viz = RemoteViz(model_path)
        except ImportError:
            pass

        self._initializing = True
        super(RobotEnv, self).__init__(
            model_path, frame_skip, camera_settings=camera_settings
        )
        self._initializing = False

    @property
    def robot(self):
        return self._robot

    @property
    def n_jnt(self):
        return self._robot.n_jnt

    @property
    def n_obj(self):
        return self._robot.n_obj

    @property
    def skip(self):
        """Alias for frame_skip. Needed for MJRL."""
        return self.frame_skip

    @property
    def initializing(self):
        return self._initializing

    def close_env(self):
        if self._robot is not None:
            self._robot.close()

    def make_robot(
        self,
        n_jnt,
        n_obj=0,
        is_hardware=False,
        device_name=None,
        legacy=False,
        **kwargs
    ):
        """Creates a new robot for the environment.

        Args:
            n_jnt: The number of joints in the robot.
            n_obj: The number of object joints in the robot environment.
            is_hardware: Whether to run on hardware or not.
            device_name: The device path for the robot hardware.
            legacy: If true, runs using direct dynamixel communication rather
              than DDS.
            kwargs: See BaseRobot for other parameters.

        Returns:
            A Robot object.
        """
        if not self.ROBOTS:
            raise NotImplementedError("Subclasses must override ROBOTS.")

        if is_hardware and not device_name:
            raise ValueError("Must provide device name if running on hardware.")

        robot_name = "dds_robot" if not legacy and is_hardware else "robot"
        if robot_name not in self.ROBOTS:
            raise KeyError(
                "Unsupported robot '{}', available: {}".format(
                    robot_name, list(self.ROBOTS.keys())
                )
            )

        cls = import_class_from_path(self.ROBOTS[robot_name])

        calibration_path = None
        if self.CALIBRATION_PATHS:
            if not device_name:
                calibration_name = "default"
            elif device_name not in self.CALIBRATION_PATHS:
                print(
                    'Device "{}" not in CALIBRATION_PATHS; using default.'.format(
                        device_name
                    )
                )
                calibration_name = "default"
            else:
                calibration_name = device_name

            calibration_path = self.CALIBRATION_PATHS[calibration_name]
            if not os.path.isfile(calibration_path):
                raise OSError(
                    "Could not find calibration file at: {}".format(calibration_path)
                )

        return cls(
            n_jnt,
            n_obj,
            is_hardware=is_hardware,
            device_name=device_name,
            calibration_path=calibration_path,
            **kwargs
        )
