import os
import random

import gradio as gr
import numpy as np
from romkan import to_roma

from ocr import Recognizer

kana_list = """
あいうえお
かきくけこ
がぎぐげご
さしすせそ
ざじずぜぞ
たちつてと
だぢづでど
なにぬねの
はひふへほ
ばびぶべぼ
ぱぴぷぺぽ
まみむめも
やゆよ
らりるれろ
わを
ん

アイウエオ
カキクケコ
ガギグゲゴ
サシスセソ
ザジズゼゾ
タチツテト
ダヂヅデド
ナニヌネノ
ハヒフヘホ
バビブベボ
パピプペポ
マミムメモ
ヤユヨ
ラリルレロ
ワヲ
ン
"""
kana_list = [c for c in kana_list if c.strip()]

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


def get_char(ch=None, assist=True, kana=True):
    if ch is None:
        img_path = random.choice(chars)
        ch = img_path[-5]
    else:
        img_path = f"chars/{ch}.png"

    if not assist:
        img_path = "bg.png"

    info = f"{ch} ({to_roma(ch)})"
    if not kana:
        info = f"{to_roma(ch)}"

    return ch, img_path, info


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

with gr.Blocks(theme=theme, title="Kana Writer") as app:
    ch, char_img, roma = get_char()
    curr_img_path = gr.State(char_img)
    curr_char = gr.State(ch)
    recog = Recognizer("model/model.xml", "model/char_list.txt")

    with gr.Tab("寫字練習"):
        brush = gr.Brush(default_color="#111", default_size=15)
        sketch = gr.Sketchpad(
            char_img,
            type="numpy",
            brush=brush,
            layers=False,
            image_mode="RGB",
            eraser=False,
            show_label=False,
        )
        with gr.Row():
            header = gr.Textbox(roma, label="目標")
            result = gr.Textbox(label="辨識結果")

        with gr.Row():
            with gr.Row():
                assist = gr.Checkbox(True, label="輔助")
                kana = gr.Checkbox(True, label="假名")
            rand_btn = gr.Button("隨機")
            check_btn = gr.Button("辨識")

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

    def valid_item(item):
        if item["prob"] < 1e-2:
            return False
        if item["char"] not in kana_list:
            return False
        return True

    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 ", ".join(parse_item(item) for items in nbest for item in items if valid_item(item))

    def rand_char(assist, kana):
        char, img, roma = get_char(None, assist, kana)
        return char, img, img, roma, None

    def clear(curr_char):
        return curr_char

    def change_opt(curr_char, assist, kana):
        char, img, roma = get_char(curr_char, assist, kana)
        return char, img, img, roma

    check_btn.click(do_recog, sketch, result, show_progress="minimal")
    rand_btn.click(
        rand_char,
        [assist, kana],
        [curr_char, sketch, curr_img_path, header, result],
        show_progress="minimal",
    )
    sketch.clear(clear, curr_img_path, sketch, show_progress="minimal")
    assist.change(
        change_opt,
        [curr_char, assist, kana],
        [curr_char, sketch, curr_img_path, header],
        show_progress="minimal",
    )
    kana.change(
        change_opt,
        [curr_char, assist, kana],
        [curr_char, sketch, curr_img_path, header],
        show_progress="minimal",
    )

app.launch()