// #include <cblas.h>
#include <cmath>
#include <chrono>
#include <iostream>
#include <limits>
#include <random>
#include <vector>
#include <algorithm>

// ---------------- utilities ----------------
template <class T>
inline void DoNotOptimize(T&& value)
{
#if defined(__clang__)
    asm volatile("" : "+r,m"(value) : : "memory");
#else
    asm volatile("" : "+m,r"(value) : : "memory");
#endif
}

// A C++ implementation to numpy.allclose()
bool allclose(const std::vector<float>& a,
              const std::vector<float>& b,
              float rtol = 1e-5f,
              float atol = 1e-6f)
{
    if (a.size() != b.size()) return false;
    for (size_t i = 0; i < a.size(); i++) {
        float diff = std::fabs(a[i] - b[i]);
        float tol = atol + rtol * std::fabs(b[i]);
        if (diff > tol) return false;
    }
    return true;
}

static inline float row_max(const float* row, int len) {
    float m = -std::numeric_limits<float>::infinity();
    for (int i = 0; i < len; ++i) m = std::max(m, row[i]);
    return m;
}
static inline float row_sum(const float* row, int len) {
    float s = 0.f;
    for (int i = 0; i < len; ++i) s += row[i];
    return s;
}
static inline void random_matrix(std::vector<float>& A) {
    std::mt19937 rng(42);
    std::uniform_real_distribution<float> dist(-1.f, 1.f);
    for (auto& x : A) x = dist(rng);
}

// ---------------- naive attention ----------------
std::vector<float> naive_attention(const float* __restrict__ Q, const float* __restrict__ Kt, const float* __restrict__ V,
                                   int N, int d, float scale)
{
    std::vector<float> O(N * d, 0.f);
    return O;
}

// --------------- FlashAttention ---------------
std::vector<float> flash_attention(const float* __restrict__ Q, const float* __restrict__ Kt, const float* __restrict__ V,
                                        int N, int d, int M_bytes, float scale)
{
    std::vector<float> O(N * d, 0.f);
    return O;
}

// --------- Main: Compare correctness + timing ----------
int main() {
    int N = 1024;   // sequence length
    int d = 64;    // head dimension
    float scale = 1.0f / std::sqrt(float(d));

    std::vector<int> M_sizes;
    for (int i = 2; (1 << i) <= 160*1024; i++) {
        int base = (1 << i) * 1024;
        M_sizes.push_back(base);
        int sqrt2 = int(base * std::sqrt(2.0));
        M_sizes.push_back(sqrt2);
    }

    std::vector<float> Q(N*d), Kt(d*N), V(N*d);
    random_matrix(Q);
    random_matrix(Kt);
    random_matrix(V);

    // --- warm-up (important) ---
    (void)naive_attention(Q.data(), Kt.data(), V.data(), N, d, scale);
    // --- time naive ---
    auto t1 = std::chrono::steady_clock::now();
    auto O_naive = naive_attention(Q.data(), Kt.data(), V.data(), N, d, scale);
    auto t2 = std::chrono::steady_clock::now();
    DoNotOptimize(O_naive.data());
    double dt_naive = std::chrono::duration<double, std::milli>(t2-t1).count();

    for (int M_bytes : M_sizes) {
        // --- warm-up (important) ---
        (void)flash_attention(Q.data(), Kt.data(), V.data(), N, d, M_bytes, scale);
        // --- time flash ---
        auto t3 = std::chrono::steady_clock::now();
        auto O_flash = flash_attention(Q.data(), Kt.data(), V.data(), N, d, M_bytes, scale);
        auto t4 = std::chrono::steady_clock::now();
        DoNotOptimize(O_flash.data());

        double dt_flash = std::chrono::duration<double, std::milli>(t4-t3).count();

        std::cout << "M=" << M_bytes/1024 << " KiB\n";
        std::cout << "\tNaive time: " << dt_naive << " ms\n";
        std::cout << "\tFlashAttention time: " << dt_flash << " ms\n";
        if (allclose(O_naive, O_flash)) {
            std::cout << "\tOutputs match within tolerance\n";
        }
        else {
            std::cout << "\tOutputs differ!\n";
        }
    }
}