use clap::{Parser, Subcommand};
use factorio_blueprint::{
    common::visualize::Visualize,
    layout::{genetic_algorithm2, GeneticAlgorithm, PathLayout},
};
use miette::{Context, IntoDiagnostic, Result};
use rand::{rngs::SmallRng, SeedableRng};
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,
}

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, &mut rng);

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

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

    Ok(())
}