#include <iostream>
#include <exception>
#include "../include/neuroflow/memory.hpp"

using namespace neuroflow;

int main() {
    try {
        std::cout << "MLA debug test..." << std::endl;
        
        LatentKVCache mla(64, 4, 16, 128);
        
        std::cout << "mla.d_model: " << mla.d_model << std::endl;
        std::cout << "mla.n_heads: " << mla.n_heads << std::endl;
        std::cout << "mla.d_latent: " << mla.d_latent << std::endl;
        std::cout << "mla.head_dim: " << mla.head_dim << std::endl;
        std::cout << "mla.cache_len: " << mla.cache_len << std::endl;
        
        Tensor input({1, 64});
        std::cout << "input shape: [" << input.shape_[0] << ", " << input.shape_[1] << "]" << std::endl;
        std::cout << "input numel: " << input.numel() << std::endl;
        
        for (size_t i = 0; i < input.numel(); ++i) input.as_fp32()[i] = 0.1f * i;
        
        std::cout << "Calling mla.forward(input, true)..." << std::endl;
        Tensor output1 = mla.forward(input, true);
        
        std::cout << "output1 shape size: " << output1.shape_.size() << std::endl;
        for (size_t i = 0; i < output1.shape_.size(); ++i) std::cout << "  dim " << i << ": " << output1.shape_[i] << std::endl;
        std::cout << "mla.cache_len after 1st forward: " << mla.cache_len << std::endl;
        
        std::cout << "Second forward..." << std::endl;
        Tensor input2({1, 64});
        for (size_t i = 0; i < input2.numel(); ++i) input2.as_fp32()[i] = 0.2f * i;
        Tensor output2 = mla.forward(input2, true);
        
        std::cout << "output2 shape size: " << output2.shape_.size() << std::endl;
        for (size_t i = 0; i < output2.shape_.size(); ++i) std::cout << "  dim " << i << ": " << output2.shape_[i] << std::endl;
        std::cout << "mla.cache_len after 2nd forward: " << mla.cache_len << std::endl;
        
        std::cout << "Success!" << std::endl;
        return 0;
    } catch (const std::exception& e) {
        std::cout << "Error: " << e.what() << std::endl;
        return 1;
    }
}