Evolution.cpp

#include "Evolution.h"
#include <algorithm>
#include <fstream>
#include <iostream>
#include <ctime>
#include <iomanip>
#include <cmath>
#include <thread>
#include <mutex>
double evaluateFitness(const std::vector<double>& genome, int gridSize);
extern const int GRID_SIZE;

// Default constructor
Evolution::Evolution()
    : config(), generation(0), logFileInitialized(false)
{
    rng.seed(static_cast<unsigned>(std::time(nullptr)));
    population.resize(config.populationSize);
    allTimeBest.fitness = -999999.0;
}

// Constructor with custom config
Evolution::Evolution(const EvolutionConfig& cfg)
    : config(cfg), generation(0), logFileInitialized(false)
{
    rng.seed(static_cast<unsigned>(std::time(nullptr)));
    population.resize(config.populationSize);
    allTimeBest.fitness = -999999.0;
}

void Evolution::initializePopulation()
{
    std::uniform_real_distribution<double> sigmaDist(0.5, 1.5);  // Initial sigma around 1.0

    for (auto& individual : population) {
        // Initialize genome weights
        for (int i = 0; i < GENOME_SIZE; i++) {
            individual.genome[i] = randomDouble(config.genomeInitMin, config.genomeInitMax);
            // Initialize each gene's sigma
            individual.sigmas[i] = sigmaDist(rng);
        }
        individual.fitness = 0.0;
    }

    // Initialize log file with headers
    initializeLogFile();
}

void Evolution::initializeLogFile()
{
    // Append mode so multiple runs can be logged to the same file
    std::ofstream file("evolution_log.txt", std::ios::app);

    if (!file.is_open()) {
        std::cerr << "Failed to initialize log file: evolution_log.txt" << std::endl;
        return;
    }

    // Write run delimiter and parameters on one line
    file << "@@@" << std::endl;
    file << "populationSize = " << config.populationSize
         << ", numParents = " << config.numParents
         << ", numOffspring = " << config.numOffspring
         << ", eliteCount = " << config.eliteCount
         << ", crossoverRate = " << config.crossoverRate
         << ", sigmaMin = " << config.sigmaMin
         << ", sigmaMax = " << config.sigmaMax
         << ", tauGlobal = " << config.tauGlobal
         << ", tauLocal = " << config.tauLocal
         << ", PER-GENE SIGMAS"
         << std::endl;
    file << std::endl;

    file.close();
    logFileInitialized = true;

    std::cout << "Log file initialized: evolution_log.txt (PER-GENE SIGMAS)" << std::endl;
}

void Evolution::logGeneration()
{
    if (!logFileInitialized) {
        initializeLogFile();
    }

    std::ofstream file("evolution_log.txt", std::ios::app);

    if (!file.is_open()) {
        std::cerr << "Failed to write to log file: evolution_log.txt" << std::endl;
        return;
    }

    // Calculate statistics
    Individual best = getBest();

    double avgFitness = 0.0;
    for (const auto& ind : population) {
        avgFitness += ind.fitness;
    }
    avgFitness /= population.size();

    // Calculate average sigma across all genes and individuals
    double avgSigma = 0.0;
    for (const auto& ind : population) {
        for (double s : ind.sigmas) {
            avgSigma += s;
        }
    }
    avgSigma /= (population.size() * GENOME_SIZE);

    // Write data line: epoch, avg fitness, best fitness, avg sigma
    file << std::fixed << std::setprecision(2);
    file << generation << ", " << avgFitness << ", " << best.fitness << ", " << avgSigma << std::endl;

    file.close();
}

void Evolution::evaluatePopulation()
{
    // This will be called from Main with threading
    // Fitness evaluation happens externally
}

void Evolution::evaluateOffspring(std::vector<Individual>& offspring)
{
    const int numOffspring = static_cast<int>(offspring.size());
    if (numOffspring == 0) return;

    // Match the style of Main.cpp: fixed worker count
    const int MAX_THREADS = 20;
    const int numThreads = std::min(MAX_THREADS, numOffspring);

    std::mutex mtx; // mirrors Main.cpp's pattern (safe but not strictly required)

    auto worker = [&offspring, &mtx](int start, int end) {
        for (int i = start; i < end; ++i) {
            double fitness = evaluateFitness(offspring[i].genome, GRID_SIZE);

            // Thread-safe write (pattern matches Main.cpp::evaluateIndividuals)
            std::lock_guard<std::mutex> lock(mtx);
            offspring[i].fitness = fitness;
        }
    };

    std::vector<std::thread> threads;
    threads.reserve(numThreads);

    int perThread = numOffspring / numThreads;
    int start = 0;

    for (int t = 0; t < numThreads; ++t) {
        int end = (t == numThreads - 1) ? numOffspring : (start + perThread);
        threads.emplace_back(worker, start, end);
        start = end;
    }

    for (auto& th : threads) {
        th.join();
    }
}

