#include "kcnf.hpp"

#include <algorithm>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <random>
#include <stdexcept>
#include <unordered_set>
#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: unrank_combination(5,3,5) -> (1,4,5)
// the full combination list would look like: (1,2,3), (1,2,4), (1,2,5), (1,3,4), (1,3,5), (1,4,5) <- this is the combination on index 5
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::unordered_set<long long> set;
    set.reserve(c);

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

    while ((long)set.size() < c) {
        set.insert(dist(gen));
    }

    return std::vector<long long>(set.begin(), set.end());
}

// 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, you may change this line for yourself if you are happy using more memory.
        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 {
    for (const auto& cl : clauses) {
        if ((long)cl.size() != k) {
            throw std::runtime_error("Not a k-CNF formula.");
        }
    }

    return CNF::write_dimacs(filename);
}

KCNF KCNF::read_dimacs(const std::string& filename) {
    CNF base = CNF::read_dimacs(filename);

    KCNF f;
    f.n = base.n;
    f.clauses = std::move(base.clauses);

    if (!f.clauses.empty()) {
        f.k = f.clauses[0].size();

        for (const auto& cl : f.clauses) {
            if ((long)cl.size() != f.k) {
                throw std::runtime_error("Input is not a k-CNF formula.");
            }
        }
    } else {
        f.k = 0;
    }

    return f;
}