#pragma once

#include <cstddef>
#include <cmath>
#include <limits>

namespace qkmeans {
namespace distance {

// SIMD capability levels
enum class SIMDLevel { SCALAR, SSE4, AVX2 };

// Detect available SIMD level at runtime
SIMDLevel detect_simd_level();

// Get string representation of SIMD level
const char* simd_level_name(SIMDLevel level);

// Squared L2 distance between two vectors
// Dispatches to best available SIMD implementation
double squared_l2(const double* a, const double* b, size_t d);

// Scalar implementation (fallback)
double squared_l2_scalar(const double* a, const double* b, size_t d);

#ifdef QKMEANS_HAS_AVX2
// AVX2 optimized implementation
double squared_l2_avx2(const double* a, const double* b, size_t d);
#endif

// Compute distance from each point to its nearest center
// Input:  X is n x d matrix (row-major), centers is k x d matrix
// Output: distances[i] = min_j ||X[i] - centers[j]||^2
// This is parallelized with OpenMP
void compute_distances_to_nearest(
    const double* X,
    size_t n,
    size_t d,
    const double* centers,
    size_t k,
    double* distances_out
);

// Compute distance from a single point to its nearest center
// Returns the index of the nearest center and the squared distance
struct NearestResult {
    size_t index;
    double distance_sq;
};

NearestResult find_nearest_center(
    const double* point,
    size_t d,
    const double* centers,
    size_t k
);

// Compute squared norms of all rows in a matrix
// Input:  X is n x d matrix (row-major)
// Output: norms[i] = ||X[i]||^2
void compute_squared_norms(
    const double* X,
    size_t n,
    size_t d,
    double* norms_out
);

// Compute mean of all rows in a matrix
// Input:  X is n x d matrix (row-major)
// Output: mean is d-dimensional vector
void compute_mean(
    const double* X,
    size_t n,
    size_t d,
    double* mean_out
);

// Center the data by subtracting mean from each row
// Input:  X is n x d matrix (row-major), mean is d-dimensional
// Output: X_centered[i] = X[i] - mean (in-place or out-of-place)
void center_data(
    const double* X,
    size_t n,
    size_t d,
    const double* mean,
    double* X_centered_out
);

} // namespace distance
} // namespace qkmeans
