#pragma once

#include <cmath>
#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() = 0;

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

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

template <typename A>
concept ParameterChangeAlgorithm = requires(A alg, float f, unsigned int i) {
    { alg(f, f, f, i, i) } -> std::same_as<vec<3>>;
};

namespace ParamChange {
    struct Const {
	vec<3> operator()(
	    float visc, float nostal, float peerp, int _1, int _2) {
	    return {visc, nostal, peerp};
	}
    };

    struct MLinAlg {
	vec<3> operator()(
	    float visc, float nostal, float peerp, int curr, int max) {
	    float pct_done = (float)curr / (float)max;
	    return {
		visc   - (0.9f - 0.4f) * pct_done,
		nostal - (2.5f - 0.5f) * pct_done,
		peerp  + (2.5f - 0.5f) * pct_done,
	    };
	}
    };

    template<auto steepness = 3>
    struct MTanhAlg {
	vec<3> operator()(
	    float visc, float nostal, float peerp, int curr, int max) {
	    float pct_done = (float)curr / (float)max;
	    float scale_factor = (std::tanh(steepness * (pct_done - 0.5)) + 1) / 2;
	    return {
		visc   - (0.9f - 0.4f) * scale_factor,
		nostal - (2.5f - 0.5f) * scale_factor,
		peerp  + (2.5f - 0.5f) * scale_factor,
	    };
	}
    };
};

template<std::size_t N, ParameterChangeAlgorithm A = ParamChange::MTanhAlg<>>
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() override {
	vec<3> params = alg(kviscosity, knostalgia, kpeer_pressure,
			    curr_iter, max_iter);

	viscosity = params.x;
	nostalgia = params.y;
	peer_pressure = params.z;

	curr_iter++;

	velocity = viscosity * velocity
	    + nostalgia * (pb_pos - position)
	    + peer_pressure * (peer.best().second - position);

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

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

    float viscosity = kviscosity;
    float nostalgia = knostalgia;
    float peer_pressure = kpeer_pressure;

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

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

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

    vec<N> pb_pos;
    float pb;

    unsigned curr_iter = 0;
    unsigned max_iter;

    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() 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();
    }

    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), max_iter)
	    );
    }

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

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

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