import json
import os
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from pdf2image import convert_from_path
import cv2
from sklearn.cluster import AgglomerativeClustering
import networkx as nx
import time

class UnionFind:
    def __init__(self, size):
        self.parent = list(range(size))

    def find(self, u):
        while self.parent[u] != u:
            self.parent[u] = self.parent[self.parent[u]]  # Path compression
            u = self.parent[u]
        return u

    def union(self, u, v):
        pu = self.find(u)
        pv = self.find(v)
        if pu != pv:
            self.parent[pu] = pv  # Merge the sets
            
def reconstruct_text(lines):
    """
    Reconstruct text from grouped lines.
    Each line is a list of dictionaries.
    """
    reconstructed_lines = []
    for line in lines:
        # Sort line by x1
        sorted_line = sorted(line, key=lambda elem: elem['x1'])
        # Join text with spaces
        line_text = ' '.join(elem['text'] for elem in sorted_line)
        reconstructed_lines.append(line_text)
   # print('\n'.join(reconstructed_lines))
    return '\n'.join(reconstructed_lines)

def visualize_with_text_and_image(image, reconstructed_text, lines, method='clustering'):
    """
    Create a side-by-side visualization of the reconstructed text and the image with bounding boxes.
    
    Parameters:
    - image: The original image as a NumPy array.
    - reconstructed_text: The reconstructed text string.
    - lines: List of lines, each line is a list of dictionaries.
    - method: The grouping method used (for title purposes).
    """
    # Create a copy of the image to draw colored bounding boxes
    img_copy = image.copy()
    
    # Define colors for different lines using a colormap
    colors = plt.colormaps['hsv'](np.linspace(0, 1, len(lines) + 1))
    
    for idx, line in enumerate(lines):
        color = tuple([int(255 * c) for c in colors[idx][:3]])  # Convert to 0-255 scale
        for elem in line:
            x1, y1, x2, y2 = elem['x1'], elem['y1'], elem['x2'], elem['y2']
            cv2.rectangle(img_copy, (x1, y1), (x2, y2), color, 2)
    
    # Convert image from BGR to RGB for matplotlib
    img_rgb = cv2.cvtColor(img_copy, cv2.COLOR_BGR2RGB)
    
    # Create a matplotlib figure with two subplots
    fig, axes = plt.subplots(1, 2, figsize=(25, 15))
    
    axes[0].axis('off')
    axes[0].text(0, 1, reconstructed_text, fontsize=12, va='top', wrap=True)
    axes[0].set_title(f'Reconstructed Text using {method.capitalize()} Method', fontsize=16)
    
    # Right subplot: Image with Bounding Boxes
    axes[1].imshow(img_rgb)
    axes[1].axis('off')
    axes[1].set_title('OCR Bounding Boxes', fontsize=16)
    
    plt.tight_layout()
    plt.show()
    
def load_ocr_data(json_filename):
    """
    Load OCR data from a JSON file into a pandas DataFrame.
    """
    with open(json_filename, 'r', encoding='utf-8') as json_file:
        bbox = json.load(json_file)
    df = pd.DataFrame(bbox)
    return df

def preprocess_df(df):
    """
    Preprocess the DataFrame by ensuring required columns are present and correctly typed.
    """
    required_columns = ['x1', 'y1', 'x2', 'y2', 'height', 'width', 'text', 'type_of_text', 'new_line', 'region']
    for col in required_columns:
        if col not in df.columns:
            raise ValueError(f"Missing required column: {col}")
    # Convert coordinates and size columns to integers
    df[['x1', 'y1', 'x2', 'y2', 'height', 'width']] = df[['x1', 'y1', 'x2', 'y2', 'height', 'width']].astype(int)
    return df
can you apply on each row in reconstructed_text a color that matches the row in right subplot with bounding boxes?