#include "fces/controller.hpp"
#include <algorithm>
#include <cmath>
#include <numeric>

namespace fces {

// Static members
std::atomic<uint64_t> FuzzyController::next_id_{1};
thread_local std::mt19937 FuzzyController::rng_{std::random_device{}()};

// ---------------------------------------------------------------
// Genome
// ---------------------------------------------------------------

void Genome::randomize(std::mt19937 &rng) {
  std::normal_distribution<float> dist(0.0f, 0.5f);
  for (auto &w : weights) {
    w = dist(rng);
  }
  gene_success.fill(0.0f);
}

Genome Genome::clone() const {
  return *this; // Copy all fields
}

// ---------------------------------------------------------------
// FuzzyController
// ---------------------------------------------------------------

FuzzyController::FuzzyController() : id(next_id_++), origin("random") {
  genome.randomize(rng_);
  // Bias output toward acceleration (V2.1 insight)
  // Set output biases (last GENOME_OUTPUT_DIM elements) to +2.0, -1.0, 0.0 with
  // noise
  constexpr int bias_start = GENOME_SIZE - GENOME_OUTPUT_DIM;
  std::normal_distribution<float> bias_noise(0.0f, 0.5f);
  genome.weights[bias_start] = 2.0f + bias_noise(rng_);
  genome.weights[bias_start + 1] = -1.0f + bias_noise(rng_);
  genome.weights[bias_start + 2] = 0.0f + bias_noise(rng_);

  // Initialize plasticity with variance to seed evolution
  std::normal_distribution<float> plast_noise(0.0f, 0.05f);
  genome.plasticity = std::max(0.01f, 0.1f + plast_noise(rng_));
  genome.sigma_gene = 0.1f;
  genome.gene_success.fill(1.0f);
}

FuzzyController::FuzzyController(Genome genome)
    : id(next_id_++), genome(std::move(genome)), origin("constructed") {}

torch::Tensor FuzzyController::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) {
  const int num_groups = static_cast<int>(layer_stats.size());
  auto actions = torch::zeros({num_groups, GENOME_OUTPUT_DIM});

  // Extract weight views for the micro-MLP
  const float *w = genome.weights.data();

  // Layer 1: input -> hidden
  const float *W1 = w; // [(GENOME_INPUT_DIM + 1) x GENOME_HIDDEN_DIM]
  // Layer 2: hidden -> output
  const float *W2 =
      w + ((GENOME_INPUT_DIM + 1) *
           GENOME_HIDDEN_DIM); // [(GENOME_HIDDEN_DIM + 1) x GENOME_OUTPUT_DIM]

  for (int g = 0; g < num_groups; ++g) {
    // One-Hot Layer Type: Clamp to 5 to avoid overflow for new categories
    float layer_type_val =
        (layer_stats[g].size() >= 3) ? layer_stats[g][2] : 5.0f;
    int type_idx = std::min(5, static_cast<int>(layer_type_val));
    std::array<float, 5> type_onehot{0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
    if (type_idx >= 0 && type_idx < 5) {
      type_onehot[type_idx] = 1.0f;
    }

    // Build input vector
    std::array<float, GENOME_INPUT_DIM> input{};
    float gn = (layer_stats[g].size() >= 1) ? layer_stats[g][0] : 0.0f;
    float sp = (layer_stats[g].size() >= 2) ? layer_stats[g][1] : 0.0f;

    // Sanitization matching nan_to_num
    if (!std::isfinite(gn) || std::isnan(gn))
      gn = 0.0f;
    if (gn > 10.0f)
      gn = 10.0f;
    if (gn < 0.0f)
      gn = 0.0f;

    if (!std::isfinite(sp) || std::isnan(sp))
      sp = 0.0f;
    if (sp > 1.0f)
      sp = 1.0f;
    if (sp < 0.0f)
      sp = 0.0f;

    input[0] = gn;
    input[1] = sp;
    input[2] = loss_trend;
    input[3] = step_pct;
    input[4] =
        (num_groups > 1) ? static_cast<float>(g) / (num_groups - 1.0f) : 0.0f;
    input[5] = rollback_rate;
    input[6] = grad_stability;
    input[7] = spectral_alpha;
    input[8] = type_onehot[0];
    input[9] = type_onehot[1];
    input[10] = type_onehot[2];
    input[11] = type_onehot[3];
    input[12] = type_onehot[4];
    input[13] = projected_drift;

    // Forward pass: hidden = tanh(W1 * [input, 1])
    std::array<float, GENOME_HIDDEN_DIM> hidden{};
    for (int h = 0; h < GENOME_HIDDEN_DIM; ++h) {
      float sum = W1[GENOME_INPUT_DIM * GENOME_HIDDEN_DIM + h]; // Bias weight
      for (int i = 0; i < GENOME_INPUT_DIM; ++i) {
        sum += input[i] * W1[i * GENOME_HIDDEN_DIM + h];
      }
      hidden[h] = std::tanh(sum);
    }

    // Output = W2 * [hidden, 1]
    std::array<float, GENOME_OUTPUT_DIM> out_layer{};
    for (int o = 0; o < GENOME_OUTPUT_DIM; ++o) {
      float sum = W2[GENOME_HIDDEN_DIM * GENOME_OUTPUT_DIM + o]; // Bias weight
      for (int h = 0; h < GENOME_HIDDEN_DIM; ++h) {
        sum += hidden[h] * W2[h * GENOME_OUTPUT_DIM + o];
      }
      out_layer[o] = sum;
    }

    // Parse Output
    float log_mult_raw = out_layer[0];
    float log_wd_raw = out_layer[2];
    float sign_logit = out_layer[1];

    if (!std::isfinite(sign_logit) || std::isnan(sign_logit)) {
      sign_logit = 0.0f;
    }

    float noise_std = 0.0f;
    if (genome.plasticity > 0.01f) {
      noise_std += genome.plasticity;
    }
    if (stagnation_intensity > 0.0f) {
      noise_std += stagnation_intensity * 0.5f;
    }

    if (kzm_damping > 0.0f) {
      noise_std *= (1.0f - kzm_damping);
      log_mult_raw *= (1.0f - kzm_damping);
    }

    if (noise_std > 0.0f) {
      std::normal_distribution<float> noise_dist(0.0f, noise_std);
      log_mult_raw += noise_dist(rng_);
      log_wd_raw += noise_dist(rng_);
    }

    if (!std::isfinite(log_mult_raw) || std::isnan(log_mult_raw)) {
      log_mult_raw = 0.0f;
    }
    log_mult_raw = std::max(-6.0f, std::min(6.0f, log_mult_raw));
    float mult = std::pow(2.0f, log_mult_raw);

    if (!std::isfinite(log_wd_raw) || std::isnan(log_wd_raw)) {
      log_wd_raw = 0.0f;
    }
    log_wd_raw = std::max(-6.0f, std::min(6.0f, log_wd_raw));
    float wd_mult = std::pow(2.0f, log_wd_raw);

    actions[g][0] = mult;
    actions[g][1] = sign_logit;
    actions[g][2] = wd_mult;
  }

  return actions;
}

FuzzyController FuzzyController::mutate(float current_loss,
                                        float sigma_scale) const {
  Genome child_genome = genome.clone();
  std::normal_distribution<float> std_normal(0.0f, 1.0f);

  float tau = 0.2f;
  float new_sigma = genome.sigma_gene * std::exp(tau * std_normal(rng_));
  new_sigma = std::max(0.001f, std::min(0.8f, new_sigma));

  float new_plast = genome.plasticity * std::exp(tau * std_normal(rng_));
  new_plast = std::max(0.0f, std::min(0.5f, new_plast));

  float loss_val = std::max(0.0f, current_loss);
  float annealing_factor = std::sqrt(loss_val + 0.1f);
  float effective_sigma = new_sigma * sigma_scale * annealing_factor;

  std::normal_distribution<float> noise(0.0f, effective_sigma);
  for (size_t i = 0; i < child_genome.weights.size(); ++i) {
    child_genome.weights[i] += noise(rng_);
    child_genome.gene_success[i] = genome.gene_success[i] * 0.95f;
  }

  child_genome.sigma_gene = new_sigma;
  child_genome.plasticity = new_plast;

  FuzzyController child(child_genome);
  child.origin = "mutation";
  return child;
}

FuzzyController FuzzyController::crossover(const FuzzyController &partner,
                                           bool /*use_alignment*/) const {
  Genome child_genome;
  std::uniform_real_distribution<float> u_dist(0.0f, 1.0f);

  for (size_t i = 0; i < child_genome.weights.size(); ++i) {
    float success_self = genome.gene_success[i];
    float success_partner = partner.genome.gene_success[i];
    float prob_a = success_self / (success_self + success_partner + 1e-9f);

    bool choose_self = false;
    if (u_dist(rng_) < 0.1f) {
      // 10% Random injection
      choose_self = (u_dist(rng_) < 0.5f);
    } else {
      // 90% Meritocratic
      choose_self = (u_dist(rng_) < prob_a);
    }

    if (choose_self) {
      child_genome.weights[i] = genome.weights[i];
      child_genome.gene_success[i] = success_self;
    } else {
      child_genome.weights[i] = partner.genome.weights[i];
      child_genome.gene_success[i] = success_partner;
    }
  }

  child_genome.sigma_gene =
      (genome.sigma_gene + partner.genome.sigma_gene) * 0.5f;
  child_genome.plasticity =
      (genome.plasticity + partner.genome.plasticity) * 0.5f;

  FuzzyController child(child_genome);
  child.origin = "crossover";
  return child;
}

FuzzyController
FuzzyController::create_orthogonal_child(float intensity) const {
  Genome child_genome;
  std::normal_distribution<float> norm_dist(0.0f, 1.0f);

  float norm_elite = 0.0f;
  for (float w : genome.weights) {
    norm_elite += w * w;
  }
  norm_elite = std::sqrt(norm_elite) + 1e-9f;

  std::array<float, GENOME_SIZE> random_vec{};
  float dot_product = 0.0f;
  for (size_t i = 0; i < GENOME_SIZE; ++i) {
    random_vec[i] = norm_dist(rng_);
    dot_product += random_vec[i] * genome.weights[i];
  }

  std::array<float, GENOME_SIZE> orthogonal_vec{};
  float norm_ortho = 0.0f;
  for (size_t i = 0; i < GENOME_SIZE; ++i) {
    float projection =
        (dot_product / (norm_elite * norm_elite)) * genome.weights[i];
    orthogonal_vec[i] = random_vec[i] - projection;
    norm_ortho += orthogonal_vec[i] * orthogonal_vec[i];
  }
  norm_ortho = std::sqrt(norm_ortho) + 1e-9f;

  std::array<float, GENOME_SIZE> scaled_vec{};
  float final_norm = 0.0f;
  for (size_t i = 0; i < GENOME_SIZE; ++i) {
    scaled_vec[i] = orthogonal_vec[i] * (norm_elite / norm_ortho) * intensity;
    final_norm += scaled_vec[i] * scaled_vec[i];
  }
  final_norm = std::sqrt(final_norm);

  float max_allowed = std::max(norm_elite, 10.0f);
  if (final_norm > max_allowed) {
    float scale = max_allowed / (final_norm + 1e-9f);
    for (size_t i = 0; i < GENOME_SIZE; ++i) {
      scaled_vec[i] *= scale;
    }
  }

  child_genome.weights = scaled_vec;
  child_genome.gene_success.fill(1.0f);
  child_genome.sigma_gene = 0.2f;
  child_genome.plasticity = 0.2f;

  FuzzyController child(child_genome);
  child.origin = "phoenix_rebirth";
  return child;
}

std::pair<FuzzyController, FuzzyController>
FuzzyController::banach_tarski_fission(float intensity) const {
  Genome plus_genome;
  Genome minus_genome;

  float norm_parent = 0.0f;
  float max_success = 0.0f;
  for (size_t i = 0; i < GENOME_SIZE; ++i) {
    norm_parent += genome.weights[i] * genome.weights[i];
    if (genome.gene_success[i] > max_success) {
      max_success = genome.gene_success[i];
    }
  }
  norm_parent = std::sqrt(norm_parent) + 1e-9f;
  max_success += 1e-9f;

  std::normal_distribution<float> norm_dist(0.0f, 1.0f);
  std::array<float, GENOME_SIZE> noise{};
  float dot_product = 0.0f;
  for (size_t i = 0; i < GENOME_SIZE; ++i) {
    float saliency = genome.gene_success[i] / max_success;
    noise[i] = norm_dist(rng_) * saliency;
    dot_product += noise[i] * genome.weights[i];
  }

  std::array<float, GENOME_SIZE> fission_vec{};
  float norm_fission = 0.0f;
  for (size_t i = 0; i < GENOME_SIZE; ++i) {
    fission_vec[i] = noise[i] - (dot_product / (norm_parent * norm_parent)) *
                                    genome.weights[i];
    norm_fission += fission_vec[i] * fission_vec[i];
  }
  norm_fission = std::sqrt(norm_fission) + 1e-9f;

  for (size_t i = 0; i < GENOME_SIZE; ++i) {
    float scaled_fission =
        fission_vec[i] * (norm_parent / norm_fission) * 0.1f * intensity;
    plus_genome.weights[i] = genome.weights[i] + scaled_fission;
    minus_genome.weights[i] = genome.weights[i] - scaled_fission;

    plus_genome.gene_success[i] = 1.0f;
    minus_genome.gene_success[i] = 1.0f;
  }

  plus_genome.sigma_gene = genome.sigma_gene * 0.9f;
  minus_genome.sigma_gene = genome.sigma_gene * 0.9f;

  plus_genome.plasticity = genome.plasticity;
  minus_genome.plasticity = genome.plasticity;

  FuzzyController child_plus(plus_genome);
  child_plus.origin = "fission_plus";

  FuzzyController child_minus(minus_genome);
  child_minus.origin = "fission_minus";

  return {child_plus, child_minus};
}

} // namespace fces