// Árbol 2-3-4: cada nodo tiene de 1 a 3 claves, y todas las hojas están a la misma profundidad.
// Sirve para comparar con el árbol rojinegro. Se elige cuándo y cómo se divide un nodo.
#ifndef TREE_234_H
#define TREE_234_H

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

enum class SplitPolicy {
    // Se divide solo si llega una clave a un nodo lleno. Sube la del medio de las 3 que tenía.
    old_middle,
    // También solo al hacer overflow, pero con las 4 claves juntas sube siempre la segunda.
    fixed_second,
    // Al bajar, se divide todo nodo lleno que se encuentre, haga falta o no.
    top_down,
};

// Invariantes: cada nodo tiene de 1 a 3 claves ordenadas; un nodo interno tiene un hijo más que
// claves; las claves de cada hijo caen entre las de su padre; todas las hojas a la misma profundidad.
template <typename K>
class Tree234 {
    struct Node {
        std::vector<K> keys;
        std::vector<Node*> children;
        bool is_leaf() const { return children.empty(); }
    };

    // Lo que un nodo dividido le entrega a su padre.
    struct Split {
        bool happened = false;
        K up{};
        Node* right = nullptr;
    };

public:
    explicit Tree234(SplitPolicy policy = SplitPolicy::old_middle) : policy_(policy) {}
    ~Tree234() { destroy(root_); }

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

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

    // Los nodos que se dividieron en la última inserción, de abajo hacia arriba.
    const std::vector<std::string>& last_splits() const { return last_splits_; }

    bool contains(const K& key) const {
        const Node* node = root_;
        while (node != nullptr) {
            std::size_t i = position(node, key);
            if (i < node->keys.size() && !(key < node->keys[i])) {
                return true;
            }
            node = node->is_leaf() ? nullptr : node->children[i];
        }
        return false;
    }

    // Devuelve false si la clave ya estaba.
    bool insert(const K& key) {
        last_splits_.clear();
        if (contains(key)) {
            return false;
        }
        size_++;
        if (root_ == nullptr) {
            root_ = new Node{{key}, {}};
        } else if (policy_ == SplitPolicy::top_down) {
            insert_top_down(key);
        } else {
            Split split = insert_at(root_, key);
            if (split.happened) {
                grow_root(split);
            }
        }
        return true;
    }

    // Cuántos niveles de nodos tiene. Vacío: 0.
    int levels() const {
        int count = 0;
        for (const Node* node = root_; node != nullptr; node = node->is_leaf() ? nullptr : node->children[0]) {
            count++;
        }
        return count;
    }

    // El árbol en una línea: [claves](hijos).
    std::string shape() const {
        std::ostringstream out;
        write_shape(root_, out);
        return out.str();
    }

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

private:
    // El primer lugar cuya clave no es menor que key.
    static std::size_t position(const Node* node, const K& key) {
        std::size_t i = 0;
        while (i < node->keys.size() && node->keys[i] < key) {
            i++;
        }
        return i;
    }

    void grow_root(const Split& split) {
        root_ = new Node{{split.up}, {root_, split.right}};
    }

    // Baja hasta una hoja, y a la vuelta recibe lo que suba de cada división.
    Split insert_at(Node* node, const K& key) {
        if (node->is_leaf()) {
            return add(node, key, nullptr);
        }
        Split below = insert_at(node->children[position(node, key)], key);
        return below.happened ? add(node, below.up, below.right) : Split{};
    }

    // Agrega a node una clave, con el hijo que va a su derecha. Si no cabe, lo divide.
    Split add(Node* node, const K& key, Node* right_child) {
        if (node->keys.size() < 3) {
            put(node, key, right_child);
            return {};
        }
        splits_++;
        last_splits_.push_back(text_of(node));
        if (policy_ == SplitPolicy::old_middle) {
            Split split = split_full(node);
            put(key < split.up ? node : split.right, key, right_child);
            return split;
        }
        // Con las 4 claves juntas: a la izquierda 1, sube la segunda, y a la derecha 2.
        put(node, key, right_child);
        Split split;
        split.happened = true;
        split.up = node->keys[1];
        split.right = new Node{{node->keys[2], node->keys[3]}, {}};
        if (!node->is_leaf()) {
            split.right->children.assign(node->children.begin() + 2, node->children.end());
            node->children.resize(2);
        }
        node->keys.resize(1);
        return split;
    }

