import gradio as gr
from PIL import Image
from openpyxl import Workbook


def gray_to_dice_1(x: int) -> int:
    if 0 <= x and x <= 41:
        return 1
    elif 41 < x and x <= 83:
        return 2
    elif 83 < x and x <= 124:
        return 3
    elif 124 < x and x <= 165:
        return 4
    elif 165 < x and x <= 206:
        return 5
    elif 206 < x and x <= 247:
        return 6
    else:
        return 6


def gray_to_dice_2(x: int) -> int:
    if 0 <= x and x <= 41:
        return 0
    elif 41 < x and x <= 83:
        return 1
    elif 83 < x and x <= 124:
        return 2
    elif 124 < x and x <= 165:
        return 3
    elif 165 < x and x <= 206:
        return 4
    elif 206 < x and x <= 247:
        return 5
    else:
        return 6


def dice_to_image(
    image: Image.Image,
    resolution_x: int,
    resolution_y: int,
    dice_size: int,
):
    resolution_x, resolution_y = int(resolution_x), int(resolution_y)
    dice_size = int(dice_size)

    image = image.convert("L")
    image = image.resize((resolution_x, resolution_y))

    dice_img_1 = Image.new("L", (image.size[0] * dice_size, image.size[1] * dice_size))
    dice_img_2 = Image.new("L", (image.size[0] * dice_size, image.size[1] * dice_size))

    dices = [Image.new("L", (dice_size, dice_size))] + [
        Image.open(f"images/dices/{i}.jpeg").resize((dice_size, dice_size)).convert("L")
        for i in range(1, 7)
    ]

    wb = Workbook()
    for sheet_name in wb.sheetnames:
        wb.remove(wb[sheet_name])
    ws1 = wb.create_sheet("Style #1")
    ws2 = wb.create_sheet("Style #2")

    for i in range(image.size[0]):
        for j in range(image.size[1]):
            pixel = int(image.getpixel((i, j)))
            dice_index_1 = gray_to_dice_1(pixel)
            dice_index_2 = gray_to_dice_2(pixel)

            ws1.cell(row=j + 1, column=i + 1).value = dice_index_1
            ws2.cell(row=j + 1, column=i + 1).value = dice_index_2

            box = (
                i * dice_size,
                j * dice_size,
                (i + 1) * dice_size,
                (j + 1) * dice_size,
            )
            dice_img_1.paste(dices[dice_index_1], box=box)
            dice_img_2.paste(dices[dice_index_2], box=box)

    wb.save("results.xlsx")

    return dice_img_1, dice_img_2, "results.xlsx"


def main():
    with gr.Blocks() as app:
        gr.Markdown("Image to Dice")

        with gr.Row(variant="panel"):
            image_input = gr.Image(type="pil")

        with gr.Row(variant="panel"):
            resolution_x = gr.Number(value=200, precision=0, label="分辨率（宽）", show_label=True)
            resolution_y = gr.Number(value=200, precision=0, label="分辨率（高）", show_label=True)
            dice_size = gr.Number(value=16, precision=0, label="骰子大小", show_label=True)

        with gr.Row(variant="panel"):
            run_button = gr.Button("生成")

        with gr.Row(variant="panel"):
            image_output1 = gr.Image(type="pil", label="不用全黑", show_label=True)
            image_output2 = gr.Image(type="pil", label="用全黑", show_label=True)

        with gr.Row(variant="panel"):
            excel_file = gr.File(label="Excel结果", show_label=True)

        run_button.click(
            dice_to_image,
            inputs=[
                image_input,
                resolution_x,
                resolution_y,
                dice_size,
            ],
            outputs=[
                image_output1,
                image_output2,
                excel_file,
            ],
        )

        app.launch()


if __name__ == "__main__":
    main()