// Pruebas públicas del laboratorio. Se corre una por proceso: lab_tests nombre.
#include <algorithm>
#include <cassert>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <queue>
#include <random>
#include <set>
#include <string>
#include <utility>
#include <vector>

#include "graph.h"
#include "maze.h"
#include "naive_union_find.h"
#include "solver.h"
#include "union_find.h"

namespace {

int depth_of(const UnionFind& sets, int x) {
    int steps = 0;
    while (sets.parent_of(x) != x) {
        x = sets.parent_of(x);
        steps++;
    }
    return steps;
}

int floor_log2(int n) {
    int result = 0;
    while (n > 1) {
        n /= 2;
        result++;
    }
    return result;
}

// ---------- Paso 1: UnionFind ----------

void new_sets_are_separate() {
    UnionFind sets(5);
    assert(sets.size() == 5);
    assert(sets.component_count() == 5);
    assert(sets.check_invariants());
    for (int x = 0; x < 5; x++) {
        assert(sets.find(x) == x);
        assert(sets.parent_of(x) == x);
        assert(sets.rank_of(x) == 0);
    }
    assert(!sets.connected(0, 1));
    assert(sets.connected(3, 3));
}

void unite_reports_and_counts() {
    UnionFind sets(4);
    assert(sets.unite(0, 1));
    assert(sets.component_count() == 3);
    assert(sets.connected(0, 1));
    assert(!sets.connected(1, 2));
    assert(!sets.unite(1, 0));
    assert(sets.component_count() == 3);
    assert(!sets.unite(2, 2));
    assert(sets.component_count() == 3);
    assert(sets.unite(1, 2));
    assert(sets.connected(0, 2));
    assert(sets.component_count() == 2);
    assert(sets.check_invariants());
}

// unir a través de un elemento que no es raíz no debe perder a nadie
void unite_uses_roots() {
    UnionFind sets(6);
    sets.unite(0, 1);
    sets.unite(2, 3);
    sets.unite(0, 2);
    sets.unite(4, 5);
    int outsider = -1;
    for (int x = 0; x < 4; x++) {
        if (sets.parent_of(x) != x) {
            outsider = x;  // cualquiera que no sea raíz
        }
    }
    assert(outsider != -1);
    sets.unite(outsider, 4);
    assert(sets.component_count() == 1);
    for (int x = 0; x < 6; x++) {
        assert(sets.connected(0, x));
    }
    assert(sets.check_invariants());
}

void rank_rises_only_on_ties() {
    UnionFind sets(8);
    sets.unite(0, 1);
    int root = sets.find(0);
    assert(sets.rank_of(root) == 1);
    assert(sets.rank_of(root == 0 ? 1 : 0) == 0);
    // Un rango 0 cuelga de un rango 1: la raíz no cambia y su rango tampoco.
    sets.unite(2, 0);
    assert(sets.find(2) == root);
    assert(sets.rank_of(root) == 1);
    sets.unite(root, 3);
    assert(sets.find(3) == root);
    assert(sets.rank_of(root) == 1);
    // Dos de rango 1: sube a 2.
    sets.unite(4, 5);
    sets.unite(0, 4);
    assert(sets.rank_of(sets.find(0)) == 2);
    assert(sets.check_invariants());
}

// Una unión que pone el árbol pequeño arriba crecería hasta n - 1.
void height_stays_logarithmic() {
    for (int n : {16, 256, 1024}) {
        UnionFind sets(n);
        for (int i = 0; i + 1 < n; i++) {
            sets.unite(i + 1, i);
            int height = 0;
            for (int x = 0; x <= i + 1; x++) {
                height = std::max(height, depth_of(sets, x));
            }
            assert(height <= floor_log2(n));
        }
        assert(sets.check_invariants());
    }
    // El peor armado: uniones de árboles iguales, sin ningún find de por medio.
    int n = 1024;
    UnionFind sets(n);
    for (int size = 1; size < n; size *= 2) {
        for (int start = 0; start + size < n; start += 2 * size) {
            sets.unite(start, start + size);
        }
    }
    assert(sets.rank_of(sets.find(0)) == 10);
    assert(sets.check_invariants());
}

// Arma un árbol alto, de abajo hacia arriba, y devuelve su nodo más hondo.
int deepest_of(const UnionFind& sets) {
    int deepest = 0;
    for (int x = 0; x < sets.size(); x++) {
        if (depth_of(sets, x) > depth_of(sets, deepest)) {
            deepest = x;
        }
    }
    return deepest;
}

int root_without_compressing(const UnionFind& sets, int x) {
    while (sets.parent_of(x) != x) {
        x = sets.parent_of(x);
    }
    return x;
}

UnionFind tall_tree(int n) {
    UnionFind sets(n);
    for (int size = 1; size < n; size *= 2) {
        for (int start = 0; start + size < n; start += 2 * size) {
            // Se une por las raíces para que la compresión no aplane nada mientras se arma.
            sets.unite(root_without_compressing(sets, start), root_without_compressing(sets, start + size));
        }
    }
    return sets;
}

void compression_flattens_the_whole_path() {
    UnionFind sets = tall_tree(64);
    int deepest = deepest_of(sets);
    int height = depth_of(sets, deepest);
    assert(height == 6);
    std::vector<int> path;
    for (int x = deepest; sets.parent_of(x) != x; x = sets.parent_of(x)) {
        path.push_back(x);
    }
    int root = sets.find(deepest);
    for (int x : path) {
        assert(sets.parent_of(x) == root);
    }
    assert(depth_of(sets, deepest) == 1);
    assert(sets.check_invariants());
}

void compression_leaves_other_branches() {
    UnionFind sets = tall_tree(64);
    int deepest = deepest_of(sets);
    std::vector<int> before;
    for (int x = 0; x < sets.size(); x++) {
        before.push_back(sets.parent_of(x));
    }
    std::set<int> on_path;
    for (int x = deepest; sets.parent_of(x) != x; x = sets.parent_of(x)) {
        on_path.insert(x);
    }
    sets.find(deepest);
    for (int x = 0; x < sets.size(); x++) {
        if (on_path.count(x) == 0) {
            assert(sets.parent_of(x) == before[x]);
        }
    }
}

void compression_keeps_ranks() {
    UnionFind sets = tall_tree(64);
    std::vector<int> before;
    for (int x = 0; x < sets.size(); x++) {
        before.push_back(sets.rank_of(x));
    }
    for (int x = 0; x < sets.size(); x++) {
        sets.find(x);
    }
    for (int x = 0; x < sets.size(); x++) {
        assert(sets.rank_of(x) == before[x]);
    }
    assert(sets.check_invariants());
}

void random_against_naive() {
    std::mt19937 rng(14);
    for (int n : {1, 2, 10, 100, 500}) {
        UnionFind sets(n);
        NaiveUnionFind naive(n);
        for (int step = 0; step < 4 * n + 20; step++) {
            int a = static_cast<int>(rng() % n);
            int b = static_cast<int>(rng() % n);
            if (rng() % 2 == 0) {
                assert(sets.unite(a, b) == naive.unite(a, b));
            } else {
                assert(sets.connected(a, b) == naive.connected(a, b));
            }
            assert(sets.component_count() == naive.component_count());
            if (step % 16 == 0) {
                assert(sets.check_invariants());
            }
        }
        for (int a = 0; a < std::min(n, 40); a++) {
            for (int b = 0; b < std::min(n, 40); b++) {
                assert(sets.connected(a, b) == naive.connected(a, b));
            }
        }
        assert(sets.check_invariants());
    }
}

// ---------- Paso 2: generar ----------

// Revisa con una búsqueda propia: no depende de lo que el estudiante escribe en el paso 3.
void assert_connected_tree(const Maze& maze) {
    int n = maze.cell_count();
    assert(static_cast<int>(maze.passages().size()) == n - 1);
    std::vector<std::vector<int>> lists(n);
    for (const Wall& wall : maze.passages()) {
        lists[wall.a].push_back(wall.b);
        lists[wall.b].push_back(wall.a);
    }
    std::vector<char> seen(n, 0);
    std::vector<int> stack{0};
    seen[0] = 1;
    int reached = 0;
    while (!stack.empty()) {
        int u = stack.back();
        stack.pop_back();
        reached++;
        for (int v : lists[u]) {
            if (!seen[v]) {
                seen[v] = 1;
                stack.push_back(v);
            }
        }
    }
    assert(reached == n);
}

void generator_makes_a_tree() {
    for (auto [width, height] : std::vector<std::pair<int, int>>{{1, 1}, {1, 8}, {8, 1}, {2, 2}, {7, 5}, {30, 20}}) {
        for (unsigned seed = 1; seed <= 5; seed++) {
            std::mt19937 rng(seed);
            Maze maze(width, height);
            generate_maze(maze, rng);
            assert_connected_tree(maze);
        }
    }
}

void generator_opens_only_real_walls() {
    std::mt19937 rng(3);
    Maze maze(9, 6);
    generate_maze(maze, rng);
    std::set<std::pair<int, int>> walls;
    for (const Wall& wall : maze.all_walls()) {
        walls.insert({wall.a, wall.b});
    }
    std::set<std::pair<int, int>> opened;
    for (const Wall& wall : maze.passages()) {
        assert(walls.count({wall.a, wall.b}) == 1);
        assert(opened.insert({wall.a, wall.b}).second);
    }
}

void generator_depends_on_seed() {
    std::set<std::string> drawings;
    for (unsigned seed = 1; seed <= 8; seed++) {
        std::mt19937 rng(seed);
        Maze maze(10, 10);
        generate_maze(maze, rng);
        drawings.insert(maze.render());
    }
    assert(drawings.size() == 8);
    // Misma semilla, mismo laberinto.
    std::mt19937 first(5);
    std::mt19937 second(5);
    Maze a(10, 10);
    Maze b(10, 10);
    generate_maze(a, first);
    generate_maze(b, second);
    assert(a.render() == b.render());
}

// ---------- Paso 3: grafo y búsquedas ----------

void assert_valid_path(const Graph& graph, const std::vector<int>& path, int source, int target) {
    assert(!path.empty());
    assert(path.front() == source);
    assert(path.back() == target);
    for (std::size_t i = 0; i + 1 < path.size(); i++) {
        assert(graph.has_edge(path[i], path[i + 1]));
    }
    std::set<int> distinct(path.begin(), path.end());
    assert(distinct.size() == path.size());
}

void assert_visited_clean(const std::vector<int>& visited, int source) {
    assert(!visited.empty());
    assert(visited.front() == source);
    std::set<int> distinct(visited.begin(), visited.end());
    assert(distinct.size() == visited.size());
}

// Referencia: la distancia en aristas, bajando valores hasta que nada cambia.
std::vector<int> brute_distances(const Graph& graph, int source) {
    int n = graph.node_count();
    const int unreached = n + 1;
    std::vector<int> distance(n, unreached);
    distance[source] = 0;
    bool changed = true;
    while (changed) {
        changed = false;
        for (int u = 0; u < n; u++) {
            for (const Edge& edge : graph.neighbors(u)) {
                if (distance[u] != unreached && distance[u] + 1 < distance[edge.to]) {
                    distance[edge.to] = distance[u] + 1;
                    changed = true;
                }
            }
        }
    }
    return distance;
}

Graph random_graph(int n, int edges, std::mt19937& rng) {
    Graph graph(n);
    std::set<std::pair<int, int>> seen;
    for (int i = 0; i < edges; i++) {
        int u = static_cast<int>(rng() % n);
        int v = static_cast<int>(rng() % n);
        if (u > v) {
            std::swap(u, v);
        }
        if (u != v && seen.insert({u, v}).second) {
            graph.add_edge(u, v);
        }
    }
    return graph;
}

// El grafo de la clase: de 0 a 6 hay caminos de 4 y de 6 aristas.
Graph class_graph() {
    Graph graph(7);
    graph.add_edge(0, 1);
    graph.add_edge(0, 2);
    graph.add_edge(1, 3);
    graph.add_edge(2, 3);
    graph.add_edge(2, 4);
    graph.add_edge(3, 5);
    graph.add_edge(4, 5);
    graph.add_edge(5, 6);
    return graph;
}

void maze_graph_has_each_passage() {
    std::mt19937 rng(2);
    Maze maze(6, 4);
    generate_maze(maze, rng);
    open_extra_walls(maze, 3, rng);
    Graph graph = maze_graph(maze);
    assert(graph.node_count() == 24);
    assert(graph.edge_count() == maze.passages().size());
    assert(!graph.directed());
    assert(graph.check_invariants());
    for (const Wall& wall : maze.passages()) {
        assert(graph.has_edge(wall.a, wall.b));
        assert(graph.has_edge(wall.b, wall.a));
    }
}

void maze_graph_keeps_passage_order() {
    Maze maze(3, 1);
    maze.open({1, 2});
    maze.open({0, 1});
    Graph graph = maze_graph(maze);
    assert(graph.neighbors(1)[0].to == 2);
    assert(graph.neighbors(1)[1].to == 0);
}

void bfs_small_case() {
    Graph graph = class_graph();
    SearchResult result = bfs_search(graph, 0, 6);
    assert_valid_path(graph, result.path, 0, 6);
    assert(result.path.size() == 5);
    assert((result.visited == std::vector<int>{0, 1, 2, 3, 4, 5, 6}));
}

void bfs_stops_when_marking_target() {
    Graph graph = class_graph();
    SearchResult result = bfs_search(graph, 0, 3);
    assert((result.path == std::vector<int>{0, 1, 3}));
    assert((result.visited == std::vector<int>{0, 1, 2, 3}));
}

void bfs_source_is_target() {
    Graph graph = class_graph();
    SearchResult result = bfs_search(graph, 4, 4);
    assert((result.path == std::vector<int>{4}));
    assert((result.visited == std::vector<int>{4}));
}

void bfs_unreachable() {
    Graph graph(5);
    graph.add_edge(0, 1);
    graph.add_edge(1, 2);
    graph.add_edge(3, 4);
    SearchResult result = bfs_search(graph, 0, 4);
    assert(result.path.empty());
    assert((result.visited == std::vector<int>{0, 1, 2}));
}

void bfs_random_against_brute_force() {
    std::mt19937 rng(14);
    for (int round = 0; round < 60; round++) {
        int n = 2 + static_cast<int>(rng() % 40);
        Graph graph = random_graph(n, n + static_cast<int>(rng() % (2 * n)), rng);
        int source = static_cast<int>(rng() % n);
        std::vector<int> distance = brute_distances(graph, source);
        for (int target = 0; target < n; target++) {
            SearchResult result = bfs_search(graph, source, target);
            assert_visited_clean(result.visited, source);
            if (distance[target] == n + 1) {
                assert(result.path.empty());
            } else {
                assert_valid_path(graph, result.path, source, target);
                assert(static_cast<int>(result.path.size()) - 1 == distance[target]);
                assert(result.visited.back() == target);
            }
        }
    }
}

// Si se marca al sacar, un nodo entra muchas veces a la cola: aquí, una cantidad que se duplica con
// cada rombo. Esta prueba no terminaría, o se quedaría sin memoria.
void bfs_marks_when_pushing() {
    const int diamonds = 40;
    Graph graph(3 * diamonds + 1);
    for (int i = 0; i < diamonds; i++) {
        int top = 3 * i;
        graph.add_edge(top, top + 1);
        graph.add_edge(top, top + 2);
        graph.add_edge(top + 1, top + 3);
        graph.add_edge(top + 2, top + 3);
    }
    SearchResult result = bfs_search(graph, 0, 3 * diamonds);
    assert(static_cast<int>(result.path.size()) == 2 * diamonds + 1);
    assert_visited_clean(result.visited, 0);
    assert(static_cast<int>(result.visited.size()) == 3 * diamonds + 1);
}

void dfs_small_case() {
    Graph graph = class_graph();
    SearchResult result = dfs_search(graph, 0, 6);
    assert((result.path == std::vector<int>{0, 1, 3, 2, 4, 5, 6}));
    assert((result.visited == std::vector<int>{0, 1, 3, 2, 4, 5, 6}));
}

// Retrocede de un callejón: 3 se visita, pero no está en el camino.
void dfs_path_skips_dead_ends() {
    Graph graph(5);
    graph.add_edge(0, 1);
    graph.add_edge(1, 3);
    graph.add_edge(1, 2);
    graph.add_edge(2, 4);
    SearchResult result = dfs_search(graph, 0, 4);
    assert((result.path == std::vector<int>{0, 1, 2, 4}));
    assert((result.visited == std::vector<int>{0, 1, 3, 2, 4}));
}

void dfs_source_is_target() {
    Graph graph = class_graph();
    SearchResult result = dfs_search(graph, 4, 4);
    assert((result.path == std::vector<int>{4}));
    assert((result.visited == std::vector<int>{4}));
}

void dfs_unreachable() {
    Graph graph(5);
    graph.add_edge(0, 1);
    graph.add_edge(1, 2);
    graph.add_edge(3, 4);
    SearchResult result = dfs_search(graph, 0, 4);
    assert(result.path.empty());
    assert((result.visited == std::vector<int>{0, 1, 2}));
}

void dfs_random_paths_are_valid() {
    std::mt19937 rng(21);
    for (int round = 0; round < 60; round++) {
        int n = 2 + static_cast<int>(rng() % 40);
        Graph graph = random_graph(n, n + static_cast<int>(rng() % (2 * n)), rng);
        int source = static_cast<int>(rng() % n);
        std::vector<int> distance = brute_distances(graph, source);
        for (int target = 0; target < n; target++) {
            SearchResult result = dfs_search(graph, source, target);
            assert_visited_clean(result.visited, source);
            if (distance[target] == n + 1) {
                assert(result.path.empty());
            } else {
                assert_valid_path(graph, result.path, source, target);
                assert(static_cast<int>(result.path.size()) - 1 >= distance[target]);
                assert(result.visited.back() == target);
            }
        }
    }
}

// Se corre con 128 KB de pila: una recursión de 6000 niveles no alcanza.
void dfs_long_corridor() {
    const int n = 6000;
    Graph graph(n);
    for (int i = 0; i + 1 < n; i++) {
        graph.add_edge(i, i + 1);
    }
    SearchResult result = dfs_search(graph, 0, n - 1);
    assert(static_cast<int>(result.path.size()) == n);
    assert(static_cast<int>(result.visited.size()) == n);
}

// ---------- Paso 3 y 4: juntos ----------

// Sin ciclos hay un solo camino: los dos recorridos deben dar exactamente el mismo.
void perfect_maze_gives_the_same_path() {
    for (unsigned seed = 1; seed <= 20; seed++) {
        std::mt19937 rng(seed);
        Maze maze(20, 12);
        generate_maze(maze, rng);
        Graph graph = maze_graph(maze);
        int target = maze.cell_count() - 1;
        SearchResult by_bfs = bfs_search(graph, 0, target);
        SearchResult by_dfs = dfs_search(graph, 0, target);
        assert_valid_path(graph, by_bfs.path, 0, target);
        assert(by_bfs.path == by_dfs.path);
    }
}

// Con ciclos, BFS nunca pierde, y en esta muestra DFS da un camino más largo varias veces.
void cycles_separate_the_paths() {
    int longer = 0;
    for (unsigned seed = 1; seed <= 20; seed++) {
        std::mt19937 rng(seed);
        Maze maze(20, 12);
        generate_maze(maze, rng);
        open_extra_walls(maze, 40, rng);
        Graph graph = maze_graph(maze);
        int target = maze.cell_count() - 1;
        SearchResult by_bfs = bfs_search(graph, 0, target);
        SearchResult by_dfs = dfs_search(graph, 0, target);
        assert_valid_path(graph, by_bfs.path, 0, target);
        assert_valid_path(graph, by_dfs.path, 0, target);
        assert(by_bfs.path.size() <= by_dfs.path.size());
        longer += by_bfs.path.size() < by_dfs.path.size() ? 1 : 0;
    }
    assert(longer >= 10);
}

struct Case {
    const char* name;
    void (*run)();
};

const Case CASES[] = {
    {"new_sets_are_separate", new_sets_are_separate},
    {"unite_reports_and_counts", unite_reports_and_counts},
    {"unite_uses_roots", unite_uses_roots},
    {"rank_rises_only_on_ties", rank_rises_only_on_ties},
    {"height_stays_logarithmic", height_stays_logarithmic},
    {"compression_flattens_the_whole_path", compression_flattens_the_whole_path},
    {"compression_leaves_other_branches", compression_leaves_other_branches},
    {"compression_keeps_ranks", compression_keeps_ranks},
    {"random_against_naive", random_against_naive},
    {"generator_makes_a_tree", generator_makes_a_tree},
    {"generator_opens_only_real_walls", generator_opens_only_real_walls},
    {"generator_depends_on_seed", generator_depends_on_seed},
    {"maze_graph_has_each_passage", maze_graph_has_each_passage},
    {"maze_graph_keeps_passage_order", maze_graph_keeps_passage_order},
    {"bfs_small_case", bfs_small_case},
    {"bfs_stops_when_marking_target", bfs_stops_when_marking_target},
    {"bfs_source_is_target", bfs_source_is_target},
    {"bfs_unreachable", bfs_unreachable},
    {"bfs_random_against_brute_force", bfs_random_against_brute_force},
    {"bfs_marks_when_pushing", bfs_marks_when_pushing},
    {"dfs_small_case", dfs_small_case},
    {"dfs_path_skips_dead_ends", dfs_path_skips_dead_ends},
    {"dfs_source_is_target", dfs_source_is_target},
    {"dfs_unreachable", dfs_unreachable},
    {"dfs_random_paths_are_valid", dfs_random_paths_are_valid},
    {"dfs_long_corridor", dfs_long_corridor},
    {"perfect_maze_gives_the_same_path", perfect_maze_gives_the_same_path},
    {"cycles_separate_the_paths", cycles_separate_the_paths},
};

}  // namespace

int main(int argc, char* argv[]) {
    if (argc == 2) {
        for (const Case& test : CASES) {
            if (std::strcmp(test.name, argv[1]) == 0) {
                test.run();
                return 0;
            }
        }
    }
    std::cerr << "uso: lab_tests nombre\n";
    return 2;
}