import os
import random

import gradio as gr
import numpy as np

from ocr import Recognizer

chars = [os.path.join(dn, fn) for dn, _, ff in os.walk("chars") for fn in ff]

font = gr.themes.GoogleFont("Noto Sans")
theme = gr.themes.Soft(font=font)

with gr.Blocks(theme=theme, title="Kana Writer") as app:
    char = random.choice(chars)
    curr_char = gr.State(char)
    recog = Recognizer("model/model.xml", "model/char_list.txt")

    brush = gr.Brush(default_color="#111", default_size=15)
    sketch = gr.Sketchpad(
        char,
        type="numpy",
        brush=brush,
        layers=False,
        image_mode="RGB",
        label="寫字練習",
    )
    result = gr.TextArea(label="辨識結果", lines=1)

    with gr.Row():
        check_btn = gr.Button("辨識")
        rand_btn = gr.Button("隨機")

    def parse_item(item):
        prob = item["prob"]
        char = item["char"]
        return f"{char}: {prob:.2%}"

    def do_recog(img: dict[str, np.ndarray]):
        img: np.ndarray = img["layers"][0]
        img[img == 0] = 255
        img[img != 255] = 0
        _, nbest = recog(img)
        return "\n".join(parse_item(i) for item in nbest for i in item if i["prob"] > 1e-2)

    def rand_char():
        char = random.choice(chars)
        return char, char

    def clear(curr_char):
        return curr_char

    check_btn.click(do_recog, sketch, result, show_progress="minimal")
    rand_btn.click(rand_char, None, [sketch, curr_char], show_progress="minimal")
    sketch.clear(clear, curr_char, sketch, show_progress="minimal")

app.launch()