import hashlib
from types import SimpleNamespace
import gradio as gr
from PIL import Image
from ultralytics import YOLO
from ultralytics.utils.plotting import save_one_box
import easyocr
import zxingcpp
import numpy as np

from utils import custom_plot


model = YOLO("YOLOV8s_Barcode_Detection.pt")
reader = easyocr.Reader(['en'])
key_hash = "6734e0268423441620ada6c66bf76c255500139962741833ab33185274b8a010"
hash = lambda x: hashlib.sha256(x.encode()).hexdigest()


def process_image(input_img, activation_key, progress=gr.Progress()):
    if input_img is None:
        return None, "請上傳圖片後再進行辨識。"
    if hash(activation_key) != key_hash:
        return None, "金鑰錯誤，請確認後再試一次。"

    # Perform object detection on an image
    result = model(input_img, imgsz=(1280))[0]

    crops = []
    for d in result.boxes:
        crops.append(save_one_box(
            d.xyxy,
            result.orig_img.copy(),
            save=False,
        ))

    texts = []

    for pr, crop in enumerate(crops):
        progress((pr+1) / len(crops), desc="辨識中")
        img = Image.fromarray(crop)
        res = zxingcpp.read_barcodes(img)

        if not res:
            # rotate and retry
            for i in range(1, 8):
                res = zxingcpp.read_barcodes(img.rotate(i, resample=2))
                if res:
                    break
                res = zxingcpp.read_barcodes(img.rotate(-i, resample=2))
                if res:
                    break
            
        if not res:
            # resort to ocr the bottom-left no.
            full_res = reader.readtext(crop, allowlist='-0123456789')
            res = sorted(
                [r for r in full_res if (
                    r[0][3][0] < (img.width/4) and r[0][3][1] > (img.height/3) and
                    r[0][0][0] < (img.width/4) and r[0][0][1] > (img.height/2)
                )],
                key=lambda x: x[2], reverse=True,
            )
            if res:
                pred_text = res[0][1]
                # sanity check if the no. is not divided into multiple box
                cur_box = res[0]
                # if len(cur_box[1]) < 13:
                other_boxes = [r for r in full_res if r[1] != cur_box[1]]
                thrs = np.linalg.norm(img.size) / 25
                while other_boxes:
                    upper = np.linalg.norm(cur_box[0][1] - np.array([b[0][0] for b in other_boxes]), axis=1)
                    lower = np.linalg.norm(cur_box[0][2] - np.array([b[0][3] for b in other_boxes]), axis=1)
                    is_same = (upper < thrs) & (lower < thrs)
                    rank = sorted(
                        [(i, dist, s) for (i, dist), s in zip(enumerate(lower + upper), is_same) if s],
                        key=lambda x: x[1]
                    )
                    if rank:
                        cur_box = other_boxes[rank[0][0]]
                        other_boxes = [r for r in other_boxes if r[1] != cur_box[1]]
                        pred_text += cur_box[1]
                    else:
                        break
                # HACK
                if len(pred_text) != 15:
                    res = []
            res = [SimpleNamespace(text=pred_text)] if res else []
        
        texts.append(res[0].text if res else None)

    output_text = '\n'.join([t for t in texts if isinstance(t, str)])
    results_img = custom_plot(
        result, font_size=40, pil=True,
        barcode_texts=texts,
    )
    return results_img, output_text



# Defining the Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Barcode")
    gr.Markdown("請用下方 example 圖片測試條碼辨識功能，或上傳您自己的圖片。")
    
    with gr.Row():
        with gr.Column():
            input_view = gr.Image(type="pil", label="輸入圖片")
            activation_key = gr.Textbox(label="金鑰")
            btn = gr.Button("辨識", variant="primary")

            gr.Examples(examples=[["example1.jpeg", ""], ["example2.jpeg", ""]], inputs=input_view)
        with gr.Column():
            output_view = gr.Image(type="pil", label="辨識結果")
        with gr.Column():
            text_output = gr.Textbox(label="條碼內容")

    # Wire up the button
    btn.click(
        fn=process_image, 
        inputs=[input_view, activation_key],
        outputs=[output_view, text_output]
    )

if __name__ == "__main__":
    demo.launch()