#pragma once

/**
 * @file optimizer.hpp
 * @brief FCESOptimizer — the main entry point. libtorch-compatible optimizer.
 */

#include <memory>
#include <string>
#include <torch/torch.h>
#include <vector>

#include "config.hpp"
#include "evolution.hpp"
#include "fitness.hpp"
#include "oscillation.hpp"
#include "population.hpp"
#include "spectral.hpp"
#include "telemetry.hpp"

namespace fces {

/**
 * FCESOptimizer — Fuzzy Controlled Evolutionary Search V49.0 (C++ Port).
 *
 * Usage:
 *   auto optimizer = FCESOptimizer(model->parameters(),
 * FCESConfig{}.set_lr(1.6e-3));
 *   // In training loop:
 *   optimizer.zero_grad();
 *   auto loss = model->forward(input);
 *   loss.backward();
 *   optimizer.step();
 *   optimizer.update_fitness(loss.item<float>());
 */
struct FCESOptimizerOptions
    : public torch::optim::OptimizerCloneableOptions<FCESOptimizerOptions> {
  explicit FCESOptimizerOptions(double lr = 0.01) : lr_(lr) {}

  double get_lr() const override { return lr_; }
  void set_lr(const double lr) override { lr_ = lr; }

  double lr_;
};

class FCESOptimizer : public torch::optim::Optimizer {
public:
  explicit FCESOptimizer(std::vector<torch::Tensor> params,
                         FCESConfig config = FCESConfig{});

  /// Perform a single optimization step
  torch::Tensor step(LossClosure closure = nullptr) override;

  /// Update evolutionary fitness with current loss
  void update_fitness(float loss);

  /// Backup model weights to CPU RAM
  void backup_to_ram();

  /// Restore model weights from CPU RAM backup
  void restore_from_ram();

  /// Get current step count
  int step_count() const { return step_counter_; }

  /// Calculate model sparsity
  float calculate_sparsity() const;

private:
  FCESConfig config_;
  Population population_;
  FitnessEngine fitness_engine_;
  FuzzyFitnessEvaluator fitness_evaluator_;
  std::unique_ptr<EvolutionManager> evolution_manager_;
  OscillationDetector oscillation_detector_;
  RunningStats grad_norm_tracker_;

  // State
  int step_counter_ = 0;
  float ema_loss_ = 0.0f;
  float last_step_loss_ = 0.0f;
  float best_loss_window_ = std::numeric_limits<float>::infinity();
  float rollback_ema_ = 0.0f;
  int stagnation_counter_ = 0;
  float last_loss_velocity_ = 0.0f;
  float last_sparsity_ = 0.0f;

  // RAM backup
  std::vector<torch::Tensor> ram_backup_;

  // Layer stats and group mappings
  std::vector<std::vector<float>> layer_stats_;
  std::vector<int> param_group_mapping_;
  std::unique_ptr<SpectralSensor> spectral_sensor_;
  SpectralController spectral_controller_;
  float last_spectral_rank_ = 0.0f;

  // Internal methods
  void gather_stats();
  void apply_parameter_updates(const torch::Tensor &actions);
  void handle_rollback();
};

} // namespace fces