# Anzahl der Klasseninstanzen berechnen:
import os
import matplotlib.pyplot as plt
from collections import defaultdict

# Pfad zum Ordner mit den .txt-Dateien im YOLO-Format
folder_path = '_MOMMERT/sorted_images_YOLO_formatted_no_conifers'

# Klassennamen und deren IDs
class_names = {
    0: 'Großer Laubbaum',
    1: 'Kleiner Laubbaum'
}

# Zähler für die Klassen
class_counts = defaultdict(int)

# Durch alle .txt-Dateien im Ordner gehen
for filename in os.listdir(folder_path):
    if filename.endswith('.txt'):
        with open(os.path.join(folder_path, filename), 'r') as file:
            for line in file:
                class_id = int(line.split()[0])  # Erste Zahl in jeder Zeile gibt die Klasse an
                class_counts[class_id] += 1

# Konvertiere den Zähler in eine Liste für die Darstellung
labels = [class_names[i] for i in range(len(class_names))]
counts = [class_counts[i] for i in range(len(class_names))]

# Balkendiagramm erstellen
plt.figure(figsize=(10, 6))
plt.bar(labels, counts, color='skyblue')
plt.xlabel('Klassen')
plt.ylabel('Anzahl der Instanzen')
plt.title('Verteilung der Klasseninstanzen im YOLO-Dataset')
plt.xticks(rotation=45)
plt.tight_layout()

# Diagramm anzeigen
plt.show()





# Möglichkeit 1: Verteilung der Bilder auf 'test, 'train' & 'val' im Verhältnis 10:70:20.
import random
from collections import Counter
import matplotlib.pyplot as plt
import shutil

# Pfade
source_folder = '_MOMMERT/sorted_images_YOLO_formatted'
base_target_folder = '_MOMMERT/sorted_images_YOLO_formatted'
train_folder = os.path.join(base_target_folder, 'train')
val_folder = os.path.join(base_target_folder, 'val')
test_folder = os.path.join(base_target_folder, 'test')

# Erstellen der Zielverzeichnisse, falls nicht vorhanden
os.makedirs(train_folder, exist_ok=True)
os.makedirs(val_folder, exist_ok=True)
os.makedirs(test_folder, exist_ok=True)

# Klassennamen und deren IDs
class_names = {
    0: 'Großer Laubbaum',
    1: 'Kleiner Laubbaum'
}

# Sammeln aller .txt-Dateien
all_txt_files = [f for f in os.listdir(source_folder) if f.endswith('.txt')]
print(f"Total .txt files found: {len(all_txt_files)}")

# Mischen der Liste
random.shuffle(all_txt_files)

# Aufteilen der Dateien in Train, Val, Test
train_ratio = 0.7
val_ratio = 0.20
test_ratio = 0.10

num_total = len(all_txt_files)
num_train = int(num_total * train_ratio)
num_val = int(num_total * val_ratio)
num_test = num_total - num_train - num_val  # Sicherstellen, dass die Summe stimmt

train_files = set(all_txt_files[:num_train])
val_files = set(all_txt_files[num_train:num_train + num_val])
test_files = set(all_txt_files[num_train + num_val:])

print(f"Train files: {len(train_files)}")
print(f"Validation files: {len(val_files)}")
print(f"Test files: {len(test_files)}")

# Funktion zum Verschieben der Dateien
def move_files(file_list, target_folder):
    for txt_file in file_list:
        txt_path = os.path.join(source_folder, txt_file)

        # Verschieben der .txt-Datei
        if os.path.exists(txt_path):
            shutil.move(txt_path, os.path.join(target_folder, txt_file))
        else:
            print(f"Textdatei {txt_file} nicht gefunden.")

        # Verschieben der zugehörigen .png-Datei
        img_file = txt_file.replace('.txt', '.png')
        img_path = os.path.join(source_folder, img_file)

        if os.path.exists(img_path):
            shutil.move(img_path, os.path.join(target_folder, img_file))
        else:
            print(f"Bilddatei {img_file} nicht gefunden.")

# Verschieben der Dateien
move_files(train_files, train_folder)
move_files(val_files, val_folder)
move_files(test_files, test_folder)

print("Dateien erfolgreich aufgeteilt und verschoben!")

# Funktion zum Erstellen eines Balkendiagramms und Ausgeben der absoluten und prozentualen Verteilung
def plot_class_distribution(folder, folder_name):
    class_count = Counter()

    for filename in os.listdir(folder):
        if filename.endswith('.txt'):
            with open(os.path.join(folder, filename), 'r') as file:
                for line in file:
                    class_id = int(line.split()[0])
                    class_count[class_id] += 1

    # Sicherstellen, dass alle Klassen vorhanden sind
    counts = [class_count.get(cid, 0) for cid in sorted(class_names.keys())]
    labels = [class_names[cid] for cid in sorted(class_names.keys())]

    # Prozentuale Verteilung berechnen
    total_instances = sum(counts)
    percentages = [(count / total_instances) * 100 if total_instances > 0 else 0 for count in counts]

    # Ausgabe der absoluten und prozentualen Verteilung
    print(f"\nVerteilung in {folder_name}:")
    for label, count, percentage in zip(labels, counts, percentages):
        print(f"{label}: {count} Instanzen ({percentage:.2f}%)")

    # Balkendiagramm erstellen
    plt.figure(figsize=(10,6))
    plt.bar(labels, counts, color='skyblue')
    plt.title(f'Klassenverteilung in {folder_name}')
    plt.xlabel('Klassen')
    plt.ylabel('Anzahl der Instanzen')
    plt.xticks(rotation=45, ha='right')
    plt.tight_layout()
    plt.show()

