#pragma once

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

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

#include "config.hpp"
#include "population.hpp"
#include "fitness.hpp"
#include "evolution.hpp"
#include "spectral.hpp"
#include "oscillation.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>());
 */
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;

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

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

}  // namespace fces