#include "kcnf.hpp"

#include <algorithm>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <random>
#include <stdexcept>
#include <vector>

#ifdef _OPENMP
#include <omp.h>
#endif

namespace fs = std::filesystem;

const fs::path CNF_DIR = "./CNF";

long binom(long n, long k) {
    return (k > n)          ? 0 :
           (k == 0 || k == n) ? 1 :
           (k == 1 || k == n - 1) ? n :
           (k + k < n) ?
                (binom(n - 1, k - 1) * n) / k :
                (binom(n - 1, k) * n) / (n - k);
}

// get the (n choose k) possible variable choices at index idx
// example: 
std::vector<int> unrank_combination(long n, long k, long long idx) {
    std::vector<int> comb;
    comb.reserve(k);
    long start = 1;
    for (long i = 0; i < k; ++i) {
        for (long v = start; v <= n; ++v) {
            long long count = binom(n - v, k - i - 1);
            if (idx < count) {
                comb.push_back((int)v);
                start = v + 1;
                break;
            } else {
                idx -= count;
            }
        }
    }
    return comb;
}

// sample c indices in [0,total] for our c = # of clauses
// this is the heart of the program as we are flattening the problem with our sampled indices and the unrank function!
std::vector<long long> sample_indices(long long total, long c) {
    std::vector<long long> idx;
    idx.reserve(c);

    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<long long> dist(0, total - 1);

    for (long i = 0; i < c; ++i) {
        idx.push_back(dist(gen));
    }

    return idx;
}

// decode one clause from global index (think of it like a hash function for better understanding!)
// idx = combo_id * 2^k + mask
Clause decode_clause(long n, long k, long long idx) {
    long long combo_id = idx >> k;
    int mask = idx & ((1 << k) - 1);

    std::vector<int> comb = unrank_combination(n, k, combo_id);

    Clause cl;
    cl.reserve(k);

    for (int i = 0; i < k; ++i) {
        int lit = comb[i];
        cl.push_back((mask & (1 << i)) ? lit : -lit);
    }

    return cl;
}

// parallel iteration to build clauses from indices
std::vector<Clause> build_clauses_parallel(
    long n,
    long k,
    const std::vector<long long>& indices) {

    std::vector<Clause> result(indices.size());

    #pragma omp parallel for schedule(dynamic)
    for (long i = 0; i < (long)indices.size(); ++i) {
        result[i] = decode_clause(n, k, indices[i]);
    }

    return result;
}

KCNF KCNF::generate_random(long n, long c, long k) {
    KCNF f;
    f.n = n;
    f.k = k;

    // test for logical limits
    long long total = binom(n,k) * (1LL << k);
    if (c > total) {
        throw std::runtime_error("Too many clauses requested.");
    }
    
    // test for hardware limits
    long long estimated_bytes =
    c * (sizeof(Clause) + k * sizeof(int));

    if (estimated_bytes > 500 * 1024 * 1024) { // 500MB
        throw std::runtime_error("Estimated memory too large");
    }

    auto indices = sample_indices(total, c);
    f.clauses = build_clauses_parallel(n, k, indices);

    return f;
}

bool KCNF::write_dimacs(const std::string& filename) const {
    fs::create_directories(CNF_DIR);

    fs::path path = CNF_DIR / filename;
    path = path.lexically_normal();

    std::ofstream out(path);
    if (!out) {
        std::cerr << "Error: Could not open file " << path << "\n";
        return false;
    }

    out << "p cnf " << n << " " << clauses.size() << "\n";

    for (const auto& cl : clauses) {
        for (int lit : cl) {
            out << lit << " ";
        }
        out << "0\n";
    }

    return true;
}