package collection;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.Objects;
import java.util.Stack;

/**
 * Arbre Rouge-Noir générique. Respecte les propriétés des arbres RB : - chaque
 * nœud est rouge ou noir - la racine est toujours noire - les feuilles null
 * sont noires - un nœud rouge ne peut pas avoir un parent rouge (pas de
 * rouge-rouge) - tous les chemins depuis un nœud jusqu’aux feuilles null
 * contiennent le même nombre de nœuds noirs
 */
public class RedBlackTree<E extends Comparable<E>> implements Collection<E> {
  private static final boolean ROUGE = true;
  private static final boolean NOIR = false;

  /** Classe interne représentant un nœud de l'arbre */
  private class Noeud {
    E cle;
    Noeud gauche;
    Noeud droit;
    Noeud parent;
    boolean color = ROUGE; // par défaut, un nouveau nœud est rouge

    Noeud(E cle, Noeud parent) {
      this.cle = cle;
      this.parent = parent;
    }
  }

  private Noeud racine;
  private int taille = 0;

  /** Rotation gauche standard */
  private void rotateLeft(Noeud x) {
    Noeud y = x.droit;
    x.droit = y.gauche;

    if (y.gauche != null) {
      y.gauche.parent = x;
    }

    y.parent = x.parent;

    if (x.parent == null) {
      racine = y;
    } else if (x == x.parent.gauche) {
      x.parent.gauche = y;
    } else {
      x.parent.droit = y;
    }

    y.gauche = x;
    x.parent = y;
  }

  /** Rotation droite standard */
  private void rotateRight(Noeud x) {
    Noeud y = x.gauche;
    x.gauche = y.droit;

    if (y.droit != null) {
      y.droit.parent = x;
    }

    y.parent = x.parent;

    if (x.parent == null) {
      racine = y;
    } else if (x == x.parent.droit) {
      x.parent.droit = y;
    } else {
      x.parent.gauche = y;
    }

    y.droit = x;
    x.parent = y;
  }

  /**
   * Correction après insertion. Assure que les propriétés RB sont conservées.
   */
  private void fixAfterInsert(Noeud z) {
    while (z.parent != null && z.parent.color == ROUGE) {
      if (z.parent == z.parent.parent.gauche) {
        Noeud y = z.parent.parent.droit; // oncle

        if (y != null && y.color == ROUGE) {
          // Cas 1 : parent rouge + oncle rouge -> recolorations
          z.parent.color = NOIR;
          y.color = NOIR;
          z.parent.parent.color = ROUGE;
          z = z.parent.parent;
        } else {
          if (z == z.parent.droit) {
            // Cas 2 : triangle -> rotation gauche
            z = z.parent;
            rotateLeft(z);
          }

          // Cas 3 : ligne -> rotation droite
          z.parent.color = NOIR;
          z.parent.parent.color = ROUGE;
          rotateRight(z.parent.parent);
        }
      } else {
        // Symétrique : parent est enfant droit
        Noeud y = z.parent.parent.gauche;

        if (y != null && y.color == ROUGE) {
          z.parent.color = NOIR;
          y.color = NOIR;
          z.parent.parent.color = ROUGE;
          z = z.parent.parent;
        } else {
          if (z == z.parent.gauche) {
            z = z.parent;
            rotateRight(z);
          }
          z.parent.color = NOIR;
          z.parent.parent.color = ROUGE;
          rotateLeft(z.parent.parent);
        }
      }
    }

    racine.color = NOIR;
  }

  /** Remplace un nœud par un autre dans l’arbre */
  private void transplant(Noeud u, Noeud v) {
    if (u.parent == null) {
      racine = v;
    } else if (u == u.parent.gauche) {
      u.parent.gauche = v;
    } else {
      u.parent.droit = v;
    }

    if (v != null) {
      v.parent = u.parent;
    }
  }

