import cv2
import numpy as np
from skimage.measure import label, regionprops
from collections import Counter
from typing import Tuple
import matplotlib.pyplot as plt
import os

class ImageRotator:
    def rotate_image(self, image: np.ndarray) -> Tuple[np.ndarray, float]:
        gray = self.convert_to_gray(image)
        threshold = self.threshold_image(gray)
        
        min_rect_angle = self.find_min_area_rect_angle(threshold)
        props_angle = self.find_regionprops_angle(threshold)
        
        combined_angles = np.concatenate(([min_rect_angle], [props_angle]))
        rotation_angle = self.calculate_mode(combined_angles)
        
        return self.rotate_image_cv2(image, rotation_angle), rotation_angle

    def convert_to_gray(self, image: np.ndarray) -> np.ndarray:
        return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    def threshold_image(self, image: np.ndarray) -> np.ndarray:
        return cv2.adaptiveThreshold(image, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 19, 5)

    def find_min_area_rect_angle(self, image: np.ndarray) -> float:
        contours, _ = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if not contours:
            return 0.0
        largest_contour = max(contours, key=cv2.contourArea)
        rect = cv2.minAreaRect(largest_contour)
        angle = rect[2]
        return angle if angle < 45 else angle - 90

    def find_regionprops_angle(self, image: np.ndarray) -> float:
        labeled_image = label(image)
        regions = regionprops(labeled_image)
        if len(regions) == 0:
            return 0.0
        angles = [region.orientation * (180 / np.pi) for region in regions]
        return np.mean(angles)

    def find_hough_lines(self, image: np.ndarray) -> np.ndarray:
        threshold = 160
        lines = cv2.HoughLinesP(image=image, rho=1, theta=np.pi / 180, threshold=threshold, minLineLength=50, maxLineGap=10)
        if lines is None:
            return np.array([0])
        angles = np.array([])
        for line in lines.squeeze():
            A = line[3] - line[1]
            B = line[0] - line[2]
            cos_theta = -B / (np.sqrt(A**2 + B**2))
            angle = np.rad2deg(np.arccos(cos_theta))
            angles = np.append(angles, angle)
        if angles.size == 0:
            return np.array([0])
        return angles

    def calculate_mode(self, angles: np.ndarray) -> float:
        return Counter(angles).most_common(1)[0][0]

    def rotate_image_cv2(self, image: np.ndarray, angle: float) -> np.ndarray:
        image_center = tuple(np.array(image.shape[1::-1]) / 2)
        rot_mat = cv2.getRotationMatrix2D(image_center, angle, 1.0)
        return cv2.warpAffine(image, rot_mat, image.shape[1::-1], flags=cv2.INTER_LINEAR, borderValue=(255, 255, 255))


def rotate_image(image: np.ndarray) -> np.ndarray:
    return ImageRotator().rotate_image(image=image)

def grayscale(image: np.ndarray) -> np.ndarray:
    return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

def prep(filename, image):
    rotated_image, rotation_angle = ImageRotator().rotate_image(image) 
    print(f"{filename[:-4].split('/')[-1]} - Rotation Angle: {rotation_angle}")

def display_images(original, rotated):
    fig, axes = plt.subplots(1, 2, figsize=(12, 6))
    axes[0].imshow(cv2.cvtColor(original, cv2.COLOR_BGR2RGB))
    axes[0].set_title("Original Image")
    axes[0].axis("off")
    axes[1].imshow(cv2.cvtColor(rotated, cv2.COLOR_BGR2RGB))
    axes[1].set_title("Rotated Image")
    axes[1].axis("off")
    plt.show()

folder = "test_rotation"

for filename in os.listdir(folder)[0:15]:
    if filename.endswith('.png'):
        image_path = os.path.join(folder, filename)
        image = cv2.imread(image_path)
        rotated_image = prep(filename, image)
        display_images(image, rotated_image)
# filename = "d79ec5e0-963932.png"
# if filename.endswith('.png'):
#     image_path = os.path.join(folder, filename)
#     image = cv2.imread(image_path)
#     rotated_image = prep(filename, image)
#     display_images(image, rotated_image)
add fourier_transform_angle