    // Divide un nodo con 3 claves: se queda con la primera, y la segunda sube.
    static Split split_full(Node* node) {
        Split split;
        split.happened = true;
        split.up = node->keys[1];
        split.right = new Node{{node->keys[2]}, {}};
        if (!node->is_leaf()) {
            split.right->children.assign(node->children.begin() + 2, node->children.end());
            node->children.resize(2);
        }
        node->keys.resize(1);
        return split;
    }

    static void put(Node* node, const K& key, Node* right_child) {
        std::size_t i = position(node, key);
        node->keys.insert(node->keys.begin() + static_cast<std::ptrdiff_t>(i), key);
        if (right_child != nullptr) {
            node->children.insert(node->children.begin() + static_cast<std::ptrdiff_t>(i) + 1, right_child);
        }
    }

    void insert_top_down(const K& key) {
        if (root_->keys.size() == 3) {
            splits_++;
            last_splits_.push_back(text_of(root_));
            grow_root(split_full(root_));
        }
        Node* node = root_;
        while (!node->is_leaf()) {
            std::size_t i = position(node, key);
            Node* child = node->children[i];
            if (child->keys.size() == 3) {
                splits_++;
                last_splits_.push_back(text_of(child));
                Split split = split_full(child);
                put(node, split.up, split.right);
                child = key < split.up ? child : split.right;
            }
            node = child;
        }
        put(node, key, nullptr);
    }

    static std::string text_of(const Node* node) {
        std::ostringstream out;
        write_keys(node, out);
        return out.str();
    }

    static void write_keys(const Node* node, std::ostringstream& out) {
        out << "[";
        for (std::size_t i = 0; i < node->keys.size(); i++) {
            out << (i > 0 ? " " : "") << node->keys[i];
        }
        out << "]";
    }

    static void write_shape(const Node* node, std::ostringstream& out) {
        if (node == nullptr) {
            out << "_";
            return;
        }
        write_keys(node, out);
        if (!node->is_leaf()) {
            out << "(";
            for (std::size_t i = 0; i < node->children.size(); i++) {
                out << (i > 0 ? "," : "");
                write_shape(node->children[i], out);
            }
            out << ")";
        }
    }

    static void check_node(const Node* node, const K* low, const K* high, int level, int levels, std::size_t& count) {
        if (node == nullptr) {
            return;
        }
        assert(node->keys.size() >= 1 && node->keys.size() <= 3);
        assert(node->is_leaf() ? level == levels : node->children.size() == node->keys.size() + 1);
        count += node->keys.size();
        for (std::size_t i = 0; i < node->keys.size(); i++) {
            assert((i == 0 ? low : &node->keys[i - 1]) == nullptr || *(i == 0 ? low : &node->keys[i - 1]) < node->keys[i]);
        }
        assert(high == nullptr || node->keys.back() < *high);
        for (std::size_t i = 0; i < node->children.size(); i++) {
            const K* child_low = i == 0 ? low : &node->keys[i - 1];
            const K* child_high = i < node->keys.size() ? &node->keys[i] : high;
            check_node(node->children[i], child_low, child_high, level + 1, levels, count);
        }
    }

    static void destroy(Node* node) {
        if (node == nullptr) {
            return;
        }
        for (Node* child : node->children) {
            destroy(child);
        }
        delete node;
    }

    SplitPolicy policy_;
    Node* root_ = nullptr;
    std::size_t size_ = 0;
    long splits_ = 0;
    std::vector<std::string> last_splits_;
};

#endif