#
#      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 pytinydiffsim as dp
#import meshcat_utils_dp
#import meshcat
import time
import copy
world = dp.TinyWorld()
mb_solver = dp.TinyMultiBodyConstraintSolver()

#vis = meshcat.Visualizer(zmq_url='tcp://127.0.0.1:6000')
#vis.delete()

urdf_parser = dp.TinyUrdfParser()

plane_urdf_data = urdf_parser.load_urdf("../../data/bunny.urdf")
#plane2vis = meshcat_utils_dp.convert_visuals(plane_urdf_data, "../../data/checker_purple.png", vis, "../../data/")
plane_mb = dp.TinyMultiBody(False)
plane2mb = dp.UrdfToMultiBody2()
res = plane2mb.convert2(plane_urdf_data, world, plane_mb)

colliders = plane_urdf_data.base_links[0].urdf_collision_shapes
print("num_colliders=", len(colliders))
ray_from = []
    def __init__(self, simulation_time_step):
        self.MPC_BODY_MASS = MPC_BODY_MASS
        self.MPC_BODY_INERTIA = MPC_BODY_INERTIA
        self.MPC_BODY_HEIGHT = MPC_BODY_HEIGHT

        self.time_step = simulation_time_step

        self.num_legs = NUM_LEGS
        self.num_motors = NUM_MOTORS

        self.skip_sync = 0
        self.world = dp.TinyWorld()
        self.world.friction = 10.0

        self.mb_solver = dp.TinyMultiBodyConstraintSolver()
        self.vis = meshcat.Visualizer(zmq_url='tcp://127.0.0.1:6000')
        self.vis.delete()

        urdf_parser = dp.TinyUrdfParser()

        plane_urdf_data = urdf_parser.load_urdf(
            "../../../data/plane_implicit.urdf")
        plane2vis = meshcat_utils_dp.convert_visuals(
            plane_urdf_data, "../../../data/checker_purple.png", self.vis,
            "../../../data/")
        self.plane_mb = dp.TinyMultiBody(False)
        plane2mb = dp.UrdfToMultiBody2()
        res = plane2mb.convert2(plane_urdf_data, self.world, self.plane_mb)

        self.urdf_data = urdf_parser.load_urdf(
            "../../../data/laikago/laikago_toes_zup_joint_order.urdf")
        print("robot_name=", self.urdf_data.robot_name)
        self.b2vis = meshcat_utils_dp.convert_visuals(
            self.urdf_data, "../../../data/laikago/laikago_tex.jpg", self.vis,
            "../../../data/laikago/")
        is_floating = True
        self.mb = dp.TinyMultiBody(is_floating)
        urdf2mb = dp.UrdfToMultiBody2()
        self.res = urdf2mb.convert2(self.urdf_data, self.world, self.mb)
        self.mb.set_base_position(dp.Vector3(0, 0, 0.6))

        knee_angle = -0.5
        abduction_angle = 0.2

        self.initial_poses = [
            abduction_angle, 0., knee_angle, abduction_angle, 0., knee_angle,
            abduction_angle, 0., knee_angle, abduction_angle, 0., knee_angle
        ]

        qcopy = self.mb.q
        print("mb.q=", self.mb.q)
        print("qcopy=", qcopy)
        for q_index in range(12):
            qcopy[q_index + 7] = self.initial_poses[q_index]
        print("qcopy=", qcopy)
        self.mb.set_q(qcopy)
        print("2 mb.q=", self.mb.q)
        print("self.mb.links=", self.mb.links)
        print("TDS:")
        for i in range(len(self.mb.links)):
            print("i=", i)
            l = self.mb.links[i]
            #print("l.link_name=", l.link_name)
            print("l.joint_name=", l.joint_name)
        dp.forward_kinematics(self.mb, self.mb.q, self.mb.qd)

        self._BuildJointNameToIdDict()
        self._BuildUrdfIds()
        self._BuildMotorIdList()
        self.ResetPose()
        self._motor_enabled_list = [True] * self.num_motors
        self._step_counter = 0
        self._state_action_counter = 0
        self._motor_offset = np.array([0.] * 12)
        self._motor_direction = np.array([1.] * 12)
        self.ReceiveObservation()
        self._kp = self.GetMotorPositionGains()
        self._kd = self.GetMotorVelocityGains()
        self._motor_model = LaikagoMotorModel(
            kp=self._kp, kd=self._kd, motor_control_mode=MOTOR_CONTROL_HYBRID)

        self.ReceiveObservation()
        self.mb.set_base_position(
            dp.Vector3(START_POS[0], START_POS[1], START_POS[2]))
        self.mb.set_base_orientation(
            dp.Quaternion(START_ORN[0], START_ORN[1], START_ORN[2],
                          START_ORN[3]))

        dp.forward_kinematics(self.mb, self.mb.q, self.mb.qd)
        meshcat_utils_dp.sync_visual_transforms(self.mb, self.b2vis, self.vis)

        self._SettleDownForReset(reset_time=1.0)