import os
from phonemizer import backend

# Konfiguracja plików
INPUT_FILE = 'train_3_filtered_dur.txt'        # Plik wejściowy
OUTPUT_FILE = 'train_3_filtered_dur_fixed.txt' # Plik wynikowy

# Zestaw znaków, których obecność w fonemach powoduje usunięcie linii
FORBIDDEN_CHARS = {'{', '}', '(', ')'}

def initialize_phonemizers():
    """
    Inicjalizuje słownik backendów espeak dla wymaganych języków.
    Zwraca słownik { 'language_id': backend_instance }
    """
    phonemizers = {}
    print("Inicjalizacja phonemizerów (może to chwilę potrwać)...")
    
    try:
        # ID 0: Polski
        phonemizers['0'] = backend.EspeakBackend(
            language='pl', 
            preserve_punctuation=True, 
            with_stress=True
        )
        
        # ID 1: Angielski Brytyjski
        phonemizers['1'] = backend.EspeakBackend(
            language='en-gb', 
            preserve_punctuation=True, 
            with_stress=True
        )
        
        # ID 2: Hindi
        phonemizers['2'] = backend.EspeakBackend(
            language='hi', 
            preserve_punctuation=True, 
            with_stress=True
        )
        
        print("Phonemizery zainicjalizowane pomyślnie.")
        return phonemizers

    except Exception as e:
        print(f"Krytyczny błąd inicjalizacji phonemizerów: {e}")
        print("Upewnij się, że masz zainstalowany 'espeak-ng'.")
        return None

def process_dataset():
    # 1. Przygotuj phonemizery
    phonemizers = initialize_phonemizers()
    if not phonemizers:
        return

    if not os.path.exists(INPUT_FILE):
        print(f"Nie znaleziono pliku: {INPUT_FILE}")
        return

    print(f"Rozpoczynam przetwarzanie pliku {INPUT_FILE}...")
    
    processed_count = 0
    skipped_count = 0
    
    with open(INPUT_FILE, 'r', encoding='utf-8') as f_in, \
         open(OUTPUT_FILE, 'w', encoding='utf-8') as f_out:
        
        for line_num, line in enumerate(f_in, 1):
            line = line.strip()
            if not line:
                continue
                
            # 2. Split z limitem 4
            parts = line.split('|', 4)
            
            # Sprawdzenie poprawności struktury (musi być 5 kolumn)
            if len(parts) < 5:
                print(f"Ostrzeżenie: Linia {line_num} ma nieprawidłowy format (pominięto).")
                continue

            # Wyciągamy potrzebne dane
            # parts[1] = phonemes (to podmieniamy i sprawdzamy)
            lang_id = parts[3]
            raw_text = parts[4]
            
            # Pobieramy odpowiedni phonemizer
            current_backend = phonemizers.get(lang_id)
            
            if current_backend:
                try:
                    # Generowanie fonemów
                    phonemes_list = current_backend.phonemize([raw_text], strip=True)
                    if phonemes_list:
                        parts[1] = phonemes_list[0]  # Aktualizacja fonemów
                except Exception as e:
                    print(f"Błąd phonemizera w linii {line_num} (ID jęz: {lang_id}): {e}")
            
            # --- NOWA LOGIKA FILTROWANIA ---
            current_phonemes = parts[1]
            
            # Sprawdzamy czy w fonemach występuje któryś z zakazanych znaków
            if any(char in current_phonemes for char in FORBIDDEN_CHARS):
                skipped_count += 1
                # Opcjonalnie: odkomentuj linię poniżej, jeśli chcesz widzieć, co jest usuwane
                # print(f"Pominięto linię {line_num}: znaleziono zakazany znak w fonemach: {current_phonemes}")
                continue

            # 3. Zapis do nowego pliku (tylko jeśli walidacja przeszła pomyślnie)
            processed_count += 1
            new_line = "|".join(parts)
            f_out.write(new_line + '\n')

            # Logowanie postępu co 100 linii
            if processed_count > 0 and processed_count % 100 == 0:
                print(f"Zapisano {processed_count} poprawnych linii...")

    print("-" * 30)
    print(f"Zakończono!")
    print(f"Poprawnie przetworzono i zapisano: {processed_count}")
    print(f"Usunięto (zawierały '{{', '}}', '(', ')'): {skipped_count}")
    print(f"Wynik zapisano w: {OUTPUT_FILE}")

if __name__ == "__main__":
    process_dataset()