// Códigos sin prefijos comunes: Huffman, ideas que fallan y un árbitro exhaustivo.
#pragma once

#include <algorithm>
#include <cstddef>
#include <functional>
#include <limits>
#include <numeric>
#include <queue>
#include <set>
#include <utility>
#include <vector>

// Costo de un código: suma de frecuencia por profundidad.
inline long long code_cost(const std::vector<long long>& freq, const std::vector<int>& depth) {
    long long total = 0;
    for (std::size_t i = 0; i < freq.size(); i++) {
        total += freq[i] * depth[i];
    }
    return total;
}

struct HuffmanResult {
    std::vector<int> depth;                                  // profundidad de cada símbolo
    std::vector<std::pair<long long, long long>> merges;     // los dos pesos de cada unión
};

// Requiere al menos dos símbolos, todos con frecuencia positiva.
// A igual peso, sale primero el árbol creado antes.
inline HuffmanResult huffman(const std::vector<long long>& freq) {
    using Entry = std::pair<long long, int>;
    std::vector<int> parent(freq.size(), -1);
    std::priority_queue<Entry, std::vector<Entry>, std::greater<Entry>> heap;
    for (std::size_t i = 0; i < freq.size(); i++) {
        heap.push({freq[i], static_cast<int>(i)});
    }
    HuffmanResult result;
    while (heap.size() >= 2) {
        Entry first = heap.top();
        heap.pop();
        Entry second = heap.top();
        heap.pop();
        int node = static_cast<int>(parent.size());
        parent.push_back(-1);
        parent[static_cast<std::size_t>(first.second)] = node;
        parent[static_cast<std::size_t>(second.second)] = node;
        result.merges.push_back({first.first, second.first});
        heap.push({first.first + second.first, node});
    }
    for (std::size_t i = 0; i < freq.size(); i++) {
        int depth = 0;
        for (int node = static_cast<int>(i); parent[static_cast<std::size_t>(node)] != -1;
             node = parent[static_cast<std::size_t>(node)]) {
            depth++;
        }
        result.depth.push_back(depth);
    }
    return result;
}

inline long long merges_total(const std::vector<std::pair<long long, long long>>& merges) {
    long long total = 0;
    for (const auto& merge : merges) {
        total += merge.first + merge.second;
    }
    return total;
}

// Un árbol lleno sobre n hojas se arma con n - 1 uniones, en algún orden, y cada unión
// baja un nivel todas sus hojas. Esta búsqueda prueba todos los pares en cada paso, sin
// ninguna regla, y llama a visit con las profundidades de cada árbol. Solo con n pequeño.
template <typename Visit>
void all_full_trees(std::vector<std::vector<int>>& groups, std::vector<int>& depth, Visit& visit) {
    if (groups.size() == 1) {
        visit(depth);
        return;
    }
    for (std::size_t i = 0; i < groups.size(); i++) {
        for (std::size_t j = i + 1; j < groups.size(); j++) {
            std::vector<std::vector<int>> next;
            std::vector<int> joined = groups[i];
            joined.insert(joined.end(), groups[j].begin(), groups[j].end());
            for (std::size_t g = 0; g < groups.size(); g++) {
                if (g != i && g != j) next.push_back(groups[g]);
            }
            next.push_back(joined);
            for (int symbol : joined) depth[static_cast<std::size_t>(symbol)]++;
            all_full_trees(next, depth, visit);
            for (int symbol : joined) depth[static_cast<std::size_t>(symbol)]--;
        }
    }
}

inline long long brute_min_cost(const std::vector<long long>& freq) {
    std::vector<std::vector<int>> groups;
    for (std::size_t i = 0; i < freq.size(); i++) {
        groups.push_back({static_cast<int>(i)});
    }
    std::vector<int> depth(freq.size(), 0);
    long long best = std::numeric_limits<long long>::max();
    auto visit = [&](const std::vector<int>& d) { best = std::min(best, code_cost(freq, d)); };
    all_full_trees(groups, depth, visit);
    return best;
}

