import os
import cv2
import numpy as np
from typing import Tuple

# Definicja kolorów i etykiet
CLASS_COLORS = {
    0: {'color_rgb': (0, 0, 0), 'label': 'background'},
    1: {'color_rgb': (255, 0, 0), 'label': 'chondroosseous border'},
    2: {'color_rgb': (0, 255, 0), 'label': 'femoral head'},
    3: {'color_rgb': (0, 0, 255), 'label': 'labrum'},
    4: {'color_rgb': (255, 255, 0), 'label': 'cartilagineous roof'},
    5: {'color_rgb': (0, 255, 255), 'label': 'bony roof'},
    6: {'color_rgb': (159, 2, 250), 'label': 'bony rim'},
    7: {'color_rgb': (255, 132, 0), 'label': 'lower limb'},
    8: {'color_rgb': (255, 0, 255), 'label': 'baseline'},
    9: {'color_rgb': (66, 135, 245), 'label': 'lower limb template'},
    10: {'color_rgb': (255, 69, 0), 'label': 'lower limb - alg. v3'}
}


def visualize_masks(input_dir, output_dir):
    # Tworzenie katalogu wyjściowego, jeśli nie istnieje
    os.makedirs(output_dir, exist_ok=True)

    # Iteracja po plikach PNG w katalogu wejściowym
    for filename in os.listdir(input_dir):
        if filename.endswith(".png"):
            # Wczytywanie maski
            mask_path = os.path.join(input_dir, filename)
            mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)

            # Tworzenie pustego obrazu RGB o tych samych wymiarach co maska
            color_image = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)

            # Zamiana etykiet na kolory RGB zgodnie z CLASS_COLORS
            for class_id, color_info in CLASS_COLORS.items():
                color_rgb = color_info['color_rgb']
                color_image[mask == class_id] = color_rgb

            # Znajdowanie punktów baseline
            right_point, left_point = find_baseline(mask)
            
            # Rysowanie baseline na obrazie, jeżeli punkty istnieją
            if right_point[0] is not None and left_point[0] is not None:
                cv2.line(color_image, right_point, left_point, (255, 255, 255), 1)

            # Zapisanie obrazu kolorowego do katalogu wyjściowego
            output_path = os.path.join(output_dir, filename)
            cv2.imwrite(output_path, color_image)
            print(f"Maska {filename} została zapisana w katalogu wyjściowym.")

# Przykładowe użycie
input_dir = "./Angles/dane/ta_jedna"
output_dir = "./Angles/dane/wizul"
visualize_masks(input_dir, output_dir)