#include <algorithm>
#include <cmath>
#include <fstream>
#include <iterator>
#include <map>
#include <print>
#include <ranges>
#include <set>
#include <vector>

#include "common/getinputpath.h"
#include "common/istest.h"

const int STEPS = is_test() ? 10 : 1000;

struct Box {
    int x;
    int y;
    int z;

    float euclidean(Box& other) {
        return std::sqrt(std::pow((other.x - this->x), 2) +
                         std::pow((other.y - this->y), 2) +
                         std::pow((other.z - this->z), 2));
    }

    bool operator==(const Box other) const {
        return this->x == other.x && this->y == other.y && this->z == other.z;
    }

    bool operator<(const Box other) const {
        return this->x < other.x || (this->x == other.x && this->y < other.y) ||
               (this->x == other.x && this->y == other.y && this->z < other.z);
    }

    std::string to_string() { return std::format("{} {} {}", x, y, z); }
};

struct Connection {
    Box   a;
    Box   b;
    float distance;

    bool operator==(const Connection& other) const {
        return (other.a == this->a && other.b == this->b) ||
               (other.a == this->b && other.b == this->a);
    }

    bool contains(Box box) { return this->a == box || this->b == box; }

    bool connected(Connection connection) {
        return this->contains(connection.a) || this->contains(connection.b);
    }

    bool operator<(const Connection other) const {
        return this->distance < other.distance;
    }
};

Connection make_connection(Box& one, Box& other) {
    return Connection{.a = one, .b = other, .distance = one.euclidean(other)};
}

// standard DSU
class DSU {
    std::vector<int> size{};
    std::vector<int> parent{};

   public:
    DSU(int n) {
        parent.resize(n);
        size.resize(n);

        for (int i = 0; i < n; ++i) {
            parent[i] = i;
            size[i]   = 1;
        }
    }

    int find(int i) {
        return (parent[i] == i) ? i : (parent[i] = find(parent[i]));
    }

    void unite(int x, int y) {
        int s1 = find(x);
        int s2 = find(y);

        if (s1 == s2) {
            return;
        }
        if (size[s1] < size[s2]) {
            parent[s1] = s2;
            size[s2] += size[s1];
        } else if (size[s1] > size[s2]) {
            parent[s2] = s1;
            size[s1] += size[s2];
        } else {
            parent[s2] = s1;
            size[s1] += size[s2];
        }
    }

    std::vector<int> sizes() { return size; }
};

int main() {
    auto          path = get_input_path(DATA_FOLDER);
    std::ifstream data_ifstream(path);
    if (!data_ifstream.is_open()) {
        std::println("Cannot open the file at {}", path.string());
        return 1;
    }

    std::vector<Box> data{};
    for (std::string t; std::getline(data_ifstream, t);) {
        int x{};
        int y{};
        int z{};

        auto first_comma_idx  = t.find_first_of(',');
        auto x_str            = t.substr(0, first_comma_idx);
        auto yz_str           = t.substr(first_comma_idx + 1);
        auto second_comma_idx = yz_str.find_first_of(',');
        auto y_str            = yz_str.substr(0, second_comma_idx);
        auto z_str            = yz_str.substr(second_comma_idx + 1);

        x = std::stoi(x_str);
        y = std::stoi(y_str);
        z = std::stoi(z_str);

        data.push_back(Box{.x = x, .y = y, .z = z});
    }

    long long estimated_conns = data.size() * data.size() - data.size();
    std::print("Generating connections... est. {}", estimated_conns);

    // replacing this with an unordered set would make things faster
    std::set<Connection> connections{};
    for (auto box1 : data) {
        for (auto box2 : data) {
            if (box1 == box2) {
                continue;
            }

            Connection conn = make_connection(box1, box2);
            if (connections.contains(conn)) {  // connection already exists
                continue;
            }
            connections.insert(conn);
            std::print("\r");
        }
    }
    std::println();

    std::println("Sorting connections...");
    std::vector<Connection> connections_vec =
        connections | std::ranges::to<std::vector<Connection>>();
    std::ranges::sort(connections_vec, [](Connection one, Connection other) {
        return one.distance < other.distance;
    });

    std::println("Calculating three largest connected components...");
    DSU dsu(data.size());

    std::map<Box, size_t> data_idxs{};
    for (auto [idx, box] : data | std::ranges::views::enumerate) {
        data_idxs[box] = idx;
    }

    long long  counter = 0;
    Connection last_connected;
    long long  password_part_1;
    for (auto conn : connections_vec) {
        counter += 1;
        if (dsu.find(data_idxs[conn.a]) != dsu.find(data_idxs[conn.b])) {
            dsu.unite(data_idxs[conn.a], data_idxs[conn.b]);
            last_connected = conn;
        }
        if (counter == STEPS) {
            auto sizes = dsu.sizes();
            std::ranges::sort(sizes);
            std::ranges::reverse(sizes);

            std::println("Three biggest: {}, {}, {}", sizes[0], sizes[1],
                         sizes[2]);
            int s1          = sizes[0];
            int s2          = sizes[1];
            int s3          = sizes[2];
            password_part_1 = s1 * s2 * s3;
        }
    }

    long long password_part_2 = last_connected.a.x * last_connected.b.x;

    std::println();
    std::println("Eureka! {} / {}", password_part_1, password_part_2);

    return 0;
}