#include "prim.hpp"
#include "rng.hpp"
#include "shake.hpp"
#include "solution.hpp"
#include "tsplib_graph.hpp"
#include "vnd1.hpp"
#include "vnd2.hpp"
#include "vnd3.hpp"
#include "vnd4.hpp"

#include <chrono>
#include <cstdint>
#include <cstdlib>
#include <cmath>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <fstream>
#include <string>
#include <vector>

int main(int argc, char* argv[]) {
    if (argc < 6) {
        std::cerr << "Usage: " << argv[0]
                  << " <d> <k> <filename> <k_max> <iters> [seed] [--no-vnd1] [--no-vnd2] [--no-vnd3] [--no-vnd4] [--vnd1bi] [--vnd2bi] [--vnd3bi] [--vnd4bi] [--shake2] [--shake3] [--vnd-order ORDER] [--max-time SECS] [--csv] [--best-known VAL]\n";
        return 1;
    }

    const int d = std::atoi(argv[1]);
    const int k = std::atoi(argv[2]);
    const std::string filename = argv[3];
    const int k_max = std::atoi(argv[4]);
    const int iters = std::atoi(argv[5]);

    uint32_t seed = 0u;
    bool use_vnd1 = true;
    bool use_vnd2 = true;
    bool use_vnd3 = true;
    bool use_vnd4 = true;
    bool bi_vnd1 = false;
    bool bi_vnd2 = false;
    bool bi_vnd3 = false;
    bool bi_vnd4 = false;
    bool output_to_file = false;
    int shake_type = 1;
    std::vector<int> vnd_order = {2, 1, 3, 4}; // default order
    bool csv_mode = false;
    double best_known = -1.0;
    double best_known_tol = 0.0;
    double max_time = -1.0;

    for (int i = 6; i < argc; ++i) {
        if (std::strcmp(argv[i], "--no-vnd1") == 0) {
            use_vnd1 = false;
        } else if (std::strcmp(argv[i], "--no-vnd2") == 0) {
            use_vnd2 = false;
        } else if (std::strcmp(argv[i], "--no-vnd3") == 0) {
            use_vnd3 = false;
        } else if (std::strcmp(argv[i], "--no-vnd4") == 0) {
            use_vnd4 = false;
        } else if (std::strcmp(argv[i], "--vnd1bi") == 0) {
            bi_vnd1 = true;
        } else if (std::strcmp(argv[i], "--vnd2bi") == 0) {
            bi_vnd2 = true;
        } else if (std::strcmp(argv[i], "--vnd3bi") == 0) {
            bi_vnd3 = true;
        } else if (std::strcmp(argv[i], "--vnd4bi") == 0) {
            bi_vnd4 = true;
        } else if (std::strcmp(argv[i], "--shake2") == 0) {
            shake_type = 2;
        } else if (std::strcmp(argv[i], "--shake3") == 0) {
            shake_type = 3;
        } else if (std::strcmp(argv[i], "--vnd-order") == 0) {
            if (i + 1 < argc) {
                const char* s = argv[++i];
                vnd_order.clear();
                for (int j = 0; s[j]; ++j)
                    if (s[j] >= '1' && s[j] <= '4')
                        vnd_order.push_back(s[j] - '0');
            }
        } else if (std::strcmp(argv[i], "--csv") == 0) {
            csv_mode = true;
        } else if (std::strcmp(argv[i], "--max-time") == 0) {
            if (i + 1 < argc) max_time = std::atof(argv[++i]);
        } else if (std::strcmp(argv[i], "--best-known") == 0) {
            if (i + 1 < argc) {
                const char* bk_str = argv[++i];
                best_known = std::atof(bk_str);
                const char* dot = std::strchr(bk_str, '.');
                int r = 0;
                if (dot) r = static_cast<int>(std::strlen(dot + 1));
                best_known_tol = 5.0 * std::pow(10.0, -(r + 1));
            }
        } else if (std::strcmp(argv[i], "--output-edges") == 0) {
            output_to_file = true;
        } else {
            seed = static_cast<uint32_t>(std::stoul(argv[i]));
        }
    }

    const bool use_txt = filename.size() >= 4 &&
                         filename.compare(filename.size() - 4, 4, ".txt") == 0;
    TSPLIBGraph g = use_txt ? TSPLIBGraph::load_txt(filename)
                            : TSPLIBGraph::load(filename);
    const int n = g.size();

    if (k > n) {
        std::cerr << "Error: need k <= n  (got k=" << k
                  << ", n=" << n << ")\n";
        return 1;
    }

    RNG rng(seed);

    Solution sol = prim_original(g, k, d, rng);
    Solution best = sol;
    int best_iter = 0;
    int64_t n_shake = 0, n_vnd1 = 0, n_vnd2 = 0, n_vnd3 = 0, n_vnd4 = 0, n_accept = 0;

    using Clock = std::chrono::steady_clock;
    auto t_start = Clock::now();
    auto t_best = t_start;

    auto time_limit = std::chrono::duration<double>(max_time);

    for (int iter = 0; iter < iters; ++iter) {
        if (max_time >= 0.0 && Clock::now() - t_start >= time_limit)
            goto done;
        int k_vns = 1;
        while (k_vns <= k_max) {
            if      (shake_type == 3) shake3(sol, k_vns, k, d, rng);
            else if (shake_type == 2) shake2(sol, k_vns, k, d, rng);
            else                      shake(sol, k_vns, rng);
            ++n_shake;

            // VND phase: loop until no VND improves
            bool vnd_improved = true;
            while (vnd_improved) {
                vnd_improved = false;
                if (max_time >= 0.0 && Clock::now() - t_start >= time_limit)
                    goto done;
                for (int vnd_id : vnd_order) {
                    bool improved = false;
                    switch (vnd_id) {
                    case 1:
                        if (!use_vnd1) break;
                        ++n_vnd1;
                        improved = bi_vnd1 ? vnd1bi(sol) : vnd1(sol);
                        break;
                    case 2:
                        if (!use_vnd2) break;
                        ++n_vnd2;
                        improved = bi_vnd2 ? vnd2bi(sol, d, rng) : vnd2(sol, d, rng);
                        break;
                    case 3:
                        if (!use_vnd3) break;
                        ++n_vnd3;
                        improved = bi_vnd3 ? vnd3bi(sol, d, rng) : vnd3(sol, d, rng);
                        break;
                    case 4:
                        if (!use_vnd4) break;
                        ++n_vnd4;
                        improved = bi_vnd4 ? vnd4bi(sol, d, rng) : vnd4(sol, d, rng);
                        break;
                    }
                    if (improved) {
                        vnd_improved = true;

                        // Acceptance
                        if (sol.cost < best.cost) {
                            best = sol;
                            best_iter = iter;
                            ++n_accept;
                            t_best = Clock::now();
                            if (best_known >= 0.0 && best.cost <= best_known + best_known_tol)
                                goto done;
                            if (max_time >= 0.0 && Clock::now() - t_start >= time_limit)
                                goto done;
                            k_vns = 1;
                        } else {
                            ++k_vns;
                            sol = best;
                        }

                        break; // restart VND sequence from first in order
                    }
                }
            }
        }
    }
    done:

    auto t_end = Clock::now();
    double time_best = std::chrono::duration<double>(t_best - t_start).count();
    double time_total = std::chrono::duration<double>(t_end - t_start).count();

    // ── Validation ─────────────────────────────────────────────────────────
    int node_count = 0;
    int max_degree = 0;
    bool connected = true;
    for (int v = 0; v < n; ++v) {
        if (best.degree[v] > 0) ++node_count;
        if (best.degree[v] > max_degree) max_degree = best.degree[v];
    }
    if (node_count > 0) {
        std::vector<bool> visited(n, false);
        int root = -1;
        for (int v = 0; v < n; ++v)
            if (best.degree[v] > 0) { root = v; break; }
        std::vector<int> stack = {root};
        visited[root] = true;
        int visit_count = 1;
        while (!stack.empty()) {
            int u = stack.back(); stack.pop_back();
            for (int w : best.adj[u]) {
                if (!visited[w]) {
                    visited[w] = true;
                    ++visit_count;
                    stack.push_back(w);
                }
            }
        }
        connected = (visit_count == node_count);
    }
    bool valid = (node_count == k) && (max_degree <= d) && connected;

    std::string vnd_order_str;
    for (int v : vnd_order) vnd_order_str += std::to_string(v);

    std::cout << std::fixed << std::setprecision(4);

    if (csv_mode) {
        std::cout << filename
                  << "," << n
                  << "," << d
                  << "," << k
                  << "," << k_max
                  << "," << iters
                  << "," << seed
                  << "," << max_time
                  << "," << best_known
                  << "," << shake_type
                  << "," << (use_vnd1 ? 1 : 0)
                  << "," << (bi_vnd1 ? 1 : 0)
                  << "," << (use_vnd2 ? 1 : 0)
                  << "," << (bi_vnd2 ? 1 : 0)
                  << "," << (use_vnd3 ? 1 : 0)
                  << "," << (bi_vnd3 ? 1 : 0)
                  << "," << (use_vnd4 ? 1 : 0)
                  << "," << (bi_vnd4 ? 1 : 0)
                  << "," << vnd_order_str
                  << "," << best_iter
                  << "," << best.edges.size()
                  << "," << best.cost
                  << "," << time_best
                  << "," << time_total
                  << "," << n_shake
                  << "," << n_vnd1
                  << "," << n_vnd2
                  << "," << n_vnd3
                  << "," << n_vnd4
                  << "," << n_accept
                  << "," << (valid ? 1 : 0)
                  << "," << node_count
                  << "," << max_degree
                  << "," << (connected ? 1 : 0)
                  << "\n";
    } else {
        std::cout << "file:       " << filename << "\n"
                  << "n:          " << n << "\n"
                  << "d:          " << d << "\n"
                  << "k:          " << k << "\n"
                  << "k_max:      " << k_max << "\n"
                  << "iters:      " << iters << "\n"
                  << "seed:       " << seed << "\n"
                  << "max_time:   " << max_time << "s\n"
                  << "best_known: " << best_known << "\n"
                  << "shake:      " << shake_type << "\n"
                  << "use_vnd1:   " << (use_vnd1 ? "yes" : "no") << "\n"
                  << "bi_vnd1:    " << (bi_vnd1 ? "yes" : "no") << "\n"
                  << "use_vnd2:   " << (use_vnd2 ? "yes" : "no") << "\n"
                  << "bi_vnd2:    " << (bi_vnd2 ? "yes" : "no") << "\n"
                  << "use_vnd3:   " << (use_vnd3 ? "yes" : "no") << "\n"
                  << "bi_vnd3:    " << (bi_vnd3 ? "yes" : "no") << "\n"
                  << "use_vnd4:   " << (use_vnd4 ? "yes" : "no") << "\n"
                  << "bi_vnd4:    " << (bi_vnd4 ? "yes" : "no") << "\n"
                  << "vnd_order:  " << vnd_order_str << "\n"
                  << "best_iter:  " << best_iter << "\n"
                  << "edges:      " << best.edges.size() << "\n"
                  << "cost:       " << best.cost << "\n"
                  << "time_best:  " << time_best << "s\n"
                  << "time_total: " << time_total << "s\n"
                  << "n_shake:    " << n_shake << "\n"
                  << "n_vnd1:     " << n_vnd1 << "\n"
                  << "n_vnd2:     " << n_vnd2 << "\n"
                  << "n_vnd3:     " << n_vnd3 << "\n"
                  << "n_vnd4:     " << n_vnd4 << "\n"
                  << "n_accept:   " << n_accept << "\n"
                  << "valid:      " << (valid ? "yes" : "NO") << "\n"
                  << "nodes:      " << node_count << "\n"
                  << "max_degree: " << max_degree << "\n"
                  << "connected:  " << (connected ? "yes" : "no") << "\n";
    }
    
    if (output_to_file) {
        std::ofstream sol_file;
        sol_file.open ("solution_output.txt");
        for (auto edge_id : best.edges) {
            auto [u, v] = g.edge_nodes(edge_id);
            auto dist = g.dist(edge_id);   
            sol_file << u << "," << v << "," << dist << std::endl;
        }
        sol_file.close();        
    }

    return 0;
}
