#include "qkmeans/sampling.hpp"
#include <algorithm>
#include <numeric>
#include <stdexcept>

namespace qkmeans {
namespace sampling {

// Alias Method implementation
// Reference: Walker's alias method for weighted sampling
void AliasTable::build(const double* weights, size_t n) {
    if (n == 0) {
        throw std::invalid_argument("Cannot build alias table with zero elements");
    }

    n_ = n;
    prob_.resize(n);
    alias_.resize(n);

    // Compute total weight
    double total = 0.0;
    for (size_t i = 0; i < n; ++i) {
        total += weights[i];
    }

    if (total <= 0.0) {
        throw std::invalid_argument("Total weight must be positive");
    }

    // Normalize weights to sum to n
    std::vector<double> scaled(n);
    for (size_t i = 0; i < n; ++i) {
        scaled[i] = weights[i] * static_cast<double>(n) / total;
    }

    // Separate into small and large
    std::vector<size_t> small, large;
    small.reserve(n);
    large.reserve(n);

    for (size_t i = 0; i < n; ++i) {
        if (scaled[i] < 1.0) {
            small.push_back(i);
        } else {
            large.push_back(i);
        }
    }

    // Build alias table
    while (!small.empty() && !large.empty()) {
        size_t s = small.back();
        small.pop_back();
        size_t l = large.back();
        large.pop_back();

        prob_[s] = scaled[s];
        alias_[s] = l;

        scaled[l] = (scaled[l] + scaled[s]) - 1.0;

        if (scaled[l] < 1.0) {
            small.push_back(l);
        } else {
            large.push_back(l);
        }
    }

    // Handle remaining elements (due to numerical precision)
    while (!large.empty()) {
        size_t l = large.back();
        large.pop_back();
        prob_[l] = 1.0;
        alias_[l] = l;
    }

    while (!small.empty()) {
        size_t s = small.back();
        small.pop_back();
        prob_[s] = 1.0;
        alias_[s] = s;
    }
}

size_t AliasTable::sample(std::mt19937_64& rng) const {
    if (n_ == 0) {
        throw std::runtime_error("Alias table not built");
    }

    // Generate random index and threshold
    std::uniform_int_distribution<size_t> idx_dist(0, n_ - 1);
    std::uniform_real_distribution<double> thresh_dist(0.0, 1.0);

    size_t i = idx_dist(rng);
    double u = thresh_dist(rng);

    // Return i or its alias based on threshold
    return (u < prob_[i]) ? i : alias_[i];
}

// ProposalSampler implementation
void ProposalSampler::initialize(
    const double* norms_sq,
    size_t n,
    double c1_norm_sq,
    double frobenius_sq
) {
    c1_norm_sq_ = c1_norm_sq;

    // Compute weights: w_i = ||x_i||² + ||c₁||²
    weights_.resize(n);
    for (size_t i = 0; i < n; ++i) {
        weights_[i] = norms_sq[i] + c1_norm_sq;
    }

    // Normalization constant: Z = ||X||²_F + n||c₁||²
    normalization_ = frobenius_sq + static_cast<double>(n) * c1_norm_sq;

    // Build alias table for sampling
    alias_.build(weights_.data(), n);
}

size_t ProposalSampler::sample(std::mt19937_64& rng) const {
    return alias_.sample(rng);
}

double ProposalSampler::get_weight(size_t idx) const {
    return weights_[idx];
}

} // namespace sampling
} // namespace qkmeans
