import numpy as np
import math
from typing import Tuple
import pandas as pd
import os
from PIL import Image
import shutil

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 detect_bump(mask: np.ndarray, skip: int = 15, threshold: int = 5) -> int:
    """
    Wykrywa garb na bony roof. Zwraca 1, jeśli garb wykryto, w przeciwnym wypadku 0.
    """
    label_bony_roof = 5
    binary_mask = (mask == label_bony_roof).astype(np.uint8)

    height, width = binary_mask.shape
    upper_contour = []

    for x in range(width):
        column = binary_mask[:, x]
        if np.any(column):
            y = np.where(column)[0][0] 
            upper_contour.append((x, y))
        else:
            upper_contour.append((x, np.nan))

    print(f"Upper contour (pierwsze 10 punktów): {upper_contour[:10]}")

    # Odfiltrowanie wartości NaN (kolumn bez bony roof)
    valid_contour = np.array([y for x, y in upper_contour if not np.isnan(y)])

    if len(valid_contour) == 0:
        print("Brak pikseli bony roof")
        return 0
    
    min_y = np.min(valid_contour)

    print(f"Najwyższy punkt (min_y): {min_y}")

    distances = valid_contour - min_y
    
    print(f"Odległości od najwyższego punktu (pierwsze 10): {distances[:10]}")

    differences = pd.Series(distances).diff(periods=skip).diff().fillna(0).abs()

    print(f"Pierwsze 3 różnice z różnic: {differences[:3].values}")
    max_diff = differences.max()
    print(f"Max różnica: {max_diff}")

    return 1 if max_diff >= threshold else 0

def process_masks_and_save_with_visualization(mask_dir, bump_dir, no_bump_dir, skip=7, threshold=6):
    """
    Przetwarza maski, wykrywa obecność garbu, zapisuje je do odpowiednich folderów 
    i tworzy ich wizualizacje.
    """
    if not os.path.exists(bump_dir):
        os.makedirs(bump_dir)
    if not os.path.exists(no_bump_dir):
        os.makedirs(no_bump_dir)

    for filename in os.listdir(mask_dir):
        file_path = os.path.join(mask_dir, filename)
        
        if os.path.isfile(file_path) and filename.endswith('.png'):
            mask = Image.open(file_path).convert('L')  # Konwertuj do odcieni szarości
            mask_np = np.array(mask)
            
            # Wykrycie garbu
            bump_present = detect_bump(mask_np, skip, threshold)
            
            # Obliczanie max_diff
            label_bony_roof = 5
            binary_mask = (mask_np == label_bony_roof).astype(np.uint8)
            height, width = binary_mask.shape
            upper_contour = []

            for x in range(width):
                column = binary_mask[:, x]
                if np.any(column):
                    y = np.where(column)[0][0]  # Najwyższy piksel w danej kolumnie
                    upper_contour.append(y)
                else:
                    upper_contour.append(height)

            upper_contour = np.array(upper_contour)
            min_y = np.min(upper_contour)
            distances = min_y - upper_contour
            differences = pd.Series(distances).diff(periods=skip).fillna(0).abs()
            max_diff = differences.max()
            
            # Wizualizacja maski
            visualized_image = np.zeros((height, width, 3), dtype=np.uint8)
            
            for label, color_info in CLASS_COLORS.items():
                color = color_info['color_rgb']
                visualized_image[mask_np == label] = color

            # Zapisanie do odpowiedniego folderu
            if bump_present:
                save_path = os.path.join(bump_dir, filename)
            else:
                save_path = os.path.join(no_bump_dir, filename)

            Image.fromarray(visualized_image).save(save_path)
            print(f'Zapisano zwizualizowaną maskę do: {save_path} - max_diff: {max_diff}')


process_masks_and_save_with_visualization('./Angles/dane/masks', './Angles/dane/garb_obecny','./Angles/dane/garb_nieobecny')