import java.util.ArrayList;

public class RBTree<E extends Comparable<E>> extends BST<E> {
  /** Tworzy domyślne drzewo czerwono-czarne */
  public RBTree() {
  }

  /** Tworzy drzewo czerwono-czarne na podstawie tablicy elementów */
  public RBTree(E[] elements) {
    super(elements);
  }

  @Override /** Przesłanianie metody createNewNode, aby tworzyła węzły RBTreeNode */
  protected RBTreeNode<E> createNewNode(E e) {
    return new RBTreeNode<E>(e);
  }

  @Override /** Przesłanianie metody insert, aby w razie potrzeby
                wyważała drzewo */
  public boolean insert(E e) {
    boolean successful = super.insert(e);
    if (!successful)
      return false; // e już znajduje się w drzewie
    else {
      ensureRBTree(e); 
    }

    return true; // e został wstawiony
  }

  /** Gwarantuje, że drzewo jest drzewem czerwono-czarnym */
  private void ensureRBTree(E e) {
    // Pobieranie ścieżki z korzenia do elementu e
    ArrayList<TreeNode<E>> path = path(e);

    int i = path.size() - 1; // Indeks bieżącego węzła w ścieżce
    
    // u jest ostatnim węzłem w ścieżce i zawiera element e
    RBTreeNode<E> u = (RBTreeNode<E>)(path.get(i));
        
    // v (jeśli istnieje) jest rodzicem węzła u 
    RBTreeNode<E> v = (u == root) ? null : 
      (RBTreeNode<E>)(path.get(i - 1));

    u.setRed(); // Można pokolorować u na czerwono
          
    if (u == root) // Jeśli e został wstawiony w korzeniu, pokoloruj korzeń na czarno
      u.setBlack();
    else if (v.isRed()) 
      fixDoubleRed(u, v, path, i); // Usuwanie problemu dwóch kolejnych czerwonych węzłów w u
  }
  
  /** Usuwa problem dwóch kolejnych czerwonych węzłów w u */
  private void fixDoubleRed(RBTreeNode<E> u, RBTreeNode<E> v, 
      ArrayList<TreeNode<E>> path, int i) {          
    // w jest dziadkiem u
    RBTreeNode<E> w = (RBTreeNode<E>)(path.get(i - 2));
    RBTreeNode<E> parentOfw = (w == root) ? null : 
      (RBTreeNode<E>)path.get(i - 3);

    // Pobieranie węzła x (brata węzła v)
    RBTreeNode<E> x = (w.left == v) ? 
      (RBTreeNode<E>)(w.right) : (RBTreeNode<E>)(w.left);
    
    if (x == null || x.isBlack()) {
      // Scenariusz 1. x (brat węzła v) jest czarny
      if (w.left == v && v.left == u) {
        // Scenariusz 1.1. u < v < w, zmiana struktury i koloru węzłów
        restructureRecolor(u, v, w, w, parentOfw);
    
        w.left = v.right; // v.right to y3 z rysunku 36.6
        v.right = w;
      }
      else if (w.left == v && v.right == u) {
        // Scenariusz 1.2. v < u < w, zmiana struktury i koloru węzłów
        restructureRecolor(v, u, w, w, parentOfw);
        v.right = u.left;
        w.left = u.right;
        u.left = v;
        u.right = w;
      }
      else if (w.right == v && v.right == u) {
        // Scenariusz 1.3. w < v < u, zmiana struktury i koloru węzłów
        restructureRecolor(w, v, u, w, parentOfw);
        w.right = v.left;
        v.left = w;
      }
      else {
        // Scenariusz 1.4. w < u < v, zmiana struktury i koloru węzłów
        restructureRecolor(w, u, v, w, parentOfw);
        w.right = u.left;
        v.left = u.right;
        u.left = w;
        u.right = v;
      }
    }
    else { // Scenariusz 2. x (brat węzła v) jest czerwony
      // Zmiana koloru węzłów
      w.setRed();
      u.setRed();
      ((RBTreeNode<E>)(w.left)).setBlack(); 
      ((RBTreeNode<E>)(w.right)).setBlack();
      
      if (w == root) {
        w.setBlack();     
      }
      else if (((RBTreeNode<E>)parentOfw).isRed()) {  
        // Przejście w górę ścieżki, aby rozwiązać nowy problem dwóch kolejnych czerwonych węzłów
        u = w;
        v = (RBTreeNode<E>)parentOfw;
        fixDoubleRed(u, v, path, i - 2); // i – 2 powoduje przejście w górę ścieżki
      }
    }
  }

