import gradio as gr
from diffusers import DiffusionPipeline
import cv2
import torch
import numpy as np
from PIL import Image
import os

model = "shadowlamer/sd-zxspectrum-model-256"
image_width = 256
image_height = 192
samples_dir = "/tmp"

pipe = DiffusionPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False)


# Borrowed from here: https://stackoverflow.com/a/73667318
def quantize_to_palette(_image, _palette):
    x_query = _image.reshape(-1, 3).astype(np.float32)
    x_index = _palette.astype(np.float32)
    knn = cv2.ml.KNearest_create()
    knn.train(x_index, cv2.ml.ROW_SAMPLE, np.arange(len(_palette)))
    ret, results, neighbours, dist = knn.findNearest(x_query, 1)
    _quantized_image = np.array([_palette[idx] for idx in neighbours.astype(int)])
    _quantized_image = _quantized_image.reshape(_image.shape)
    return Image.fromarray(cv2.cvtColor(np.array(_quantized_image, dtype=np.uint8), cv2.COLOR_BGR2RGB))


def collect_char_colors(image, _x, _y):
    _colors = {}
    for _char_y in range(8):
        for _char_x in range(8):
            _pixel = image.getpixel((_x + _char_x, _y + _char_y))
            _colors[_pixel] = 1 if _pixel not in _colors else _colors[_pixel] + 1
    _colors = sorted(_colors.items(), key=lambda _v: _v[1], reverse=True)
    return [list(_tuple[0]) for _tuple in list(_colors)]


def palette_to_attr(_palette):
    if len(_palette) == 0:
        return 0x00
    _attr = 0x40
    _paper = _palette[0]
    if _paper[0] > 0:
        _attr = _attr + 0x10  # r
    if _paper[1] > 0:
        _attr = _attr + 0x20  # g
    if _paper[2] > 0:
        _attr = _attr + 0x08  # b
    if len(_palette) == 1:
        return _attr
    _ink = _palette[1]
    if _ink[0] > 0:
        _attr = _attr + 0x02  # r
    if _ink[1] > 0:
        _attr = _attr + 0x04  # g
    if _ink[2] > 0:
        _attr = _attr + 0x01  # b
    return _attr


def generate(prompt, seed, steps):
    generator = torch.Generator("cpu").manual_seed(int(seed))
    raw_image = \
    pipe(prompt, height=image_height, width=image_width, num_inference_steps=int(steps), generator=generator).images[0]
    palette = np.array(
        [[0, 0, 0], [0, 0, 255], [0, 255, 0], [0, 255, 255], [255, 0, 0], [255, 0, 255], [255, 255, 0],
         [255, 255, 255]])
    input_image = np.array(raw_image)
    input_image = input_image[:, :, ::-1].copy()
    image = quantize_to_palette(_image=input_image, _palette=palette)

    out = samples_dir + "/" + prompt.replace(" ", "_") + "_" + str(seed) + ".scr"

    if not os.path.exists(out):
        byte_buffer = [0] * 0x1800
        attr_buffer = [0b00111000] * 0x300

        for y in range(0, image_height, 8):
            for x in range(0, image_width, 8):
                px = int(x / 8)
                py = int(y / 8)
                palette = collect_char_colors(image, x, y)
                byte_index = int(py / 8) * 0x800 + (py % 8) * 32 + px
                for cy in range(8):
                    byte = 0
                    for cx in range(8):
                        byte = byte * 2
                        pixel = list(image.getpixel((x + cx, y + cy)))
                        if palette[0] != pixel:
                            byte = byte + 1
                    byte_buffer[byte_index] = byte
                    byte_index = byte_index + 0x100
                attr = palette_to_attr(palette)
                attr_buffer[py * 32 + px] = attr

        scr = open(out, 'wb')
        scr.write(bytearray(byte_buffer))
        scr.write(bytearray(attr_buffer))
        scr.close()

    return [image, out]


iface = gr.Interface(fn=generate,
                     title="ZX-Spectrum inspired images generator ",
                     inputs=["text", "number", "number"],
                     outputs=["image", "file"],
                     examples=[["Cute cat", 123, 20], ["Solar system", 123, 20], ["Disco ball", 123, 20]])
iface.launch()