use clap::Parser;
use rand::prelude::*;
use std::time::Instant;
use triadic_memory::{
    sdr::{self, Sdr},
    tri_mem::*,
};

const RMAX: u128 = 999_999_999;

#[derive(Debug, Parser)]
#[clap(author, version, about)]
struct Cli {
    /// Size of the SDRs
    #[clap(long, short = 'N', default_value_t = sdr::SIZE)]
    size: usize,

    /// Number of active connections in the SDRs
    #[clap(long, short = 'P', default_value_t = sdr::POP)]
    pop: usize,

    /// Number of triples in reporting interval, defaults to <TRIPLES>
    #[clap(long, short = 'i')]
    reporting_interval: Option<usize>,

    /// Number of triples to store, defaults to <SIZE> * 100
    #[clap(long, short)]
    triples: Option<usize>,

    /// Only do the writing benchmarks
    #[clap(long = "only-writes", action)]
    only_writes: bool,

    /// Don't benchmark fetch_y
    #[clap(long = "skip-y", action)]
    no_y: bool,

    /// Check recall errors on first <NOSTALGIA> triples once all writes are done
    #[clap(long)]
    nostalgia: Option<usize>,
}

fn main() {
    let cli = Cli::parse();

    let size = cli.size;
    let pop = cli.pop;
    let triples = cli.triples.unwrap_or(size * 100);
    let interval = cli.reporting_interval.unwrap_or(triples);
    let num_sdrs = triples * 3;

    let mut rng = StdRng::from_entropy();

    let mut mem = TriadicMemory::new(size, pop);

    let mut sdrs: Vec<Sdr> = Vec::with_capacity(num_sdrs);

    for _ in 0..num_sdrs {
        sdrs.push(Sdr::random(size, pop, &mut rng));
    }

    let span = interval * 3;
    let mut writes = 0;
    while writes < triples {
        let start = writes * 3;
        let end = (start + span).min(num_sdrs);
        let sdrs = &sdrs[start..end];
        let prev_writes = writes;

        let storing = Instant::now();
        for triple in sdrs.chunks(3) {
            mem.store(&triple[0], &triple[1], &triple[2]);
            writes += 1;
        }
        let elapsed = storing.elapsed();

        let interval = writes - prev_writes;

        // "triples,N,P,writes/s,x/s,y/s,z/s,max x-dist,max y-dist,max z-dist,max x-pop,max y-pop,max z-pop,max mem weight"
        let mut report = format!("{},{},{},", writes, size, pop);

        let write_rate = ((interval * 1000) as u128)
            .checked_div(elapsed.as_millis())
            .unwrap_or(RMAX);
        report.push_str(&write_rate.to_string());
        report.push(',');

        if !cli.only_writes {
            // X
            let mut x_overlap = usize::MAX;
            let mut x_pop = 0;
            let timer = Instant::now();
            for triple in sdrs.chunks(3) {
                let s = mem.fetch(None, Some(&triple[1]), Some(&triple[2])).unwrap();
                x_overlap = x_overlap.min(s.overlap(&triple[0]));
                x_pop = x_pop.max(s.pop());
            }
            let elapsed = timer.elapsed();
            let x_rate = ((interval * 1000) as u128)
                .checked_div(elapsed.as_millis())
                .unwrap_or(RMAX);

            // Y
            let mut y_overlap = usize::MAX;
            let mut y_pop = 0;
            //
            let y_rate = if cli.no_y {
                0
            } else {
                let timer = Instant::now();
                for triple in sdrs.chunks(3) {
                    let s = mem.fetch(Some(&triple[0]), None, Some(&triple[2])).unwrap();
                    y_overlap = y_overlap.min(s.overlap(&triple[1]));
                    y_pop = y_pop.max(s.pop());
                }
                let elapsed = timer.elapsed();
                ((interval * 1000) as u128)
                    .checked_div(elapsed.as_millis())
                    .unwrap_or(RMAX)
            };

            // Z
            let mut z_overlap = usize::MAX;
            let mut z_pop = 0;
            let timer = Instant::now();
            for triple in sdrs.chunks(3) {
                let s = mem.fetch(Some(&triple[0]), Some(&triple[1]), None).unwrap();
                z_overlap = z_overlap.min(s.overlap(&triple[2]));
                z_pop = z_pop.max(s.pop());
            }
            let elapsed = timer.elapsed();
            let z_rate = ((interval * 1000) as u128)
                .checked_div(elapsed.as_millis())
                .unwrap_or(RMAX);

            // OK, add what we have to our report
            let (y_rate, y_overlap, y_pop) = if cli.no_y {
                ("*".to_string(), "*".to_string(), "*".to_string())
            } else {
                (y_rate.to_string(), y_overlap.to_string(), y_pop.to_string())
            };
            let s = format!(
                "{x_rate},{y_rate},{z_rate},{x_overlap},{y_overlap},{z_overlap},{x_pop},{y_pop},{z_pop},"
            );
            report.push_str(&s);
        } else {
            report.push_str(",,,,,,,,,");
        }

        println!("{}", report);
    }

    if let Some(t) = cli.nostalgia {
        let n = t * 3;
        let mut x_dist = 0;
        let mut y_dist = 0;
        let mut z_dist = 0;
        let mut x_over = size;
        let mut y_over = size;
        let mut z_over = size;
        for triple in sdrs[0..n].chunks(3) {
            let (x, y, z) = (&triple[0], &triple[1], &triple[2]);
            let rx = mem.fetch(None, Some(y), Some(z)).unwrap();
            let ry = mem.fetch(Some(x), None, Some(z)).unwrap();
            let rz = mem.fetch(Some(x), Some(y), None).unwrap();
            x_over = x.overlap(&rx).min(x_over);
            x_dist = x.hamming(&rx).max(x_dist);
            y_over = y.overlap(&ry).min(y_over);
            y_dist = y.hamming(&ry).max(y_dist);
            z_over = z.overlap(&rz).min(z_over);
            z_dist = z.hamming(&rz).max(z_dist);
        }
        println!(
            "NOSTALGIA, first {} triples: x_max_dist: {}, x_min_overlap: {}",
            t, x_dist, x_over
        );
        println!(
            "NOSTALGIA, first {} triples: y_max_dist: {}, y_min_overlap: {}",
            t, y_dist, y_over
        );
        println!(
            "NOSTALGIA, first {} triples: z_max_dist: {}, z_min_overlap: {}",
            t, z_dist, z_over
        );
    }
}