#!/usr/bin/env python3 

#  #####  ##   ##   ####   #######   ####    ####
#    #    # # # #  #    #     #     #    #  #    #
#    #    #  #  #  #          #     #    #  #
#    #    #     #  #  ###     #     #    #  #
#    #    #     #  #    #     #     #    #  #    #
#  #####  #     #   ####      #      ####    ####
#
#  Converts a bitmap image to c source.


import os
import typer
from PIL import Image


class PixelList:
    byte_align_rows = 0
    align_pixels = 0
    depth = 0
    current_row = []
    rows = []
    current_byte = 0
    current_byte_length = 0
    num_row_bytes = 0

    def __init__(self, depth: int, align_pixels: int, byte_align_rows: bool):
        if depth not in [1, 2, 4, 8]:
            raise 'depth must be 1, 2, 4 or 8'
        elif align_pixels % depth != 0:
            raise 'align_pixels must be a multiple of depth'
        elif align_pixels not in [1, 2, 4, 8]:
            raise 'align_pixels must be 1, 2, 4 or 8'


        self.depth = depth
        self.byte_align_rows = byte_align_rows
        self.align_pixels = align_pixels

    def append_pixel(self, pixel: int):
        self.current_byte = (self.current_byte << self.align_pixels) | (pixel & ((1 << self.depth) - 1))
        self.current_byte_length += self.align_pixels
        if self.current_byte_length == 8:
            self.current_row.append(self.current_byte)
            self.num_row_bytes += 1
            self.current_byte = 0
            self.current_byte_length = 0
        pass

    def end_row(self):
        if self.current_byte_length > 0 and self.byte_align_rows:
            byte = self.current_byte << (8 - self.current_byte_length)
            self.current_row.append(byte)
            self.num_row_bytes += 1
            self.current_byte = 0
            self.current_byte_length = 0
        self.rows.append(self.current_row)
        self.current_row = []

    def end_image(self):
        if self.current_byte_length > 0:
            byte = self.current_byte << (8 - self.current_byte_length)
            self.rows.append([byte])

    def get_rows(self):
        return self.rows


def img_to_c(im: Image, fmt: str, invert: bool, depth: int, align_pixels: int, byte_align_rows: bool):
    pixels = PixelList(depth, align_pixels, byte_align_rows)

    for iy in range(0, im.size[1]):  # iterate rows
        for ix in range(0, im.size[0]):  # iterate columns
            pix = im.getpixel((ix, iy))
            pixels.append_pixel(pix)
        pixels.end_row()
    pixels.end_image()

    src_c = ''
    for row in pixels.get_rows():  # iterate rows
        src_c += '        '
        for byte in row:
            if invert:
                byte = 0xff - byte
            if fmt == 'hex':
                src_c += '0x' + format(byte, '02x')
            elif fmt == 'bin':
                src_c += 'B' + format(byte, '08b')
            elif fmt == 'dec':
                src_c += f'{byte:>3}'
            src_c += ', '
        src_c += '\n'
    return src_c


def make_palette(depth: int):
    num_colors = 1 << depth
    colors = []
    color_step = 255 / (num_colors - 1)
    for i in range(0, num_colors - 1):
        colors.append(int(color_step * i))
    colors.append(255)
    return colors


def find_closest_palette_color(color: int, palette: [int]) -> (int, int):
    return min(enumerate(palette), key=lambda x: abs(x[1] - color))


def dither(original: Image, palette: [int]):
    im = Image.new('L', original.size)
    im.paste(original)
    for iy in range(0, im.size[1]):  # iterate rows
        for ix in range(0, im.size[0]):  # iterate columns
            oldpix = im.getpixel((ix, iy))
            closest = find_closest_palette_color(oldpix, palette)
            im.putpixel((ix, iy), closest[0])
            quant_error = oldpix - closest[1]
            if ix < im.size[0] - 1:
                im.putpixel((ix + 1, iy), int(im.getpixel((ix + 1, iy)) + quant_error * 7 / 16))
            if ix > 0 and iy < im.size[1] - 1:
                im.putpixel((ix - 1, iy + 1), int(im.getpixel((ix - 1, iy + 1)) + quant_error * 3 / 16))
            if iy < im.size[1] - 1:
                im.putpixel((ix, iy + 1), int(im.getpixel((ix, iy + 1)) + quant_error * 5 / 16))
            if ix < im.size[0] - 1 and iy < im.size[1] - 1:
                im.putpixel((ix + 1, iy + 1), int(im.getpixel((ix + 1, iy + 1)) + quant_error * 1 / 16))

    return im


def main(input_file: str,
         out: str = typer.Option(None, help='File to write to, defaults to standard output'),
         name: str = typer.Option(None, help='Name of the bitmap variable in generated source, defaults to the '
                                             'filename without extension.'),
         fmt: str = typer.Option('hex', help='Format to generate image data in (hex or bin).'),
         invert: bool = typer.Option(False, help='Invert brightness.'),
         depth: int = typer.Option(1, help='Bit depth of the output, can be 1, 2, 4 or 8.'),
         pixel_align: int = typer.Option(None, help='Bit-alignment of the pixels, defaults to the depth, can be 1, 2, 4 or 8.'),
         no_row_align: bool = typer.Option(False, help='Disable byte alignment for rows.'),
         ):
    im: Image = Image.open(input_file)
    if im.mode != 'L':
        im = im.convert(mode='L')

    im = dither(im, make_palette(depth))

    if name is None:
        name = os.path.splitext(os.path.basename(input_file))[0]

    if pixel_align is None:
        pixel_align = depth

    row_align_val = 0 if no_row_align else 1

    ctext = f'// Image definition of: {input_file}\n'
    ctext += 'static PROGMEM bitmap_t ' + name + ' { \n'
    ctext += f'    .width  = {im.size[0]},\n'
    ctext += f'    .height = {im.size[1]},\n'
    ctext += f'    .depth = {depth},\n'
    ctext += f'    .pixel_align = {pixel_align},\n'
    ctext += f'    .row_align = {row_align_val},\n'
    ctext += '     .data = (PROGMEM uint8_t[]) {\n'

    ctext += img_to_c(im, fmt, invert, depth, pixel_align, not no_row_align)

    ctext += '    }\n};\n'

    if out:
        with open(out, mode='w') as f:
            f.write(ctext)
            f.close()
    else:
        print(ctext)


if __name__ == "__main__":
    typer.run(main)