# -*- encoding: utf-8 -*-
# @Author: SWHL
# @Contact: liekkaskono@163.com
from pathlib import Path

import numpy as np
import streamlit as st
from PIL import Image
from rapid_latex_ocr import LatexOCR
from streamlit_cropper import st_cropper
from streamlit_image_select import image_select

st.set_option("deprecation.showfileUploaderEncoding", False)


class RecEquation:
    def __init__(self, model_dir: str):
        model_dir = Path(model_dir)

        image_resizer_path = model_dir / "image_resizer.onnx"
        encoder_path = model_dir / "encoder.onnx"
        decoder_path = model_dir / "decoder.onnx"
        tokenizer_json = model_dir / "tokenizer.json"
        self.model = LatexOCR(
            image_resizer_path=str(image_resizer_path),
            encoder_path=str(encoder_path),
            decoder_path=str(decoder_path),
            tokenizer_json=str(tokenizer_json),
        )

    def __call__(self, img: np.ndarray):
        result, elapse = self.model(img)
        return result, elapse


if __name__ == "__main__":
    st.markdown(
        "<h1 style='text-align: center;'><a href='https://github.com/RapidAI/RapidLatexOCR' style='text-decoration: none'>Rapid Latex OCR</a></h1>",
        unsafe_allow_html=True,
    )
    st.markdown(
        """
    <p align="left">
        <a href=""><img src="https://img.shields.io/badge/Python->=3.6,<3.12-aff.svg"></a>
        <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>
        <a href="https://pepy.tech/project/rapid_latex_ocr"><img src="https://static.pepy.tech/personalized-badge/rapid_latex_ocr?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads"></a>
        <a href="https://pypi.org/project/rapid_latex_ocr/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapid_latex_ocr"></a>
        <a href="https://semver.org/"><img alt="SemVer2.0" src="https://img.shields.io/badge/SemVer-2.0-brightgreen"></a>
        <a href="https://github.com/psf/black"><img src="https://img.shields.io/badge/code%20style-black-000000.svg"></a>
    </p>
    """,
        unsafe_allow_html=True,
    )

    # Upload an image and set some options for demo purposes
    img_file = st.sidebar.file_uploader(label="Upload a file", type=["png", "jpg"])
    realtime_update = st.sidebar.checkbox(label="Update in Real Time", value=False)
    box_color = st.sidebar.color_picker(label="Box Color", value="#0000FF")
    aspect_choice = st.sidebar.radio(
        label="Aspect Ratio", options=["Free", "1:1", "16:9", "4:3", "2:3"]
    )
    aspect_dict = {
        "Free": None,
        "1:1": (1, 1),
        "16:9": (16, 9),
        "4:3": (4, 3),
        "2:3": (2, 3),
    }
    aspect_ratio = aspect_dict[aspect_choice]
    with st.sidebar.container():
        img = image_select(
            label="Examples（click to select）:",
            images=[
                "images/equation.png",
                "images/eq_2.png",
                "images/eq_3.png",
                "images/eq_4.png",
            ],
            key="equation_default",
            use_container_width=False,
        )

    rec_eq_sys = RecEquation(model_dir="models")

    select_img_container = st.container()

    st.markdown("#### Select image:")
    img_empty = st.empty()

    img_empty.image(img, use_column_width=False)
    rec_res, elapse = rec_eq_sys(img)

    if img_file:
        img = Image.open(img_file)

        # Get a cropped image from the frontend
        with select_img_container:
            if not realtime_update:
                select_img_container.markdown("#### Double click to save crop")

            img = st_cropper(
                img,
                realtime_update=realtime_update,
                box_color=box_color,
                aspect_ratio=aspect_ratio,
            )

        img_empty.image(img, use_column_width=False)
        rec_res, elapse = rec_eq_sys(np.array(img))

    st.markdown(f"#### Rec result (cost: {elapse:.4f}s):")
    st.latex(rec_res)

    st.markdown("#### Latex source code:")
    st.code(rec_res)