Main.cpp

#include <iostream>
#include <thread>
#include <vector>
#include <mutex>
#include <algorithm>
#include <fstream>
#include <cstdlib>
#include "AI.h"
#include "Evolution.h"
#include "snake.h"
#include "food.h"
#include "HitCheck.h"
#include "Observe.h"


float SCREEN_WIDTH = 625.0f;
float SCREEN_HEIGHT = 625.0f;
int amt = 0;

extern const int GRID_SIZE = 25;
float wallthickness = GRID_SIZE / 2.0f;
extern const float CELL_SIZE = (SCREEN_WIDTH - 2.0f * wallthickness) / GRID_SIZE;

bool gameOver = false;
bool specialmessage = false;
bool ContinueGame = true;
bool AIMODE = true;


// Simulation parameters
const int MAX_STEPS = 10000; //Mostly legacy, kept just in case it decides to wander forever
const int THREADS = 20;
const int GENERATIONS = 300;

// Fitness evaluation function (runs headless)
double evaluateFitness(const std::vector<double>& genome, int gridSize)
{
    Snake snake;
    Food food;
    sf::Vector2i appleXY = food.apple(snake);

    int steps = 0;
    int stepsSinceLastFood = 0;
    int foodEaten = 0;
    bool gameOver = false;

    while (!gameOver && steps < MAX_STEPS) {
        // Get sensor inputs
        std::vector<double> inputs = AI::getSensorInputs(snake, appleXY, gridSize);

        // Run neural network
        std::vector<double> outputs = AI::forward(inputs, genome);

        // Get direction decision
        sf::Vector2i newDir = AI::getDirection(outputs, snake);
        snake.SetDirection(newDir.x, newDir.y);

        // Move snake
        snake.move();
        steps++;
        stepsSinceLastFood++;

        // Check wall collision
        if (DidItHitWall(snake)) {
            gameOver = true;
            break;
        }

        // Check self collision
        if (snake.checkSelfCollision()) {
            gameOver = true;
            break;
        }

        // Check if ate food
        sf::Vector2i head = snake.body.back();
        if (head.x == appleXY.x && head.y == appleXY.y) {
            snake.grow();
            foodEaten++;
            stepsSinceLastFood = 0;
            appleXY = food.apple(snake);
        }

        // Timeout if stuck wandering
        if (stepsSinceLastFood > 100) {
            gameOver = true;
            break;
        }
    }

    // Fitness function: heavily reward food, reward survival, penalize wandering
    double fitness = (foodEaten);

    return fitness;
}

// Thread worker function
void evaluateIndividuals(std::vector<Individual>& individuals, int start, int end, int gridSize, std::mutex& mtx)
{
    for (int i = start; i < end; i++) {
        double fitness = evaluateFitness(individuals[i].genome, gridSize);

        // Thread-safe update
        std::lock_guard<std::mutex> lock(mtx);
        individuals[i].fitness = fitness;
    }
}

