#pragma once

#include "anns/base.hpp"
#include "sampling.hpp"
#include <vector>
#include <random>
#include <memory>

namespace qkmeans {

// Configuration for QKMeans algorithm
struct QKMeansConfig {
    size_t k = 8;                               // Number of clusters
    size_t max_rejections = 0;                  // 0 = auto (m·ln(k))
    anns::ANNSType anns_type = anns::ANNSType::LSH;
    uint64_t random_seed = 0;                   // 0 = random seed
    int n_threads = -1;                         // -1 = all available

    // LSH-specific parameters
    size_t lsh_num_tables = 10;
    size_t lsh_num_hashes = 8;
    double lsh_rho = 0.5;

    // HNSW-specific parameters
    size_t hnsw_M = 16;
    size_t hnsw_ef_construction = 200;
    size_t hnsw_ef_search = 50;

    // FAISS-specific parameters
    int faiss_index_type = 0;           // 0=FLAT, 1=HNSW, 2=LSH, 3=IVF
    size_t faiss_hnsw_M = 32;           // HNSW: number of links per node
    size_t faiss_hnsw_efConstruction = 40;  // HNSW: construction-time search depth
    size_t faiss_hnsw_efSearch = 16;    // HNSW: search-time depth
    size_t faiss_lsh_nbits = 0;         // LSH: number of bits (0 = auto: 2*dim)
    size_t faiss_nlist = 100;           // IVF: number of inverted lists
    size_t faiss_nprobe = 10;           // IVF: number of lists to probe
};

// Result of QKMeans seeding
struct QKMeansResult {
    std::vector<double> centers;            // k x d flattened (row-major)
    std::vector<size_t> center_indices;     // Indices of chosen data points
    double total_cost = 0.0;                // Sum of D² values
    size_t total_rejections = 0;            // Number of rejected samples
    size_t total_accepted = 0;              // Number of accepted samples
    size_t fallback_samples = 0;            // Number of uniform fallback samples
    double seeding_time_ms = 0.0;           // Time for seeding
    double preprocess_time_ms = 0.0;        // Time for preprocessing
};

// Main QKMeans algorithm class
// Implements Algorithm 2 from "Faster k-means Seeding Under The Manifold Hypothesis"
class QKMeans {
public:
    explicit QKMeans(const QKMeansConfig& config);
    ~QKMeans();

    // Run the seeding algorithm
    // X: n x d data matrix (row-major, contiguous)
    // Returns: QKMeansResult with k centers
    QKMeansResult fit(const double* X, size_t n, size_t d);

    // Get the configuration
    const QKMeansConfig& config() const { return config_; }

private:
    // PREPROCESS(X): Center data and compute norms
    void preprocess(const double* X, size_t n, size_t d);

    // SAMPLE(L, C, m): Rejection sampling for one center
    // Returns index of the sampled point
    size_t sample_center();

    // Initialize ANNS data structure
    void init_anns();

    // Compute beta (oversampling factor) dynamically
    double compute_beta() const;

    QKMeansConfig config_;
    std::mt19937_64 rng_;

    // Dataset properties
    size_t n_ = 0;                          // Number of points
    size_t d_ = 0;                          // Dimension

    // Preprocessed data
    std::vector<double> X_centered_;        // n x d centered data
    std::vector<double> norms_sq_;          // ||x_i||² for each point
    std::vector<double> mean_;              // Mean vector
    double frobenius_sq_ = 0.0;             // ||X||²_F

    // First center info
    double c1_norm_sq_ = 0.0;               // ||c₁||²
    size_t c1_idx_ = 0;                     // Index of first center

    // ANNS data structure
    std::unique_ptr<anns::ANNSBase> anns_;

    // Sampling
    sampling::ProposalSampler proposal_;
    std::uniform_real_distribution<double> uniform_01_;

    // Current state
    std::vector<size_t> center_indices_;
    size_t current_k_ = 0;                  // Number of centers chosen so far

    // Statistics
    size_t total_rejections_ = 0;
    size_t total_accepted_ = 0;
    size_t fallback_samples_ = 0;
};

} // namespace qkmeans