// Todas las formas de correr Huffman: en cada paso, cualquier par cuyos pesos sean los dos
// menores. Devuelve las profundidades distintas que se pueden obtener.
inline void huffman_ties_rec(const std::vector<long long>& weight, const std::vector<std::vector<int>>& groups,
                             std::vector<int>& depth, std::set<std::vector<int>>& out) {
    if (groups.size() == 1) {
        out.insert(depth);
        return;
    }
    std::vector<long long> sorted = weight;
    std::sort(sorted.begin(), sorted.end());
    for (std::size_t i = 0; i < groups.size(); i++) {
        for (std::size_t j = i + 1; j < groups.size(); j++) {
            long long low = std::min(weight[i], weight[j]);
            long long high = std::max(weight[i], weight[j]);
            if (low != sorted[0] || high != sorted[1]) continue;
            std::vector<long long> next_weight;
            std::vector<std::vector<int>> next_groups;
            for (std::size_t g = 0; g < groups.size(); g++) {
                if (g != i && g != j) {
                    next_weight.push_back(weight[g]);
                    next_groups.push_back(groups[g]);
                }
            }
            std::vector<int> joined = groups[i];
            joined.insert(joined.end(), groups[j].begin(), groups[j].end());
            next_weight.push_back(low + high);
            next_groups.push_back(joined);
            for (int symbol : joined) depth[static_cast<std::size_t>(symbol)]++;
            huffman_ties_rec(next_weight, next_groups, depth, out);
            for (int symbol : joined) depth[static_cast<std::size_t>(symbol)]--;
        }
    }
}

inline std::set<std::vector<int>> huffman_all_ties(const std::vector<long long>& freq) {
    std::vector<std::vector<int>> groups;
    for (std::size_t i = 0; i < freq.size(); i++) {
        groups.push_back({static_cast<int>(i)});
    }
    std::vector<int> depth(freq.size(), 0);
    std::set<std::vector<int>> out;
    huffman_ties_rec(freq, groups, depth, out);
    return out;
}

// Idea que falla: dividir de arriba hacia abajo. Se ordena de mayor a menor y se parte la
// lista donde las dos sumas queden más parejas (a igual diferencia, el primer corte).
inline void split_rec(const std::vector<long long>& freq, const std::vector<int>& ids, std::size_t lo, std::size_t hi,
                      int level, std::vector<int>& depth) {
    if (hi - lo == 1) {
        depth[static_cast<std::size_t>(ids[lo])] = level;
        return;
    }
    long long total = 0;
    for (std::size_t i = lo; i < hi; i++) total += freq[static_cast<std::size_t>(ids[i])];
    long long left = 0;
    long long best_gap = std::numeric_limits<long long>::max();
    std::size_t cut = lo + 1;
    for (std::size_t i = lo; i + 1 < hi; i++) {
        left += freq[static_cast<std::size_t>(ids[i])];
        long long gap = std::max(left, total - left) - std::min(left, total - left);
        if (gap < best_gap) {
            best_gap = gap;
            cut = i + 1;
        }
    }
    split_rec(freq, ids, lo, cut, level + 1, depth);
    split_rec(freq, ids, cut, hi, level + 1, depth);
}

inline std::vector<int> top_down_split(const std::vector<long long>& freq) {
    std::vector<int> ids(freq.size());
    std::iota(ids.begin(), ids.end(), 0);
    std::stable_sort(ids.begin(), ids.end(), [&](int a, int b) {
        return freq[static_cast<std::size_t>(a)] > freq[static_cast<std::size_t>(b)];
    });
    std::vector<int> depth(freq.size(), 0);
    split_rec(freq, ids, 0, freq.size(), 0, depth);
    return depth;
}

// Idea que falla: un código en cadena, 0, 10, 110, ..., con los más frecuentes arriba.
inline std::vector<int> chain_code(const std::vector<long long>& freq) {
    std::vector<int> ids(freq.size());
    std::iota(ids.begin(), ids.end(), 0);
    std::stable_sort(ids.begin(), ids.end(), [&](int a, int b) {
        return freq[static_cast<std::size_t>(a)] > freq[static_cast<std::size_t>(b)];
    });
    std::vector<int> depth(freq.size(), 0);
    for (std::size_t i = 0; i < ids.size(); i++) {
        depth[static_cast<std::size_t>(ids[i])] = static_cast<int>(std::min(i + 1, ids.size() - 1));
    }
    return depth;
}
