// Árbol AVL sin valores: el conjunto de claves de los apuntes.
// Es recursivo y no tiene puntero al padre: la vuelta de la recursión es la subida hacia la raíz.
// El laboratorio construye un mapa con puntero al padre, sobre el BstMap de la semana 10.
#ifndef AVL_SET_H
#define AVL_SET_H

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

// Invariantes:
// - En cada nodo, todas las claves del subárbol izquierdo son menores que la suya,
//   y todas las del derecho son mayores.
// - En cada nodo, height es la altura de su subárbol. Un nodo nulo tiene altura -1.
// - En cada nodo, las alturas de sus dos subárboles difieren en 1 como máximo.
// - El árbol tiene exactamente size_ nodos, y root_ es nulo si y solo si size_ == 0.
//
// K solo necesita operator<, y se copia al borrar una clave con dos hijos.
// La recursión es segura aquí: gasta pila proporcional a la altura, y la altura es O(log n).
template <typename K>
class AvlSet {
    struct Node {
        K key;
        Node* left = nullptr;
        Node* right = nullptr;
        int height = 0;
    };

    friend struct AvlSetAccess;

public:
    // Una reparación: en qué clave estaba el nodo desbalanceado, y qué caso fue.
    struct Repair {
        K key;
        std::string kind;
    };

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

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

    std::size_t size() const { return size_; }
    bool empty() const { return size_ == 0; }

    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_repairs_.clear();
        bool added = false;
        root_ = insert_at(root_, key, added);
        if (added) {
            size_++;
        }
        return added;
    }

    // Devuelve false si la clave no estaba.
    bool erase(const K& key) {
        last_repairs_.clear();
        bool erased = false;
        root_ = erase_at(root_, key, erased);
        if (erased) {
            size_--;
        }
        return erased;
    }

    int height() const { return height_of(root_); }

    // Las reparaciones que hizo la última operación, de abajo hacia arriba.
    const std::vector<Repair>& last_repairs() const { return last_repairs_; }

    // Rotaciones desde que se creó: una reparación doble cuenta como dos.
    long rotations() const { return rotations_; }

    // 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.
    std::string shape() const {
        std::ostringstream out;
        write_shape(root_, out);
        return out.str();
    }

    void check_invariants() const {
        assert((root_ == nullptr) == (size_ == 0));
        assert(count_checked(root_, nullptr, nullptr) == size_);
    }

private:
    static int height_of(const Node* node) { return node != nullptr ? node->height : -1; }

    static void update_height(Node* node) {
        node->height = 1 + std::max(height_of(node->left), height_of(node->right));
    }

    // Positivo: pesa más la izquierda. Negativo: pesa más la derecha.
    static int balance_of(const Node* node) { return height_of(node->left) - height_of(node->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;
    }

    // Actualiza la altura de node, lo repara si hace falta, y devuelve la raíz de su subárbol.
    // Los dos hijos de node ya tienen su altura correcta.
    Node* rebalance(Node* node) {
        update_height(node);
        int balance = balance_of(node);
        if (balance > 1) {
            return repair_left_heavy(node);
        }
        if (balance < -1) {
            return repair_right_heavy(node);
        }
        return node;
    }

    Node* repair_left_heavy(Node* x) {
        Node* y = x->left;
        // Con balance 0 en y, que solo pasa al borrar, basta la rotación simple.
        if (balance_of(y) < 0) {
            last_repairs_.push_back({x->key, "izquierda-derecha"});
            x->left = rotate_left(y);
            update_height(y);
            update_height(x->left);
        } else {
            last_repairs_.push_back({x->key, "izquierda-izquierda"});
        }
        Node* top = rotate_right(x);
        update_height(x);
        update_height(top);
        return top;
    }

    Node* repair_right_heavy(Node* x) {
        Node* y = x->right;
        if (balance_of(y) > 0) {
            last_repairs_.push_back({x->key, "derecha-izquierda"});
            x->right = rotate_right(y);
            update_height(y);
            update_height(x->right);
        } else {
            last_repairs_.push_back({x->key, "derecha-derecha"});
        }
        Node* top = rotate_left(x);
        update_height(x);
        update_height(top);
        return top;
    }

    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 rebalance(node);
    }

    Node* erase_at(Node* node, const K& key, bool& erased) {
        if (node == nullptr) {
            return nullptr;
        }
        if (key < node->key) {
            node->left = erase_at(node->left, key, erased);
        } else if (node->key < key) {
            node->right = erase_at(node->right, key, erased);
        } else {
            erased = true;
            if (node->left == nullptr || node->right == nullptr) {
                Node* child = node->left != nullptr ? node->left : node->right;
                delete node;
                return child;
            }
            // Dos hijos: se copia el sucesor, y se quita su nodo del subárbol derecho.
            const Node* next = node->right;
            while (next->left != nullptr) {
                next = next->left;
            }
            node->key = next->key;
            bool ignored = false;
            node->right = erase_at(node->right, node->key, ignored);
        }
        return rebalance(node);
    }

    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->left != nullptr || node->right != nullptr) {
            out << "(";
            write_shape(node->left, out);
            out << ",";
            write_shape(node->right, out);
            out << ")";
        }
    }

    // Cuenta los nodos, y comprueba el orden, cada altura guardada y la condición AVL.
    static std::size_t count_checked(const Node* node, const K* low, const K* high) {
        if (node == nullptr) {
            return 0;
        }
        assert(low == nullptr || *low < node->key);
        assert(high == nullptr || node->key < *high);
        std::size_t count = 1 + count_checked(node->left, low, &node->key)
                              + count_checked(node->right, &node->key, high);
        assert(node->height == 1 + std::max(height_of(node->left), height_of(node->right)));
        assert(std::abs(balance_of(node)) <= 1);
        return count;
    }

    Node* root_ = nullptr;
    std::size_t size_ = 0;
    long rotations_ = 0;
    std::vector<Repair> last_repairs_;
};

#endif
