#pragma once
#include <vector>
#include <cstddef>
#include "raylib.h"
#include "raymath.h"

class NervousSystem {
 public:
  struct Neuron {
   private:
    friend class NervousSystem;
    float potential = 0.0f;
    float recovery = 0.92f;
    float threshold = 1.0f;
    bool fired = false;

    bool update(float input) {
      potential = (potential * recovery) + input;
      if (potential < 0.0f) potential = 0.0f;
      if (potential >= threshold) {
        potential = 0.0f;
        fired = true;
        return true;
      }
      fired = false;
      return false;
    }

   public:
    float getCharge() const { return potential / threshold; }
    bool hasFired() const { return fired; }
  };

 private:
  struct Synapse {
    int from;
    int to;
    float weight;
  };

  std::vector<Neuron> units;
  std::vector<Synapse> connections;
  std::vector<float> incomingSignals;

 public:
  int addUnit() {
    units.push_back(Neuron());
    incomingSignals.push_back(0.0f);
    return static_cast<int>(units.size()) - 1;
  }

  void connect(int fromIdx, int toIdx, float weight) {
    if (fromIdx < (int)units.size() && toIdx < (int)units.size()) {
      connections.push_back({fromIdx, toIdx, weight});
    }
  }

  void tick() {
    for (std::size_t i = 0; i < units.size(); ++i) {
      units[i].update(incomingSignals[i]);
      incomingSignals[i] = 0.0f;
    }

    for (const auto& syn : connections) {
      if (units[syn.from].fired) {
        incomingSignals[syn.to] += syn.weight;
      }
    }
  }

  void stimulate(int idx, float power) {
    if (idx >= 0 && idx < (int)incomingSignals.size()) {
      incomingSignals[idx] += power;
    }
  }

  void draw(Vector2 pos, int rows, float spacing) const {
    for (const auto& syn : connections) {
      Vector2 p1 = {pos.x + (syn.from / rows) * spacing * 2, pos.y + (syn.from % rows) * spacing};
      Vector2 p2 = {pos.x + (syn.to / rows) * spacing * 2, pos.y + (syn.to % rows) * spacing};

      Color synColor = units[syn.from].fired ? YELLOW : Fade(GRAY, 0.3f);
      DrawLineEx(p1, p2, 1.0f, synColor);
    }

    for (std::size_t i = 0; i < units.size(); ++i) {
      Vector2 nodePos = {pos.x + (i / rows) * spacing * 2, pos.y + (i % rows) * spacing};

      float charge = units[i].getCharge();
      Color nodeColor = ColorLerp(DARKGRAY, LIME, charge);
      if (units[i].fired) nodeColor = WHITE;

      DrawCircleV(nodePos, 8, nodeColor);
      DrawCircleLinesV(nodePos, 8, units[i].fired ? YELLOW : DARKGRAY);
    }
  }

  const Neuron& getUnit(int idx) const { return units[idx]; }
};