#include "agents/two_way_engine.hpp"

#include <numeric>
#include <random>

namespace agents {

TwoWayEngine::TwoWayEngine() : cfg_() {
  init_();
}

TwoWayEngine::TwoWayEngine(const Cfg& cfg) : cfg_(cfg) {
  init_();
}

void TwoWayEngine::init_() {
  params_.assign(cfg_.dim, 0.0f);
  target_.assign(cfg_.dim, 0.5f);  // neutral mid target
  std::mt19937 rng(42);
  std::uniform_real_distribution<float> dist(-0.5f, 0.5f);
  for (auto& p : params_) p = dist(rng);
  update_error_();
}

void TwoWayEngine::set_target(std::span<const float> t) {
  if (target_.size() != cfg_.dim) target_.assign(cfg_.dim, 0.0f);
  const std::size_t n = std::min<std::size_t>(cfg_.dim, t.size());
  std::copy_n(t.begin(), n, target_.begin());
  // If provided target is shorter than dim, keep remaining values as-is.
}

std::vector<float> TwoWayEngine::behavior() const {
  std::vector<float> z(cfg_.dim);
  for (std::size_t i = 0; i < cfg_.dim; ++i) {
    // Simple bounded nonlinearity; can be replaced by your exact equations.
    // Using sigmoid of param gives [0,1] range naturally.
    z[i] = sigmoid(params_[i]);
  }
  return z;
}

void TwoWayEngine::gradient_descent_step_(float lr) {
  // Minimize: L = sum_i (sigmoid(s_i) - target_i)^2
  // dL/ds_i = 2*(f_i - t_i) * f_i*(1 - f_i), where f_i = sigmoid(s_i)
  for (std::size_t i = 0; i < cfg_.dim; ++i) {
    float f = sigmoid(params_[i]);
    float diff = (f - target_[i]);
    float grad = 2.0f * diff * f * (1.0f - f);
    params_[i] -= lr * grad;
  }
}

void TwoWayEngine::update_error_() {
  auto z = behavior();
  double acc = 0.0;
  for (std::size_t i = 0; i < cfg_.dim; ++i) {
    double d = static_cast<double>(target_[i]) - static_cast<double>(z[i]);
    acc += d * d;
  }
  current_error_ = std::sqrt(acc / static_cast<double>(cfg_.dim));
}

void TwoWayEngine::step() {
  // A tiny alternating schedule: gradient step -> joint "blend"
  const float lr0 = 0.5f;
  const float lr1 = 0.25f;
  for (int iter = 0; iter < cfg_.iters; ++iter) {
    gradient_descent_step_(lr0);

    // Joint tweak towards target in parameter space (proxy for reverse step)
    for (std::size_t i = 0; i < cfg_.dim; ++i) {
      // Move params subtly to favor convergence robustness.
      float f = sigmoid(params_[i]);
      float delta = (target_[i] - f);
      params_[i] += lr1 * delta;  // gentle coupling
    }

    if (cfg_.use_gpu) {
#ifdef __APPLE__
      metal_step_();
#else
      simd_step_(params_);
#endif
    }
  }
  update_error_();
}

void TwoWayEngine::simd_step_(std::span<float> s) {
  // Placeholder: no-op to keep portability by default.
  (void)s;
}

#ifdef __APPLE__
void TwoWayEngine::metal_step_() {
  // Placeholder: tie in a .metal kernel here; fallback to simd_step_ for now.
  simd_step_(params_);
}
#endif

}  // namespace agents