#include "knn_opencl.h"

#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <string>
#include <vector>

#include <cstdio>
#include <unistd.h>

namespace {


class ScopedStdoutToStderr {
public:
    ScopedStdoutToStderr()
    {
        std::cout.flush();
        std::fflush(stdout);

        saved_stdout_ =
            ::dup(STDOUT_FILENO);

        if (saved_stdout_ < 0) {
            throw std::runtime_error(
                "Impossibile salvare stdout"
            );
        }

        if (
            ::dup2(
                STDERR_FILENO,
                STDOUT_FILENO
            ) < 0
        ) {
            ::close(saved_stdout_);
            saved_stdout_ = -1;

            throw std::runtime_error(
                "Impossibile deviare stdout"
            );
        }
    }

    ~ScopedStdoutToStderr()
    {
        if (saved_stdout_ < 0) {
            return;
        }

        std::cout.flush();
        std::fflush(stdout);

        ::dup2(
            saved_stdout_,
            STDOUT_FILENO
        );

        ::close(saved_stdout_);
    }

    ScopedStdoutToStderr(
        const ScopedStdoutToStderr&
    ) = delete;

    ScopedStdoutToStderr& operator=(
        const ScopedStdoutToStderr&
    ) = delete;

private:
    int saved_stdout_ = -1;
};


struct Arguments {
    std::string queries_path;

    std::string database_path;

    std::string self_map_path;

    std::string backend = "cpu";

    std::size_t query_count = 0;

    std::size_t database_count = 0;

    std::uint32_t dimension = 128;

    std::uint32_t top_k = 1;