# Plotten der Klassenverteilung und Ausgabe der Werte
plot_class_distribution(train_folder, 'Train')
plot_class_distribution(val_folder, 'Validation')
plot_class_distribution(test_folder, 'Test')







# Möglichkeit 2: Stratifizierte Aufteilung
import random
from collections import Counter
import matplotlib.pyplot as plt
from iterstrat.ml_stratifiers import MultilabelStratifiedShuffleSplit
import numpy as np
import shutil

# Pfade
source_folder = '_MOMMERT/sorted_images_YOLO_formatted'
base_target_folder = '_MOMMERT/sorted_images_YOLO_formatted'
train_folder = os.path.join(base_target_folder, 'data/train')
val_folder = os.path.join(base_target_folder, 'data/val')
test_folder = os.path.join(base_target_folder, 'data/test')

# Erstellen der Zielverzeichnisse, falls nicht vorhanden
os.makedirs(train_folder, exist_ok=True)
os.makedirs(val_folder, exist_ok=True)
os.makedirs(test_folder, exist_ok=True)

# Klassennamen und deren IDs
class_names = {
    0: 'Großer Laubbaum',
    1: 'Kleiner Laubbaum'
}

# Sammeln aller .txt-Dateien
all_txt_files = [f for f in os.listdir(source_folder) if f.endswith('.txt')]
print(f"Total .txt files found: {len(all_txt_files)}")

# Sammeln aller Dateien und deren Klassen
file_classes = {}
for txt_file in all_txt_files:
    with open(os.path.join(source_folder, txt_file), 'r') as file:
        classes = set(int(line.split()[0]) for line in file)
        # Erstelle eine Binärliste, die angibt, ob jede Klasse in der Datei vorhanden ist oder nicht
        class_vector = [1 if i in classes else 0 for i in range(len(class_names))]
        file_classes[txt_file] = class_vector

# Verwenden von MultilabelStratifiedShuffleSplit für stratifizierte Aufteilung
msss = MultilabelStratifiedShuffleSplit(n_splits=1, test_size=0.3, random_state=42)
X = list(file_classes.keys())
y = np.array(list(file_classes.values()))

for train_val_index, test_index in msss.split(X, y):
    X_train_val = [X[i] for i in train_val_index]
    X_test = [X[i] for i in test_index]

msss_val = MultilabelStratifiedShuffleSplit(n_splits=1, test_size=0.1765, random_state=42)  # 0.1765 ≈ 0.15 / 0.85
y_train_val = [file_classes[f] for f in X_train_val]

for train_index, val_index in msss_val.split(X_train_val, y_train_val):
    X_train = [X_train_val[i] for i in train_index]
    X_val = [X_train_val[i] for i in val_index]

train_files = set(X_train)
val_files = set(X_val)
test_files = set(X_test)

print(f"Train files: {len(train_files)}")
print(f"Validation files: {len(val_files)}")
print(f"Test files: {len(test_files)}")

# Funktion zum Verschieben der Dateien
def move_files(file_list, target_folder):
    for txt_file in file_list:
        txt_path = os.path.join(source_folder, txt_file)

        # Verschieben der .txt-Datei
        if os.path.exists(txt_path):
            shutil.move(txt_path, os.path.join(target_folder, txt_file))
        else:
            print(f"Textdatei {txt_file} nicht gefunden.")

        # Verschieben der zugehörigen .png-Datei
        img_file = txt_file.replace('.txt', '.png')
        img_path = os.path.join(source_folder, img_file)

        if os.path.exists(img_path):
            shutil.move(img_path, os.path.join(target_folder, img_file))
        else:
            print(f"Bilddatei {img_file} nicht gefunden.")

# Verschieben der Dateien
move_files(train_files, train_folder)
move_files(val_files, val_folder)
move_files(test_files, test_folder)

print("Dateien erfolgreich aufgeteilt und verschoben!")

# Funktion zum Erstellen eines Balkendiagramms und Ausgeben der absoluten und prozentualen Verteilung
def plot_class_distribution(folder, folder_name):
    class_count = Counter()

    for filename in os.listdir(folder):
        if filename.endswith('.txt'):
            with open(os.path.join(folder, filename), 'r') as file:
                for line in file:
                    class_id = int(line.split()[0])
                    class_count[class_id] += 1

    # Sicherstellen, dass alle Klassen vorhanden sind
    counts = [class_count.get(cid, 0) for cid in sorted(class_names.keys())]
    labels = [class_names[cid] for cid in sorted(class_names.keys())]

    # Prozentuale Verteilung berechnen
    total_instances = sum(counts)
    percentages = [(count / total_instances) * 100 if total_instances > 0 else 0 for count in counts]

    # Ausgabe der absoluten und prozentualen Verteilung
    print(f"\nVerteilung in {folder_name}:")
    for label, count, percentage in zip(labels, counts, percentages):
        print(f"{label}: {count} Instanzen ({percentage:.2f}%)")

    # Balkendiagramm erstellen
    plt.figure(figsize=(10,6))
    plt.bar(labels, counts, color='skyblue')
    plt.title(f'Klassenverteilung in {folder_name}')
    plt.xlabel('Klassen')
    plt.ylabel('Anzahl der Instanzen')
    plt.xticks(rotation=45, ha='right')
    plt.tight_layout()
    plt.show()

# Plotten der Klassenverteilung und Ausgabe der Werte
plot_class_distribution(train_folder, 'Train')
plot_class_distribution(val_folder, 'Validation')
plot_class_distribution(test_folder, 'Test')
