#pragma once

#include <cstddef>

#include <functional>
#include <utility>
#include <vec.hpp>
#include <vector>

template<typename T>
struct Agent {
    using F = std::function<float(T)>;

    virtual void move(float dt) = 0;
    virtual void step(float dt) = 0;

    virtual std::pair<float, T> best() const = 0;

    Agent(F f) : f(f) {}
protected:
    F f;
};

template<std::size_t N>
struct Particle : public Agent<vec<N>> {
    std::pair<float, vec<N>> best() const override { return {pb, pb_pos}; };

    void move(float dt=1) override { position = position + velocity * dt; }

    void step(float dt=1) override {
	velocity = kviscosity * velocity
	    + knostalgia * (pb_pos - position)
	    + kpeer_pressure * (peer.best().second - position);

	float y = f(position);
	if(y < pb) {
	    pb = y;
	    pb_pos = position;
	}
    };

    float kviscosity = 0.9f;
    float knostalgia = 2.5f;
    float kpeer_pressure = 0.5f;

    Particle(Agent<vec<N>>::F f, const Agent<vec<N>>& peer,
	     vec<N> position, vec<N> velocity)
	: Agent<vec<N>>(f), position(position), velocity(velocity),
	  pb_pos(position), peer(peer) {
	pb = f(position);
    }

    const vec<N> &get_position() const { return position; };

private:
    vec<N> position;
    vec<N> velocity;

    vec<N> pb_pos;
    float pb;

    const Agent<vec<N>>& peer;
};

template <std::size_t N>
struct Swarm : public Agent<vec<N>> {
    std::pair<float, vec<N>> best() const override { return {best_val, best_pos}; }

    void move(float dt=1) override {
	for(auto &p : particles)
	    p.move(dt);
    }

    void step(float dt=1) override {
	if(particles.empty()) return;

	// find best before step() on each particle
	// otherwise, if this were in best(), each p.step() would update it
	const Particle<N> *b = &particles[0];
	for(const auto &p : particles) {
	    if(b->best().first > p.best().first) b = &p;
	}

	best_val = b->best().first;
	best_pos = b->best().second;

	for(auto &p : particles)
	    p.step(dt);
    }

    void add_particle(std::size_t n=1) {
	for(std::size_t i=0; i <= n; ++i)
	    particles.push_back(
		Particle<N>(this->_f, *this, vec<N>(0), vec<N>(10))
	    );
    }

    void add_particle(const vec<N> &pos, const vec<N> &vel) {
	particles.push_back(
	    Particle<N>(this->f, *this, pos, vel)
	);
    }

    const std::vector<Particle<N>>& get_particles() { return particles; };

    Swarm(Agent<vec<N>>::F f) : Agent<vec<N>>(f) {}
private:
    vec<N> best_pos;
    float best_val;
    std::vector<Particle<N>> particles;
};