evolution/src/apps/SnakeAI/SnakeAI.tsx

import { useState, useEffect, useRef, useCallback } from 'react';
import AppContainer from '../../components/AppContainer';
import SnakeGrid from './SnakeGrid';
import Controls from './Controls';
import Stats from './Stats';
import Tips from './Tips';
import BestSnakeDisplay from './BestSnakeDisplay';
import {
    createPopulation,
} from '../../lib/snakeAI/evolution';
import type { EvolutionConfig } from '../../lib/snakeAI/types';
import './SnakeAI.css';

const DEFAULT_CONFIG: EvolutionConfig = {
    populationSize: 36,
    mutationRate: 0.1,
    eliteCount: 5,
    gridSize: 20,
    maxGameSteps: 20000,
};

import { WorkerPool } from '../../lib/snakeAI/workerPool';
import { evolveGeneration, updateBestStats, type Population } from '../../lib/snakeAI/evolution';

export default function SnakeAI() {
    const [population, setPopulation] = useState<Population>(() =>
        createPopulation(DEFAULT_CONFIG)
    );
    const [config, setConfig] = useState<EvolutionConfig>(DEFAULT_CONFIG);
    const [isRunning, setIsRunning] = useState(false);
    const [speed, setSpeed] = useState(5);
    const [gamesPlayed, setGamesPlayed] = useState(0);
    const [fitnessHistory, setFitnessHistory] = useState<Array<{ generation: number, best: number, average: number }>>([]);

    // Keep a ref to population for the worker
    const populationRef = useRef(population);
    useEffect(() => {
        populationRef.current = population;
    }, [population]);

    const animationFrameRef = useRef<number>();
    const lastUpdateRef = useRef<number>(0);

    // Compute derived values for display
    const currentBestFitness = population.lastGenerationStats?.bestFitness || 0;
    const currentAverageFitness = population.lastGenerationStats?.averageFitness || 0;

    const workerPoolRef = useRef<WorkerPool | null>(null);
    const isProcessingRef = useRef(false);

    useEffect(() => {
        // Initialize Worker Pool with logical cores (default)
        workerPoolRef.current = new WorkerPool();

        return () => {
            workerPoolRef.current?.terminate();
        };
    }, []);

    const runGeneration = useCallback(async (generations: number = 1) => {
        if (isProcessingRef.current || !workerPoolRef.current) return;

        isProcessingRef.current = true;
        let currentPop = populationRef.current;

        try {
            for (let i = 0; i < generations; i++) {
                // 1. Evaluate in parallel
                let evaluatedPop = await workerPoolRef.current.evaluateParallel(currentPop, config);

                // 1.5 Update Best Stats (Critical for UI)
                evaluatedPop = updateBestStats(evaluatedPop);

                // 2. Evolve on main thread (fast)
                currentPop = evolveGeneration(evaluatedPop, config);
            }

            // Update state
            populationRef.current = currentPop;
            setPopulation(currentPop);

            // Update history
            if (currentPop.lastGenerationStats) {
                setFitnessHistory(prev => {
                    const newEntry = {
                        generation: currentPop.generation - 1,
                        best: currentPop.lastGenerationStats!.bestFitness,
                        average: currentPop.lastGenerationStats!.averageFitness
                    };
                    const newHistory = [...prev, newEntry];
                    if (newHistory.length > 100) return newHistory.slice(newHistory.length - 100);
                    return newHistory;
                });
            }

        } catch (err) {
            console.error("Evolution error:", err);
            setIsRunning(false);
        } finally {
            isProcessingRef.current = false;
        }
    }, [config]);

    // Update stats when generation changes
    useEffect(() => {
        if (population.generation > 1) {
            setGamesPlayed((count) => count + config.populationSize);
        }
    }, [population.generation, config.populationSize]);

    useEffect(() => {
        if (!isRunning) {
            if (animationFrameRef.current) {
                cancelAnimationFrame(animationFrameRef.current);
            }
            return;
        }

        const animate = (timestamp: number) => {
            const elapsed = timestamp - lastUpdateRef.current;

            // Speed handling logic:
            // 1-20: Visual speeds (delay between generations)
            // 21-100: Turbo speeds (multiple generations per frame)

            if (speed <= 20) {
                // Standard visual mode
                let updateInterval;
                if (speed <= 5) {
                    // Speeds 1-5: Very slow observation
                    updateInterval = 12000 / speed;
                } else {
                    // Speeds 6-20: 62.5ms to 1000ms
                    // speed 20 -> 1000/16 = 62.5ms
                    // speed 6 -> 1000/2 = 500ms
                    updateInterval = 1000 / (speed - 4);
                }

                if (elapsed >= updateInterval) {
                    runGeneration(1);
                    lastUpdateRef.current = timestamp;
                }
            } else {
                // Turbo mode: Run multiple generations per frame
                // Speed 21 -> 1 gen per frame (~60 eps)
                // Speed 100 -> 10 gens per frame (~600 eps)
                const gensPerFrame = Math.floor((speed - 10) / 10);

                // For turbo mode, we just fire once per frame (or whenever the worker is ready)
                // asking for multiple generations
                runGeneration(gensPerFrame);
                lastUpdateRef.current = timestamp;
            }

            animationFrameRef.current = requestAnimationFrame(animate);
        };

        animationFrameRef.current = requestAnimationFrame(animate);

        return () => {
            if (animationFrameRef.current) {
                cancelAnimationFrame(animationFrameRef.current);
            }
        };
    }, [isRunning, speed, runGeneration]);

    const handleReset = () => {
        setIsRunning(false);
        const newPop = createPopulation(config);
        populationRef.current = newPop;
        setPopulation(newPop);
        setGamesPlayed(0);
    };

    const handleMutationRateChange = (rate: number) => {
        setConfig((prev) => ({ ...prev, mutationRate: rate }));
    };

    return (
        <AppContainer title="Snake Evolution">
            <div className="snake-ai-layout">
                <div className="left-panel">
                    <BestSnakeDisplay
                        network={population.bestNetworkEver}
                        gridSize={config.gridSize}
                        fitness={population.bestFitnessEver}
                    />
                    <SnakeGrid
                        individuals={population.individuals}
                        gridSize={config.gridSize}
                        count={config.populationSize}
                        columns={6}
                    />

                </div>

                <div className="right-panel">
                    <Controls
                        isRunning={isRunning}
                        onToggleRunning={() => setIsRunning(!isRunning)}
                        onReset={handleReset}
                        speed={speed}
                        onSpeedChange={setSpeed}
                        mutationRate={config.mutationRate}
                        onMutationRateChange={handleMutationRateChange}
                        populationSize={config.populationSize}
                    />

                    <Stats
                        generation={population.generation}
                        bestFitness={currentBestFitness}
                        bestFitnessEver={population.bestFitnessEver}
                        averageFitness={currentAverageFitness}
                        gamesPlayed={gamesPlayed}
                        history={fitnessHistory}
                    />

                    <Tips />

                </div>
            </div>
        </AppContainer>
    );
}