import os
import pybullet
from pybullet_envs.robot_bases import MJCFBasedRobot, URDFBasedRobot

# Use our custom data
from rlberry.envs.bullet3 import data


class MJCFBasedRobot2(MJCFBasedRobot):
    def reset(self, bullet_client):
        self._p = bullet_client
        # print("Created bullet_client with id=", self._p._client)
        if self.doneLoading == 0:
            self.ordered_joints = []
            self.doneLoading = 1
            if self.self_collision:
                self.objects = self._p.loadMJCF(
                    os.path.join(data.getDataPath(), "mjcf", self.model_xml),
                    flags=pybullet.URDF_USE_SELF_COLLISION
                    | pybullet.URDF_USE_SELF_COLLISION_EXCLUDE_ALL_PARENTS
                    | pybullet.URDF_GOOGLEY_UNDEFINED_COLORS,
                )
                (
                    self.parts,
                    self.jdict,
                    self.ordered_joints,
                    self.robot_body,
                ) = self.addToScene(self._p, self.objects)
            else:
                self.objects = self._p.loadMJCF(
                    os.path.join(
                        data.getDataPath(),
                        "mjcf",
                        self.model_xml,
                        flags=pybullet.URDF_GOOGLEY_UNDEFINED_COLORS,
                    )
                )
                (
                    self.parts,
                    self.jdict,
                    self.ordered_joints,
                    self.robot_body,
                ) = self.addToScene(self._p, self.objects)
        self.robot_specific_reset(self._p)

        s = (
            self.calc_state()
        )  # optimization: calc_state() can calculate something in self.* for calc_potential() to use

        return s


class URDFBasedRobot2(URDFBasedRobot):
    def __init__(
        self,
        model_urdf,
        robot_name,
        action_dim,
        obs_dim,
        basePosition=[0, 0, 0],
        baseOrientation=[0, 0, 0, 1],
        fixed_base=False,
        self_collision=False,
    ):
        super().__init__(
            model_urdf,
            robot_name,
            action_dim,
            obs_dim,
            basePosition,
            baseOrientation,
            fixed_base,
            self_collision,
        )
        self.doneLoading = 0

    def reset(self, bullet_client):
        self._p = bullet_client
        if self.doneLoading == 0:
            self.ordered_joints = []
            self.doneLoading = 1
            if self.self_collision:
                (
                    self.parts,
                    self.jdict,
                    self.ordered_joints,
                    self.robot_body,
                ) = self.addToScene(
                    self._p,
                    self._p.loadURDF(
                        os.path.join(data.getDataPath(), self.model_urdf),
                        basePosition=self.basePosition,
                        baseOrientation=self.baseOrientation,
                        useFixedBase=self.fixed_base,
                        flags=pybullet.URDF_USE_SELF_COLLISION
                        | pybullet.URDF_GOOGLEY_UNDEFINED_COLORS,
                    ),
                )
            else:
                (
                    self.parts,
                    self.jdict,
                    self.ordered_joints,
                    self.robot_body,
                ) = self.addToScene(
                    self._p,
                    self._p.loadURDF(
                        os.path.join(data.getDataPath(), self.model_urdf),
                        basePosition=self.basePosition,
                        baseOrientation=self.baseOrientation,
                        useFixedBase=self.fixed_base,
                        flags=pybullet.URDF_GOOGLEY_UNDEFINED_COLORS,
                    ),
                )

        self.robot_specific_reset(self._p)

        s = (
            self.calc_state()
        )  # optimization: calc_state() can calculate something in self.* for calc_potential() to use
        self.potential = self.calc_potential()

        return s
