#include "qkmeans/distance.hpp"
#include <algorithm>
#include <cstring>
#include <vector>

#ifdef QKMEANS_HAS_AVX2
#include <immintrin.h>
#endif

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

namespace qkmeans {
namespace distance {

// Detect SIMD level at runtime
SIMDLevel detect_simd_level() {
#ifdef QKMEANS_HAS_AVX2
    // Check for AVX2 support
    #if defined(__GNUC__) || defined(__clang__)
    __builtin_cpu_init();
    if (__builtin_cpu_supports("avx2")) {
        return SIMDLevel::AVX2;
    }
    #elif defined(_MSC_VER)
    int cpuInfo[4];
    __cpuid(cpuInfo, 7);
    if (cpuInfo[1] & (1 << 5)) {  // AVX2 bit
        return SIMDLevel::AVX2;
    }
    #endif
#endif
    return SIMDLevel::SCALAR;
}

const char* simd_level_name(SIMDLevel level) {
    switch (level) {
        case SIMDLevel::SCALAR: return "SCALAR";
        case SIMDLevel::SSE4: return "SSE4";
        case SIMDLevel::AVX2: return "AVX2";
        default: return "UNKNOWN";
    }
}

// Scalar implementation
double squared_l2_scalar(const double* a, const double* b, size_t d) {
    double sum = 0.0;
    for (size_t i = 0; i < d; ++i) {
        double diff = a[i] - b[i];
        sum += diff * diff;
    }
    return sum;
}

#ifdef QKMEANS_HAS_AVX2
// AVX2 implementation - processes 4 doubles at a time
double squared_l2_avx2(const double* a, const double* b, size_t d) {
    __m256d sum = _mm256_setzero_pd();

    size_t i = 0;
    // Process 4 elements at a time
    for (; i + 4 <= d; i += 4) {
        __m256d va = _mm256_loadu_pd(a + i);
        __m256d vb = _mm256_loadu_pd(b + i);
        __m256d diff = _mm256_sub_pd(va, vb);
        sum = _mm256_fmadd_pd(diff, diff, sum);
    }

    // Horizontal sum of the 4 doubles
    __m128d low = _mm256_castpd256_pd128(sum);
    __m128d high = _mm256_extractf128_pd(sum, 1);
    __m128d sum128 = _mm_add_pd(low, high);
    sum128 = _mm_hadd_pd(sum128, sum128);
    double result = _mm_cvtsd_f64(sum128);

    // Handle remaining elements
    for (; i < d; ++i) {
        double diff = a[i] - b[i];
        result += diff * diff;
    }

    return result;
}
#endif

// Main dispatch function
double squared_l2(const double* a, const double* b, size_t d) {
    static SIMDLevel level = detect_simd_level();

#ifdef QKMEANS_HAS_AVX2
    if (level == SIMDLevel::AVX2) {
        return squared_l2_avx2(a, b, d);
    }
#endif
    return squared_l2_scalar(a, b, d);
}

// Find nearest center for a single point
NearestResult find_nearest_center(
    const double* point,
    size_t d,
    const double* centers,
    size_t k
) {
    NearestResult result{0, std::numeric_limits<double>::max()};

    for (size_t j = 0; j < k; ++j) {
        double dist_sq = squared_l2(point, centers + j * d, d);
        if (dist_sq < result.distance_sq) {
            result.index = j;
            result.distance_sq = dist_sq;
        }
    }

    return result;
}

// Compute distances to nearest center for all points
void compute_distances_to_nearest(
    const double* X,
    size_t n,
    size_t d,
    const double* centers,
    size_t k,
    double* distances_out
) {
    #pragma omp parallel for 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 j = 0; j < k; ++j) {
            double dist_sq = squared_l2(point, centers + j * d, d);
            min_dist_sq = std::min(min_dist_sq, dist_sq);
        }

        distances_out[i] = min_dist_sq;
    }
}

// Compute squared norms
void compute_squared_norms(
    const double* X,
    size_t n,
    size_t d,
    double* norms_out
) {
    #pragma omp parallel for schedule(static)
    for (size_t i = 0; i < n; ++i) {
        const double* point = X + i * d;
        double norm_sq = 0.0;

        for (size_t j = 0; j < d; ++j) {
            norm_sq += point[j] * point[j];
        }

        norms_out[i] = norm_sq;
    }
}

// Compute mean vector
void compute_mean(
    const double* X,
    size_t n,
    size_t d,
    double* mean_out
) {
    // Initialize mean to zero
    std::memset(mean_out, 0, d * sizeof(double));

    // Accumulate sum (with OpenMP reduction)
    #pragma omp parallel
    {
        std::vector<double> local_sum(d, 0.0);

        #pragma omp for schedule(static) nowait
        for (size_t i = 0; i < n; ++i) {
            const double* point = X + i * d;
            for (size_t j = 0; j < d; ++j) {
                local_sum[j] += point[j];
            }
        }

        #pragma omp critical
        {
            for (size_t j = 0; j < d; ++j) {
                mean_out[j] += local_sum[j];
            }
        }
    }

    // Divide by n
    for (size_t j = 0; j < d; ++j) {
        mean_out[j] /= static_cast<double>(n);
    }
}

// Center data
void center_data(
    const double* X,
    size_t n,
    size_t d,
    const double* mean,
    double* X_centered_out
) {
    #pragma omp parallel for schedule(static)
    for (size_t i = 0; i < n; ++i) {
        const double* point = X + i * d;
        double* centered = X_centered_out + i * d;

        for (size_t j = 0; j < d; ++j) {
            centered[j] = point[j] - mean[j];
        }
    }
}

} // namespace distance
} // namespace qkmeans
