import math
import os
import random
import numpy as np
import pybullet as p
import pybullet_data
import gymnasium as gym
from gymnasium import spaces
import config as cfg
import robot_init as ri


class JackBotEnv(gym.Env):
    """Gymnasium environment for joint-command learning with the JackBot URDF."""

    metadata = {"render_modes": ["human", "rgb_array"]}

    def __init__(
        self,
        urdf_path: str | None = None,
        use_gui: bool = True,
        frame_skip: int = 4,
        random_command: bool = True,
        max_episode_steps: int = 2000,
        num_robots: int = 1,
        robot_spacing: float = 0.5,
        start_pose: str = "init_deg",
    ):
        """
        Initialize the JackBot environment.

        Args:
            urdf_path: Path to the robot's URDF file. Defaults to config value.
            use_gui: If True, uses p.GUI (shows window), else p.DIRECT (headless).
            frame_skip: Number of physics simulation steps per agent step. 
                Higher values mean the agent takes actions less frequently (e.g., 4 steps at 240Hz = 60Hz control).
            random_command: If True, assigns a random velocity command (vx, vy, omega) to each robot every reset.
            max_episode_steps: Maximum number of steps (interactions) before the episode is truncated.
            num_robots: Number of robots to spawn for parallel training in the same scene.
            robot_spacing: Distance in meters between robots when spawning multiple.
            start_pose: The initial joint configuration name (e.g., 'init_deg' or 'init90_deg').
        """
        self.urdf_path = urdf_path or cfg.urdf_path
        self.use_gui = use_gui
        self.frame_skip = frame_skip
        self.random_command = random_command
        self.max_episode_steps = max_episode_steps
        self.num_robots = max(1, num_robots)
        self.robot_spacing = robot_spacing
        self.start_pose = start_pose
        self.action_scale = math.radians(8.0)
        self.physics_client = None
        self.robots = []
        self.plane = None
        self.robot_joint_indices = []
        self.joint_lower = None
        self.joint_upper = None
        self.initial_angles = None
        self.commands = None
        self.step_count = 0
        self.last_action = None
        self._first_reset = True
        self._min_steps_before_done = 150

        self._connect_sim()

        total_joints = self.num_robots * len(self.robot_joint_indices[0])
        observation_dim = total_joints + self.num_robots * 4
        action_dim = total_joints

        self.observation_space = spaces.Box(
            low=-np.inf,
            high=np.inf,
            shape=(observation_dim,),
            dtype=np.float32,
        )
        self.action_space = spaces.Box(
            low=-1.0,
            high=1.0,
            shape=(action_dim,),
            dtype=np.float32,
        )

    def _connect_sim(self):
        if self.physics_client is not None:
            p.disconnect(self.physics_client)

        flags = p.GUI if self.use_gui else p.DIRECT
        self.physics_client = p.connect(flags)
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        p.setGravity(0, 0, -9.81)

        self.plane = p.loadURDF("plane.urdf")
        self.robots = []
        self.robot_joint_indices = []

        for robot_id in range(self.num_robots):
            base_pos = self._robot_base_position(robot_id)
            robot = p.loadURDF(self.urdf_path, basePosition=base_pos, useFixedBase=False)
            self.robots.append(robot)
            joint_indices = [
                i
                for i in range(p.getNumJoints(robot))
                if p.getJointInfo(robot, i)[2] == p.JOINT_REVOLUTE
            ]
            self.robot_joint_indices.append(joint_indices)

        joint_count = len(self.robot_joint_indices[0])
        if any(len(indices) != joint_count for indices in self.robot_joint_indices):
            raise ValueError("All robots must have the same number of revolute joints.")

        lower_limits = []
        upper_limits = []
        for joint_index in self.robot_joint_indices[0]:
            info = p.getJointInfo(self.robots[0], joint_index)
            lower = info[8]
            upper = info[9]
            if lower >= upper:
                lower = -math.pi
                upper = math.pi
            lower_limits.append(lower)
            upper_limits.append(upper)

        limits = np.array([lower_limits, upper_limits], dtype=np.float32)
        self.joint_lower = np.tile(limits[0], (self.num_robots, 1))
        self.joint_upper = np.tile(limits[1], (self.num_robots, 1))

        pose = ri.init_deg if self.start_pose == "init_deg" else ri.init90_deg
        initial_pose = pose.to_rad().data.flatten()
        self.initial_angles = np.tile(initial_pose, (self.num_robots, 1)).astype(np.float32)
        self.initial_angles = np.clip(
            self.initial_angles,
            self.joint_lower,
            self.joint_upper,
        )

        self.last_action = np.zeros(self.num_robots * joint_count, dtype=np.float32)
        self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)

    def reset(self, command: np.ndarray | None = None, seed: int | None = None, options: dict | None = None):
        if self._first_reset:
            self.step_count = 0
            self.last_action = np.zeros(self.num_robots * len(self.robot_joint_indices[0]), dtype=np.float32)

            if seed is not None:
                self.seed(seed)

            if command is not None:
                self.commands = np.array(command, dtype=np.float32).reshape(self.num_robots, 4)
            elif self.random_command:
                self.commands = np.stack([self.sample_command() for _ in range(self.num_robots)])
            else:
                self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)

            self._reset_pose()
            self._first_reset = False
            return self._get_obs(), {}

        p.resetSimulation()
        p.setGravity(0, 0, -9.81)
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        self.plane = p.loadURDF("plane.urdf")
        self.robots = []
        self.robot_joint_indices = []

        for robot_id in range(self.num_robots):
            base_pos = self._robot_base_position(robot_id)
            robot = p.loadURDF(self.urdf_path, basePosition=base_pos, useFixedBase=False)
            self.robots.append(robot)
            joint_indices = [
                i
                for i in range(p.getNumJoints(robot))
                if p.getJointInfo(robot, i)[2] == p.JOINT_REVOLUTE
            ]
            self.robot_joint_indices.append(joint_indices)

        self.step_count = 0
        self.last_action = np.zeros(self.num_robots * len(self.robot_joint_indices[0]), dtype=np.float32)

        if seed is not None:
            self.seed(seed)

        if command is not None:
            self.commands = np.array(command, dtype=np.float32).reshape(self.num_robots, 4)
        elif self.random_command:
            self.commands = np.stack([self.sample_command() for _ in range(self.num_robots)])
        else:
            self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)

        self._reset_pose()
        return self._get_obs(), {}

    def seed(self, seed: int | None = None):
        """Set RNG seed for reproducibility (Gym compatibility)."""
        if seed is None:
            seed = int(np.random.randint(0, 2**31 - 1))
        random.seed(seed)
        np.random.seed(seed)
        self._seed = seed
        return [seed]

    def _reset_pose(self):
        for robot, joint_indices, init_angles in zip(self.robots, self.robot_joint_indices, self.initial_angles):
            for joint_index, target_angle in zip(joint_indices, init_angles):
                p.resetJointState(robot, joint_index, target_angle)
                p.setJointMotorControl2(
                    bodyIndex=robot,
                    jointIndex=joint_index,
                    controlMode=p.POSITION_CONTROL,
                    targetPosition=target_angle,
                    force=250,
                )

        for _ in range(400):
            p.stepSimulation()

        self.joint_angles = self.initial_angles.copy()

    def sample_command(self) -> np.ndarray:
        vx = np.random.uniform(-1.0, 1.0)
        vy = np.random.uniform(-0.5, 0.5)
        vz = 0.0
        omega = np.random.uniform(-1.0, 1.0)
        return np.array([vx, vy, vz, omega], dtype=np.float32)

    def _get_obs(self) -> np.ndarray:
        return np.concatenate([self.joint_angles.flatten(), self.commands.flatten()]).astype(np.float32)

    def step(self, action: np.ndarray):
        action = np.clip(action, self.action_space.low, self.action_space.high).astype(np.float32)
        self.last_action = action
        action_matrix = action.reshape(self.num_robots, -1)
        self.joint_angles = np.clip(
            self.joint_angles + action_matrix * self.action_scale,
            self.joint_lower,
            self.joint_upper,
        )

        for robot, joint_indices, angles in zip(self.robots, self.robot_joint_indices, self.joint_angles):
            for joint_index, target_angle in zip(joint_indices, angles):
                p.setJointMotorControl2(
                    bodyIndex=robot,
                    jointIndex=joint_index,
                    controlMode=p.POSITION_CONTROL,
                    targetPosition=target_angle,
                    force=250,
                )

        for _ in range(self.frame_skip):
            p.stepSimulation()

        self.step_count += 1
        observation = self._get_obs()
        reward = self._compute_reward()
        done = self._is_done()
        if self.step_count <= self._min_steps_before_done:
            done = False
        info = {"step": self.step_count}

        # Gymnasium-style return: (obs, reward, terminated, truncated, info)
        if done:
            terminated = True
            truncated = False
        else:
            terminated = False
            truncated = False

        return observation, float(reward), terminated, truncated, info

    def _compute_reward(self) -> float:
        rewards = []
        for robot_id, robot in enumerate(self.robots):
            linear_vel, angular_vel = p.getBaseVelocity(robot)
            _, orientation = p.getBasePositionAndOrientation(robot)
            roll, pitch, _ = p.getEulerFromQuaternion(orientation)
            command = self.commands[robot_id]

            forward_reward = command[0] * linear_vel[0] + command[1] * linear_vel[1]
            rotation_reward = command[3] * angular_vel[2]
            stability_penalty = abs(roll) + abs(pitch)
            action_penalty = float(np.sum(np.square(self.last_action.reshape(self.num_robots, -1)[robot_id]))) * 0.01

            rewards.append(0.1 + forward_reward + rotation_reward - 0.2 * stability_penalty - action_penalty)

        return float(np.sum(rewards))

    def _is_done(self) -> bool:
        for robot in self.robots:
            _, orientation = p.getBasePositionAndOrientation(robot)
            roll, pitch, _ = p.getEulerFromQuaternion(orientation)
            if abs(roll) > 0.7 or abs(pitch) > 0.7:
                return True
            position, _ = p.getBasePositionAndOrientation(robot)
            if position[2] < 0.02:
                return True

        if self.step_count >= self.max_episode_steps:
            return True
        return False

    def render(self, mode="human"):
        if self.use_gui:
            return None
        raise NotImplementedError("Render is only supported with use_gui=True")

    def _robot_base_position(self, robot_id: int) -> list[float]:
        cols = int(math.sqrt(self.num_robots - 1)) + 1
        row = robot_id // cols
        col = robot_id % cols
        x = (col - (cols - 1) / 2.0) * self.robot_spacing
        y = (row - (cols - 1) / 2.0) * self.robot_spacing
        return [x, y, 0.1]

    def close(self):
        if self.physics_client is not None:
            p.disconnect(self.physics_client)
            self.physics_client = None