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

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

    double score;
    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(1), money(money), stock(stock), n(n),
        q([](std::shared_ptr<trained<T>> x, const dataset& input, double money, unsigned stock) { 
            return money + input.back() * stock; 
        }), 
        random_engine(std::time(0)) 
    {
        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()
    {
        sort();
        filter();
        breed();

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

    void train(const dataset& input, std::shared_ptr<trained<T>> trainee) 
    {
        trainee->decider.start_money = money;
        trainee->decider.start_stock = stock;
        trainee->decider.reset();

        double   money = this->money;
        unsigned stock = this->stock;

        for (double price : input) {
            auto decision   = trainee->decider.decide(price, money, stock);
            auto current    = price * stock + money;
            auto max_credit = std::max(current * 0.05, -1e4);

            if (decision < 0) {
                decision = std::max<int>(decision, -stock); // cannot sell more than we actually have
            } else if (decision > 0) {
                decision = std::min<int>(floor((money + max_credit) / price), decision);   
            }

            money -= price * decision;
            stock += decision;
        }

        trainee->score += q(trainee, input, money, stock);

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

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

        std::cout 
            << "#" << trainee->id << ": " << wealth 
            
            << std::showpos
            << " H: " << wealth - hodl << " (" << (wealth - hodl) / hodl * 100 << "%)" 
            << " S: " << wealth - start << " (" << (wealth - start) / start * 100 << "%) "
            << std::noshowpos

            << stock << " akcji, " 
            << money << " gelda w banku. " 
            << std::endl; 
    }

    void train(const dataset& input, const std::string& name)
    {
        std::cout << "Zestaw " << name
            << " GEN #" << this->generation 
            << " start: " << money + input.front() * stock
            << " HODL: " << money + input.back() * stock
            << std::endl;

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

        std::cout << "Przy życiu pozostają: ";
        for (auto trainee : trainees) {
            std::cout << "#" << trainee->id << " (" << trainee->position << ") ";
        }
    }

    void breed()
    {   
        std::size_t diff = n - trainees.size();
        std::cout << "W populacji brakuje " << diff << " sieci, aktualnie " << trainees.size() << "." << std::endl;

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

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

        auto combiner = [=](const double& a, const double& b){
            auto mutation = (rand() % 2 ? -1 : 1) * exponential(random_engine);

            if (exponential(random_engine) < 1) {
                auto r = ratio(random_engine);
                return r * a + (1 - r) * b + mutation; 
            }

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

        unsigned first, second;
        for (int i = 0; i < diff - 4; i++) {
            first = distribution(random_engine);
            do { second = distribution(random_engine); } while (first == second);

            auto combined = trainees[first]->decider.combine(trainees[second]->decider, combiner);
            std::cout << "Łączenie #" << trainees[first]->id << " z #" << trainees[second]->id << " dało #" << (id+1) << std::endl;

            add(combined);
        }

        add(4); // some random things

        generation++;
    }

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