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 RED = true;
    private static final boolean BLACK = false;

    /** Classe interne représentant un nœud de l'arbre */
    private class Node {
	E key;
	Node left, right, parent;
	boolean color = RED; // par défaut, un nouveau nœud est rouge

	Node(E key, Node parent) {
	    this.key = key;
	    this.parent = parent;
	}
    }

    private Node root;
    private int size = 0;

    /** Rotation gauche standard */
    private void rotateLeft(Node x) {
	Node y = x.right;
	x.right = y.left;

	if (y.left != null)
	    y.left.parent = x;

	y.parent = x.parent;

	if (x.parent == null)
	    root = y;
	else if (x == x.parent.left)
	    x.parent.left = y;
	else
	    x.parent.right = y;

	y.left = x;
	x.parent = y;
    }

    /** Rotation droite standard */
    private void rotateRight(Node x) {
	Node y = x.left;
	x.left = y.right;

	if (y.right != null)
	    y.right.parent = x;

	y.parent = x.parent;

	if (x.parent == null)
	    root = y;
	else if (x == x.parent.right)
	    x.parent.right = y;
	else
	    x.parent.left = y;

	y.right = x;
	x.parent = y;
    }

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

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

		    // Cas 3 : ligne -> rotation droite
		    z.parent.color = BLACK;
		    z.parent.parent.color = RED;
		    rotateRight(z.parent.parent);
		}
	    } else {
		// Symétrique : parent est enfant droit
		Node y = z.parent.parent.left;

		if (y != null && y.color == RED) {
		    z.parent.color = BLACK;
		    y.color = BLACK;
		    z.parent.parent.color = RED;
		    z = z.parent.parent;
		} else {
		    if (z == z.parent.left) {
			z = z.parent;
			rotateRight(z);
		    }
		    z.parent.color = BLACK;
		    z.parent.parent.color = RED;
		    rotateLeft(z.parent.parent);
		}
	    }
	}

	root.color = BLACK;
    }

    /** Remplace un nœud par un autre dans l’arbre */
    private void transplant(Node u, Node v) {
	if (u.parent == null)
	    root = v;
	else if (u == u.parent.left)
	    u.parent.left = v;
	else
	    u.parent.right = v;

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

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

    /**
     * Correction après suppression. Restaure les propriétés Rouge-Noir.
     */
    private void fixAfterDelete(Node x, Node parent) {
	while ((x != root) && (x == null || x.color == BLACK)) {
	    if (parent != null && x == parent.left) {
		Node w = parent.right;

		if (w != null && w.color == RED) {
		    w.color = BLACK;
		    parent.color = RED;
		    rotateLeft(parent);
		    w = parent.right;
		}

		if ((w == null)
			|| ((w.left == null || w.left.color == BLACK) && (w.right == null || w.right.color == BLACK))) {
		    if (w != null)
			w.color = RED;

		    x = parent;
		    parent = x.parent;
		} else {
		    if (w.right == null || w.right.color == BLACK) {
			if (w.left != null)
			    w.left.color = BLACK;

			w.color = RED;
			rotateRight(w);
			w = parent.right;
		    }

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

		    parent.color = BLACK;

		    if (w != null && w.right != null)
			w.right.color = BLACK;

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

		Node w = parent.left;

		if (w != null && w.color == RED) {
		    w.color = BLACK;
		    parent.color = RED;
		    rotateRight(parent);
		    w = parent.left;
		}

		if ((w == null)
			|| ((w.right == null || w.right.color == BLACK) && (w.left == null || w.left.color == BLACK))) {
		    if (w != null)
			w.color = RED;

		    x = parent;
		    parent = x.parent;
		} else {
		    if (w.left == null || w.left.color == BLACK) {
			if (w.right != null)
			    w.right.color = BLACK;

			w.color = RED;
			rotateLeft(w);
			w = parent.left;
		    }

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

		    parent.color = BLACK;

		    if (w != null && w.left != null)
			w.left.color = BLACK;

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

	if (x != null)
	    x.color = BLACK;
    }

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

	// Arbre vide : nouvelle racine noire
	if (root == null) {
	    root = new Node(e, null);
	    root.color = BLACK;
	    size = 1;
	    return true;
	}

	Node cur = root, parent = null;
	int cmp = 0;

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

	    if (cmp < 0)
		cur = cur.left;
	    else if (cmp > 0)
		cur = cur.right;
	    else
		return false; // pas de doublons
	}

	Node node = new Node(e, parent);

	if (cmp < 0)
	    parent.left = node;
	else
	    parent.right = node;

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

	return true;
    }

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

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

	    if (cmp == 0)
		return cur;
	    if (cmp < 0)
		cur = cur.left;
	    else
		cur = cur.right;
	}

	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;
	}

	Node z = findNode(e);

	if (z == null)
	    return false;

	Node y = z;
	boolean yOriginalColor = y.color;
	Node x;
	Node xParent;

	// Cas 1 : un seul enfant
	if (z.left == null) {
	    x = z.right;
	    xParent = z.parent;
	    transplant(z, z.right);
	} else if (z.right == null) {
	    x = z.left;
	    xParent = z.parent;
	    transplant(z, z.left);
	}

	// Cas 2 : deux enfants → successeur
	else {
	    y = minimum(z.right);
	    yOriginalColor = y.color;
	    x = y.right;

	    if (y.parent == z) {
		if (x != null)
		    x.parent = y;
		xParent = y;
	    } else {
		transplant(y, y.right);
		y.right = z.right;

		if (y.right != null)
		    y.right.parent = y;

		xParent = y.parent;
	    }

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

	    if (y.left != null)
		y.left.parent = y;

	    y.color = z.color;
	}

	if (yOriginalColor == BLACK)
	    fixAfterDelete(x, xParent);

	size--;
	return true;
    }

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

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

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

		return s;
	    }

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

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

		Node n = stack.pop();
		E res = n.key;
		Node cur = n.right;

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

		return res;
	    }

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

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

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

    @Override
    public void clear() {
	root = null;
	size = 0;
    }

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

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

	return arr;
    }

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

	int i = 0;

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

	if (a.length > size)
	    a[size] = 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();
    }
}