import pybullet as p
import numpy as np
import RobotState as rs
import GlobalVariables as gv
import config as cfg
import DataTypes as dt
import time
import os
import math


class Simulation:
    def __init__(self):
        self.physics_client = p.connect(p.GUI)
        self.robot = p.loadURDF(cfg.urdf_path, useFixedBase=True)

        self.revolute_joints = [
            i
            for i in range(p.getNumJoints(self.robot))
            if p.getJointInfo(self.robot, i)[2] == p.JOINT_REVOLUTE
        ]

        self.set_all_joints_to_90()
        p.resetDebugVisualizerCamera(
            cameraDistance=1.0,
            cameraYaw=50,
            cameraPitch=-35,
            cameraTargetPosition=[0, 0, 0],
        )

    def set_all_joints_to_90(self):
        for joint_index in range(p.getNumJoints(self.robot)):
            joint_info = p.getJointInfo(self.robot, joint_index)
            joint_type = joint_info[2]
            if joint_type == p.JOINT_REVOLUTE:
                p.resetJointState(self.robot, joint_index, math.radians(90))

    def updatePos(self, current_rad: dt.RadArray):
        radflat = current_rad.data.flatten()
        for joint_index, target_angle in zip(self.revolute_joints, radflat):
            p.setJointMotorControl2(
                bodyIndex=self.robot,
                jointIndex=joint_index,
                controlMode=p.POSITION_CONTROL,
                targetPosition=target_angle,
                force=500,
            )

    def step(self):
        p.stepSimulation()

    def disconnect(self):
        p.disconnect(self.physics_client)


if __name__ == "__main__":
    count = 0
    current_rad: dt.RadArray = gv.init_deg.to_rad()
    gv.shared_sim.updatePos(current_rad)

    rs.walking()
    while True:
        count = +1
        gv.shared_sim.step()
        time.sleep(1 / 240)
        if count > 50:
            rs.walking()
            count = 0