#include <iostream>
#include "../include/neuroflow/model.hpp"

using namespace neuroflow;

int main() {
    std::cout << "Simple forward test..." << std::endl;
    
    NeuroFlowModel::Config cfg;
    cfg.input_dim = 64;
    cfg.hidden_dim = 32;
    cfg.output_dim = 5;
    cfg.memory_slots = 8;
    cfg.memory_dim = 16;
    cfg.num_layers = 1;
    cfg.num_associations = 2;
    cfg.use_mla = false;
    
    NeuroFlowModel model(cfg);
    
    Tensor input({2, cfg.input_dim});
    for (size_t i = 0; i < input.numel(); ++i) input.as_fp32()[i] = 0.1f * i;
    
    std::cout << "Calling forward..." << std::endl;
    auto output = model.forward(input);
    
    std::cout << "Success!" << std::endl;
    return 0;
}