#include "ijk.h"

#include <algorithm>
#include <cassert>
#include <chrono>
#include <iomanip>
#include <iostream>
#include <random>
#include <stdfloat>
#include <type_traits>

template<class T, class F> void test_type(F &f)
{
    T a[N][N], b[N][N], c[N][N];

    for (size_t i = 0; i < N; ++i)
        for (size_t j = 0; j < N; ++j)
        {
            a[i][j] = 0;
            b[i][j] = B;
            c[i][j] = C;
        }

    const size_t R = 33;
    std::chrono::duration<std::float64_t, std::milli> rep[R];

    for (auto &r : rep)
    {
        auto start = std::chrono::high_resolution_clock::now();
        f(a, b, c);
        auto stop = std::chrono::high_resolution_clock::now();
        r = stop - start;
    }

    std::nth_element(rep, rep + R / 2, rep + R);
    std::cout << std::fixed << std::setprecision(2) << std::setw(8)
              << rep[R / 2].count();

    if constexpr (std::is_integral_v<T>)
        assert(a[N - 1][N - 1] == T(B * C * N * R));
};

#define test_function(f)                                                  \
    std::cout << std::setw(8) << #f;                                      \
    test_type<int8_t>(f<int8_t>);                                         \
    test_type<int16_t>(f<int16_t>);                                       \
    test_type<int32_t>(f<int32_t>);                                       \
    test_type<int64_t>(f<int64_t>);                                       \
    test_type<std::bfloat16_t>(f<std::bfloat16_t>);                       \
    test_type<std::float16_t>(f<std::float16_t>);                         \
    test_type<std::float32_t>(f<std::float32_t>);                         \
    test_type<std::float64_t>(f<std::float64_t>);                         \
    test_type<std::float128_t>(f<std::float128_t>);                       \
    std::cout << std::endl;

int32_t main()
{
    const char *types[] = {"(ms)",
                           "int8",
                           "int16",
                           "int32",
                           "int64",
                           "bfloa16",
                           "float16",
                           "float32",
                           "float64",
                           "floa128"};

    for (auto &i : types)
        std::cout << std::setw(8) << i;
    std::cout << std::endl;

    for (size_t i = 0; i < sizeof(types) / sizeof(char *); ++i)
        std::cout << std::setw(8) << "-------";
    std::cout << std::endl;

    test_function(ijk);
    test_function(ikj);
    test_function(jik);
    test_function(jki);
    test_function(kij);
    test_function(kji);
}