// Árbol rojinegro clásico sin valores: el conjunto de claves de los apuntes.
// Es recursivo y no tiene puntero al padre: se repara desde el abuelo, a la vuelta de la recursión.
// Toma las mismas decisiones que el algoritmo con puntero al padre, que construye el laboratorio.
#ifndef RB_SET_H
#define RB_SET_H

#include <cassert>
#include <cstddef>
#include <sstream>
#include <string>
#include <vector>

// Invariantes, además del orden de un árbol de búsqueda:
// 1. Cada nodo es rojo o negro.
// 2. La raíz es negra.
// 3. Las hojas nulas son negras: aquí un puntero nulo cuenta como negro, sin ocupar memoria.
// 4. Un nodo rojo solo tiene hijos negros.
// 5. Todo camino desde un nodo hasta sus hojas nulas pasa por la misma cantidad de negros.
// Además, el árbol tiene exactamente size_ nodos, y root_ es nulo si y solo si size_ == 0.
//
// K solo necesita operator<, y se copia al insertar.
template <typename K>
class RedBlackSet {
    enum class Color { red, black };

    struct Node {
        K key;
        Node* left = nullptr;
        Node* right = nullptr;
        Color color = Color::red;
    };

public:
    // Un paso de la reparación: en qué clave estaba el abuelo, y qué se hizo.
    struct Step {
        K key;
        std::string kind;
    };

    RedBlackSet() = default;
    ~RedBlackSet() { destroy(root_); }

    // Este conjunto es para estudiar las operaciones: no se copia ni se mueve.
    RedBlackSet(const RedBlackSet&) = delete;
    RedBlackSet& operator=(const RedBlackSet&) = delete;

    std::size_t size() const { return size_; }

    bool contains(const K& key) const {
        const Node* node = root_;
        while (node != nullptr) {
            if (key < node->key) {
                node = node->left;
            } else if (node->key < key) {
                node = node->right;
            } else {
                return true;
            }
        }
        return false;
    }

    // Devuelve false si la clave ya estaba.
    bool insert(const K& key) {
        last_steps_.clear();
        bool added = false;
        root_ = insert_at(root_, key, added);
        if (added) {
            size_++;
        }
        if (root_->color == Color::red) {
            last_steps_.push_back({root_->key, "raiz roja: pintar de negro"});
            root_->color = Color::black;
        }
        return added;
    }

    // Altura: aristas del camino más largo desde la raíz. Vacío: -1. Un solo nodo: 0.
    int height() const { return height_of(root_); }

    // Cuántos negros hay en un camino desde la raíz hasta una hoja nula. La hoja nula no se cuenta.
    int black_height() const {
        int blacks = 0;
        for (const Node* node = root_; node != nullptr; node = node->left) {
            blacks += is_red(node) ? 0 : 1;
        }
        return blacks;
    }

    // Lo que hizo la última inserción, de abajo hacia arriba.
    const std::vector<Step>& last_steps() const { return last_steps_; }

    long rotations() const { return rotations_; }

    // Cuántas veces se aplicó el caso del tío rojo. No cuenta los otros cambios de color.
    long recolors() const { return recolors_; }

    // La suma de las profundidades de todos los nodos.
    long total_depth() const { return depth_sum(root_, 0); }

    std::vector<K> in_order() const {
        std::vector<K> keys;
        walk_in(root_, keys);
        return keys;
    }

    // El árbol en una línea: clave(izquierdo,derecho), con _ si falta un hijo y r si el nodo es rojo.
    std::string shape() const {
        std::ostringstream out;
        write_shape(root_, out);
        return out.str();
    }

    // El árbol 2-3-4 que representa: cada negro con sus hijos rojos forma un nodo.
    std::string collapsed() const {
        std::ostringstream out;
        write_group(root_, out);
        return out.str();
    }

    // Cuántos nodos de 2 claves tiene el árbol 2-3-4, según hacia dónde quedó el rojo.
    long left_leaning_groups() const { return count_leaning(root_, true); }
    long right_leaning_groups() const { return count_leaning(root_, false); }

    void check_invariants() const {
        assert((root_ == nullptr) == (size_ == 0));
        assert(!is_red(root_));
        std::size_t count = 0;
        check_node(root_, nullptr, nullptr, count);
        assert(count == size_);
    }

private:
    static bool is_red(const Node* node) { return node != nullptr && node->color == Color::red; }

    static int height_of(const Node* node) {
        if (node == nullptr) {
            return -1;
        }
        int left = height_of(node->left);
        int right = height_of(node->right);
        return 1 + (left > right ? left : right);
    }

    // Sube al hijo derecho de x, y devuelve la nueva raíz del subárbol. Solo cambia enlaces.
    Node* rotate_left(Node* x) {
        Node* y = x->right;
        x->right = y->left;
        y->left = x;
        rotations_++;
        return y;
    }

    // Sube al hijo izquierdo de x, y devuelve la nueva raíz del subárbol. Solo cambia enlaces.
    Node* rotate_right(Node* x) {
        Node* y = x->left;
        x->left = y->right;
        y->right = x;
        rotations_++;
        return y;
    }

    Node* insert_at(Node* node, const K& key, bool& added) {
        if (node == nullptr) {
            added = true;
            return new Node{key};
        }
        if (key < node->key) {
            node->left = insert_at(node->left, key, added);
        } else if (node->key < key) {
            node->right = insert_at(node->right, key, added);
        } else {
            return node;
        }
        return repair_from_grandparent(node);
    }

