use clap::{Parser, Subcommand};
use factorio_core::visualize::Visualize;
use factorio_layout::layout::{GeneticAlgorithm, PathLayout, genetic_algorithm2};
use miette::{Context, IntoDiagnostic, Result};
use rand::{SeedableRng, rngs::SmallRng};
use std::path::PathBuf;

#[derive(Debug, Parser)]
struct Args {
    #[clap(short, long, default_value_t = 0)]
    seed: u64,

    problem: PathBuf,

    #[command(subcommand)]
    subcommand: Commands,
}

#[derive(Debug, Subcommand)]
enum Commands {
    V1,
    V2,
    Bench {
        #[clap(short, long, default_value_t = 100)]
        runs: usize,

        #[clap(short, long, default_value_t = 100)]
        mutations: usize,
    },
}

fn main() -> Result<()> {
    let args = Args::parse();

    let mut rng = SmallRng::seed_from_u64(args.seed);
    let file = std::fs::File::open(args.problem)
        .into_diagnostic()
        .context("Failed to open problem file.")?;
    let p = serde_yaml::from_reader(file)
        .into_diagnostic()
        .context("Failed to parse yaml.")?;

    match args.subcommand {
        Commands::V1 => {
            let mut g = GeneticAlgorithm::new(&p, 20, 2, 0, &mut rng);

            for i in 0..100 {
                println!("Generatrion {i}");
                g.generation(&mut rng);
            }

            g.output_population();
        }
        Commands::V2 => {
            let mut m: Option<PathLayout> = None;
            for _ in 0..20 {
                let g = genetic_algorithm2(&p, 10, 320, 10000, &mut rng);

                g.print_visualization();
                if m.as_ref().is_none_or(|m| g.score() < m.score()) {
                    m = Some(g);
                }
            }

            m.unwrap().print_visualization();
        }
        Commands::Bench { runs, mutations } => {
            let mut map = Vec::new();

            let mut m: Option<PathLayout> = None;
            let start = std::time::Instant::now();
            for _ in 0..runs {
                let g = genetic_algorithm2(&p, 1, mutations, mutations, &mut rng);

                map.push(g.score());
                if m.as_ref().is_none_or(|m| g.score() < m.score()) {
                    m = Some(g);
                }
            }

            println!("Time: {:.2}", start.elapsed().as_secs_f32());

            let mean = map.iter().sum::<usize>() / runs;
            println!("Mean: {}", mean);
            let min = map.iter().min().unwrap();
            println!("Min: {}", min);
            let max = map.iter().max().unwrap();
            println!("Max: {}", max);
            let stddev =
                ((map.iter().map(|v| (v - mean) * (v - mean)).sum::<usize>() / runs) as f64).sqrt();
            println!("Stddev: {:.1}", stddev);

            m.unwrap().print_visualization();
        }
    }

    Ok(())
}