import math
from collections import namedtuple

import pygame

RGB_COLOR_CODES = [
    (0, 0, 0),  # black
    (70, 70, 70),  # grey
    (0, 0, 255),  # blue
    (0, 255, 0),  # green
    (255, 0, 0),  # red
    (255, 255, 255)  # white
]

Colors = namedtuple("Colors", "black,grey,blue,green,red,white", defaults=RGB_COLOR_CODES)


class Environment:

    def __init__(self, map_path: str, map_height: int, map_width: int, window_name: str = "RRT path planning",
                 colors: Colors = None):
        self.point_cloud = []
        self.external_map = pygame.image.load(map_path)
        self.map_height = map_height
        self.map_width = map_width
        self.window_name = window_name
        self.colors = colors if colors else Colors()
        self._initialize_pygame()

    def _initialize_pygame(self):
        pygame.init()
        pygame.display.set_caption(self.window_name)
        self.map = pygame.display.set_mode((self.map_width, self.map_height))
        # overlay external map
        self.map.blit(self.external_map, (0, 0))

    def calc_2d_position(self, distance, angle, robot_position):
        """Calculates a 2D point from a robots position and a Lidar sensor reading."""
        x = distance * math.cos(angle) + robot_position[0]
        y = -distance * math.sin(angle) + robot_position[1]
        return (int(x), int(y))

    # TODO add type hints
    def store_points(self, sensor_sweep_data):
        """Stores sensor reading in the environment's point cloud."""
        for element in sensor_sweep_data:
            # TODO this could be refactored into a LidarReading & RobotPosition named tuples
            point = self.calc_2d_position(element[0], element[1], element[2])
            # check for dupes
            # TODO ok an optimization could probably be added here too
            if point not in self.point_cloud:
                self.point_cloud.append(point)

    def show_sensor_data(self):
        # Draw data on new map
        self.info_map = self.map.copy()
        for point in self.point_cloud:
            self.info_map.set_at((int(point[0]), int(point[1])), self.colors.red)