void Evolution::runTheGreatTournament(std::vector<Individual>& offspring,
                                      const std::vector<Individual>& eliteParents)
{
    // Build μ+λ tournament pool: all children (λ) + elite parents (μ_elite)
    std::vector<Individual> tournamentPool;
    tournamentPool.reserve(offspring.size() + eliteParents.size());

    // std::cout << "\n=== TheGreatTournament ===\n"
    //           << " Children (λ): " << offspring.size()
    //           << " | Elite parents (μ_elite): " << eliteParents.size()
    //           << " | Target survivors: " << config.populationSize
    //           << "\n";

    tournamentPool.insert(tournamentPool.end(), offspring.begin(), offspring.end());
    tournamentPool.insert(tournamentPool.end(), eliteParents.begin(), eliteParents.end());

    // Sort by fitness (DESCENDING) to pick the best μ = populationSize
    std::sort(tournamentPool.begin(), tournamentPool.end(),
        [](const Individual& a, const Individual& b) {
            return a.fitness > b.fitness;
        });

    if (static_cast<int>(tournamentPool.size()) < config.populationSize) {
        std::cerr << "[TheGreatTournament] WARNING: tournament pool smaller "
                  << "than populationSize; copying everyone into next generation.\n";
        population = std::move(tournamentPool);
    } else {
        population.assign(tournamentPool.begin(),
                          tournamentPool.begin() + config.populationSize);
    }

    // Optional: brief log of best contestant in TheGreatTournament
    // if (!population.empty()) {
    //     std::cout << "TheGreatTournament winner fitness: "
    //               << population.front().fitness << "\n\n";
    // }
}

void Evolution::evolve()
{
    // ============================================================
    // 1) Sort current population by fitness (DESCENDING)
    //    (population has just been evaluated in Main.cpp)
    // ============================================================
    std::sort(population.begin(), population.end(),
        [](const Individual& a, const Individual& b) {
            return a.fitness > b.fitness;
        });

    // Safety: clamp numParents and eliteCount to population size
    const int effectiveNumParents =
        std::min(config.numParents, static_cast<int>(population.size()));
    const int effectiveEliteCount =
        std::min(config.eliteCount, effectiveNumParents);

    // ============================================================
    // 2) Select parents (μ) and elite parents (top eliteCount)
    // ============================================================
    std::vector<Individual> parents(population.begin(),
                                    population.begin() + effectiveNumParents);

    std::vector<Individual> eliteParents(parents.begin(),
                                         parents.begin() + effectiveEliteCount);

    // ============================================================
    // 3) Create children (λ) from parents
    //    * Crossover path: recombination ONLY (NO mutation)
    //    * ES path: clone + self-adaptive Gaussian mutation
    // ============================================================
    std::vector<Individual> offspring;
    offspring.reserve(config.numOffspring);

    while (static_cast<int>(offspring.size()) < config.numOffspring) {
        // Select two random parents from the parent pool
        std::pair<int, int> parentIndices = selectTwoParents(effectiveNumParents);
        int idx1 = parentIndices.first;
        int idx2 = parentIndices.second;

        // Decide between crossover or ES (mutation) branch
        double roll = randomDouble(0.0, 1.0);

        if (roll < config.crossoverRate) {
            // ----------------------------------------------------
            // CROSSOVER PATH:
            //   - 1 child
            //   - recombination ONLY (NO mutation)
            // ----------------------------------------------------
            Individual child;
            child.genome = crossover(
                parents[idx1].genome,
                parents[idx2].genome
            );

            child.sigmas = crossoverSigmas(
                parents[idx1].sigmas,
                parents[idx2].sigmas
            );

            // NO mutate() call here
            child.fitness = 0.0;
            offspring.push_back(std::move(child));
        } else {
            // ----------------------------------------------------
            // ES PATH:
            //   - up to 2 children
            //   - clone parents, then self-adaptive mutation
            // ----------------------------------------------------
            Individual child1 = parents[idx1];  // clone genome + sigmas
            Individual child2 = parents[idx2];  // clone genome + sigmas

            mutate(child1.genome, child1.sigmas);
            mutate(child2.genome, child2.sigmas);

            child1.fitness = 0.0;
            child2.fitness = 0.0;

            int childrenNeeded = config.numOffspring - static_cast<int>(offspring.size());

            if (childrenNeeded >= 2) {
                offspring.push_back(std::move(child1));
                offspring.push_back(std::move(child2));
            } else if (childrenNeeded == 1) {
                // Randomly pick one of the two to finish filling λ
                if (randomDouble(0.0, 1.0) < 0.5) {
                    offspring.push_back(std::move(child1));
                } else {
                    offspring.push_back(std::move(child2));
                }
            }
        }
    }

    // ============================================================
    // 4) Evaluate ALL children in parallel (TheGreatTournament prep)
    //    Parents already have fitness from Main.cpp evaluation.
    // ============================================================
    evaluateOffspring(offspring);

    // ============================================================
    // 5) μ+λ selection: run TheGreatTournament
    //    Tournament pool = all λ children + elite parents (μ_elite)
    //    Keep the best populationSize as next generation
    // ============================================================
    runTheGreatTournament(offspring, eliteParents);

    generation++;
}

