import math
from PIL import Image


class ImageAssembler:
    def __init__(self, metadata, atlas: Image, cell_size: int, padding: int):
        self.atlas = atlas
        self.atlas_width_cells = int(atlas.size[0] / cell_size)
        self.atlas_height_cells = int(atlas.size[1] / cell_size)

        self.cell_size = cell_size
        self.padding = padding
        self.metadata = metadata

        self.output = Image.new('RGBA', (metadata['width'], metadata['height']), (255, 255, 255, 0))

        # TODO double check this logic
        self.output_width_cells = math.ceil(float(metadata['width']) / cell_size)

    def atlas_get_cell(self, cell_no: int):
        cell_x_pos = cell_no % self.atlas_width_cells
        cell_y_pos = self.atlas_height_cells - math.floor(cell_no / self.atlas_width_cells) - 1  # From top
        return self.atlas.crop((
            cell_x_pos * self.cell_size,
            cell_y_pos * self.cell_size,
            (cell_x_pos + 1) * self.cell_size,
            (cell_y_pos + 1) * self.cell_size
        ))

    def paint(self):
        for i in range(len(self.metadata['cellIndexList'])):
            source_cell_number = self.metadata['cellIndexList'][i]
            source_cell = self.atlas_get_cell(source_cell_number)

            output_x_pos = i % self.output_width_cells
            output_y_pos = math.floor(i / self.output_width_cells)  # From bottom
            self.output.paste(source_cell, (
                output_x_pos * (self.cell_size - (2*self.padding)),
                self.output.size[1] - self.cell_size - (output_y_pos * (self.cell_size - (2*self.padding)))
            ))

    def get_output(self):
        return self.output