from PIL import Image
import numpy as np

DEFAULT_FILL_COLOR = [128, 128, 128]
UV_PIXEL_CENTER_OFFSET = 0.5


def build_atlas(colors, width, height):
    atlas_pixel_array = create_empty_atlas_array(width, height)
    fill_atlas_with_palette(atlas_pixel_array, colors, width, height)
    return Image.fromarray(atlas_pixel_array, "RGB")


def create_empty_atlas_array(width, height):
    return np.zeros((height, width, 3), dtype=np.uint8)


def fill_atlas_with_palette(atlas_array, palette_colors, width, height):
    palette_index = 0

    for row in range(height):
        for column in range(width):
            if palette_index < len(palette_colors):
                atlas_array[row, column] = palette_colors[palette_index]
                palette_index += 1
            else:
                atlas_array[row, column] = DEFAULT_FILL_COLOR


def palette_to_atlas(palette, atlas_size=16):
    atlas_image = build_atlas(palette, atlas_size, atlas_size)
    uv_mapping = create_color_to_uv_mapping(palette, atlas_size)
    return atlas_image, uv_mapping


def create_color_to_uv_mapping(palette, atlas_dimensions):
    color_to_uv_coordinates = {}
    palette_index = 0

    for row in range(atlas_dimensions):
        for column in range(atlas_dimensions):
            if palette_index < len(palette):
                current_color = palette[palette_index]
                rgb_tuple = convert_to_rgb_tuple(current_color)
                uv_position = calculate_uv_for_pixel(column, row, atlas_dimensions)
                color_to_uv_coordinates[rgb_tuple] = uv_position
                palette_index += 1

    return color_to_uv_coordinates


def convert_to_rgb_tuple(color_array):
    return tuple(int(channel) for channel in color_array)


def calculate_uv_for_pixel(pixel_x, pixel_y, atlas_size):
    u_coordinate = (pixel_x + UV_PIXEL_CENTER_OFFSET) / atlas_size
    v_coordinate = 1.0 - (pixel_y + UV_PIXEL_CENTER_OFFSET) / atlas_size
    return (u_coordinate, v_coordinate)