    float threshold = 0.50f;
};


void print_usage(
    const char* program)
{
    std::cerr
        << "Uso:\n"
        << "  "
        << program
        << " --queries FILE"
        << " --database FILE"
        << " --query-count N"
        << " --database-count N"
        << " --dimension N"
        << " --top-k N"
        << " --backend cpu"
        << " [--threshold SCORE]"
        << " [--self-map FILE]"
        << "\n";
}


std::string require_value(
    int argc,
    char** argv,
    int& index)
{
    if (index + 1 >= argc) {
        throw std::invalid_argument(
            std::string(
                "Valore mancante per "
            )
            + argv[index]
        );
    }

    index += 1;

    return argv[index];
}


std::size_t parse_size(
    const std::string& text,
    const char* name)
{
    std::size_t consumed = 0;

    const auto value =
        std::stoull(
            text,
            &consumed,
            10
        );

    if (
        consumed != text.size()
        || value == 0
    ) {
        throw std::invalid_argument(
            std::string(name)
            + " non valido: "
            + text
        );
    }

    return static_cast<std::size_t>(
        value
    );
}


std::uint32_t parse_u32(
    const std::string& text,
    const char* name)
{
    const std::size_t value =
        parse_size(
            text,
            name
        );

    if (
        value
        > std::numeric_limits<
            std::uint32_t
        >::max()
    ) {
        throw std::invalid_argument(
            std::string(name)
            + " supera uint32"
        );
    }

    return static_cast<std::uint32_t>(
        value
    );
}


float parse_float(
    const std::string& text,
    const char* name)
{
    std::size_t consumed = 0;

    const float value =
        std::stof(
            text,
            &consumed
        );

    if (consumed != text.size()) {
        throw std::invalid_argument(
            std::string(name)
            + " non valido: "
            + text
        );
    }

    return value;
}


Arguments parse_arguments(
    int argc,
    char** argv)
{
    Arguments arguments;

    for (
        int index = 1;
        index < argc;
        ++index
    ) {
        const std::string option =
            argv[index];

        if (
            option == "--help"
            || option == "-h"
        ) {
            print_usage(argv[0]);
            std::exit(0);
        }

        if (option == "--queries") {
            arguments.queries_path =
                require_value(
                    argc,
                    argv,
                    index
                );
        } else if (
            option == "--database"
        ) {
            arguments.database_path =
                require_value(
                    argc,
                    argv,
                    index
                );
        } else if (
            option == "--self-map"
        ) {
            arguments.self_map_path =
                require_value(
                    argc,
                    argv,
                    index
                );
        } else if (
            option == "--backend"
        ) {
            arguments.backend =
                require_value(
                    argc,
                    argv,
                    index
                );
        } else if (
            option == "--query-count"
        ) {
            arguments.query_count =
                parse_size(
                    require_value(
                        argc,
                        argv,
                        index
                    ),
                    "query-count"
                );
        } else if (
            option == "--database-count"
        ) {
            arguments.database_count =
                parse_size(
                    require_value(
                        argc,
                        argv,
                        index
                    ),
                    "database-count"
                );
        } else if (
            option == "--dimension"
        ) {
            arguments.dimension =
                parse_u32(
                    require_value(
                        argc,
                        argv,
                        index
                    ),
                    "dimension"
                );
        } else if (
            option == "--top-k"
        ) {
            arguments.top_k =
                parse_u32(
                    require_value(
                        argc,
                        argv,
                        index
                    ),
                    "top-k"
                );
        } else if (
            option == "--threshold"
        ) {
            arguments.threshold =
                parse_float(
                    require_value(
                        argc,
                        argv,
                        index
                    ),
                    "threshold"
                );
        } else {
            throw std::invalid_argument(
                "Opzione sconosciuta: "
                + option
            );
        }
    }

    if (
        arguments.queries_path.empty()
        || arguments.database_path.empty()
        || arguments.query_count == 0
        || arguments.database_count == 0
    ) {
        throw std::invalid_argument(
            "Argomenti obbligatori mancanti"
        );
    }

    if (
        arguments.backend != "cpu"
        && arguments.backend != "opencl"
    ) {
        throw std::invalid_argument(
            "Backend non supportato: "
            + arguments.backend
        );
    }

    return arguments;
}


std::vector<float> read_float_file(
    const std::string& path,
    std::size_t expected_values)
{
    std::ifstream input(
        path,
        std::ios::binary
    );

    if (!input) {
        throw std::runtime_error(
            "Impossibile aprire: "
            + path
        );
    }

    input.seekg(
        0,
        std::ios::end
    );

    const std::streamoff byte_count =
        input.tellg();

    input.seekg(
        0,
        std::ios::beg
    );

    const auto expected_bytes =
        static_cast<std::streamoff>(
            expected_values
            * sizeof(float)
        );

    if (byte_count != expected_bytes) {
        throw std::runtime_error(
            "Dimensione file incoerente: "
            + path
            + ", byte="
            + std::to_string(byte_count)
            + ", attesi="
            + std::to_string(
                expected_bytes
            )
        );
    }

    std::vector<float> values(
        expected_values
    );

    input.read(
        reinterpret_cast<char*>(
            values.data()
        ),
        byte_count
    );

    if (!input) {
        throw std::runtime_error(
            "Lettura incompleta: "
            + path
        );
    }

    return values;
}


std::vector<std::int64_t> read_self_map(
    const std::string& path,
    std::size_t query_count)
{
    std::vector<std::int64_t> values(
        query_count,
        -1
    );

    if (path.empty()) {
        return values;
    }

    std::ifstream input(path);

    if (!input) {
        throw std::runtime_error(
            "Impossibile aprire self-map: "
            + path
        );
    }

    std::size_t index = 0;
    std::int64_t value = -1;

    while (input >> value) {
        if (index >= query_count) {
            throw std::runtime_error(
                "Self-map contiene "
                "troppe righe"
            );
        }

        values[index] = value;
        index += 1;
    }

    if (index != query_count) {
        throw std::runtime_error(
            "Self-map contiene "
            + std::to_string(index)
            + " righe, attese "
            + std::to_string(
                query_count
            )
        );
    }

    return values;
}


void write_json(
    const face_knn::SearchOutput& output,
    const std::vector<std::int64_t>& self_map,
    float threshold)
{
    std::size_t recognized_count = 0;

    std::cout
        << std::setprecision(9)
        << "{\n"
        << "  \"backend\": \""
        << output.backend
        << "\",\n"
        << "  \"device_name\": \""
        << output.device_name
        << "\",\n"
        << "  \"query_count\": "
        << output.query_count
        << ",\n"
        << "  \"database_count\": "
        << output.database_count
        << ",\n"
        << "  \"dimension\": "
        << output.dimension
        << ",\n"
        << "  \"top_k\": "
        << output.top_k
        << ",\n"
        << "  \"threshold\": "
        << threshold
        << ",\n"
        << "  \"leave_one_out\": "
        << (
            self_map.empty()
                ? "false"
                : "true"
        )
        << ",\n"
        << "  \"elapsed_ms\": "
        << output.elapsed_ms
        << ",\n"
        << "  \"results\": [\n";

    for (
        std::size_t query_index = 0;
        query_index < output.query_count;
        ++query_index
    ) {
        auto neighbors =
            face_knn::query_neighbors(
                output,
                query_index
            );

        const std::int64_t excluded =
            self_map[query_index];

        if (excluded >= 0) {
            neighbors =
                face_knn::exclude_neighbor(
                    neighbors,
                    static_cast<
                        std::uint32_t
                    >(excluded)
                );
        }

        if (neighbors.empty()) {
            throw std::runtime_error(
                "Nessun vicino disponibile "
                "dopo leave-one-out"
            );
        }

        const auto& best =
            neighbors.front();

        const bool recognized =
            best.score >= threshold;

        if (recognized) {
            recognized_count += 1;
        }

        std::cout
            << "    {\n"
            << "      \"query_index\": "
            << query_index
            << ",\n"
            << "      \"excluded_index\": ";

        if (excluded < 0) {
            std::cout << "null";
        } else {
            std::cout << excluded;
        }

        std::cout
            << ",\n"
            << "      \"recognized\": "
            << (
                recognized
                    ? "true"
                    : "false"
            )
            << ",\n"
            << "      \"best_index\": "
            << best.index
            << ",\n"
            << "      \"best_score\": "
            << best.score
            << ",\n"
            << "      \"neighbors\": [\n";

        for (
            std::size_t position = 0;
            position < neighbors.size();
            ++position
        ) {
            const auto& neighbor =
                neighbors[position];

            std::cout
                << "        {"
                << "\"index\": "
                << neighbor.index
                << ", \"score\": "
                << neighbor.score
                << "}";

            if (
                position + 1
                < neighbors.size()
            ) {
                std::cout << ",";
            }

            std::cout << "\n";
        }

        std::cout
            << "      ]\n"
            << "    }";

        if (
            query_index + 1
            < output.query_count
        ) {
            std::cout << ",";
        }

        std::cout << "\n";
    }

    const std::size_t below_count =
        output.query_count
        - recognized_count;

    const double coverage =
        output.query_count > 0
            ? (
                static_cast<double>(
                    recognized_count
                )
                / static_cast<double>(
                    output.query_count
                )
                * 100.0
            )
            : 0.0;

    std::cout
        << "  ],\n"
        << "  \"recognized_count\": "
        << recognized_count
        << ",\n"
        << "  \"below_threshold_count\": "
        << below_count
        << ",\n"
        << "  \"coverage_percent\": "
        << coverage
        << "\n"
        << "}\n";
}


}  // namespace