// Run a single experiment with given config, returns the best individual
Individual runExperiment(EvolutionConfig config, int experimentNum, int totalExperiments)
{
    std::cout << "\n========================================" << std::endl;
    std::cout << "EXPERIMENT " << experimentNum << " / " << totalExperiments << std::endl;
    std::cout << "========================================" << std::endl;
    std::cout << "Population: " << config.populationSize << std::endl;
    std::cout << "Parents (μ): " << config.numParents << std::endl;
    std::cout << "Offspring (λ): " << config.numOffspring << std::endl;
    std::cout << "Crossover rate: " << config.crossoverRate << std::endl;
    std::cout << "Tau global: " << config.tauGlobal << std::endl;
    std::cout << "Tau local: " << config.tauLocal << std::endl;
    std::cout << "========================================\n" << std::endl;

    Evolution evolution(config);
    evolution.initializePopulation();

    std::mutex mtx;

    for (int gen = 0; gen < GENERATIONS; gen++) {
        // Get population
        std::vector<Individual>& population = evolution.getPopulation();

        // Evaluate fitness using multiple threads
        std::vector<std::thread> threads;
        int individualsPerThread = config.populationSize / THREADS;

        for (int t = 0; t < THREADS; t++) {
            int start = t * individualsPerThread;
            int end = (t == THREADS - 1) ? config.populationSize : (t + 1) * individualsPerThread;

            threads.emplace_back(evaluateIndividuals, std::ref(population), start, end, GRID_SIZE, std::ref(mtx));
        }

        // Wait for all threads to complete
        for (auto& thread : threads) {
            thread.join();
        }

        // Update all-time best tracker
        evolution.updateAllTimeBest();

        // Log generation stats to file
        evolution.logGeneration();

        // Get statistics
        Individual best = evolution.getBest();
        double avgFitness = 0.0;
        // double avgSigma = 0.0;
        for (const auto& ind : population) {
            avgFitness += ind.fitness;
            // Calculate average sigma across all genes for this individual
            // for (double s : ind.sigmas) {
            //     avgSigma += s;
            // }
        }
        avgFitness /= config.populationSize;
        // avgSigma /= (config.populationSize * GENOME_SIZE);  // Average of all sigmas across all individuals

        // Print progress
        Individual allTimeBest = evolution.getAllTimeBest();
        std::cout << "Generation " << gen + 1
                  << " | Best: " << best.fitness
                  << " | Avg Fitness: " << avgFitness
                  // << " | Average Sigma: " << avgSigma
                  << " | All-Time Best: " << allTimeBest.fitness << '\n';

        // // Check for population collapse and inject diversity if needed
        // if (best.fitness < 0) {
        //     std::cout << "*** WARNING: Population collapse detected! Injecting diversity... ***" << std::endl;
        //
        //     std::sort(population.begin(), population.end(),
        //         [](const Individual& a, const Individual& b) {
        //             return a.fitness > b.fitness;
        //         });
        //
        //     int keepTop = config.populationSize / 10;  // Keep top 10%
        //     for (int i = keepTop; i < config.populationSize; i++) {
        //         // Reinitialize genome
        //         for (int j = 0; j < GENOME_SIZE; j++) {
        //             population[i].genome[j] = (rand() / (double)RAND_MAX) * 2.0 - 1.0;
        //         }
        //         // Reinitialize ALL sigmas (per-gene sigmas!)
        //         for (int j = 0; j < GENOME_SIZE; j++) {
        //             population[i].sigmas[j] = 0.5 + (rand() / (double)RAND_MAX);  // sigmas[j] not sigma
        //         }
        //         population[i].fitness = 0.0;
        //     }
        // }

        // Evolve to next generation
        evolution.evolve();
    }

    Individual allTimeBest = evolution.getAllTimeBest();
    std::cout << "\nExperiment " << experimentNum << " complete. All-time best: " << allTimeBest.fitness << std::endl;

    return allTimeBest;
}

