import os
import cv2
import numpy as np
import prediction_utils  # Zależności od twojej struktury projektu

# Funkcja przetwarzająca obrazy i maski
def process_images_and_masks(images_directory: str, masks_directory: str, output_directory: str):
    """
    Przetwarza obrazy i maski przez kolejne funkcje, zapisując wyniki do plików
    oraz wyświetlając tylko kąty.
    """
    os.makedirs(output_directory, exist_ok=True)
    images_output_dir = os.path.join(output_directory, 'images')
    masks_output_dir = os.path.join(output_directory, 'masks')
    masks_with_lines_output_dir = os.path.join(output_directory, 'masks_with_lines')

    os.makedirs(images_output_dir, exist_ok=True)
    os.makedirs(masks_output_dir, exist_ok=True)
    os.makedirs(masks_with_lines_output_dir, exist_ok=True)
    images = [cv2.imread(os.path.join(images_directory, f)) for f in os.listdir(images_directory) if f.endswith('.png')]
    masks = [cv2.imread(os.path.join(masks_directory, f), 0) for f in os.listdir(masks_directory) if f.endswith('.png')]

    images = images[700:800]
    masks = masks[700:800]

    data = {
        "images": images,
        "masks": masks
    }

    print(f"Initial number of images: {len(data['images'])}")
    print(f"Initial number of masks: {len(data['masks'])}")

    print("1. Przepuszczanie danych przez `choose_frame_1`...")
    data = prediction_utils.choose_frame_1(data)
    print(f"After `choose_frame_1`, remaining images: {len(data['images'])}")
    print(f"After `choose_frame_1`, remaining masks: {len(data['masks'])}")

    print("2. Wyznaczanie punktów i bazowych linii...")
    data = prediction_utils.calculate_points_and_baseline_8class(data)
    print(f"After `calculate_points_and_baseline`, remaining images: {len(data['images'])}")
    print(f"After `calculate_points_and_baseline`, remaining masks: {len(data['masks'])}")

    print("3. Przepuszczanie danych przez `choose_frame_2`...")
    data = prediction_utils.choose_frame_2(data)
    print(f"After `choose_frame_2`, remaining images: {len(data['images'])}")
    print(f"After `choose_frame_2`, remaining masks: {len(data['masks'])}")

    print("4. Obliczanie kąta alpha...")
    try:
        result = prediction_utils.identify_alpha_beta_angle_new(data)

        print("Wynik z `identify_alpha_beta_angle_new`:")
        print(result)

        image_filename = "image_with_max_angle.png"
        image_path = os.path.join(images_output_dir, image_filename)
        cv2.imwrite(image_path, result["image"])

        mask_filename = "mask_with_max_angle.png"
        mask_path = os.path.join(masks_output_dir, mask_filename)
        cv2.imwrite(mask_path, result["mask"])

        mask_with_lines_filename = "mask_with_lines.png"
        mask_with_lines_path = os.path.join(masks_with_lines_output_dir, mask_with_lines_filename)
        cv2.imwrite(mask_with_lines_path, result["angle_lines_mask"])

    except ValueError as e:
        print(f"Błąd: {e}")

images_directory = './app/angle_utils_5class/images'
masks_directory = './app/angle_utils_5class/masks'
output_directory = './app/angle_utils_5class/output'

process_images_and_masks(images_directory, masks_directory, output_directory)