#pragma once

#include <cstddef>
#include <vector>
#include <random>

namespace qkmeans {
namespace sampling {

// Alias method for O(1) weighted sampling
// Build time: O(n), Sample time: O(1)
class AliasTable {
public:
    AliasTable() = default;

    // Build the alias table from weights (need not be normalized)
    // Time complexity: O(n)
    void build(const double* weights, size_t n);

    // Sample an index according to the weights
    // Time complexity: O(1)
    size_t sample(std::mt19937_64& rng) const;

    // Get the number of elements
    size_t size() const { return n_; }

    // Check if the table is built
    bool is_built() const { return n_ > 0; }

private:
    size_t n_ = 0;
    std::vector<double> prob_;      // Probability table
    std::vector<size_t> alias_;     // Alias table
};

// Sampler for the proposal distribution κ(x|C)
// κ(x|C) = (||x||² + ||c₁||²) / (||X||²_F + n||c₁||²)
class ProposalSampler {
public:
    ProposalSampler() = default;

    // Initialize with squared norms and first center's squared norm
    // norms_sq[i] = ||x_i||² (after centering)
    // c1_norm_sq = ||c₁||²
    // frobenius_sq = ||X||²_F = sum of norms_sq
    void initialize(
        const double* norms_sq,
        size_t n,
        double c1_norm_sq,
        double frobenius_sq
    );

    // Sample a point index from the proposal distribution
    size_t sample(std::mt19937_64& rng) const;

    // Get the unnormalized weight for a given index
    // weight(i) = ||x_i||² + ||c₁||²
    double get_weight(size_t idx) const;

    // Get the normalization constant
    // Z = ||X||²_F + n||c₁||²
    double get_normalization() const { return normalization_; }

private:
    AliasTable alias_;
    std::vector<double> weights_;   // Unnormalized weights
    double normalization_ = 0.0;    // Sum of weights
    double c1_norm_sq_ = 0.0;
};

// Uniform sampler (for fallback in rejection sampling)
class UniformSampler {
public:
    explicit UniformSampler(size_t n) : n_(n), dist_(0, n - 1) {}

    size_t sample(std::mt19937_64& rng) {
        return dist_(rng);
    }

private:
    size_t n_;
    std::uniform_int_distribution<size_t> dist_;
};

} // namespace sampling
} // namespace qkmeans
