/**
 * @file simple_optimization.cpp
 * @brief Minimal example: optimize a quadratic function with FCES.
 */

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

int main() {
    // Target: minimize f(x) = ||x - target||^2
    auto target = torch::tensor({1.0f, 2.0f, 3.0f, 4.0f, 5.0f});
    auto x = torch::randn({5}, torch::requires_grad());

    std::vector<torch::Tensor> params = {x};
    fces::FCESOptimizer optimizer(params, fces::FCESConfig{}.set_lr(1e-2f));

    for (int step = 0; step < 500; ++step) {
        optimizer.zero_grad();
        auto loss = (x - target).pow(2).sum();
        loss.backward();
        optimizer.step();
        optimizer.update_fitness(loss.item<float>());

        if (step % 50 == 0) {
            std::cout << "Step " << step
                      << " | Loss: " << loss.item<float>()
                      << " | x: " << x << std::endl;
        }
    }

    std::cout << "\nFinal x: " << x << std::endl;
    std::cout << "Target:  " << target << std::endl;

    return 0;
}