  /** Łączenie b z parentOfw i zmiana koloru a, b, c na podstawie zależności a < b < c */
  private void restructureRecolor(RBTreeNode<E> a, RBTreeNode<E> b, 
      RBTreeNode<E> c, RBTreeNode<E> w, RBTreeNode<E> parentOfw) {
    if (parentOfw == null)
      root = b;
    else if (parentOfw.left == w)
      parentOfw.left = b;
    else
      parentOfw.right = b;

    b.setBlack(); // b staje się korzeniem w poddrzewie
    a.setRed(); // a staje się lewym dzieckiem węzła b
    c.setRed(); // c staje się prawym dzieckiem węzła b
  }      
  
  @Override /** Usuwa element z drzewa RBTree.
              * Zwraca true po udanym usunięciu elementu.
              * Zwraca false, jeśli element nie znajduje się w drzewie */
  public boolean delete(E e) {
    // Znajdowanie usuwanego węzła
    TreeNode<E> current = root;
    while (current != null) {
      if (e.compareTo(current.element) < 0) {
        current = current.left;
      }
      else if (e.compareTo(current.element) > 0) {
        current = current.right;
      }
      else
        break; // Element is znajduje się w drzewie wskazywanym przez current
    }

    if (current == null)
      return false; // Element nie występuje w drzewie

    java.util.ArrayList<TreeNode<E>> path;

    // Węzeł current jest węzłem wewnętrznym
    if (current.left != null && current.right != null) {
      // Znajdowanie pierwszego od prawej węzła w lewym poddrzewie węzła current
      TreeNode<E> rightMost = current.left;
      while (rightMost.right != null) {
        rightMost = rightMost.right; // Należy przechodzić w prawo
      }

      path = path(rightMost.element); // Pobieranie ścieżki przed zastępowaniem elementów

      // Zastępowanie elementu z current elementem z rightMost
      current.element = rightMost.element;
    }
    else
      path = path(e); // Pobieranie ścieżki do węzła current

    // Usuwanie ostatniego węzła w ścieżce i w razie potrzeby przejście w górę ścieżki
    deleteLastNodeInPath(path);
    
    size--; // Po usunięciu 1 elementu
    return true; // Element został usunięty
  }

  /** Usuwa ostatni węzeł w ścieżce */
  public void deleteLastNodeInPath(ArrayList<TreeNode<E>> path) {
    int i = path.size() - 1; // Indeks do węzła w ścieżce
    // u jest ostatnim węzłem w ścieżce
    RBTreeNode<E> u = (RBTreeNode<E>)(path.get(i));
    RBTreeNode<E> parentOfu = (u == root) ? null :
      (RBTreeNode<E>)(path.get(i - 1));
    RBTreeNode<E> grandparentOfu = (parentOfu == null ||
      parentOfu == root) ? null :
      (RBTreeNode<E>)(path.get(i - 2));
    RBTreeNode<E> childOfu = (u.left == null) ?
      (RBTreeNode<E>)(u.right) : (RBTreeNode<E>)(u.left);

    // Usuwanie węzła u; łączenie childOfu z parentOfu
    connectNewParent(parentOfu, u, childOfu);
    
    // Zmiana koloru węzłów i usunięcie problemu podwójnie czarnego węzła, jeśli jest to konieczne
    if (childOfu == root || u.isRed())
      return; // Jeśli childOfu jest korzeniem lub u jest czerwony, można zakończyć pracę
    else if (childOfu != null && childOfu.isRed()) 
      childOfu.setBlack(); // Kolorowanie childOfu na czarno — gotowe
    else // u jest czarny, childOfu jest czarny lub równy null
      // Usunięcie problemu podwójnie czarnego węzła w parentOfu
      fixDoubleBlack(grandparentOfu, parentOfu, childOfu, path, i);
  }

