import cv2
import numpy as np
import streamlit as st
from run_yolo import get_layout_results
from order_text_blocks import get_ordered_data
from run_ocr import OCR
from main import driver
import json
import pandas as pd

colors = {
    'Articles': [0, 0, 0],        # Red
    'Advertisement': [0, 255, 0],   # Green
    'Headlines': [0, 0, 255],       # Blue
    'Sub-headlines': [255, 255, 0],  # Yellow
    'Graphics': [255, 0, 255],      # Magenta
    'Images': [128, 0, 128],        # Purple
    'Tables': [0, 255, 255],        # Teal
    'Header': [0, 0, 0],             # Black
    'Text Block': [255, 0, 0]
}

try:
    st.set_page_config(layout="wide", page_title="Newspaper Layout Detection and OCR Demo")
    st.markdown("<h1 style='text-align: center; color: #333;'>Newspaper Layout Detection and OCR Demo</h1>", unsafe_allow_html=True)
    
    # Streamlit UI for user input
    uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
    language_name = st.text_input("Enter the language name: (hin, en, mal, tel, tam, kan))")
    submit_button = st.button("Submit")

    # Check if the user clicked the submit button
    if submit_button:
        # Check if image and language are provided
        if uploaded_image is not None and language_name:
            # Convert Streamlit file uploader to OpenCV image
            image_bytes = uploaded_image.read()
            nparr = np.frombuffer(image_bytes, np.uint8)
            img_ori = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            img = img_ori.copy()
            st.markdown("<p style='text-align: center; color: red'>Image Uploaded Successfully!</p>", unsafe_allow_html=True)
            # Continue with the rest of the code...
            output_dict, article_ocr_dict = driver(img_ori, language_name, st)
        
        # Create a list to store dictionaries for OCR results
        image_data = []

        # Visualizing Results
        itr = 1
        for art_key in article_ocr_dict:
            art_coords = art_key.split('_')
            art_x1, art_y1, art_x2, art_y2 = int(art_coords[0]), int(art_coords[1]), int(art_coords[2]), int(art_coords[3])
            
            # Mark the article bounding box with dark green color
            img_ori = cv2.rectangle(img, (art_x1, art_y1), (art_x2, art_y2), (0, 0, 0), 4)
            # Put the article number on the image in large font
            img_ori = cv2.putText(img_ori, str(itr), (art_x1, art_y1), cv2.FONT_HERSHEY_SIMPLEX, 3, (0, 0, 255), 3, cv2.LINE_AA)            
            ocr_dict = article_ocr_dict[art_key]

            # Initialize variables to store OCR text for each type
            headlines_text = ''
            subheadlines_text = ''
            textblocks_text = ''

            for obj_key in ocr_dict:
                # obj_key is either of Headlines, Sub-headlines, Text Block
                obj_list = ocr_dict[obj_key]
                for obj_dict in obj_list:
                    for key in obj_dict:
                        coords = key.split('_')
                        x1, y1, x2, y2 = int(coords[0]), int(coords[1]), int(coords[2]), int(coords[3])
                        # Mark the bounding box with color corresponding to the object type
                        img = cv2.rectangle(img, (x1, y1), (x2, y2), colors[obj_key], 2)

                        # Add the OCR text to the corresponding variable
                        if obj_key == 'Headlines':
                            headlines_text += obj_dict[key] + '\n'
                        elif obj_key == 'Sub-headlines':
                            subheadlines_text += obj_dict[key] + '\n'
                        elif obj_key == 'Text Block':
                            textblocks_text += obj_dict[key] + '\n'

            # Add a dictionary to the list for the current article
            image_data.append({'Article': itr, 'Headlines': headlines_text, 'Subheadlines': subheadlines_text, 'Textblocks': textblocks_text})

            itr += 1

        # Create a DataFrame from the list of dictionaries
        image_df = pd.DataFrame(image_data)

        # Use Streamlit columns to display the image and results side by side
        col1, col2 = st.columns(2)

        # Display the image with marked bounding boxes in the left column
        col1.image(img_ori, use_column_width=True, channels="BGR", caption="Image with Marked Bounding Boxes")

        # Display the DataFrame for OCR results for the whole image in the right column
        col2.table(image_df.set_index('Article').style.set_table_styles([{'selector': 'th', 'props': [('text-align', 'center')]}]))

        
except Exception as e:
    st.exception(e)