use ndarray::Array2;
use ndarray_npy::read_npy;
use std::{env, iter::zip, path::PathBuf};

use muscat::distinguishers::cpa::CpaProcessor;
use muscat::leakage_model::aes::sbox;

fn leakage_model(plaintext_byte: usize, guess: usize) -> usize {
    sbox((plaintext_byte ^ guess) as u8) as usize
}

fn main() {
    let traces_dir =
        PathBuf::from(env::var("TRACES_DIR").expect("Missing TRACES_DIR environment variable"));

    let traces: Array2<f64> =
          read_npy(traces_dir.join("traces.npy")).expect("Failed to read traces.npy");
    let plaintexts: Array2<u8> =
          read_npy(traces_dir.join("plaintexts.npy")).expect("Failed to read plaintexts.npy");
    assert_eq!(traces.shape()[0], plaintexts.shape()[0]);

    let mut processor = CpaProcessor::new(traces.shape()[1], 256);
    for (trace, plaintext) in zip(traces.rows(), plaintexts.rows()) {
        processor.update(trace.view(), plaintext[0], leakage_model);
    }
    let cpa = processor.finalize(leakage_model);
    let best_guess = cpa.best_guess();
    println!("Best subkey guess: {best_guess:?}");
}