import os
import cv2
import pandas as pd
from fer import FER
from concurrent.futures import ThreadPoolExecutor, as_completed
import csv
from collections import Counter

def remove_file_if_exists(file_path):
    """Löscht eine Datei, wenn sie existiert."""
    if os.path.exists(file_path):
        os.remove(file_path)

def analyze_image(png_file, detector, input_folder, nd_log_file):
    """
    Analysiert ein einzelnes Bild und gibt ein Dictionary mit Informationen zur
    dominanten Emotion und Emotionen zurück.
    """
    img_path = os.path.join(input_folder, png_file)
    img = cv2.imread(img_path)

    if img is None:
        print(f"Could not load image {png_file}")
        return {
            'image_name': png_file,
            'dominant_emotion': 'neutral',
            'emotion': {'neutral': 100}
        }

    try:
        results = detector.detect_emotions(img)
        if results:
            emotions = results[0]['emotions']
            dominant_emotion = max(emotions, key=emotions.get)

            # Falls die dominante Emotion 'neutral' ist und mindestens noch eine weitere Emotion > 0 vorliegt,
            # wähle die zweithöchste Emotion.
            if dominant_emotion == 'neutral' and len(emotions) > 1:
                sorted_emotions = sorted(emotions.items(), key=lambda x: x[1], reverse=True)
                if sorted_emotions[1][1] > 0:
                    dominant_emotion = sorted_emotions[1][0]

            result = {
                'image_name': png_file,
                'dominant_emotion': dominant_emotion,
                'emotion': emotions
            }
        else:
            raise ValueError("No face detected")

    except Exception as e:
        with open(nd_log_file, 'a') as log_file:
            log_file.write(f"{png_file}: {str(e)}\n")
        result = {
            'image_name': png_file,
            'dominant_emotion': 'neutral',
            'emotion': {'neutral': 100}
        }

    return result

def count_dominant_emotions(csv_file, output_txt):
    """
    Zählt die dominanten Emotionen aus einer CSV-Datei und speichert die Ergebnisse in einer Textdatei.
    """
    emotion_counter = Counter()

    try:
        # CSV-Datei lesen und Emotionen zählen
        with open(csv_file, mode='r', encoding='utf-8') as file:
            reader = csv.DictReader(file)
            for row in reader:
                dominant_emotion = row.get('dominant_emotion', '').strip()
                if dominant_emotion:
                    emotion_counter[dominant_emotion] += 1

        # Ergebnisse in Textdatei speichern
        with open(output_txt, mode='w', encoding='utf-8') as output_file:
            output_file.write("Dominante Emotionen:\n")
            for emotion, count in emotion_counter.items():
                output_file.write(f"{emotion}: {count}\n")

        print(f"Die Ergebnisse wurden in '{output_txt}' gespeichert.")

    except FileNotFoundError:
        print(f"Die Datei '{csv_file}' wurde nicht gefunden.")
    except Exception as e:
        print(f"Fehler beim Verarbeiten der Datei: {e}")

if __name__ == "__main__":
    # Ordner und Dateinamen definieren
    input_folder = 'output_crop'
    csv_file = 'results-internalset-fer.csv'
    nd_log_file = 'nd-log-internalset-fer.txt'
    output_txt = 'emotion_counts.txt'

    # Vorhandene Dateien löschen, um sie neu zu erstellen
    remove_file_if_exists(csv_file)
    remove_file_if_exists(nd_log_file)

    # PNG-Dateien im Ordner sammeln (ohne versteckte Dateien wie ._*)
    absolute_path = os.path.abspath(input_folder)
    png_files = sorted([
        f for f in os.listdir(input_folder)
        if f.lower().endswith('.png') and not f.startswith('._')
    ])
    
    total_files = len(png_files)
    print(f"{total_files} PNG-Datei(en) gefunden.")

    # FER-Detector initialisieren
    detector = FER(mtcnn=True)

    results_list = []

    # Bilder parallel analysieren
    with ThreadPoolExecutor(max_workers=8) as executor:
        future_to_file = {
            executor.submit(analyze_image, png_file, detector, input_folder, nd_log_file): png_file
            for png_file in png_files
        }
        for index, future in enumerate(as_completed(future_to_file), start=1):
            result = future.result()
            results_list.append(result)
            progress = int((index / total_files) * 100)
            print(f'Processing {index}/{total_files} [{progress}%]', end='\r')

    # CSV-Datei schreiben, wenn Ergebnisse vorhanden sind
    if results_list:
        result_df = pd.DataFrame(results_list)
        result_df.to_csv(csv_file, index=False)
        print("\nCSV-Datei wurde erstellt und gespeichert.")
    else:
        print("\nKeine Ergebnisse vorhanden. CSV-Datei wird nicht erstellt.")

    # Dominante Emotionen aus der CSV zählen und in .txt speichern
    count_dominant_emotions(csv_file, output_txt)