  /** Renvoie le minimum du sous-arbre */
  private Noeud minimum(Noeud x) {
    while (x.gauche != null) {
      x = x.gauche;
    }

    return x;
  }

  /**
   * Correction après suppression. Restaure les propriétés Rouge-Noir.
   */
  private void fixAfterDelete(Noeud x, Noeud parent) {
    while ((x != racine) && (x == null || x.color == NOIR)) {
      if (parent != null && x == parent.gauche) {
        Noeud w = parent.droit;

        if (w != null && w.color == ROUGE) {
          w.color = NOIR;
          parent.color = ROUGE;
          rotateLeft(parent);
          w = parent.droit;
        }

        if ((w == null) || ((w.gauche == null || w.gauche.color == NOIR)
            && (w.droit == null || w.droit.color == NOIR))) {
          if (w != null) {
            w.color = ROUGE;
          }

          x = parent;
          parent = x.parent;
        } else {
          if (w.droit == null || w.droit.color == NOIR) {
            if (w.gauche != null) {
              w.gauche.color = NOIR;
            }

            w.color = ROUGE;
            rotateRight(w);
            w = parent.droit;
          }

          if (w != null) {
            w.color = parent.color;
          }

          parent.color = NOIR;

          if (w != null && w.droit != null) {
            w.droit.color = NOIR;
          }

          rotateLeft(parent);
          x = racine;
          break;
        }
      } else {
        if (parent == null) {
          break;
        }

        Noeud w = parent.gauche;

        if (w != null && w.color == ROUGE) {
          w.color = NOIR;
          parent.color = ROUGE;
          rotateRight(parent);
          w = parent.gauche;
        }

        if ((w == null) || ((w.droit == null || w.droit.color == NOIR)
            && (w.gauche == null || w.gauche.color == NOIR))) {
          if (w != null) {
            w.color = ROUGE;
          }

          x = parent;
          parent = x.parent;
        } else {
          if (w.gauche == null || w.gauche.color == NOIR) {
            if (w.droit != null) {
              w.droit.color = NOIR;
            }

            w.color = ROUGE;
            rotateLeft(w);
            w = parent.gauche;
          }

          if (w != null) {
            w.color = parent.color;
          }

          parent.color = NOIR;

          if (w != null && w.gauche != null) {
            w.gauche.color = NOIR;
          }

          rotateRight(parent);
          x = racine;
          break;
        }
      }
    }

    if (x != null) {
      x.color = NOIR;
    }
  }

  @Override
  public boolean add(E e) {
    Objects.requireNonNull(e);

    // Arbre vide : nouvelle racine noire
    if (racine == null) {
      racine = new Noeud(e, null);
      racine.color = NOIR;
      taille = 1;
      return true;
    }

    Noeud cur = racine;
    Noeud parent = null;
    int cmp = 0;

    // Recherche de la position d'insertion
    while (cur != null) {
      parent = cur;
      cmp = e.compareTo(cur.cle);

      if (cmp < 0) {
        cur = cur.gauche;
      } else if (cmp > 0) {
        cur = cur.droit;
      } else {
        return false; // pas de doublons
      }
    }

    Noeud node = new Noeud(e, parent);

    if (cmp < 0) {
      parent.gauche = node;
    } else {
      parent.droit = node;
    }

    // Correction RB après insertion
    fixAfterInsert(node);
    taille++;

    return true;
  }

  /** Recherche un nœud contenant la clé */
  private Noeud findNode(E e) {
    Noeud cur = racine;

    while (cur != null) {
      int cmp = e.compareTo(cur.cle);

      if (cmp == 0) {
        return cur;
      }
      if (cmp < 0) {
        cur = cur.gauche;
      } else {
        cur = cur.droit;
      }
    }

    return null;
  }

  @Override
  public boolean contains(Object o) {
    if (o == null) {
      return false;
    }

    try {
      @SuppressWarnings("unchecked")
      E e = (E) o;
      return findNode(e) != null;
    } catch (ClassCastException ex) {
      return false;
    }
  }