Individual Evolution::getBest() const
{
    auto best = std::max_element(population.begin(), population.end(),
        [](const Individual& a, const Individual& b) {
            return a.fitness < b.fitness;
        });
    return *best;
}

void Evolution::updateAllTimeBest()
{
    Individual currentBest = getBest();
    if (currentBest.fitness > allTimeBest.fitness) {
        allTimeBest = currentBest;
        std::cout << "*** NEW ALL-TIME BEST: " << allTimeBest.fitness << " ***" << std::endl;
    }
}

std::pair<int, int> Evolution::selectTwoParents(int numParents)
{
    int parent1 = rng() % numParents;
    int parent2 = rng() % numParents;

    // Make sure they're different
    while (parent2 == parent1) {
        parent2 = rng() % numParents;
    }

    return {parent1, parent2};
}

std::vector<double> Evolution::crossover(const std::vector<double>& parent1, const std::vector<double>& parent2)
{
    std::vector<double> child(GENOME_SIZE);

    // Uniform crossover: 50/50 chance per gene
    for (int i = 0; i < GENOME_SIZE; i++) {
        if (randomDouble(0.0, 1.0) < 0.5) {
            child[i] = parent1[i];
        } else {
            child[i] = parent2[i];
        }
    }

    return child;
}

std::vector<double> Evolution::crossoverSigmas(const std::vector<double>& sigma1, const std::vector<double>& sigma2)
{
    std::vector<double> childSigmas(GENOME_SIZE);

    // Uniform crossover: 50/50 chance per sigma
    for (int i = 0; i < GENOME_SIZE; i++) {
        if (randomDouble(0.0, 1.0) < 0.5) {
            childSigmas[i] = sigma1[i];
        } else {
            childSigmas[i] = sigma2[i];
        }
    }

    return childSigmas;
}

void Evolution::mutate(std::vector<double>& genome, std::vector<double>& sigmas)
{
    // Self-Adaptive Mutation with PER-GENE sigmas
    // Classic ES formula:
    // τ' (tau_global) = 1 / √(2n)
    // τ  (tau_local)  = 1 / √(2√n)
    // σ'_i = σ_i * exp(τ' * N_global + τ * N_local,i)
    // x'_i = x_i + σ'_i * N_i(0,1)

    double n = static_cast<double>(GENOME_SIZE);
    double tau_global = config.tauGlobal / std::sqrt(2.0 * n);
    double tau_local = config.tauLocal / std::sqrt(2.0 * std::sqrt(n));

    // Generate ONE global noise value (shared by all genes)
    std::normal_distribution<double> globalNoise(0.0, 1.0);
    double N_global = globalNoise(rng);

    // STEP 1: Mutate ALL sigmas using self-adaptive formula
    for (int i = 0; i < GENOME_SIZE; i++) {
        // Generate local noise for this gene
        std::normal_distribution<double> localNoise(0.0, 1.0);
        double N_local = localNoise(rng);

        // Update sigma for this gene
        // σ'_i = σ_i * exp(τ' * N_global + τ * N_local,i)
        sigmas[i] = sigmas[i] * std::exp(tau_global * N_global + tau_local * N_local);

        // Clamp sigma to configured bounds
        sigmas[i] = std::max(config.sigmaMin, sigmas[i]);
        sigmas[i] = std::min(config.sigmaMax, sigmas[i]);
    }

    // STEP 2: Mutate ALL genome values using their respective sigmas
    // x'_i = x_i + σ'_i * N(0,1)
    for (int i = 0; i < GENOME_SIZE; i++) {
        std::normal_distribution<double> mutationNoise(0.0, 1.0);
        double noise = mutationNoise(rng);

        genome[i] += sigmas[i] * noise;

        // Clamp to reasonable range
        genome[i] = std::max(-5.0, std::min(5.0, genome[i]));
    }
}

double Evolution::randomDouble(double min, double max)
{
    std::uniform_real_distribution<double> dist(min, max);
    return dist(rng);
}

void Evolution::saveBest(const std::string& filename)
{
    Individual best = getBest();
    std::ofstream file(filename);

    if (!file.is_open()) {
        std::cerr << "Failed to save genome to " << filename << std::endl;
        return;
    }

    for (double gene : best.genome) {
        file << gene << "\n";
    }

    file.close();
    std::cout << "Saved best genome (fitness: " << best.fitness << ") to " << filename << std::endl;
}

void Evolution::loadBest(const std::string& filename, std::vector<double>& genome)
{
    std::ifstream file(filename);

    if (!file.is_open()) {
        std::cerr << "Failed to load genome from " << filename << std::endl;
        return;
    }

    genome.clear();
    double gene;
    while (file >> gene) {
        genome.push_back(gene);
    }

    file.close();

    if (genome.size() != GENOME_SIZE) {
        std::cerr << "Warning: Loaded genome size (" << genome.size()
                  << ") doesn't match expected size (" << GENOME_SIZE << ")" << std::endl;
    }

    std::cout << "Loaded genome from " << filename << std::endl;
}