    // Tío rojo: el abuelo pasa a rojo, y el padre y el tío a negro. El problema sube.
    void recolor(Node* grandparent) {
        last_steps_.push_back({grandparent->key, "tio rojo: recolorear"});
        grandparent->color = Color::red;
        grandparent->left->color = Color::black;
        grandparent->right->color = Color::black;
        recolors_++;
    }

    // Si hay un rojo con hijo rojo debajo de grandparent, lo repara y devuelve la raíz del subárbol.
    Node* repair_from_grandparent(Node* grandparent) {
        Node* left = grandparent->left;
        Node* right = grandparent->right;
        if (is_red(left) && (is_red(left->left) || is_red(left->right))) {
            if (is_red(right)) {
                recolor(grandparent);
                return grandparent;
            }
            if (is_red(left->right)) {
                last_steps_.push_back({grandparent->key, "tio negro, zigzag: dos rotaciones"});
                grandparent->left = rotate_left(left);
            } else {
                last_steps_.push_back({grandparent->key, "tio negro, en linea: una rotacion"});
            }
            Node* top = rotate_right(grandparent);
            top->color = Color::black;
            grandparent->color = Color::red;
            return top;
        }
        if (is_red(right) && (is_red(right->left) || is_red(right->right))) {
            if (is_red(left)) {
                recolor(grandparent);
                return grandparent;
            }
            if (is_red(right->left)) {
                last_steps_.push_back({grandparent->key, "tio negro, zigzag: dos rotaciones"});
                grandparent->right = rotate_right(right);
            } else {
                last_steps_.push_back({grandparent->key, "tio negro, en linea: una rotacion"});
            }
            Node* top = rotate_left(grandparent);
            top->color = Color::black;
            grandparent->color = Color::red;
            return top;
        }
        return grandparent;
    }

    static void destroy(Node* node) {
        if (node == nullptr) {
            return;
        }
        destroy(node->left);
        destroy(node->right);
        delete node;
    }

    static long depth_sum(const Node* node, long depth) {
        if (node == nullptr) {
            return 0;
        }
        return depth + depth_sum(node->left, depth + 1) + depth_sum(node->right, depth + 1);
    }

    static void walk_in(const Node* node, std::vector<K>& keys) {
        if (node == nullptr) {
            return;
        }
        walk_in(node->left, keys);
        keys.push_back(node->key);
        walk_in(node->right, keys);
    }

    static void write_shape(const Node* node, std::ostringstream& out) {
        if (node == nullptr) {
            out << "_";
            return;
        }
        out << node->key;
        if (node->color == Color::red) {
            out << "r";
        }
        if (node->left != nullptr || node->right != nullptr) {
            out << "(";
            write_shape(node->left, out);
            out << ",";
            write_shape(node->right, out);
            out << ")";
        }
    }

    // Escribe el nodo 2-3-4 cuya parte de arriba es node, que es negro: [claves](hijos).
    static void write_group(const Node* node, std::ostringstream& out) {
        if (node == nullptr) {
            out << "_";
            return;
        }
        std::vector<const Node*> keys;
        std::vector<const Node*> children;
        if (is_red(node->left)) {
            keys.push_back(node->left);
            children.push_back(node->left->left);
            children.push_back(node->left->right);
        } else {
            children.push_back(node->left);
        }
        keys.push_back(node);
        if (is_red(node->right)) {
            keys.push_back(node->right);
            children.push_back(node->right->left);
            children.push_back(node->right->right);
        } else {
            children.push_back(node->right);
        }
        out << "[";
        for (std::size_t i = 0; i < keys.size(); i++) {
            out << (i > 0 ? " " : "") << keys[i]->key;
        }
        out << "]";
        bool is_leaf = true;
        for (const Node* child : children) {
            is_leaf = is_leaf && child == nullptr;
        }
        if (!is_leaf) {
            out << "(";
            for (std::size_t i = 0; i < children.size(); i++) {
                out << (i > 0 ? "," : "");
                write_group(children[i], out);
            }
            out << ")";
        }
    }

    static long count_leaning(const Node* node, bool to_left) {
        if (node == nullptr) {
            return 0;
        }
        long own = 0;
        if (!is_red(node) && is_red(node->left) != is_red(node->right) && is_red(node->left) == to_left) {
            own = 1;
        }
        return own + count_leaning(node->left, to_left) + count_leaning(node->right, to_left);
    }

    // Comprueba todo el subárbol y devuelve sus negros hasta una hoja nula.
    static int check_node(const Node* node, const K* low, const K* high, std::size_t& count) {
        if (node == nullptr) {
            return 0;
        }
        count++;
        assert(low == nullptr || *low < node->key);
        assert(high == nullptr || node->key < *high);
        assert(node->color == Color::red || node->color == Color::black);
        if (is_red(node)) {
            assert(!is_red(node->left) && !is_red(node->right));
        }
        int left = check_node(node->left, low, &node->key, count);
        int right = check_node(node->right, &node->key, high, count);
        assert(left == right);
        return left + (is_red(node) ? 0 : 1);
    }

    Node* root_ = nullptr;
    std::size_t size_ = 0;
    long rotations_ = 0;
    long recolors_ = 0;
    std::vector<Step> last_steps_;
};

#endif