#pragma once

#include "base.hpp"
#include <vector>
#include <unordered_map>
#include <random>

namespace qkmeans {
namespace anns {

// Locality-Sensitive Hashing (LSH) implementation for approximate nearest neighbor
// Based on random hyperplane projections (SimHash variant for Euclidean distance)
//
// For Euclidean distance, we use the following hash family:
// h(x) = sign(<x, r>) where r is a random unit vector
//
// We use L hash tables, each with K hash functions, giving O(n^ρ) query time
// where ρ depends on the approximation ratio c: ρ = 1/c
class LSHIndex : public ANNSBase {
public:
    // Constructor
    // dim: dimension of the points
    // num_tables: number of hash tables (L)
    // num_hashes: number of hash functions per table (K)
    // rho: approximation parameter (affects accuracy vs speed tradeoff)
    LSHIndex(size_t dim, size_t num_tables, size_t num_hashes, double rho);

    // Create from config
    explicit LSHIndex(const ANNSConfig& config);

    ~LSHIndex() override = default;

    // Insert a point
    void insert(const double* point, size_t idx) override;

    // Query nearest neighbor
    QueryResult query(const double* point) const override;

    // Get number of points
    size_t size() const override { return num_points_; }

    // Get dimension
    size_t dimension() const override { return dim_; }

    // Clear the index
    void clear() override;

    // Reserve space
    void reserve(size_t n) override;

private:
    // Hash a point using the hash functions for a specific table
    uint64_t compute_hash(const double* point, size_t table_idx) const;

    // Compute squared distance between a point and a stored point
    double compute_distance_sq(const double* point, size_t stored_idx) const;

    // Initialize random hyperplanes
    void init_hyperplanes();

    size_t dim_;
    size_t num_tables_;
    size_t num_hashes_;
    double rho_;
    size_t num_points_ = 0;

    // Random hyperplanes for hashing: [num_tables][num_hashes][dim]
    std::vector<std::vector<std::vector<double>>> hyperplanes_;

    // Hash tables: [num_tables] -> {hash_value -> list of point indices}
    std::vector<std::unordered_map<uint64_t, std::vector<size_t>>> tables_;

    // Stored points for distance computation: [point_idx][dim]
    std::vector<std::vector<double>> points_;

    // Random number generator
    mutable std::mt19937_64 rng_;
};

// Brute-force nearest neighbor (exact, for comparison/debugging)
class BruteForceIndex : public ANNSBase {
public:
    explicit BruteForceIndex(size_t dim);
    explicit BruteForceIndex(const ANNSConfig& config);

    ~BruteForceIndex() override = default;

    void insert(const double* point, size_t idx) override;
    QueryResult query(const double* point) const override;
    size_t size() const override { return num_points_; }
    size_t dimension() const override { return dim_; }
    void clear() override;
    void reserve(size_t n) override;

private:
    size_t dim_;
    size_t num_points_ = 0;
    std::vector<std::vector<double>> points_;
    std::vector<size_t> indices_;
};

} // namespace anns
} // namespace qkmeans
