import bz2
import espeak_ng
import json
import sys

# Zbiór samogłosek IPA (bez oznaczeń akcentu/nosowości)
SAMOGLOSKI_BAZOWE = set("aeiouyɛɔɨʉæøɪʊəɐɑɒœɜɞʌ")
# Znaki modyfikujące, które nie zmieniają statusu samogłoski
MODYFIKATORY = set("ˈˌːˑ̩̯̝̞̘̙̃̈̽͡")


def czy_samogloska(fonem):
    """Sprawdza, czy fonem jest samogłoską (zawiera przynajmniej jeden znak samogłoskowy)."""
    for ch in fonem:
        if ch in SAMOGLOSKI_BAZOWE:
            return True
    return False


def normalizuj(wymowa):
    """Normalizacja fonetyczna wymowy zgodnie z regułami rymów dokładnych.
    
    Reguły:
    1. ˈɨ -> ˈi  (y akcentowane = i akcentowane)
    2. Końcowe ɔ̃ -> ɔ
    3. Akcentowane ɔɲ / ɛɲ przed spółgłoską -> ɔn / ɛn (i utożsamione z ɔ̃ / ɛ̃)
    """
    wynik = list(wymowa)
    
    # Reguła 1: ˈɨ -> ˈi
    for i, f in enumerate(wynik):
        if f == "ˈɨ":
            wynik[i] = "ˈi"
    
    # Reguła 2: końcowe ɔ̃ -> ɔ
    if wynik and wynik[-1] == "ɔ̃":
        wynik[-1] = "ɔ"
    
    # Reguła 3: akcentowane oń/eń przed spółgłoską
    # Szukam sekwencji: samogłoska akcentowana (ˈɔ lub ˈɛ), potem ɲ, potem spółgłoska
    i = 0
    while i < len(wynik) - 2:
        if wynik[i] in ("ˈɔ", "ˌɔ", "ˈɛ", "ˌɛ"):
            if wynik[i + 1] == "ɲ" and i + 2 < len(wynik) and not czy_samogloska(wynik[i + 2]):
                wynik[i + 1] = "n"
        i += 1
    
    return wynik


def znajdz_koncowke_rymowa(wymowa):
    """Znajduje końcówkę rymową: od ostatniej samogłoski z akcentem głównym ˈ do końca."""
    ostatni_akcent = -1
    for i, fonem in enumerate(wymowa):
        if (fonem.startswith("ˈ") or fonem.startswith("ˌ")) and czy_samogloska(fonem):
            ostatni_akcent = i
    
    if ostatni_akcent == -1:
        return None
    
    return wymowa[ostatni_akcent:]


def wyciagnij_samogloski(koncowka):
    """Wyciąga same samogłoski z końcówki (rdzeń samogłoskowy).
    Usuwa oznaczenia akcentu z fonemów samogłoskowych."""
    rdzen = []
    for fonem in koncowka:
        if czy_samogloska(fonem):
            # Usuń oznaczenia akcentu
            czysty = fonem.replace("ˈ", "").replace("ˌ", "")
            rdzen.append(czysty)
    return rdzen


def licz_samogloski(koncowka):
    """Liczy samogłoski w końcówce."""
    return sum(1 for f in koncowka if czy_samogloska(f))


def klucz_a_tergo(seq):
    """Klucz sortowania a tergo: odwrócona sekwencja."""
    if isinstance(seq, list):
        return list(reversed(seq))
    return seq[::-1]


def main():
    espeak_ng.initialize()
    espeak_ng.set_voice_by_properties(name="pl")

    # Etap 1: Zbierz czwórki (słowo, wymowa, końcówka, rdzeń samogłoskowy)
    czworki = []

    for idx, line in enumerate(bz2.BZ2File("slowa.txt.bz2", "r")):
        slowo = line.decode("utf-8").rstrip()
        wymowa_raw = [
            fonem
            for fonem in espeak_ng.text_to_phonemes(
                slowo, phonememode=(0x02 | (ord("_") << 8))
            ).split("_")
            if fonem != ""
        ]

        wymowa = normalizuj(wymowa_raw)
        koncowka = znajdz_koncowke_rymowa(wymowa)

        if koncowka is None:
            continue

        # Odrzuć rymy męskie (jedna samogłoska lub mniej w końcówce)
        if licz_samogloski(koncowka) <= 1:
            continue

        rdzen = wyciagnij_samogloski(koncowka)

        czworki.append({
            "slowo": slowo,
            "wymowa": wymowa,
            "koncowka": koncowka,
            "rdzen": rdzen,
        })

    # Etap 2: Grupowanie

    # Poziom 1: klucz = rdzeń samogłoskowy (jako tuple dla hashowalności)
    from collections import defaultdict

    grupy_rdzen = defaultdict(list)
    for cz in czworki:
        klucz = tuple(cz["rdzen"])
        grupy_rdzen[klucz].append(cz)

    # Sortuj rdzenie a tergo
    #rdzenie_sorted = sorted(grupy_rdzen.keys(), key=lambda r: list(reversed(r)))
    # Sortuj rdzenie a tergo
    rdzenie_sorted = sorted(grupy_rdzen.keys())

    # Buduj wynikową strukturę
    wynik = {}

    for rdzen in rdzenie_sorted:
        rdzen_str = " ".join(rdzen)
        trojki = grupy_rdzen[rdzen]

        # Poziom 2: klucz = końcówka rymowa
        grupy_koncowka = defaultdict(list)
        for tr in trojki:
            klucz_konc = tuple(tr["koncowka"])
            grupy_koncowka[klucz_konc].append(tr)

        # Sortuj końcówki a tergo
        #koncowki_sorted = sorted(grupy_koncowka.keys(), key=lambda k: list(reversed(k)))
        # Sortuj końcówki
        koncowki_sorted = sorted(grupy_koncowka.keys())

        slownik_koncowek = {}
        for konc in koncowki_sorted:
            konc_str = " ".join(konc)
            pary = grupy_koncowka[konc]

            # Poziom 3: pary (słowo, wymowa) sortowane a tergo po wymowie
            pary_sorted = sorted(pary, key=lambda p: list(reversed(p["wymowa"])))

            slownik_koncowek[konc_str] = [
                {"slowo": p["slowo"], "wymowa": p["wymowa"]}
                for p in pary_sorted
            ]

        wynik[rdzen_str] = slownik_koncowek

    # Zapisz
    with open("rymy.json", "w", encoding="utf-8") as f:
        json.dump(wynik, f, ensure_ascii=False, indent=None)

    print(f"Zapisano {len(wynik)} rdzeni samogłoskowych do rymy.json", file=sys.stderr)


if __name__ == "__main__":
    main()