  @Override
  public boolean remove(Object o) {
    if (o == null) {
      return false;
    }

    E e;

    try {
      e = (E) o;
    } catch (ClassCastException ex) {
      return false;
    }

    Noeud z = findNode(e);

    if (z == null) {
      return false;
    }

    Noeud y = z;
    boolean yorigColor = y.color;
    Noeud x;
    Noeud xparent;

    // Cas 1 : un seul enfant
    if (z.gauche == null) {
      x = z.droit;
      xparent = z.parent;
      transplant(z, z.droit);
    } else if (z.droit == null) {
      x = z.gauche;
      xparent = z.parent;
      transplant(z, z.gauche);
    } else { // Cas 2 : deux enfants → successeur
      y = minimum(z.droit);
      yorigColor = y.color;
      x = y.droit;

      if (y.parent == z) {
        if (x != null) {
          x.parent = y;
        }

        xparent = y;
      } else {
        transplant(y, y.droit);
        y.droit = z.droit;

        if (y.droit != null) {
          y.droit.parent = y;
        }

        xparent = y.parent;
      }

      transplant(z, y);
      y.gauche = z.gauche;

      if (y.gauche != null) {
        y.gauche.parent = y;
      }

      y.color = z.color;
    }

    if (yorigColor == NOIR) {
      fixAfterDelete(x, xparent);
    }

    taille--;
    return true;
  }

  @Override
  public Iterator<E> iterator() {
    return new Iterator<E>() {
      private final Stack<Noeud> stack = initStack(racine);

      private Stack<Noeud> initStack(Noeud r) {
        Stack<Noeud> s = new Stack<>();
        Noeud cur = r;

        while (cur != null) {
          s.push(cur);
          cur = cur.gauche;
        }

        return s;
      }

      @Override
      public boolean hasNext() {
        return !stack.isEmpty();
      }

      @Override
      public E next() {
        if (!hasNext()) {
          throw new NoSuchElementException();
        }

        Noeud n = stack.pop();
        E res = n.cle;
        Noeud cur = n.droit;

        // Descente maximale à gauche
        while (cur != null) {
          stack.push(cur);
          cur = cur.gauche;
        }

        return res;
      }

      @Override
      public void remove() {
        throw new UnsupportedOperationException();
      }
    };
  }

  @Override
  public int size() {
    return taille;
  }

  @Override
  public boolean isEmpty() {
    return taille == 0;
  }

  @Override
  public void clear() {
    racine = null;
    taille = 0;
  }

  @Override
  public Object[] toArray() {
    Object[] arr = new Object[taille];
    int i = 0;

    for (E e : this) {
      arr[i++] = e;
    }

    return arr;
  }

  @Override
  public <T> T[] toArray(T[] a) {
    if (a.length < taille) {
      a = (T[]) java.lang.reflect.Array.newInstance(a.getClass().getComponentType(), taille);
    }

    int i = 0;

    for (E e : this) {
      a[i++] = (T) e;
    }

    if (a.length > taille) {
      a[taille] = null;
    }

    return a;
  }

  @Override
  public boolean containsAll(Collection<?> c) {
    for (Object o : c) {
      if (!contains(o)) {
        return false;
      }
    }

    return true;
  }

  @Override
  public boolean addAll(Collection<? extends E> c) {
    boolean changed = false;
    for (E e : c) {
      changed |= add(e);
    }

    return changed;
  }

  @Override
  public boolean removeAll(Collection<?> c) {
    boolean changed = false;
    for (Object o : c) {
      changed |= remove(o);
    }

    return changed;
  }

  @Override
  public boolean retainAll(Collection<?> c) {
    List<E> toRemove = new ArrayList<>();

    for (E e : this) {
      if (!c.contains(e)) {
        toRemove.add(e);
      }
    }

    for (E e : toRemove) {
      remove(e);
    }

    return !toRemove.isEmpty();
  }
}