#include "qkmeans/core.hpp"
#include "qkmeans/distance.hpp"
#include "qkmeans/utils.hpp"
#include <cmath>
#include <algorithm>
#include <stdexcept>

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

namespace qkmeans {

QKMeans::QKMeans(const QKMeansConfig& config)
    : config_(config)
    , uniform_01_(0.0, 1.0)
{
    // Initialize RNG
    if (config_.random_seed == 0) {
        std::random_device rd;
        rng_.seed(rd());
    } else {
        rng_.seed(config_.random_seed);
    }

    // Set number of threads
#ifdef _OPENMP
    if (config_.n_threads > 0) {
        omp_set_num_threads(config_.n_threads);
    }
#endif
}

QKMeans::~QKMeans() = default;

void QKMeans::init_anns() {
    anns::ANNSConfig anns_config;
    anns_config.dimension = d_;
    anns_config.expected_size = config_.k;

    // Copy ANNS-specific parameters
    anns_config.lsh_num_tables = config_.lsh_num_tables;
    anns_config.lsh_num_hashes = config_.lsh_num_hashes;
    anns_config.lsh_rho = config_.lsh_rho;
    anns_config.hnsw_M = config_.hnsw_M;
    anns_config.hnsw_ef_construction = config_.hnsw_ef_construction;
    anns_config.hnsw_ef_search = config_.hnsw_ef_search;
    anns_config.faiss_index_type = config_.faiss_index_type;
    anns_config.faiss_hnsw_M = config_.faiss_hnsw_M;
    anns_config.faiss_hnsw_efConstruction = config_.faiss_hnsw_efConstruction;
    anns_config.faiss_hnsw_efSearch = config_.faiss_hnsw_efSearch;
    anns_config.faiss_lsh_nbits = config_.faiss_lsh_nbits;
    anns_config.faiss_nlist = config_.faiss_nlist;
    anns_config.faiss_nprobe = config_.faiss_nprobe;

    anns_ = anns::ANNSBase::create(config_.anns_type, anns_config);
    anns_->reserve(config_.k);
}

void QKMeans::preprocess(const double* X, size_t n, size_t d) {
    n_ = n;
    d_ = d;

    // Allocate space
    X_centered_.resize(n * d);
    norms_sq_.resize(n);
    mean_.resize(d);

    // Compute mean
    distance::compute_mean(X, n, d, mean_.data());

    // Center data and compute norms
    distance::center_data(X, n, d, mean_.data(), X_centered_.data());
    distance::compute_squared_norms(X_centered_.data(), n, d, norms_sq_.data());

    // Compute Frobenius norm squared
    frobenius_sq_ = 0.0;
    for (size_t i = 0; i < n; ++i) {
        frobenius_sq_ += norms_sq_[i];
    }
}

double QKMeans::compute_beta() const {
    // Beta is the oversampling factor from Lemma D.18
    // τ = 2ε⁻¹(||X||²_F + n||c₁||²) / cost(X, S)
    // For simplicity, we use a conservative estimate based on the scaling laws
    // β ≈ O(k^ε) where ε = 2/d (quantization exponent)

    // Start with a baseline
    double beta = 4.0 * (frobenius_sq_ + n_ * c1_norm_sq_) /
                  std::max(1.0, frobenius_sq_);

    // Clamp to reasonable range
    return std::clamp(beta, 2.0, 100.0);
}

size_t QKMeans::sample_center() {
    // Maximum rejection iterations: m·ln(k) where m = β
    size_t max_iter = config_.max_rejections;
    if (max_iter == 0) {
        double beta = compute_beta();
        max_iter = static_cast<size_t>(std::ceil(beta * std::log(config_.k + 1)));
        max_iter = std::max(max_iter, size_t(10));  // At least 10 iterations
    }

    for (size_t iter = 0; iter < max_iter; ++iter) {
        // Sample from proposal distribution κ(·|C)
        size_t x_idx = proposal_.sample(rng_);

        // Query approximate nearest center from ANNS
        const double* point = X_centered_.data() + x_idx * d_;
        auto [nn_internal_idx, nn_dist_sq] = anns_->query(point);

        // Compute acceptance ratio
        // r = cost(x, QUERY(L,x)) / (2(||x||² + ||c₁||²))
        double weight = proposal_.get_weight(x_idx);
        double r = nn_dist_sq / (2.0 * weight);

        // Clamp r to [0, 1] for numerical stability
        r = std::clamp(r, 0.0, 1.0);

        // Accept with probability r
        double u = uniform_01_(rng_);
        if (u <= r) {
            ++total_accepted_;
            return x_idx;
        }

        ++total_rejections_;
    }

    // Fallback: uniform sample
    ++fallback_samples_;
    std::uniform_int_distribution<size_t> uniform_idx(0, n_ - 1);
    return uniform_idx(rng_);
}

QKMeansResult QKMeans::fit(const double* X, size_t n, size_t d) {
    if (n == 0 || d == 0) {
        throw std::invalid_argument("Data dimensions must be positive");
    }
    if (config_.k == 0 || config_.k > n) {
        throw std::invalid_argument("k must be between 1 and n");
    }

    QKMeansResult result;
    Timer timer;

    // Reset statistics
    total_rejections_ = 0;
    total_accepted_ = 0;
    fallback_samples_ = 0;
    center_indices_.clear();
    center_indices_.reserve(config_.k);

    // Step 1: Preprocessing
    timer.start();
    preprocess(X, n, d);
    result.preprocess_time_ms = timer.elapsed_ms();

    // Step 2: Initialize ANNS
    init_anns();

    // Step 3: Start seeding
    timer.start();

    // Choose first center uniformly at random
    std::uniform_int_distribution<size_t> uniform_idx(0, n - 1);
    c1_idx_ = uniform_idx(rng_);
    c1_norm_sq_ = norms_sq_[c1_idx_];

    center_indices_.push_back(c1_idx_);

    // Insert first center into ANNS
    const double* c1_ptr = X_centered_.data() + c1_idx_ * d_;
    anns_->insert(c1_ptr, 0);

    // Initialize proposal sampler
    proposal_.initialize(norms_sq_.data(), n, c1_norm_sq_, frobenius_sq_);

    current_k_ = 1;

    // Choose remaining k-1 centers using rejection sampling
    for (size_t t = 1; t < config_.k; ++t) {
        // Sample next center
        size_t ct_idx = sample_center();
        center_indices_.push_back(ct_idx);

        // Insert into ANNS
        const double* ct_ptr = X_centered_.data() + ct_idx * d_;
        anns_->insert(ct_ptr, t);

        ++current_k_;
    }

    result.seeding_time_ms = timer.elapsed_ms();

    // Step 4: Build result
    result.centers.resize(config_.k * d);
    result.center_indices = center_indices_;

    // Copy centers (original coordinates, not centered)
    for (size_t t = 0; t < config_.k; ++t) {
        size_t idx = center_indices_[t];
        const double* src = X + idx * d;
        double* dst = result.centers.data() + t * d;
        std::copy(src, src + d, dst);
    }

    // Compute total cost (sum of D² values)
    double total_cost = 0.0;
    #pragma omp parallel for reduction(+:total_cost) schedule(static)
    for (size_t i = 0; i < n; ++i) {
        const double* point = X + i * d;
        double min_dist_sq = std::numeric_limits<double>::max();

        for (size_t t = 0; t < config_.k; ++t) {
            const double* center = result.centers.data() + t * d;
            double dist_sq = distance::squared_l2(point, center, d);
            min_dist_sq = std::min(min_dist_sq, dist_sq);
        }

        total_cost += min_dist_sq;
    }
    result.total_cost = total_cost;

    result.total_rejections = total_rejections_;
    result.total_accepted = total_accepted_;
    result.fallback_samples = fallback_samples_;

    return result;
}

} // namespace qkmeans
