use clap::Parser;
use indicatif::{ProgressBar, ProgressStyle};
use ldpc::code::{CodeTopology, GenerationMethod, LdpcCode, LdpcParams};
use rayon::prelude::*;
use std::fs::File;
use std::io::Write;
use std::sync::{
    atomic::{AtomicBool, AtomicU64, Ordering},
    Arc,
};
use std::time::{Duration, Instant};

#[derive(Parser, Debug)]
#[command(author, version, about = "Generateur LDPC Haute Performance")]
struct Args {
    #[arg(short, long)]
    n: usize,
    #[arg(short, long)]
    k: usize,
    #[arg(long, default_value_t = 3)]
    wc: usize,
    #[arg(long, default_value_t = 6)]
    wr: usize,
    #[arg(short, long, default_value = "mackay")]
    method: String,
}

fn main() -> ldpc::Result<()> {
    let args = Args::parse();
    let m = args.n - args.k;

    if (args.n * args.wc) % m != 0 || (args.n * args.wc) / m != args.wr {
        println!("Erreur : Parametres impossibles (n*wc != m*wr)");
        std::process::exit(1);
    }

    let found = Arc::new(AtomicBool::new(false));
    let attempts = Arc::new(AtomicU64::new(0));
    let start_global = Instant::now();

    let pb = ProgressBar::new_spinner();
    pb.enable_steady_tick(Duration::from_millis(80));
    pb.set_style(
        ProgressStyle::default_spinner()
            .template(
                "{spinner:.cyan} [{elapsed_precise}] {msg} | Tentatives: {pos} | {per_sec} mats/s",
            )
            .unwrap(),
    );
    pb.set_message("Recherche d'une matrice valide...");

    let method = if args.method == "mackay" {
        GenerationMethod::MacKayNeal { max_attempts: 1000 }
    } else {
        GenerationMethod::Gallager
    };

    let pb_clone = pb.clone();
    let found_clone = Arc::clone(&found);
    let attempts_clone = Arc::clone(&attempts);

    let result = (0..num_cpus::get())
        .into_par_iter()
        .map(|_| {
            while !found_clone.load(Ordering::Relaxed) {
                let current_attempt = attempts_clone.fetch_add(1, Ordering::SeqCst);
                pb_clone.set_position(current_attempt);

                let params = LdpcParams {
                    n: args.n,
                    k: args.k,
                    topology: CodeTopology::Regular {
                        wc: args.wc,
                        wr: args.wr,
                    },
                    generation: method.clone(),
                    seed: Some(rand::random()),
                };

                if let Ok(mut code) = LdpcCode::new(params) {
                    // Cette etape est la plus longue (Gauss-Jordan)
                    if code.compute_systematic_form().is_ok() {
                        if !found_clone.swap(true, Ordering::SeqCst) {
                            return Some(code);
                        }
                    }
                }
            }
            None
        })
        .find_any(|res| res.is_some())
        .flatten();

    if let Some(code) = result {
        pb.finish_with_message(format!("Succes en {:.2?}", start_global.elapsed()));

        let filename = format!(
            "cache_ldpc_{}_n{}_k{}.bin",
            args.method.to_uppercase(),
            args.n,
            args.k
        );
        let encoded = bincode::serialize(&code).expect("Erreur serialisation");
        let mut file = File::create(&filename).expect("Erreur creation");
        file.write_all(&encoded).expect("Erreur ecriture");

        println!("\nFichier genere : {}", filename);
        println!(
            "Girth : {} | Densite : {:.2}%",
            code.girth(),
            code.h.density() * 100.0
        );
    } else {
        pb.abandon_with_message("Recherche arretee.");
    }

    Ok(())
}
