#include <cstdio>
#include <string>
#include <fstream>
#include <numeric>
#include <stdexcept>

#include <migraphx/onnx.hpp>

#include <migraphx/gpu/target.hpp>
#include <migraphx/gpu/hip.hpp>
#include <migraphx/generate.hpp>

#include "softmax.hpp"

auto reverse_int(unsigned int i)
{
    unsigned char c1;
    unsigned char c2;
    unsigned char c3;
    unsigned char c4;
    c1 = i & 255u;
    c2 = (i >> 8u) & 255u;
    c3 = (i >> 16u) & 255u;
    c4 = (i >> 24u) & 255u;
    return (static_cast<unsigned int>(c1) << 24u) + (static_cast<unsigned int>(c2) << 16u) +
           (static_cast<unsigned int>(c3) << 8u) + c4;
};

std::vector<float>
read_mnist_images(const std::string& full_path, int& number_of_images, int& image_size)
{
    using uchar = unsigned char;

    std::ifstream file(full_path, std::ios::binary);

    if(file.is_open())
    {
        int magic_number = 0;
        int n_rows       = 0;
        int n_cols       = 0;

        file.read(reinterpret_cast<char*>(&magic_number), sizeof(magic_number));
        magic_number = reverse_int(magic_number);

        if(magic_number != 2051)
            throw std::runtime_error("Invalid MNIST image file!");

        file.read(reinterpret_cast<char*>(&number_of_images), sizeof(number_of_images));
        number_of_images = reverse_int(number_of_images);
        file.read(reinterpret_cast<char*>(&n_rows), sizeof(n_rows));
        n_rows = reverse_int(n_rows);
        file.read(reinterpret_cast<char*>(&n_cols), sizeof(n_cols));
        n_cols = reverse_int(n_cols);

        image_size = n_rows * n_cols;

        std::vector<float> result(number_of_images * image_size);
        for(int i = 0; i < number_of_images; i++)
        {
            for(int j = 0; j < image_size; j++)
            {
                uchar tmp;
                file.read(reinterpret_cast<char*>(&tmp), 1);
                result[i * image_size + j] = tmp / 255.0;
            }
        }
        return result;
    }
    else
    {
        throw std::runtime_error("Cannot open file `" + full_path + "`!");
    }
}

std::vector<int32_t> read_mnist_labels(const std::string& full_path, int& number_of_labels)
{
    using uchar = unsigned char;

    std::ifstream file(full_path, std::ios::binary);

    if(file.is_open())
    {
        int magic_number = 0;
        file.read(reinterpret_cast<char*>(&magic_number), sizeof(magic_number));
        magic_number = reverse_int(magic_number);

        if(magic_number != 2049)
            throw std::runtime_error("Invalid MNIST label file!");

        file.read(reinterpret_cast<char*>(&number_of_labels), sizeof(number_of_labels));
        number_of_labels = reverse_int(number_of_labels);

        std::vector<int32_t> result(number_of_labels);
        for(int i = 0; i < number_of_labels; i++)
        {
            uchar tmp;
            file.read(reinterpret_cast<char*>(&tmp), 1);
            result[i] = tmp;
        }
        return result;
    }
    else
    {
        throw std::runtime_error("Unable to open file `" + full_path + "`!");
    }
}

int main(int argc, char const* argv[])
{
    if(argc > 3)
    {
        std::string datafile        = argv[2];
        std::string labelfile       = argv[3];
        int nimages                 = -1;
        int image_size              = -1;
        int nlabels                 = -1;
        std::vector<float> input    = read_mnist_images(datafile, nimages, image_size);
        std::vector<int32_t> labels = read_mnist_labels(labelfile, nlabels);

        std::string file = argv[1];
        auto prog        = migraphx::parse_onnx(file);
        std::cout << prog << std::endl << std::endl;
        prog.compile(migraphx::gpu::target{});
        auto s = migraphx::shape{migraphx::shape::float_type, {1, 1, 28, 28}};
        std::cout << s << std::endl;
        auto* ptr = input.data();
        migraphx::program::parameter_map m;
        m["output"] =
            migraphx::gpu::to_gpu(migraphx::generate_argument(prog.get_parameter_shape("output")));
        for(int i = 0; i < 20; i++)
        {
            std::cout << "label: " << labels[i] << "  ---->  ";
            m["0"]       = migraphx::gpu::to_gpu(migraphx::argument{s, &ptr[784 * i]});
            auto results = prog.eval(m).back();
            auto result  = migraphx::gpu::from_gpu(results);
            std::vector<float> logits;
            result.visit([&](auto output) { logits.assign(output.begin(), output.end()); });
            std::vector<float> probs = softmax(logits);
            for(auto x : probs)
                std::cout << x << "  ";
            std::cout << std::endl;
        }
        std::cout << std::endl;
    }
}