// Author: Simon-Pierre Boucher — contact@spboucher.ai
#include <Foundation/Foundation.hpp>
#include <Metal/Metal.hpp>
#include "core/device.h"
#include "core/tensor.h"
#include "ops/cpu/cpu_ops.h"
#include "ops/metal/metal_ops.h"
#include <cmath>
#include <cstdio>
#include <random>
#include <string>
using namespace forge;
struct Params { uint32_t M, N, K; };
int main() {
    NS::AutoreleasePool* pool = NS::AutoreleasePool::alloc()->init();
    const int64_t M = 128, N = 64, K = 96;   // exact multiples of the 64x32 tile
    Tensor a = Tensor::empty({M, K}), b = Tensor::empty({K, N});
    Tensor c = Tensor::zeros({M, N}), ref = Tensor::empty({M, N});
    std::mt19937 rng(5);
    std::uniform_real_distribution<float> d(-1.f, 1.f);
    for (int64_t i = 0; i < a.numel(); ++i) a.data<float>()[i] = d(rng);
    for (int64_t i = 0; i < b.numel(); ++i) b.data<float>()[i] = d(rng);
    cpu::matmul(a, b, ref, false, false);

    Device& dev = Device::get();
    MTL::ComputePipelineState* pso = dev.pipeline("matmul_mpp_f32");
    Params p{uint32_t(M), uint32_t(N), uint32_t(K)};
    MTL::ComputeCommandEncoder* enc = metal::Stream::get().encoder();
    enc->setComputePipelineState(pso);
    enc->setBuffer(a.buffer(), a.buffer_offset(), 0);
    enc->setBuffer(b.buffer(), b.buffer_offset(), 1);
    enc->setBuffer(c.buffer(), c.buffer_offset(), 2);
    enc->setBytes(&p, sizeof(p), 3);
    enc->dispatchThreadgroups(MTL::Size(N/32, M/64, 1), MTL::Size(128,1,1));
    metal::sync();

    float worst = 0.f; int64_t nz = 0;
    for (int64_t i = 0; i < c.numel(); ++i) {
        worst = std::max(worst, std::fabs(c.data<float>()[i] - ref.data<float>()[i]));
        if (c.data<float>()[i] != 0.f) ++nz;
    }
    std::printf("MPP matmul2d f32 vs CPU: max abs err %.3e  (nonzero outputs %lld/%lld)\n",
                double(worst), (long long)nz, (long long)c.numel());
    const bool ok32 = worst <= 1e-4f && nz == c.numel();
    std::printf("%s\n", ok32 ? "CORRECT" : "WRONG");

    // f16 operands, f32 accumulator: the headline throughput number, so check
    // it computes the same product within f16 input precision.
    Tensor ah = Tensor::empty({M, K}, DType::F16);
    Tensor bh = Tensor::empty({K, N}, DType::F16);
    Tensor ch = Tensor::zeros({M, N});
    Tensor refh = Tensor::empty({M, N});
    for (int64_t i = 0; i < a.numel(); ++i) ah.set_item(i, a.data<float>()[i]);
    for (int64_t i = 0; i < b.numel(); ++i) bh.set_item(i, b.data<float>()[i]);
    // reference from the ROUNDED f16 values, so we compare arithmetic not rounding
    Tensor ar = Tensor::empty({M, K}), br = Tensor::empty({K, N});
    for (int64_t i = 0; i < a.numel(); ++i) ar.data<float>()[i] = ah.item_at(i);
    for (int64_t i = 0; i < b.numel(); ++i) br.data<float>()[i] = bh.item_at(i);
    cpu::matmul(ar, br, refh, false, false);

    MTL::ComputePipelineState* pso16 = dev.pipeline("matmul_mpp_f16");
    MTL::ComputeCommandEncoder* e2 = metal::Stream::get().encoder();
    e2->setComputePipelineState(pso16);
    e2->setBuffer(ah.buffer(), ah.buffer_offset(), 0);
    e2->setBuffer(bh.buffer(), bh.buffer_offset(), 1);
    e2->setBuffer(ch.buffer(), ch.buffer_offset(), 2);
    e2->setBytes(&p, sizeof(p), 3);
    e2->dispatchThreadgroups(MTL::Size(N/32, M/64, 1), MTL::Size(128,1,1));
    metal::sync();

    float w16 = 0.f; int64_t nz16 = 0;
    for (int64_t i = 0; i < ch.numel(); ++i) {
        w16 = std::max(w16, std::fabs(ch.data<float>()[i] - refh.data<float>()[i]));
        if (ch.data<float>()[i] != 0.f) ++nz16;
    }
    const bool ok16 = w16 <= 1e-3f && nz16 == ch.numel();
    std::printf("MPP matmul2d f16 vs CPU(f16 inputs): max abs err %.3e  "
                "(nonzero %lld/%lld)\n%s\n", double(w16), (long long)nz16,
                (long long)ch.numel(), ok16 ? "CORRECT" : "WRONG");
    // Transposed variants — training needs nt (forward X.W^T) and tn (dW = dY^T.X).
    bool okT = true;
    struct TCase { const char* suffix; bool ta, tb; const char* tag; };
    const TCase tcases[] = {{"_nt", false, true, "nt (fwd X.Wt)"},
                            {"_tn", true, false, "tn (dW = dYt.X)"}};
    for (const TCase& tc : tcases) {
        Tensor at = tc.ta ? Tensor::empty({K, M}) : Tensor::empty({M, K});
        Tensor bt = tc.tb ? Tensor::empty({N, K}) : Tensor::empty({K, N});
        Tensor ct = Tensor::zeros({M, N}), rt = Tensor::empty({M, N});
        for (int64_t i = 0; i < at.numel(); ++i) at.data<float>()[i] = d(rng);
        for (int64_t i = 0; i < bt.numel(); ++i) bt.data<float>()[i] = d(rng);
        cpu::matmul(at, bt, rt, tc.ta, tc.tb);

        std::string kn = std::string("matmul_mpp_f32") + tc.suffix;
        MTL::ComputePipelineState* ps = dev.pipeline(kn);
        MTL::ComputeCommandEncoder* e3 = metal::Stream::get().encoder();
        e3->setComputePipelineState(ps);
        e3->setBuffer(at.buffer(), at.buffer_offset(), 0);
        e3->setBuffer(bt.buffer(), bt.buffer_offset(), 1);
        e3->setBuffer(ct.buffer(), ct.buffer_offset(), 2);
        e3->setBytes(&p, sizeof(p), 3);
        e3->dispatchThreadgroups(MTL::Size(N/32, M/64, 1), MTL::Size(128,1,1));
        metal::sync();

        float w = 0.f; int64_t nzt = 0;
        for (int64_t i = 0; i < ct.numel(); ++i) {
            w = std::max(w, std::fabs(ct.data<float>()[i] - rt.data<float>()[i]));
            if (ct.data<float>()[i] != 0.f) ++nzt;
        }
        const bool ok = w <= 1e-4f && nzt == ct.numel();
        okT = okT && ok;
        std::printf("MPP matmul2d f32 %-16s max abs err %.3e  (nonzero %lld/%lld)  %s\n",
                    tc.tag, double(w), (long long)nzt, (long long)ct.numel(),
                    ok ? "CORRECT" : "WRONG");
    }

    pool->drain();
    return (ok32 && ok16 && okT) ? 0 : 1;
}