/**
 * @file telemetry_and_inference.cpp
 * @brief Example showcasing telemetry instrumentation and model inference.
 */

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

// Define a simple neural network for nonlinear regression: y = x^2
struct RegressionNet : torch::nn::Module {
  torch::nn::Linear fc1{nullptr}, fc2{nullptr};

  RegressionNet() {
    fc1 = register_module("fc1", torch::nn::Linear(1, 16));
    fc2 = register_module("fc2", torch::nn::Linear(16, 1));
  }

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

int main() {
  fces::Telemetry::get().info("app_start",
                              "Telemetry and Inference demo initialized.");

  // 1. Create Model and Data
  auto model = std::make_shared<RegressionNet>();

  // Generate training data: x in [-2, 2], y = x^2 + noise
  auto x_train = torch::linspace(-2.0, 2.0, 100).unsqueeze(1);
  auto y_train = x_train.pow(2) + 0.1 * torch::randn({100, 1});

  // 2. Configure Optimizer
  std::vector<torch::Tensor> params;
  for (auto &p : model->parameters()) {
    params.push_back(p);
  }

  fces::FCESOptimizer optimizer(
      params,
      fces::FCESConfig{}.set_lr(2e-3f).set_population_size(150).set_total_steps(
          100));

  fces::Telemetry::get().info("training_start",
                              "Beginning neural net optimization with FCES.");

  auto start_train = std::chrono::high_resolution_clock::now();

  // 3. Optimization Loop
  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 % 20 == 0) {
      fces::Telemetry::get().info(
          "epoch_checkpoint", "Epoch " + std::to_string(epoch) + " | Loss: " +
                                  std::to_string(loss.item<float>()));
    }
  }

  auto end_train = std::chrono::high_resolution_clock::now();
  double train_duration =
      std::chrono::duration<double, std::milli>(end_train - start_train)
          .count();

  fces::Telemetry::get().info("training_complete",
                              "Duration: " + std::to_string(train_duration) +
                                  " ms");

  // 4. Inference Phase
  fces::Telemetry::get().info("inference_phase_start",
                              "Evaluating model on new test inputs.");

  // Generate test inputs
  auto x_test = torch::tensor({-1.5f, -0.5f, 0.0f, 0.5f, 1.5f}).unsqueeze(1);
  auto y_expected = x_test.pow(2);

  // Switch model to evaluation mode
  model->eval();

  // Run inference and measure latency
  auto start_inf = std::chrono::high_resolution_clock::now();
  torch::Tensor y_pred;
  {
    torch::NoGradGuard no_grad;
    y_pred = model->forward(x_test);
  }
  auto end_inf = std::chrono::high_resolution_clock::now();
  double inf_duration =
      std::chrono::duration<double, std::milli>(end_inf - start_inf).count();

  // Log telemetry for inference performance
  fces::Telemetry::get().info(
      "inference_perf", "Inputs: " + std::to_string(x_test.size(0)) +
                            " | Latency: " + std::to_string(inf_duration) +
                            " ms");

  // Print predictions and expected values side-by-side
  std::cout << "\n================ INFERENCE RESULTS ================"
            << std::endl;
  std::cout << "Input (x)  |  Predicted (y_pred)  |  Expected (y_expected)"
            << std::endl;
  std::cout << "----------------------------------------------------"
            << std::endl;
  for (int i = 0; i < x_test.size(0); ++i) {
    float x_val = x_test[i][0].item<float>();
    float pred_val = y_pred[i][0].item<float>();
    float exp_val = y_expected[i][0].item<float>();
    std::printf("  %7.2f  |        %7.4f       |        %7.4f\n", x_val,
                pred_val, exp_val);
  }
  std::cout << "====================================================\n"
            << std::endl;

  fces::Telemetry::get().info("app_finish", "Exiting demo successfully.");
  return 0;
}