use clap::Parser;
use rand::prelude::*;
use triadic_memory::{
    sdr::{self, Sdr},
    sequence_mem::SequenceMemory,
};

#[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 sequences to store
    #[clap(long, short, default_value_t = 10)]
    sequences: usize,

    /// Length of each stored sequence
    #[clap(long, short, default_value_t = 100)]
    length: usize,

    /// How many items to feed each sequence in before trying to recall
    #[clap(long, short = 'p', default_value_t = 2)]
    pre_roll: usize,

    /// Check recall on earliest sequence at the end of the test.
    #[clap(long, action)]
    nostalgia: bool,

    /// Don't reuse the same `LENGTH` SDRs in different orders for each sequence
    #[clap(long, short, action)]
    unique: bool,
}

fn main() {
    let cli = Cli::parse();
    let size = cli.size;
    let pop = cli.pop;
    let seqs = cli.sequences;
    let len = cli.length;
    let pre_roll = cli.pre_roll;

    let mut rng = StdRng::from_entropy();

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

    let mut raw_sdrs: Vec<Sdr> = Vec::with_capacity(len);
    for _ in 0..len {
        raw_sdrs.push(Sdr::random(size, pop, &mut rng));
    }

    let mut sdrs = Vec::with_capacity(seqs);
    for _ in 1..seqs {
        if cli.unique {
            let mut v = Vec::with_capacity(len);
            for _ in 0..len {
                v.push(Sdr::random(size, pop, &mut rng));
            }
            sdrs.push(v);
        } else {
            let mut v = raw_sdrs.clone();
            v.shuffle(&mut rng);
            sdrs.push(v);
        }
    }
    sdrs.push(raw_sdrs);

    for (sequence, sdrs) in sdrs.iter().enumerate() {
        // first, add the sequence
        mem.new_seq();
        for sdr in sdrs {
            mem.add_to_sequence(sdr);
        }

        // now, check recall
        mem.new_seq();
        for sdr in sdrs[0..pre_roll].iter() {
            mem.predict_next(sdr);
        }
        let mut prev = sdrs[pre_roll].clone();
        for (i, sdr) in sdrs[pre_roll..len].iter().enumerate() {
            let pop = prev.pop();
            let overlap = sdr.overlap(&prev);
            prev = mem.predict_next(&prev);
            println!(
                "Sequence {sequence}, SDR {idx}: pop: {pop}, overlap: {overlap}",
                idx = i + pre_roll,
            );
        }
    }

    if cli.nostalgia {
        mem.new_seq();
        let sdrs = &sdrs[0];
        for sdr in sdrs[0..pre_roll].iter() {
            mem.predict_next(sdr);
        }
        let mut prev = sdrs[pre_roll].clone();
        for (i, sdr) in sdrs[pre_roll..len].iter().enumerate() {
            let dist = sdr.hamming(&prev);
            let overlap = sdr.overlap(&prev);
            prev = mem.predict_next(&prev);
            println!(
                "NOSTALGIA SDR {}: distance: {}, overlap: {}",
                i + pre_roll,
                dist,
                overlap
            );
        }
    }
}