package zaj4;

public class AvlTree {
    public Wezel korzen;

    int liczWysokosc(Wezel wezel) {
        if (wezel == null)
            return 0;
        return wezel.wysokosc;
    }

    Wezel rotacjaPrawo(Wezel wezelA) {
        Wezel wezelB = wezelA.lewy;
        Wezel choinkaII = wezelB.prawy;

        wezelB.prawy = wezelA;
        wezelA.lewy = choinkaII;

        aktualizujWysokosci(wezelA, wezelB);

        return wezelB;
    }

    Wezel rotacjaLewo(Wezel wezelA) {
        Wezel wezelB = wezelA.prawy;
        Wezel choinkaII = wezelB.lewy;

        wezelB.lewy = wezelA;
        wezelA.prawy = choinkaII;

        aktualizujWysokosci(wezelA, wezelB);

        return wezelB;
    }

    private void aktualizujWysokosci(Wezel wezelA, Wezel wezelB) {
        wezelA.wysokosc = Math.max(liczWysokosc(wezelA.lewy), liczWysokosc(wezelA.prawy)) + 1;
        wezelB.wysokosc = Math.max(liczWysokosc(wezelB.lewy), liczWysokosc(wezelB.prawy)) + 1;
    }

    int liczBalans(Wezel wezel) {
        if (wezel == null) return 0;
        return liczWysokosc(wezel.lewy) - liczWysokosc(wezel.prawy);
    }

    Wezel dodaj(Wezel wezel, String slowo) {
        // jesli nie mamy zadnego elementu
        if (wezel == null)
            return (new Wezel(slowo));
        // idziemy w lewo
        if (slowo.compareTo(wezel.slowo) < 0)
            wezel.lewy = dodaj(wezel.lewy, slowo);
            // idziemy w prawo
        else if (slowo.compareTo(wezel.slowo) > 0)
            wezel.prawy = dodaj(wezel.prawy, slowo);
        else
            // gdy takie same to wraca
            return wezel;

        wezel.wysokosc = Math.max(liczWysokosc(wezel.lewy), liczWysokosc(wezel.prawy)) + 1;

        wezel = balansujDrzewo(wezel);

        return wezel;
    }

    private Wezel balansujDrzewo(Wezel wezel) {
        int aktualnyBalans = liczBalans(wezel);

        // rotuj w prawo
        if (aktualnyBalans == 2 && liczBalans(wezel.lewy) == 1)
            return rotacjaPrawo(wezel);

        // rotuj w lewo
        if (aktualnyBalans == -2 && liczBalans(wezel.prawy) == -1)
            return rotacjaLewo(wezel);

        // rotuj w lewo prawo
        if (aktualnyBalans == 2 && liczBalans(wezel.lewy) == -1) {
            // rotacja w lewo, wezełem B staje sie wezel C
            wezel.lewy = rotacjaLewo(wezel.lewy);
            // wezel C staje sie korzeniem
            return rotacjaPrawo(wezel);
        }

        // rotuj w prawo lewo
        if (aktualnyBalans == -2 && liczBalans(wezel.prawy) == 1) {
            // wezel C staje sie wezlem B
            wezel.prawy = rotacjaPrawo(wezel.prawy);
            // korzeniem staje sie wezel C
            return rotacjaLewo(wezel);
        }
        return wezel;
    }

    public boolean czyZawiera(Wezel wezel, String word) {
        if (wezel == null) return false;
        int wynikPorownania = wezel.slowo.compareTo(word);
        if (wynikPorownania == 0) return true;
        else if (wynikPorownania > 0) return czyZawiera(wezel.lewy, word);
        else return czyZawiera(wezel.prawy, word);
    }

    public Wezel usunWezel(Wezel wezel, String slowo) {
        // jeśli nie ma węzłów
        if (wezel == null) return null;
        // compareTo zwraca ujemną wartość gdy slowo po lewej jest pierwsze

        // slowo jest mniejsze
        if (slowo.compareTo(wezel.slowo) < 0)
            // idziemy w lewo
            wezel.lewy = usunWezel(wezel.lewy, slowo);
        // slowo jest wieksze
        if (slowo.compareTo(wezel.slowo) > 0)
            // idziemy w prawo
            wezel.prawy = usunWezel(wezel.prawy, slowo);

        // jesli trafilismy
        if (slowo.compareTo(wezel.slowo) == 0)
            // jesli nie ma dzieci albo jedno dziecko
            if (maMaksymalnieJednoDziecko(wezel)) {
                // jesli lewe dziecko jest nullem
                if (maLeweDziecko(wezel)) wezel = wezel.lewy;
                else wezel = wezel.prawy;
            } else {
                // Szukamy następnika
                Wezel temp = wybierzNajmniejszy(wezel.prawy);
                // wskaznik zostawiamy, wartosc kopiujemy
                wezel.slowo = temp.slowo;
                // usuwamy wskaznik na nast
                wezel.prawy = usunWezel(wezel.prawy, temp.slowo);
            }

        // jeśli jest liściem
        if (wezel == null) return null;

        wezel.wysokosc = Math.max(liczWysokosc(wezel.lewy), liczWysokosc(wezel.prawy)) + 1;

        wezel = balansujDrzewo(wezel);

        return wezel;
    }

    private boolean maMaksymalnieJednoDziecko(Wezel wezel) {
        if (wezel == null) return true;
        return wezel.lewy == null || wezel.prawy == null;
    }

    private boolean maLeweDziecko(Wezel wezel) {
        if (wezel == null) return false;
        return wezel.lewy != null;
    }

    public Wezel wybierzNajmniejszy(Wezel wezel) {
        Wezel current = wezel;
        while (current.lewy != null)
            current = current.lewy;
        return current;
    }

    public int liczZaczynajaceSie(Wezel wezel, String poczatek) {
        if (wezel == null) return 0;
        if (wezel.slowo.startsWith(poczatek))
            return 1
                    + liczZaczynajaceSie(wezel.lewy, poczatek)
                    + liczZaczynajaceSie(wezel.prawy, poczatek);
        return liczZaczynajaceSie(wezel.lewy, poczatek)
                + liczZaczynajaceSie(wezel.prawy, poczatek);
    }
}