/*
 * Copyright (C) 2019 Emeric Poupon
 *
 * This file is part of LMS.
 *
 * LMS is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * LMS is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with LMS.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <numeric>

#include "core/Random.hpp"

#include "ParallelFor.hpp"

template<typename Individual>
class GeneticAlgorithm
{
public:
    using Score = float;

    using BreedFunction = std::function<Individual(const Individual&, const Individual&)>;
    using MutateFunction = std::function<void(Individual&)>;
    using ScoreFunction = std::function<Score(const Individual&)>;

    struct Params
    {
        std::size_t nbWorkers{ 1 };
        std::size_t nbGenerations;
        float crossoverRatio{ 0.5 };
        float mutationProbability{ 0.05 };
        BreedFunction breedFunction;
        MutateFunction mutateFunction;
        ScoreFunction scoreFunction;
    };

    GeneticAlgorithm(const Params& params);

    // Returns the individual that has the maximum score after processing the requested generations
    Individual simulate(const std::vector<Individual>& initialPopulation);

private:
    struct ScoredIndividual
    {
        Individual individual;
        std::optional<Score> score{};
    };

    void scoreAndSortPopulation(std::vector<ScoredIndividual>& population);
    Score getTotalScore(const std::vector<ScoredIndividual>& population) const;
    typename std::vector<ScoredIndividual>::const_iterator pickRandomRouletteWheel(const std::vector<ScoredIndividual>& population, Score totalScore);

    Params _params;
};

template<typename Individual>
GeneticAlgorithm<Individual>::GeneticAlgorithm(const Params& params)
    : _params{ params }
{
}

template<typename Individual>
Individual
GeneticAlgorithm<Individual>::simulate(const std::vector<Individual>& initialPopulation)
{
    const std::size_t childrenCountPerGeneration{ static_cast<std::size_t>(initialPopulation.size() * _params.crossoverRatio) };
    if (initialPopulation.size() < 10)
        throw std::runtime_error("Initial population must has at least 10 elements");

    std::vector<ScoredIndividual> scoredPopulation;
    scoredPopulation.reserve(initialPopulation.size());

    std::transform(std::cbegin(initialPopulation), std::cend(initialPopulation), std::back_inserter(scoredPopulation),
                   [](const Individual& individual) { return ScoredIndividual{ individual }; });

    scoreAndSortPopulation(scoredPopulation);

    for (std::size_t currentGeneration{}; currentGeneration < _params.nbGenerations; ++currentGeneration)
    {
        assert(scoredPopulation.size() == initialPopulation.size());
        std::cout << "Processing generation " << currentGeneration << "..." << std::endl;
        std::cout << "Need to create " << childrenCountPerGeneration << " new children" << std::endl;

        // breed
        const Score populationTotalScore{ getTotalScore(scoredPopulation) };
        std::vector<ScoredIndividual> children;
        children.reserve(childrenCountPerGeneration);

        while (children.size() < childrenCountPerGeneration)
        {
            // Select two random parents using their score as weight
            const auto itParent1{ pickRandomRouletteWheel(scoredPopulation, populationTotalScore) };
            const auto itParent2{ pickRandomRouletteWheel(scoredPopulation, populationTotalScore) };

            if (itParent1 == itParent2)
                continue;

            ScoredIndividual child{ _params.breedFunction(itParent1->individual, itParent2->individual) };

            if (core::random::getRealRandom(float{}, float{ 1 }) <= _params.mutationProbability)
                _params.mutateFunction(child.individual);

            children.emplace_back(std::move(child));
        }

        // Elitist selection
        scoredPopulation.resize(initialPopulation.size() - childrenCountPerGeneration);

        scoredPopulation.insert(std::end(scoredPopulation), std::make_move_iterator(std::begin(children)), std::make_move_iterator(std::end(children)));
        assert(scoredPopulation.size() == initialPopulation.size());

        scoreAndSortPopulation(scoredPopulation);

        std::cout << "Mean score = " << getTotalScore(scoredPopulation) / scoredPopulation.size() << std::endl;
        std::cout << "Current best score = " << *scoredPopulation.front().score << std::endl;
    }

    std::cout << "Best score = " << *scoredPopulation.front().score << std::endl;
    return scoredPopulation.front().individual;
}

template<typename Individual>
void GeneticAlgorithm<Individual>::scoreAndSortPopulation(std::vector<ScoredIndividual>& scoredPopulation)
{
    parallel_foreach(_params.nbWorkers, std::begin(scoredPopulation), std::end(scoredPopulation),
                     [&](ScoredIndividual& scoredIndividual) {
                         if (!scoredIndividual.score)
                             scoredIndividual.score = _params.scoreFunction(scoredIndividual.individual);
                     });

    std::sort(std::begin(scoredPopulation), std::end(scoredPopulation), [](const ScoredIndividual& a, const ScoredIndividual& b) { return a.score > b.score; });
}

template<typename Individual>
typename GeneticAlgorithm<Individual>::Score
GeneticAlgorithm<Individual>::getTotalScore(const std::vector<ScoredIndividual>& scoredPopulation) const
{
    return std::accumulate(std::cbegin(scoredPopulation), std::cend(scoredPopulation), Score{}, [](Score score, const ScoredIndividual& individual) { return score + *individual.score; });
}

template<typename Individual>
typename std::vector<typename GeneticAlgorithm<Individual>::ScoredIndividual>::const_iterator
GeneticAlgorithm<Individual>::pickRandomRouletteWheel(const std::vector<ScoredIndividual>& population, Score totalScore)
{
    const Score randomScore{ core::random::getRealRandom(Score{}, totalScore) };

    Score curScore{};
    for (auto itScoredIndividual{ std::cbegin(population) }; itScoredIndividual != std::cend(population); ++itScoredIndividual)
    {
        if (curScore + *itScoredIndividual->score > randomScore)
            return itScoredIndividual;

        curScore += *itScoredIndividual->score;
    }

    throw std::runtime_error("bad random or empty population");
}
