#include <vector>
#include <random>
#include <ctime>
#include <functional>
#include <iostream>
#include <algorithm>
#include <memory>
#include <string>

#include "simulator.h"

template<class T>
struct trained {
    unsigned id;
    unsigned position;

    double score, profit;
    T      decider;
};

template <class T>
class trainer {
    std::vector<std::shared_ptr<trained<T>>> trainees;
    std::function<T()> factory;

    unsigned int id;

    std::default_random_engine random_engine;
    std::size_t n;

public:
    using dataset = std::vector<double>;

    double money;
    unsigned stock;
    unsigned int generation;

    std::function<double(std::shared_ptr<trained<T>>, const dataset&, double, unsigned)> q;

    trainer(double money, unsigned stock, std::size_t n, std::function<T()> factory) 
        : factory(factory), id(0), generation(0), money(money), stock(stock), n(n), random_engine(std::time(0)) 
    {
        this->q = [=](std::shared_ptr<trained<T>> x, const dataset& input, double money, unsigned stock) { 
            auto current = input.back() * stock + money;

            auto start   = input.front() * this->stock + this->money;
            auto hodl    = input.back() * this->stock + this->money;

            auto result = std::min((current - hodl)/hodl, (current - start)/start);
            /* if (result < 0) result *= 4; */
            return result / (1 + abs(result));
        };
        add(n);
    }

    void add(std::size_t n, std::function<T()> factory)
    {
        for(std::size_t i = 0; i < n; i++) {
            add(factory());
        }
    }

    void add(std::size_t n) 
    {
        add(n, factory);
    }

    void add(const T& decider)
    {
        auto trainee = std::make_shared<trained<T>>();
        trainee->id = ++id;
        trainee->decider = decider;

        trainees.push_back(trainee);
    }

    void evolve()
    {
        if (generation) {
            sort();
            filter();
            breed();
        }

        // cleanup before next training sessions
        for (auto t : trainees) {
            t->score = t->profit = 0;
        }

        generation++;
    }

    void train(const dataset& input, std::shared_ptr<trained<T>> trainee) 
    {
        simulator sim(&(trainee->decider), this->money, this->stock);
        sim.proceed(input);

        trainee->score = std::min(trainee->score, q(trainee, input, sim.money, sim.stock));

        auto last   = input.back();
        auto first  = input.front();
        auto wealth = sim.money + sim.stock * last;

        auto start  = this->money + this->stock * first;

        trainee->profit += (wealth - start) / start;
    }

    void train(const dataset& input, const std::string& name)
    {
        for (auto trainee : trainees) {
            train(input, trainee);
        }
    }

    void sort()
    {
        std::sort(trainees.begin(), trainees.end(), [=](std::shared_ptr<trained<T>> a, std::shared_ptr<trained<T>> b){
            return a->score > b->score;
        });

        unsigned i = 0;
        for (auto t : trainees) {
            t->position = i++;
        }
    }

    void filter()
    {
        static std::uniform_real_distribution<double> distribution(0.0, 1.0);
        auto random = [=](){ return distribution(random_engine); };
        
        auto iterator = std::remove_if(trainees.begin(), trainees.end(), [&](std::shared_ptr<trained<T>> t) {
            return random() < (double)t->position / n;
        });

        trainees.erase(iterator, std::end(trainees));
    }

    void breed()
    {   
        std::size_t diff = n - trainees.size();

        std::vector<double> probability;
        for (auto t : trainees) {
            probability.push_back(t->position);
        }

        std::discrete_distribution<unsigned>   distribution(probability.begin(), probability.end());
        std::exponential_distribution<double>  exponential(2.5);
        std::uniform_real_distribution<double> ratio(0.0, 1.0);

        auto combiner = [=](const double& a, const double& b){
            if (exponential(random_engine) < 1) {
                auto r = ratio(random_engine);
                return r * a + (1 - r) * b; 
            }

            return rand() % 2 ? a : b;
        };

        auto mutator = [=](const double& a) {
            if(ratio(random_engine) < .25) return a;

            auto mutation = (rand() % 2 ? -1 : 1) * exponential(random_engine);
            return a + mutation;
        };

        std::size_t to_combine = diff * 0.5, to_mutate = diff * 0.3;

        unsigned first, second;
        for (int i = 0; i < to_combine; i++) {
            first = distribution(random_engine);
            do { second = distribution(random_engine); } while (first == second);
            
            add(trainees[first]->decider.combine(trainees[second]->decider, combiner));
        }

        for (int i = 0; i < to_mutate; i++) {
            first = distribution(random_engine);
            add(trainees[first]->decider.mutate(mutator));
        }

        add(diff - to_combine - to_mutate); // some random things
    }

    std::vector<std::shared_ptr<trained<T>>> population() {
        return trainees;
    }
};