int main(
    int argc,
    char** argv)
{
    try {
        const Arguments arguments =
            parse_arguments(
                argc,
                argv
            );

        const std::size_t query_values =
            arguments.query_count
            * arguments.dimension;

        const std::size_t database_values =
            arguments.database_count
            * arguments.dimension;

        const auto queries =
            read_float_file(
                arguments.queries_path,
                query_values
            );

        const auto database =
            read_float_file(
                arguments.database_path,
                database_values
            );

        const auto self_map =
            read_self_map(
                arguments.self_map_path,
                arguments.query_count
            );

        face_knn::SearchOptions options;

        options.dimension =
            arguments.dimension;

        options.top_k =
            arguments.top_k;

        options.kernel_path =
            "kernels/knn_float4_topk.cl";

        face_knn::SearchOutput output;

        if (
            arguments.backend == "opencl"
        ) {
            /*
             * Il runtime Mali scrive la propria versione
             * direttamente su stdout. Durante OpenCL
             * deviamo stdout su stderr, affinché l'output
             * JSON resti puro.
             */
            ScopedStdoutToStderr redirect;

            output =
                face_knn::search_opencl(
                    database,
                    queries,
                    arguments.database_count,
                    arguments.query_count,
                    options
                );
        } else {
            output =
                face_knn::search_cpu(
                    database,
                    queries,
                    arguments.database_count,
                    arguments.query_count,
                    options
                );
        }

        write_json(
            output,
            self_map,
            arguments.threshold
        );

        return 0;
    } catch (const std::exception& error) {
        std::cerr
            << "[face_knn] errore: "
            << error.what()
            << "\n";

        print_usage(argv[0]);

        return 1;
    }
}