#include "knn_opencl.h"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <limits>
#include <stdexcept>
#include <string>
#include <vector>

namespace face_knn {

SearchOutput search_opencl_impl(
    const std::vector<float>& database,
    const std::vector<float>& queries,
    std::size_t database_count,
    std::size_t query_count,
    const SearchOptions& options);


namespace {

bool better_neighbor(
    const Neighbor& left,
    const Neighbor& right)
{
    if (left.score != right.score) {
        return left.score > right.score;
    }

    return left.index < right.index;
}


void validate_finite_values(
    const std::vector<float>& values,
    const char* name)
{
    for (
        std::size_t index = 0;
        index < values.size();
        ++index
    ) {
        if (!std::isfinite(values[index])) {
            throw std::invalid_argument(
                std::string(name)
                + " contiene un valore non finito"
            );
        }
    }
}


std::uint32_t effective_top_k(
    std::size_t database_count,
    std::uint32_t requested_top_k)
{
    return std::min(
        requested_top_k,
        static_cast<std::uint32_t>(
            database_count
        )
    );
}

}  // namespace


void validate_input(
    const std::vector<float>& database,
    const std::vector<float>& queries,
    std::size_t database_count,
    std::size_t query_count,
    const SearchOptions& options)
{
    if (options.dimension == 0) {
        throw std::invalid_argument(
            "dimension deve essere maggiore di zero"
        );
    }

    if (options.top_k == 0) {
        throw std::invalid_argument(
            "top_k deve essere maggiore di zero"
        );
    }

    if (options.lanes_per_query == 0) {
        throw std::invalid_argument(
            "lanes_per_query deve essere maggiore di zero"
        );
    }

    if (options.chunk_size == 0) {
        throw std::invalid_argument(
            "chunk_size deve essere maggiore di zero"
        );
    }

    if (database_count == 0) {
        throw std::invalid_argument(
            "database_count deve essere maggiore di zero"
        );
    }

    if (query_count == 0) {
        throw std::invalid_argument(
            "query_count deve essere maggiore di zero"
        );
    }

    if (
        database_count
        > std::numeric_limits<
            std::uint32_t
        >::max()
    ) {
        throw std::invalid_argument(
            "database_count supera uint32"
        );
    }

    const std::size_t expected_database =
        database_count
        * static_cast<std::size_t>(
            options.dimension
        );

    const std::size_t expected_queries =
        query_count
        * static_cast<std::size_t>(
            options.dimension
        );

    if (database.size() != expected_database) {
        throw std::invalid_argument(
            "dimensione database incoerente"
        );
    }

    if (queries.size() != expected_queries) {
        throw std::invalid_argument(
            "dimensione query incoerente"
        );
    }

    validate_finite_values(
        database,
        "database"
    );

    validate_finite_values(
        queries,
        "queries"
    );
}


SearchOutput search_cpu(
    const std::vector<float>& database,
    const std::vector<float>& queries,
    std::size_t database_count,
    std::size_t query_count,
    const SearchOptions& options)
{
    validate_input(
        database,
        queries,
        database_count,
        query_count,
        options
    );

    const std::uint32_t result_top_k =
        effective_top_k(
            database_count,
            options.top_k
        );

    SearchOutput output;

    output.query_count = query_count;
    output.database_count = database_count;
    output.dimension = options.dimension;
    output.top_k = result_top_k;
    output.backend = "cpu";
    output.device_name = "CPU";

    output.neighbors.resize(
        query_count
        * static_cast<std::size_t>(
            result_top_k
        )
    );

    const auto started_at =
        std::chrono::steady_clock::now();

    for (
        std::size_t query_index = 0;
        query_index < query_count;
        ++query_index
    ) {
        const float* query =
            &queries[
                query_index
                * options.dimension
            ];

        std::vector<Neighbor> candidates;

        candidates.reserve(database_count);

        for (
            std::size_t database_index = 0;
            database_index < database_count;
            ++database_index
        ) {
            const float* database_vector =
                &database[
                    database_index
                    * options.dimension
                ];

            float score = 0.0f;

            for (
                std::uint32_t component = 0;
                component < options.dimension;
                ++component
            ) {
                score +=
                    database_vector[component]
                    * query[component];
            }

            candidates.push_back({
                score,
                static_cast<std::uint32_t>(
                    database_index
                ),
            });
        }

        if (
            result_top_k
            < candidates.size()
        ) {
            std::nth_element(
                candidates.begin(),
                candidates.begin()
                    + result_top_k,
                candidates.end(),
                better_neighbor
            );
        }

        std::sort(
            candidates.begin(),
            candidates.begin()
                + result_top_k,
            better_neighbor
        );

        const std::size_t output_offset =
            query_index
            * result_top_k;

        for (
            std::uint32_t position = 0;
            position < result_top_k;
            ++position
        ) {
            output.neighbors[
                output_offset + position
            ] = candidates[position];
        }
    }

    const auto completed_at =
        std::chrono::steady_clock::now();

    output.elapsed_ms =
        std::chrono::duration<
            double,
            std::milli
        >(
            completed_at - started_at
        ).count();

    return output;
}


SearchOutput search_opencl(
    const std::vector<float>& database,
    const std::vector<float>& queries,
    std::size_t database_count,
    std::size_t query_count,
    const SearchOptions& options)
{
    return search_opencl_impl(
        database,
        queries,
        database_count,
        query_count,
        options
    );
}


std::vector<Neighbor> query_neighbors(
    const SearchOutput& output,
    std::size_t query_index)
{
    if (query_index >= output.query_count) {
        throw std::out_of_range(
            "query_index fuori intervallo"
        );
    }

    const std::size_t begin =
        query_index
        * output.top_k;

    const std::size_t end =
        begin
        + output.top_k;

    if (end > output.neighbors.size()) {
        throw std::runtime_error(
            "layout SearchOutput incoerente"
        );
    }

    return {
        output.neighbors.begin() + begin,
        output.neighbors.begin() + end,
    };
}


std::vector<Neighbor> exclude_neighbor(
    const std::vector<Neighbor>& neighbors,
    std::uint32_t excluded_index)
{
    std::vector<Neighbor> filtered;

    filtered.reserve(neighbors.size());

    for (const Neighbor& neighbor : neighbors) {
        if (neighbor.index == invalid_index) {
            continue;
        }

        if (neighbor.index == excluded_index) {
            continue;
        }

        filtered.push_back(neighbor);
    }

    return filtered;
}


ValidationResult compare_outputs(
    const SearchOutput& expected,
    const SearchOutput& actual,
    float score_tolerance,
    float recognition_threshold)
{
    if (score_tolerance < 0.0f) {
        throw std::invalid_argument(
            "score_tolerance non valida"
        );
    }

    if (
        expected.query_count
            != actual.query_count
        || expected.database_count
            != actual.database_count
        || expected.dimension
            != actual.dimension
        || expected.top_k
            != actual.top_k
        || expected.neighbors.size()
            != actual.neighbors.size()
    ) {
        throw std::invalid_argument(
            "Output kNN non confrontabili"
        );
    }

    ValidationResult result;

    result.positions =
        expected.neighbors.size();

    for (
        std::size_t index = 0;
        index < result.positions;
        ++index
    ) {
        const Neighbor& expected_neighbor =
            expected.neighbors[index];

        const Neighbor& actual_neighbor =
            actual.neighbors[index];

        if (
            expected_neighbor.index
            == actual_neighbor.index
        ) {
            result.exact_index_matches += 1;
        }

        const float error = std::fabs(
            expected_neighbor.score
            - actual_neighbor.score
        );

        result.maximum_score_error =
            std::max(
                result.maximum_score_error,
                error
            );
    }

    for (
        std::size_t query_index = 0;
        query_index < expected.query_count;
        ++query_index
    ) {
        const auto expected_neighbors =
            query_neighbors(
                expected,
                query_index
            );

        const auto actual_neighbors =
            query_neighbors(
                actual,
                query_index
            );

        const bool expected_recognized =
            !expected_neighbors.empty()
            && expected_neighbors.front().score
                >= recognition_threshold;

        const bool actual_recognized =
            !actual_neighbors.empty()
            && actual_neighbors.front().score
                >= recognition_threshold;

        if (
            expected_recognized
            == actual_recognized
        ) {
            result.decision_matches += 1;
        }
    }

    result.exact_indices =
        result.exact_index_matches
        == result.positions;

    result.decisions_equal =
        result.decision_matches
        == expected.query_count;

    if (
        result.maximum_score_error
        > score_tolerance
    ) {
        result.decisions_equal = false;
    }

    return result;
}

}  // namespace face_knn