#pragma once

/**
 * @file controller.hpp
 * @brief FuzzyController and Genome — the decision-making units of FCES.
 *
 * Each controller contains a Genome (neural network weights) that maps
 * layer statistics to update decisions (multiplier, sign_gate, wd_mult).
 *
 * Port of: packages/fces/core/controller.py
 */

#include <array>
#include <cstdint>
#include <memory>
#include <random>
#include <string>
#include <vector>

#include <torch/torch.h>

namespace fces {

// Controller input dimension (layer stats features)
constexpr int GENOME_INPUT_DIM = 14;
// Controller hidden dimension
constexpr int GENOME_HIDDEN_DIM = 8;
// Controller output dimension: [multiplier, sign_gate, wd_mult]
constexpr int GENOME_OUTPUT_DIM = 3;
// Total genome size: input->hidden weights + hidden biases + hidden->output
// weights + output biases
constexpr int GENOME_SIZE =
    (GENOME_INPUT_DIM * GENOME_HIDDEN_DIM) +  // input -> hidden weights
    GENOME_HIDDEN_DIM +                       // hidden biases
    (GENOME_HIDDEN_DIM * GENOME_OUTPUT_DIM) + // hidden -> output weights
    GENOME_OUTPUT_DIM;                        // output biases

/**
 * Genome — the "DNA" of a fuzzy controller.
 * A flat array of floats encoding a micro-MLP.
 */
struct Genome {
  std::array<float, GENOME_SIZE> weights{};
  std::array<float, GENOME_SIZE> gene_success{};
  float sigma_gene = 0.1f;
  float plasticity = 1.0f;

  /// Initialize with random weights from a normal distribution
  void randomize(std::mt19937 &rng);

  /// Deep copy
  Genome clone() const;
};

/**
 * FuzzyController — a single agent in the evolutionary population.
 *
 * Lifecycle:
 * 1. Created via random initialization or crossover/mutation
 * 2. Activated for `selection_interval` steps
 * 3. Evaluated based on loss improvement during its tenure
 * 4. Evolved (crossover/mutation) or culled based on fitness
 */
class FuzzyController {
public:
  /// Unique identifier
  uint64_t id;

  /// The neural genome
  Genome genome;

  /// Fitness scores
  float fitness = 0.0f;
  float lifetime_fitness = 0.0f;
  float ema_fitness = 0.0f;
  int evaluation_count = 0;
  int age = 0;

  /// Origin tracking
  std::string origin = "random";

  /// Trust region violation counter
  int trust_violations = 0;

  /// Rolling fitness history (for Phase 23 strategies)
  std::vector<float> fitness_history;

  // ---------------------------------------------------------------
  // Construction
  // ---------------------------------------------------------------

  FuzzyController();
  explicit FuzzyController(Genome genome);

  // ---------------------------------------------------------------
  // Core Operations
  // ---------------------------------------------------------------

  /**
   * Forward pass through the micro-MLP to produce update decisions.
   *
   * @param layer_stats  Vector of per-layer feature maps
   * @param loss_trend   Current loss velocity
   * @param step_pct     Training progress [0, 1]
   * @param rollback_rate Rolling average rollback frequency
   * @param grad_stability Gradient coefficient of variation
   * @param spectral_alpha Log spectral rank
   * @param stagnation_intensity Stagnation counter / 500
   * @param kzm_damping Kibble-Zurek damping factor
   * @param projected_drift Projected loss drift
   * @return Tensor of shape [num_groups, 3] — (mult, sign_gate, wd_mult)
   */
  torch::Tensor
  decide_update(const std::vector<std::vector<float>> &layer_stats,
                float loss_trend, float step_pct, float rollback_rate,
                float grad_stability, float spectral_alpha,
                float stagnation_intensity, float kzm_damping,
                float projected_drift);

  // ---------------------------------------------------------------
  // Evolutionary Operators
  // ---------------------------------------------------------------

  /// Create a mutated child
  FuzzyController mutate(float current_loss, float sigma_scale = 1.0f) const;

  /// Crossover with another controller
  FuzzyController crossover(const FuzzyController &partner,
                            bool use_alignment = true) const;

  /// Create an orthogonal counter-strategy (Phoenix Rebirth)
  FuzzyController create_orthogonal_child(float intensity = 1.0f) const;

  /// Banach-Tarski fission: split into two complementary children
  std::pair<FuzzyController, FuzzyController>
  banach_tarski_fission(float intensity = 1.0f) const;

private:
  static std::atomic<uint64_t> next_id_;
  static thread_local std::mt19937 rng_;
};

} // namespace fces