/**
 * @file pytorch_integration.cpp
 * @brief Example: train a small neural network with FCES via libtorch.
 */

#include "fces/optimizer.hpp"
#include <iostream>
#include <torch/torch.h>

struct TinyNet : torch::nn::Module {
  torch::nn::Linear fc1{nullptr}, fc2{nullptr};

  TinyNet()
      : fc1(register_module("fc1", torch::nn::Linear(10, 32))),
        fc2(register_module("fc2", torch::nn::Linear(32, 1))) {}

  torch::Tensor forward(torch::Tensor x) {
    x = torch::relu(fc1->forward(x));
    return fc2->forward(x);
  }
};

int main() {
  auto model = std::make_shared<TinyNet>();

  std::vector<torch::Tensor> params;
  for (auto &p : model->parameters())
    params.push_back(p);

  fces::FCESOptimizer optimizer(params, fces::FCESConfig{}
                                            .set_lr(1.6e-3f)
                                            .set_population_size(200)
                                            .set_total_steps(1000));

  // Generate synthetic regression data
  auto x_train = torch::randn({100, 10});
  auto y_train = torch::sin(x_train.sum(1, true));

  for (int epoch = 0; epoch < 100; ++epoch) {
    optimizer.zero_grad();
    auto pred = model->forward(x_train);
    auto loss = torch::mse_loss(pred, y_train);
    loss.backward();
    optimizer.step();
    optimizer.update_fitness(loss.item<float>());

    if (epoch % 10 == 0) {
      std::cout << "Epoch " << epoch << " | Loss: " << loss.item<float>()
                << std::endl;
    }
  }

  std::cout << "\nTraining complete. Final loss: "
            << torch::mse_loss(model->forward(x_train), y_train).item<float>()
            << std::endl;

  return 0;
}