  /** Usunięcie problemu w rodzicu */
  private void fixDoubleBlack(
      RBTreeNode<E> grandparent, RBTreeNode<E> parent, 
      RBTreeNode<E> db, ArrayList<TreeNode<E>> path, int i) {
    // Pobieranie y, y1 i y2
    RBTreeNode<E> y = (parent.right == db) ? 
      (RBTreeNode<E>)(parent.left) : (RBTreeNode<E>)(parent.right);
    RBTreeNode<E> y1 = (RBTreeNode<E>)(y.left);
    RBTreeNode<E> y2 = (RBTreeNode<E>)(y.right);

    if (y.isBlack() && y1 != null && y1.isRed()) {
      if (parent.right == db) {
        // Scenariusz 1.1. y jest czarnym lewym bratem, a y1 jest czerwony
        connectNewParent(grandparent, parent, y);
        recolor(parent, y, y1); // Dostosowanie kolorów

        // Dostosowanie wskaźników do dzieci
        parent.left = y.right;
        y.right = parent;
      }
      else {
        // Scenariusz 1.3. y jest czarnym lewym bratem, a y1 jest czerwony
        connectNewParent(grandparent, parent, y1);
        recolor(parent, y1, y); // Dostosowanie kolorów

        // Dostosowanie wskaźników do dzieci
        parent.right = y1.left;
        y.left = y1.right;
        y1.left = parent;
        y1.right = y;
      }
    }
    else if (y.isBlack() && y2 != null && y2.isRed()) {
      if (parent.right == db) {
        // Scenariusz 1.2. y jest czarnym lewym bratem, a y2 jest czerwony
        connectNewParent(grandparent, parent, y2);
        recolor(parent, y2, y); // Dostosowanie kolorów

        // Dostosowanie wskaźników do dzieci
        y.right = y2.left;
        parent.left = y2.right;
        y2.left = y;
        y2.right = parent;
      }
      else {
        // Scenariusz 1.4. y jest czarnym prawym bratem, a y2 jest czerwony
        connectNewParent(grandparent, parent, y);
        recolor(parent, y, y2); // Dostosowanie kolorów

        // Dostosowywanie wskaźników do dzieci
        y.left = parent;
        parent.right = y1;
      }
    }
    else if (y.isBlack()) { 
      // Scenariusz 2. y jest czarny, a jego dzieci są czarne lub równe null
      y.setRed(); // Kolorowanie y na czerwono
      if (parent.isRed())
        parent.setBlack(); // Gotowe
      else if (parent != root) {
        // Przenoszenie problemu podwójnie czarnego węzła do rodzica;
        // rekurencyjne eliminowanie nowych wystąpień problemu podwójnie czarnego węzła
        db = parent;
        parent = grandparent;
        grandparent = 
          (i >= 3) ? (RBTreeNode<E>)(path.get(i - 3)) : null;
        fixDoubleBlack(grandparent, parent, db, path, i - 1);
      }
    }
    else { // y.isRed()
      if (parent.right == db) {       
        // Scenariusz 3.1. y jest czerwonym lewym dzieckiem węzła parent
        parent.left = y2;
        y.right = parent;
      }
      else {
        // Scenariusz 3.2. y jest czerwonym prawym dzieckiem węzła parent
        parent.right = y.left;
        y.left = parent;
      } 
      
      parent.setRed(); // Kolorowanie węzła parent na czerwono
      y.setBlack(); // Kolorowanie węzła y na czarno
      connectNewParent(grandparent, parent, y); // y jest nowym węzłem parent
      fixDoubleBlack(y, parent, db, path, i - 1); 
    }
  }

  /** Zmiana koloru węzłów parent, newParent i c. Scenariusz 1. usuwania */
  private void recolor(RBTreeNode<E> parent, 
      RBTreeNode<E> newParent, RBTreeNode<E> c) {
    // Zachowanie koloru węzła parent w węźle newParent
    if (parent.isRed())
      newParent.setRed(); 
    else
      newParent.setBlack();
    
    // c i parent stają się czarnymi dziećmi węzła newParen
    parent.setBlack();  
    c.setBlack();
  }

  /** Łączenie newParent z grandParent */
  private void connectNewParent(RBTreeNode<E> grandparent,
      RBTreeNode<E> parent, RBTreeNode<E> newParent) {
    if (parent == root) {
      root = newParent;
      if (root != null)
        newParent.setBlack();
    }
    else if (grandparent.left == parent)
      grandparent.left = newParent;
    else
      grandparent.right = newParent;
  }

  @Override /** Odwiedzanie węzłów w porządku preorder od poddrzewa */
  protected void preorder(TreeNode<E> root) {
    if (root == null) return;
    System.out.print(root.element +
      (((RBTreeNode<E>)root).isRed() ? " (czerwony) " : " (czarny) "));
    preorder(root.left);
    preorder(root.right);
  }

  /** RBTreeNode to klasa TreeNode plus pole reprezentujące kolor */
  protected static class RBTreeNode<E extends Comparable<E>> extends
      BST.TreeNode<E> {
    private boolean red = true; // Określa kolor węzła

    public RBTreeNode(E e) {
      super(e);
    }

    public boolean isRed() {
      return red;
    }

    public boolean isBlack() {
      return!red;
    }

    public void setBlack() {
      red = false;
    }

    public void setRed() {
      red = true;
    }

    int blackHeight;
  }
}