#include "qkmeans/anns/lsh.hpp"
#include "qkmeans/anns/hnsw_backend.hpp"
#include "qkmeans/anns/faiss_backend.hpp"
#include "qkmeans/distance.hpp"
#include <cmath>
#include <algorithm>
#include <limits>
#include <stdexcept>

namespace qkmeans {
namespace anns {

// LSH Implementation
LSHIndex::LSHIndex(size_t dim, size_t num_tables, size_t num_hashes, double rho)
    : dim_(dim)
    , num_tables_(num_tables)
    , num_hashes_(num_hashes)
    , rho_(rho)
    , rng_(std::random_device{}())
{
    if (dim == 0) {
        throw std::invalid_argument("Dimension must be positive");
    }
    if (num_tables == 0 || num_hashes == 0) {
        throw std::invalid_argument("Number of tables and hashes must be positive");
    }

    init_hyperplanes();
    tables_.resize(num_tables_);
}

LSHIndex::LSHIndex(const ANNSConfig& config)
    : LSHIndex(config.dimension, config.lsh_num_tables, config.lsh_num_hashes, config.lsh_rho)
{
}

void LSHIndex::init_hyperplanes() {
    std::normal_distribution<double> normal(0.0, 1.0);

    hyperplanes_.resize(num_tables_);
    for (size_t t = 0; t < num_tables_; ++t) {
        hyperplanes_[t].resize(num_hashes_);
        for (size_t h = 0; h < num_hashes_; ++h) {
            hyperplanes_[t][h].resize(dim_);

            // Generate random unit vector
            double norm = 0.0;
            for (size_t i = 0; i < dim_; ++i) {
                hyperplanes_[t][h][i] = normal(rng_);
                norm += hyperplanes_[t][h][i] * hyperplanes_[t][h][i];
            }
            norm = std::sqrt(norm);
            for (size_t i = 0; i < dim_; ++i) {
                hyperplanes_[t][h][i] /= norm;
            }
        }
    }
}

uint64_t LSHIndex::compute_hash(const double* point, size_t table_idx) const {
    uint64_t hash = 0;
    for (size_t h = 0; h < num_hashes_; ++h) {
        // Compute dot product with hyperplane
        double dot = 0.0;
        const auto& hyperplane = hyperplanes_[table_idx][h];
        for (size_t i = 0; i < dim_; ++i) {
            dot += point[i] * hyperplane[i];
        }
        // Set bit based on sign
        if (dot >= 0) {
            hash |= (1ULL << h);
        }
    }
    return hash;
}

double LSHIndex::compute_distance_sq(const double* point, size_t stored_idx) const {
    return distance::squared_l2(point, points_[stored_idx].data(), dim_);
}

void LSHIndex::insert(const double* point, size_t idx) {
    // Store the point
    points_.emplace_back(point, point + dim_);

    // Insert into each hash table
    for (size_t t = 0; t < num_tables_; ++t) {
        uint64_t hash = compute_hash(point, t);
        tables_[t][hash].push_back(num_points_);
    }

    ++num_points_;
}

QueryResult LSHIndex::query(const double* point) const {
    if (num_points_ == 0) {
        return {0, std::numeric_limits<double>::max()};
    }

    QueryResult best{0, std::numeric_limits<double>::max()};

    // Collect candidate points from all hash tables
    std::vector<bool> seen(num_points_, false);

    for (size_t t = 0; t < num_tables_; ++t) {
        uint64_t hash = compute_hash(point, t);
        auto it = tables_[t].find(hash);
        if (it != tables_[t].end()) {
            for (size_t stored_idx : it->second) {
                if (!seen[stored_idx]) {
                    seen[stored_idx] = true;
                    double dist_sq = compute_distance_sq(point, stored_idx);
                    if (dist_sq < best.distance_sq) {
                        best.index = stored_idx;
                        best.distance_sq = dist_sq;
                    }
                }
            }
        }
    }

    // If no candidates found, fall back to brute force
    // (this shouldn't happen often with well-tuned parameters)
    if (best.distance_sq == std::numeric_limits<double>::max()) {
        for (size_t i = 0; i < num_points_; ++i) {
            double dist_sq = compute_distance_sq(point, i);
            if (dist_sq < best.distance_sq) {
                best.index = i;
                best.distance_sq = dist_sq;
            }
        }
    }

    return best;
}

void LSHIndex::clear() {
    points_.clear();
    for (auto& table : tables_) {
        table.clear();
    }
    num_points_ = 0;
}

void LSHIndex::reserve(size_t n) {
    points_.reserve(n);
}

// Brute Force Implementation
BruteForceIndex::BruteForceIndex(size_t dim)
    : dim_(dim)
{
    if (dim == 0) {
        throw std::invalid_argument("Dimension must be positive");
    }
}

BruteForceIndex::BruteForceIndex(const ANNSConfig& config)
    : BruteForceIndex(config.dimension)
{
}

void BruteForceIndex::insert(const double* point, size_t idx) {
    points_.emplace_back(point, point + dim_);
    indices_.push_back(idx);
    ++num_points_;
}

QueryResult BruteForceIndex::query(const double* point) const {
    if (num_points_ == 0) {
        return {0, std::numeric_limits<double>::max()};
    }

    QueryResult best{0, std::numeric_limits<double>::max()};

    for (size_t i = 0; i < num_points_; ++i) {
        double dist_sq = distance::squared_l2(point, points_[i].data(), dim_);
        if (dist_sq < best.distance_sq) {
            best.index = i;
            best.distance_sq = dist_sq;
        }
    }

    return best;
}

void BruteForceIndex::clear() {
    points_.clear();
    indices_.clear();
    num_points_ = 0;
}

void BruteForceIndex::reserve(size_t n) {
    points_.reserve(n);
    indices_.reserve(n);
}

// Base class batch_query default implementation
void ANNSBase::batch_query(
    const double* points,
    size_t n,
    size_t* indices_out,
    double* distances_out
) const {
    size_t d = dimension();
    for (size_t i = 0; i < n; ++i) {
        auto result = query(points + i * d);
        indices_out[i] = result.index;
        distances_out[i] = result.distance_sq;
    }
}

// Factory method
std::unique_ptr<ANNSBase> ANNSBase::create(
    ANNSType type,
    const ANNSConfig& config
) {
    switch (type) {
        case ANNSType::LSH:
            return std::make_unique<LSHIndex>(config);
        case ANNSType::BRUTE_FORCE:
            return std::make_unique<BruteForceIndex>(config);
        case ANNSType::FAISS:
#ifdef QKMEANS_HAS_FAISS
            return std::make_unique<FAISSIndex>(config);
#else
            throw std::runtime_error("FAISS support not compiled in. Install faiss-cpu or faiss-gpu and rebuild.");
#endif
        case ANNSType::HNSW:
#ifdef QKMEANS_HAS_HNSWLIB
            return std::make_unique<HNSWIndex>(config);
#else
            throw std::runtime_error("hnswlib support not compiled in. Install hnswlib and rebuild.");
#endif
        default:
            throw std::invalid_argument("Unknown ANNS type");
    }
}

bool ANNSBase::is_available(ANNSType type) {
    switch (type) {
        case ANNSType::LSH:
        case ANNSType::BRUTE_FORCE:
            return true;
        case ANNSType::FAISS:
#ifdef QKMEANS_HAS_FAISS
            return true;
#else
            return false;
#endif
        case ANNSType::HNSW:
#ifdef QKMEANS_HAS_HNSWLIB
            return true;
#else
            return false;
#endif
        default:
            return false;
    }
}

} // namespace anns
} // namespace qkmeans
