import streamlit as st
from PIL import Image
import numpy as np
from streamlit_cropper import st_cropper
from annotated_text import annotated_text
from simpletransformers.ner import NERModel
import pytesseract
import os
def process_image(img_file):
    st.write("Crop ส่วนที่ต้องการ (คุณลักษณะ)")
    aspect_ratio = None

    col1, col2 = st.sidebar.columns(2)
    with col1:
        box_color = st.color_picker(label="Box Color", value='#0000FF')
    return_type = "image"
    realtime_update = st.sidebar.checkbox(label="Update in Real Time", value=True)

    img = Image.open(img_file)
    width, height = img.size

    if not realtime_update:
        st.write("Double click to save crop")

    if return_type == 'box':
        rect = st_cropper(
            img,
            realtime_update=realtime_update,
            box_color=box_color,
            aspect_ratio=aspect_ratio,
            return_type=return_type
        )
        raw_image = np.asarray(img).astype('uint8')
        left, top, width, height = tuple(map(int, rect.values()))
        st.write(rect)
        masked_image = np.zeros(raw_image.shape, dtype='uint8')
        masked_image[top:top + height, left:left + width] = raw_image[top:top + height, left:left + width]
        st.image(Image.fromarray(masked_image))
    else:
        # Get a cropped image from the frontend
        cropped_img = st_cropper(
            img,
            realtime_update=realtime_update,
            box_color=box_color,
            aspect_ratio=aspect_ratio,
            return_type=return_type
        )

        # Manipulate cropped image at will
        st.write("Preview")
        _ = cropped_img.thumbnail((1000, 1000))

        desired_dpi = 300

        # Resize the image
        resized_image = resize_image(cropped_img, desired_dpi)

        # Display the image
        st.image(resized_image)
        # Perform OCR
        # extracted_text = 

        # Display extracted text
        txt_input = st.text_area("โปรดตรวจสอบความถูกต้อง", perform_ocr(cropped_img))

        word_output = ""
        word_tuple = ()
        check = st.button("Check", type="primary", use_container_width=True)
        if check:
            st.header("สรุปผล")
            predictions, raw_outputs = model.predict([txt_input])
            wrongword_lst = []

            for i in predictions[0]:
                word = i.values()
                label = list(word)[0]
                if label == "1":
                    wrong_word = list(i.keys())[0]
                    wrongword_lst.append(wrong_word)

            word_tuple = ()
            for i in predictions[0]:
                label = i.values()
                word = list(i.keys())[0]
                label = list(label)[0]
                if label == "1":
                    word_tuple = (*word_tuple, (f'{word} ', 'มีการล็อคสเปค', '#F41B15'))
                else:
                    word_tuple = (*word_tuple, f'{word} ')
            word_output = wrongword_lst

        if word_output != "":
            annotated_text(list(word_tuple))

            col1, col2 = st.columns(2)
            with col1:
                st.header("คำที่สุ่มเสี่ยง")
                st.write(word_output)






def get_sample_images(folder):
    sample_images = []
    for filename in os.listdir(folder):
        if filename.endswith(".jpg") or filename.endswith(".png"):
            sample_images.append(os.path.join(folder, filename))
    return sample_images
# Function to resize the image based on desired DPI
def resize_image(image, desired_dpi):
    # Calculate the current DPI
    current_dpi = image.info.get('dpi', (72, 72))

    # Calculate the scale factor based on the desired and current DPI
    scale_factor = desired_dpi / current_dpi[0]

    # Resize the image based on the scale factor
    new_size = (int(image.width * scale_factor), int(image.height * scale_factor))
    resized_image = image.resize(new_size)

    return resized_image

# Function to perform OCR on the image
def perform_ocr(image):
    extracted_text =pytesseract.image_to_string(image=image, lang='tha')

    return extracted_text

# Main function
def main():
 

    st.info('''
        การใช้งาน
        กรอกลักษณะของครุภัณฑ์ที่ต้องการตรวจสอบ โดยในขณะนี้รองรับเฉพาะ เครื่องปรับอากาศเเบบเเยกส่วน ตู้เหล็ก เครื่องทำลายเอกสาร เครื่องถ่ายเอกสารระบบดิจิตอล 1.) เลือกไฟล์รูป 2.) Crop ส่วนที่ต้องการตรวจสอบ 3.) ตรวจสอบข้อความ 4.) กดปุ่มตรวจสอบ
    ''', icon="ℹ️")
    
 
    st.set_option('deprecation.showfileUploaderEncoding', False)

    img_files = st.sidebar.file_uploader(label='Upload a file', type=['png', 'jpg'], accept_multiple_files=False)
    sample_folder = "test_img"
    sample_images = get_sample_images(sample_folder)
    
    # Display the sample images in the sidebar
    selection = st.sidebar.selectbox("Select a sample", ["Browse files Mode"] + sample_images, format_func=lambda x: os.path.basename(x))
    if selection and selection != "Browse files Mode":
        img_file = selection
        process_image(img_file)
        
        
        

    elif img_files:
        img_file = img_files
        process_image(img_file)
        
    

if __name__ == '__main__':
    st.title("TEEOR")
    teeor_logo = Image.open(r'TEEOR.png')
    logo = st.sidebar.image(teeor_logo)
    st.sidebar.write("[AI BUILDER 3] [ภัทรดนัย อัคราช สังกัด Arcane Whales] - ตรวจสอบการล็อคสเปคในหนังสือจัดซื้อจัดจ้าง (ครุภัณฑ์)")
    genre = st.sidebar.radio(
    "Hardware",
    ('Cpu', 'Cuda (GPU)'))
    if genre == 'Cuda':
        Hardware_Type = True
    elif genre == 'Cpu':
        Hardware_Type = False
    # Model initialization
    _NER_TAGS = ['0', '1']
    # try:
    model = NERModel(
        model_name=r"model",
        model_type="camembert",
        labels=_NER_TAGS,
        use_cuda=Hardware_Type
    )
    
    # except:
    #     st.warning("Cuda not available, Please checking your hardware")
        # Get the sample images

    
    main()
    # Display the selected image

    aib_logo = Image.open(r'aib_logo.png')
    add_logo = st.sidebar.image(aib_logo)