int main()
{
    // Menu system
    std::cout << "=== Snake AI Experiments ===" << std::endl;
    std::cout << "Press 0 to RUN EXPERIMENTS" << std::endl;
    std::cout << "Press 1 to OBSERVE best snake" << std::endl;
    std::cout << "Your choice: ";

    int choice;
    std::cin >> choice;
    std::cin.ignore();
    std::cout << std::endl;

    if (choice == 1) {
        // OBSERVE MODE
        std::cout << "=== Observation Mode ===" << std::endl;
        std::cout << "Loading best genome from file..." << std::endl;

        std::vector<double> bestGenome;

        std::ifstream file("Best Snake Final.txt");
        if (!file.is_open()) {
            std::cerr << "Error: Could not load 'Best Snake Final.txt'" << std::endl;
            std::cerr << "Please run experiments first." << std::endl;
            return 1;
        }

        double gene;
        while (file >> gene) {
            bestGenome.push_back(gene);
        }
        file.close();

        std::cout << "Genome loaded successfully!" << std::endl;
        std::cout << "Starting observation..." << std::endl;

        Observe observer(bestGenome);
        observer.Run();

        return 0;
    }
    else if (choice != 0) {
        std::cout << "Invalid choice. Please restart and select 0 or 1." << std::endl;
        return 1;
    }

    // EXPERIMENT MODE
    std::cout << "=== Snake AI Parameter Experiments ===" << std::endl;
    std::cout << "This will run 12 experiments testing:" << std::endl;
    std::cout << "  - populationSize: [200, 500, 700]" << std::endl;
    std::cout << "  - crossoverRate: [0, 0.5, 0.8]" << std::endl;
    std::cout << "  - tauGlobal: [0.5, 1.0, 1.5]" << std::endl;
    std::cout << "  - eliteCount: [25, 50, 100]" << std::endl;
    std::cout << "\nResults will be logged to evolution_log.txt" << std::endl;
    std::cout << "\nPress Enter to start experiments...";
    std::cin.get();

    // ============================================================
    // DEFAULT VALUES (used when parameter is held constant)
    // ============================================================
    const int DEFAULT_POP_SIZE = 500;
    const int DEFAULT_NUM_PARENTS = 155;
    const int DEFAULT_NUM_OFFSPRING = 995;
    const double DEFAULT_CROSSOVER_RATE = 0.1;
    const double DEFAULT_TAU_GLOBAL = 1.0;
    const double DEFAULT_TAU_LOCAL = 1.0;

    // ============================================================
    // EXPERIMENT VALUES
    // ============================================================
    int popSizes[] = {200, 500, 700};
    double crossoverRates[] = {0.0, 0.5, 0.8};
    double tauGlobalValues[] = {0.5, 1.0, 1.5};
    int eliteCounts[] = {25, 50, 100};

    int experimentNum = 0;
    int totalExperiments = 12;

    // Track absolute best across ALL experiments
    Individual absoluteBest;
    absoluteBest.fitness = -999999.0;

    // ============================================================
    // EXPERIMENT SET 1: Vary populationSize
    // Hold crossoverRate=0.1, tauGlobal=1.0, tauLocal=1.0
    // ============================================================
    std::cout << "\n\n*** EXPERIMENT SET 1: Varying Population Size ***" << std::endl;

    for (int popSize : popSizes) {
        experimentNum++;

        EvolutionConfig config;
        config.populationSize = popSize;
        // Scale parents and offspring proportionally
        config.numParents = static_cast<int>(popSize * 0.31);      // ~31% of pop
        config.numOffspring = static_cast<int>(popSize * 1.99);    // ~199% of pop
        config.eliteCount = 25;
        config.crossoverRate = DEFAULT_CROSSOVER_RATE;
        config.tauGlobal = DEFAULT_TAU_GLOBAL;
        config.tauLocal = DEFAULT_TAU_LOCAL;
        config.sigmaMin = 0.02;
        config.sigmaMax = 2.0;
        config.genomeInitMin = -1.0;
        config.genomeInitMax = 1.0;

        Individual expBest = runExperiment(config, experimentNum, totalExperiments);

        if (expBest.fitness > absoluteBest.fitness) {
            absoluteBest = expBest;
            std::cout << "*** NEW ABSOLUTE BEST: " << absoluteBest.fitness << " ***" << std::endl;
        }
    }

    // ============================================================
    // EXPERIMENT SET 2: Vary crossoverRate
    // Hold populationSize=500, tauGlobal=1.0, tauLocal=1.0
    // ============================================================
    std::cout << "\n\n*** EXPERIMENT SET 2: Varying Crossover Rate ***" << std::endl;

    for (double crossRate : crossoverRates) {
        experimentNum++;

        EvolutionConfig config;
        config.populationSize = DEFAULT_POP_SIZE;
        config.numParents = DEFAULT_NUM_PARENTS;
        config.numOffspring = DEFAULT_NUM_OFFSPRING;
        config.eliteCount = 25;
        config.crossoverRate = crossRate;
        config.tauGlobal = DEFAULT_TAU_GLOBAL;
        config.tauLocal = DEFAULT_TAU_LOCAL;
        config.sigmaMin = 0.02;
        config.sigmaMax = 2.0;
        config.genomeInitMin = -1.0;
        config.genomeInitMax = 1.0;

        Individual expBest = runExperiment(config, experimentNum, totalExperiments);

        if (expBest.fitness > absoluteBest.fitness) {
            absoluteBest = expBest;
            std::cout << "*** NEW ABSOLUTE BEST: " << absoluteBest.fitness << " ***" << std::endl;
        }
    }

    // ============================================================
    // EXPERIMENT SET 3: Vary tauGlobal
    // Hold populationSize=500, crossoverRate=0.1, tauLocal=1.0
    // ============================================================
    std::cout << "\n\n*** EXPERIMENT SET 3: Varying Tau Global ***" << std::endl;

    for (double tau : tauGlobalValues) {
        experimentNum++;

        EvolutionConfig config;
        config.populationSize = DEFAULT_POP_SIZE;
        config.numParents = DEFAULT_NUM_PARENTS;
        config.numOffspring = DEFAULT_NUM_OFFSPRING;
        config.eliteCount = 25;
        config.crossoverRate = DEFAULT_CROSSOVER_RATE;
        config.tauGlobal = tau;
        config.tauLocal = DEFAULT_TAU_LOCAL;
        config.sigmaMin = 0.02;
        config.sigmaMax = 2.0;
        config.genomeInitMin = -1.0;
        config.genomeInitMax = 1.0;

        Individual expBest = runExperiment(config, experimentNum, totalExperiments);

        if (expBest.fitness > absoluteBest.fitness) {
            absoluteBest = expBest;
            std::cout << "*** NEW ABSOLUTE BEST: " << absoluteBest.fitness << " ***" << std::endl;
        }
    }

    // ============================================================
    // EXPERIMENT SET 4: Vary eliteCount
    // Hold populationSize=500, crossoverRate=0.1, tauGlobal=1.0, tauLocal=1.0
    // ============================================================
    std::cout << "\n\n*** EXPERIMENT SET 4: Varying Elite Count ***" << std::endl;

    for (int elite : eliteCounts) {
        experimentNum++;

        EvolutionConfig config;
        config.populationSize = DEFAULT_POP_SIZE;
        config.numParents = DEFAULT_NUM_PARENTS;
        config.numOffspring = DEFAULT_NUM_OFFSPRING;
        config.eliteCount = elite;
        config.crossoverRate = DEFAULT_CROSSOVER_RATE;
        config.tauGlobal = DEFAULT_TAU_GLOBAL;
        config.tauLocal = DEFAULT_TAU_LOCAL;
        config.sigmaMin = 0.02;
        config.sigmaMax = 2.0;
        config.genomeInitMin = -1.0;
        config.genomeInitMax = 1.0;

        Individual expBest = runExperiment(config, experimentNum, totalExperiments);

        if (expBest.fitness > absoluteBest.fitness) {
            absoluteBest = expBest;
            std::cout << "*** NEW ABSOLUTE BEST: " << absoluteBest.fitness << " ***" << std::endl;
        }
    }

    // ============================================================
    // SAVE ABSOLUTE BEST
    // ============================================================
    std::cout << "\n========================================" << std::endl;
    std::cout << "ALL EXPERIMENTS COMPLETE!" << std::endl;
    std::cout << "Results saved to evolution_log.txt" << std::endl;
    std::cout << "Absolute best fitness: " << absoluteBest.fitness << std::endl;
    std::cout << "========================================" << std::endl;

    // Save the absolute best genome
    std::ofstream bestFile("Best Snake Final.txt");
    for (double gene : absoluteBest.genome) {
        bestFile << gene << "\n";
    }
    bestFile.close();
    std::cout << "Absolute best genome saved to Best Snake Final.txt" << std::endl;

    return 0;
}