import streamlit as st
import cv2
import numpy as np
from PIL import Image, ImageDraw
# import imutils
# import easyocr
# import os
# import pathlib
# import platform
# from xyxy_converter import yolov5_to_image_coordinates
# import shutil
from models import get_odometer_xy, get_digit

# system_platform = platform.system()
# if system_platform == 'Windows': pathlib.PosixPath = pathlib.WindowsPath

# CUR_DIR = os.getcwd()
# YOLO_PATH = f"{CUR_DIR}/yolov5"
# MODEL_PATH = "runs/train/exp/weights/best.pt"

def main():
    st.title("Odometer value extractor with Streamlit")

    # Use st.camera to capture images from the user's camera
    img_file_buffer = st.camera_input(label='Please, take a photo of odometer', key="odometer")

    # Check if an image is captured
    if img_file_buffer is not None:
        # Convert the image to a NumPy array
        image = Image.open(img_file_buffer)
        image_np = np.array(image)
        resized_image = cv2.resize(image_np, (640, 640))
        resized_image = resized_image.astype(np.uint8)
        cv2.imwrite('odometer_image.jpg', resized_image)

        # original_img = cv2.imread('odometer_image.jpg')
        gray = cv2.cvtColor(resized_image, cv2.COLOR_BGR2GRAY)

        x1, y1, x2, y2, odo_confidence = get_odometer_xy(
            model_path='odo_detector.tflite',
            image_path='odometer_image.jpg'
        )

        st.write(odo_confidence)

        if odo_confidence == 0:
            display_text = "An odometer is not detected in the image!!!"
            st.image('odometer_image.jpg', caption=f"{display_text}", use_column_width=True)
        else:
            # cropped_image = gray[y1:y2, x1:x2]
            cropped_image = resized_image[y1:y2, x1:x2]
            cropped_image = cv2.resize(cropped_image, (640, 640))
            cv2.imwrite('odometer_number_image.jpg', cropped_image)
            
            extracted_digit = get_digit(
                model_path="digit_yolov8_best_float16.tflite", 
                image_path='odometer_number_image.jpg', 
                threshold=0.4
            )

            display_text = f'Here is the zoomed odometer value: {extracted_digit}'
            st.image('odometer_number_image.jpg', caption=f"{display_text}", use_column_width=True)

            image = Image.open('odometer_image.jpg')
            image_resized = image.resize((640, 640))
            draw = ImageDraw.Draw(image_resized)
            draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
            class_name = 'odometer'
            text = f"Class: {class_name}, Confidence: {odo_confidence:.2f}"
            draw.text((x1, y1), text, fill="red")
            # Saving Images
            image_resized.save('odometer_highlighted_image.jpg')
            display_text = 'Here is the odometer on the image.'
            st.image('odometer_highlighted_image.jpg', caption=f"{display_text}", use_column_width=True)

        # detect(
        #     weights='yolov5\runs\train\exp\weights\best.pt',
        #     source='odometer_image.jpg',
        #     img=640,
        #     conf=0.4,
        #     name='temp_exp',
        #     hide_labels=True,
        #     hide_conf=True,
        #     save_txt=True,
        #     exist_ok=True
        # )

        # # os.system('wandb disabled')

        # os.chdir(YOLO_PATH)

        # # try:
        # #     shutil.rmtree('runs/detect/temp_exp')
        # # except:
        # #     pass

        # image_path = "../odometer_image.jpg"
        # # command = f"python detect.py --weights {MODEL_PATH} --source {image_path} --img 640 --conf 0.4 --name 'temp_exp' --hide-labels --hide-conf --save-txt --exist-ok"
        # command = f'''
        # python detect.py \
        #     --weights {MODEL_PATH} \
        #     --source {image_path} \
        #     --img 640 \
        #     --conf 0.4 \
        #     --name temp_exp \
        #     --hide-labels \
        #     --hide-conf \
        #     --save-txt \
        #     --exist-ok \
        #     --save-conf
        # '''

        # # Run the command
        # os.system(command)

        # # st.write('The detection is completed!!!')

        # os.chdir(CUR_DIR)

        # # st.write(os.path.exists('yolov5/runs/detect/temp_exp'))

        # if os.path.exists('yolov5/runs/detect/temp_exp'):
        #     processed_image = cv2.imread('yolov5/runs/detect/temp_exp/odometer_image.jpg')
        #     # st.write('Image boxed and loaded')
        #     text_files = os.listdir('yolov5/runs/detect/temp_exp/labels')
        #     original_img = cv2.imread('odometer_image.jpg')
        #     gray = cv2.cvtColor(original_img, cv2.COLOR_BGR2GRAY)

        #     if len(text_files) == 0:
        #         display_text = "An odometer is not detected in the image!!!"
        #     else:
        #         text_file_path = f'yolov5/runs/detect/temp_exp/labels/{text_files[0]}'
        #         x1, y1, x2, y2 = yolov5_to_image_coordinates(text_file_path)
        #         # st.write(x1, y1, x2, y2)
        #         cropped_image = gray[x1:x2, y1:y2]

        #         reader = easyocr.Reader(['en'])
        #         result = reader.readtext(cropped_image)

        #         if len(result) != 0:
        #             odometer_value = sorted(result, key=lambda x: x[2], reverse=True)[0][1]
        #             display_text = f"Odometer value: {odometer_value}"
        #         else:
        #             odometer_value = 'not detected'
        #             display_text = f"The odometer value is {odometer_value}!!!"
        # else:
        #     display_text = "An odometer is not detected in the image!!!"
        #     processed_image = cv2.imread('odometer_image.jpg')

        # try:
        #     shutil.rmtree('odometer_image.jpg')
        #     shutil.rmtree('runs/detect/temp_exp')
        # except:
        #     pass

        # # Resize or preprocess the image as needed for your model
        # # For example, resizing to a specific input size
        # # processed_image = cv2.resize(image_np, (224, 224))
        
        # # Process the image using your deep learning model
        # # processed_image = process_image(image_np)

        # # Display the processed image
        # st.image(processed_image, caption=f"{display_text}", use_column_width=True)

        st.session_state.pop("odometer")

if __name__ == "__main__":
    main()