from PIL import Image, ImageDraw
from typing import List
from xmlrpc.client import Boolean
from gymnasium import spaces
import gymnasium as gym
import numpy as np


class MyGymPointsEnv(gym.Env):
    metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 4}
    max_size = 10
    epsilon = 0.01

    def __init__(
        self,
        a=10,
        b=10,
        max_size=10,
        verbose=False,
    ):
        self.verbose = verbose
        self.max_size = max_size
        self.size = (a, b)
        self.step_num = 0
        self.points = self._generate_points()
        while self._check_collisions():
            self.points = self._generate_points()
        # self.low = np.array([-a / 2, -b / 2, -max(a, b), -np.pi]).astype(np.float32)
        # self.high = np.array([a / 2, b / 2, max(a, b), np.pi]).astype(np.float32)
        self.observation_space = spaces.Box(-1.0, 1.0, (self.max_size, 4))
        # self.action_space = spaces.Discrete(
        #     5
        # )  # 0 -- nop, 1 -- slow down, 2 -- speed up, 3 -- rotate left, 4 -- rotate right
        self.action_space = spaces.Box(-1.0, 1.0, shape=(3,), dtype=np.float32)
        # 0 -- pos_x
        # 1 -- pos_y
        # 2 -- velocity
        # 3 -- angle(vector)
        #     np.reshape(
        #         np.repeat(np.array([0, -np.pi, -max(a, b)]), self.max_size),
        #         (self.max_size, 3),
        #     ),
        #     np.reshape(
        #         np.repeat(np.array([self.max_size, np.pi, max(a, b)]), self.max_size),
        #         (self.max_size, 3),
        #     ),
        # )

    def step(self, action):
        reward = 0.0
        if self.verbose:
            print(action)
        target = int(action[0] * (self.max_size - 1) / 2)
        change_velocity = action[1] * max(self.size)
        change_angle = action[2] * np.pi
        if self.verbose:
            print(
                f"target: {target}",
                f"change_velocity: {change_velocity}",
                f"change_angle: {change_angle}",
            )
        self.points[target][2] += change_velocity
        self.points[target][2] = np.clip(
            self.points[target][2], -max(self.size), max(self.size)
        )
        self.points[target][3] += change_angle
        self.points[target][3] = np.clip(self.points[target][3], -np.pi, np.pi)
        if np.abs(action[1]) < self.epsilon and np.abs(action[2]) < self.epsilon:
            reward += 1.0
        else:
            reward -= 1.0
        # observation, reward, terminated, False, info
        self._move_forward()
        self._draw_state_image(f"{self.step_num}")
        terminated = False
        if not (self._check_collisions()):
            reward += 3.0
        else:
            terminated = True
            reward -= 5.0

        self.step_num += 1
        return (
            self._get_obs(),
            reward,
            terminated,
            False,
            {},
        )

    def _move_forward(self, dt=0.1):
        # Update the positions of points based on their velocity and angle
        # Using Euler integration scheme
        self.points[:, 0] += (
            self.points[:, 2] * np.cos(self.points[:, 3]) * dt
        )  # Update x positions
        self.points[:, 1] += (
            self.points[:, 2] * np.sin(self.points[:, 3]) * dt
        )  # Update y positions

    def _check_collisions(self) -> Boolean:
        """
        Check for collisions between points and return a boolean indicating if any collisions were found.
        Return: True if any collisions were found, False otherwise.
        """
        for i in range(len(self.points)):
            for j in range(i + 1, len(self.points)):
                distance = (
                    (self.points[i][0] - self.points[j][0]) ** 2
                    + (self.points[i][1] - self.points[j][1]) ** 2
                ) ** 0.5
                if distance < self.epsilon:
                    return True
        return False

    def _get_obs(self):
        return np.array(
            [
                self.points[:, 0] / (self.size[0]),
                self.points[:, 1] / (self.size[1]),
                self.points[:, 2] / (max(self.size)),
                self.points[:, 3] / (np.pi),
            ],
            dtype=np.float32,
        ).T

    def _generate_points(self):

        x_values = np.random.randint(0, self.size[0], size=self.max_size)
        y_values = np.random.randint(0, self.size[1], size=self.max_size)
        velocity_values = np.random.uniform(0, max(self.size), size=self.max_size)
        angle_values = np.random.uniform(-np.pi, np.pi, size=self.max_size)

        # Combine the values into self.points
        return np.column_stack((x_values, y_values, velocity_values, angle_values))

    def _draw_state_image(self, image_name: str):
        image_size = self.size  # Replace this with your desired image size
        image = Image.new(
            "L", image_size, color=0
        )  # 'L' mode represents a grayscale image

        # Create a draw object
        draw = ImageDraw.Draw(image)

        # Assuming self.points is a NumPy array with shape [10, 4]
        for point in self.points:
            x, y = (
                point[0],
                point[1],
            )  # Adjust for positive coordinates
            draw.point((x, y), fill=255)  # Set pixel to white

        # Save or display the image
        image.save(f"{image_name}.png")  # Display the image

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        self.points = self._generate_points()
        while self._check_collisions():
            self.points = self._generate_points()
        observation = self._get_obs()

        return np.array(observation, dtype=np.float32), {}
