// MIT License

// Copyright (c) 2025 Tobias Raayoni Last (@randogoth)

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

//! Builds the null calibration tables in `src/calibration/table.rs`, or checks them.
//!
//!     cargo run --release --example calibrate -- [--reps N] [--seed S] [--urandom] [--out PATH] [--dump PATH]
//!     cargo run --release --example calibrate -- --from-dump PATH [--out PATH]
//!     cargo run --release --example calibrate -- --check N [--sizes A,B,..] [--seed S] [--urandom]
//!
//! Generation runs every test on N random samples at each calibrated size and stores the
//! quantiles of the z-scores, then the quantiles of the Cauchy combination statistic. The
//! raw z-scores are dumped so the tables can be rebuilt with `--from-dump` without
//! simulating again. `--check` scores N fresh samples per size with the compiled-in
//! tables and reports how often the calibrated p-values fall below 0.05 ... 0.0001.

use std::fs::File;
use std::io::{self, Read, Write};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Mutex;

use onod3000::calibration::{self, LEVELS, MAX_LOG2, MIN_LOG2, SIZES, TWO_SIDED};
use onod3000::score::acat;
use onod3000::{Onod, TESTS};

const DEFAULT_REPS: usize = 20_000;
const DEFAULT_SEED: u64 = 0x0d3000;

/// xoshiro256** seeded through SplitMix64.
struct Xoshiro([u64; 4]);

impl Xoshiro {
    fn new(seed: u64) -> Self {
        let mut x = seed;
        Xoshiro(std::array::from_fn(|_| {
            x = x.wrapping_add(0x9e3779b97f4a7c15);
            let mut z = x;
            z = (z ^ (z >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
            z = (z ^ (z >> 27)).wrapping_mul(0x94d049bb133111eb);
            z ^ (z >> 31)
        }))
    }

    fn next(&mut self) -> u64 {
        let s = &mut self.0;
        let result = s[1].wrapping_mul(5).rotate_left(7).wrapping_mul(9);
        let t = s[1] << 17;
        s[2] ^= s[0];
        s[3] ^= s[1];
        s[1] ^= s[2];
        s[0] ^= s[3];
        s[2] ^= t;
        s[3] = s[3].rotate_left(45);
        result
    }

    fn fill(&mut self, buf: &mut [u8]) {
        for chunk in buf.chunks_mut(8) {
            let bytes = self.next().to_le_bytes();
            chunk.copy_from_slice(&bytes[..chunk.len()]);
        }
    }
}

struct Args {
    reps: usize,
    seed: u64,
    out: String,
    dump: String,
    from_dump: Option<String>,
    urandom: bool,
    check: Option<usize>,
    sizes: Option<Vec<usize>>,
}

fn parse_args() -> Args {
    let mut args = Args {
        reps: DEFAULT_REPS,
        seed: DEFAULT_SEED,
        out: String::from("src/calibration/table.rs"),
        dump: String::from("target/calibration-runs.bin"),
        from_dump: None,
        urandom: false,
        check: None,
        sizes: None,
    };
    let mut it = std::env::args().skip(1);
    while let Some(arg) = it.next() {
        let mut value = || it.next().unwrap_or_else(|| usage());
        match arg.as_str() {
            "--reps" => args.reps = value().parse().unwrap_or_else(|_| usage()),
            "--seed" => args.seed = value().parse().unwrap_or_else(|_| usage()),
            "--out" => args.out = value(),
            "--dump" => args.dump = value(),
            "--from-dump" => args.from_dump = Some(value()),
            "--check" => args.check = Some(value().parse().unwrap_or_else(|_| usage())),
            "--sizes" => {
                args.sizes = Some(value().split(',').map(|n| n.parse().unwrap_or_else(|_| usage())).collect())
            }
            "--urandom" => args.urandom = true,
            _ => usage(),
        }
    }
    args
}

fn usage() -> ! {
    eprintln!("Usage: calibrate [--reps N] [--seed S] [--urandom] [--out PATH] [--dump PATH]");
    eprintln!("       calibrate --from-dump PATH [--out PATH]");
    eprintln!("       calibrate --check N [--sizes A,B,..] [--seed S] [--urandom]");
    std::process::exit(1);
}

/// Fills `buf` with the random sample for repetition `rep` at size index `size`.
fn sample(buf: &mut [u8], urandom: &Option<Mutex<File>>, seed: u64, size: usize, rep: usize) {
    match urandom {
        Some(file) => file.lock().unwrap().read_exact(buf).expect("reading /dev/urandom"),
        None => Xoshiro::new(seed ^ ((size as u64) << 40) ^ rep as u64).fill(buf),
    }
}

/// Runs `job(rep, buf)` for every repetition on all cores, collecting the results in order.
fn parallel<T: Send>(reps: usize, n: usize, job: impl Fn(usize, &mut [u8]) -> T + Sync) -> Vec<T> {
    let next = AtomicUsize::new(0);
    let results: Mutex<Vec<Option<T>>> = Mutex::new((0..reps).map(|_| None).collect());
    let threads = std::thread::available_parallelism().map_or(4, |n| n.get());
    std::thread::scope(|scope| {
        for _ in 0..threads {
            scope.spawn(|| {
                let mut buf = vec![0u8; n];
                loop {
                    let rep = next.fetch_add(1, Ordering::Relaxed);
                    if rep >= reps {
                        break;
                    }
                    let result = job(rep, &mut buf);
                    results.lock().unwrap()[rep] = Some(result);
                }
            });
        }
    });
    results.into_inner().unwrap().into_iter().map(Option::unwrap).collect()
}

/// Quantiles of `values` at the calibration levels (linear between order statistics).
fn quantiles(values: &mut [f64]) -> [f64; LEVELS] {
    values.sort_by(f64::total_cmp);
    let last = (values.len() - 1) as f64;
    calibration::levels().map(|p| {
        let h = p * last;
        let (i, w) = (h.floor() as usize, h.fract());
        let next = values[(i + 1).min(values.len() - 1)];
        if w == 0.0 || next == values[i] { values[i] } else { values[i] + w * (next - values[i]) }
    })
}

/// z-score and raw p of every test per repetition and size, NaN where the test skipped.
type Runs = Vec<Vec<[(f64, f64); 15]>>;

fn simulate(args: &Args, urandom: &Option<Mutex<File>>) -> Runs {
    (0..SIZES)
        .map(|s| {
            let n = 1usize << (MIN_LOG2 as usize + s);
            let started = std::time::Instant::now();
            let runs = parallel(args.reps, n, |rep, buf| {
                sample(buf, urandom, args.seed, s, rep);
                std::array::from_fn(|t| {
                    let (stat, z, p) = Onod::run(TESTS[t].1, buf);
                    calibration::calibration_z(t, stat, z, n).map_or((f64::NAN, f64::NAN), |z| (z, p))
                })
            });
            eprintln!("{:>8} B  {} reps in {:.1} s", n, args.reps, started.elapsed().as_secs_f64());
            runs
        })
        .collect()
}

/// Dump format: header line "<source>\n", then per size a u64 count followed by
/// count * 15 * (z, p) little-endian f64 values.
fn save_runs(path: &str, source: &str, runs: &Runs) -> io::Result<()> {
    let mut out = io::BufWriter::new(File::create(path)?);
    writeln!(out, "{}", source)?;
    for size in runs {
        out.write_all(&(size.len() as u64).to_le_bytes())?;
        for run in size {
            for &(z, p) in run {
                out.write_all(&z.to_le_bytes())?;
                out.write_all(&p.to_le_bytes())?;
            }
        }
    }
    eprintln!("Saved raw runs to {}", path);
    Ok(())
}

fn load_runs(path: &str) -> io::Result<(String, Runs)> {
    let mut data = Vec::new();
    File::open(path)?.read_to_end(&mut data)?;
    let newline = data.iter().position(|&b| b == b'\n').expect("dump header");
    let source = String::from_utf8_lossy(&data[..newline]).into_owned();
    let mut pos = newline + 1;
    let mut next = || {
        let v = u64::from_le_bytes(data[pos..pos + 8].try_into().unwrap());
        pos += 8;
        v
    };
    let runs = (0..SIZES)
        .map(|_| {
            let count = next() as usize;
            (0..count)
                .map(|_| std::array::from_fn(|_| (f64::from_bits(next()), f64::from_bits(next()))))
                .collect()
        })
        .collect();
    Ok((source, runs))
}

fn generate(args: &Args, urandom: &Option<Mutex<File>>) -> io::Result<()> {
    let (source, runs) = match &args.from_dump {
        Some(path) => load_runs(path)?,
        None => {
            let source = if args.urandom {
                String::from("/dev/urandom")
            } else {
                format!("xoshiro256** seed {:#x}", args.seed)
            };
            let runs = simulate(args, urandom);
            save_runs(&args.dump, &source, &runs)?;
            (source, runs)
        }
    };
    let reps = runs[0].len();

    let mut z_rows = vec![[[f64::NAN; LEVELS]; 15]; SIZES];
    let mut t_rows = vec![[f64::NAN; LEVELS]; SIZES];
    let mut raw_fail = vec![[f64::NAN; 15]; SIZES];
    let mut notes = Vec::new();

    for (s, runs) in runs.iter().enumerate() {
        let n = 1usize << (MIN_LOG2 as usize + s);
        for t in 0..15 {
            let mut zs: Vec<f64> = runs.iter().map(|r| r[t].0).filter(|z| !z.is_nan()).collect();
            let skipped = 1.0 - zs.len() as f64 / runs.len() as f64;
            if skipped > 0.5 {
                notes.push(format!("{:>8} B  {:<14} skipped in {:.0}% of runs, not calibrated", n, TESTS[t].0, 100.0 * skipped));
                continue;
            }
            raw_fail[s][t] = runs.iter().filter(|r| r[t].1 < 0.01).count() as f64 / zs.len() as f64;
            z_rows[s][t] = quantiles(&mut zs);
            if z_rows[s][t][LEVELS - 1] - z_rows[s][t][0] <= 1e-9 * (1.0 + z_rows[s][t][0].abs()) {
                notes.push(format!("{:>8} B  {:<14} constant on random data (z = {})", n, TESTS[t].0, z_rows[s][t][0]));
            }
        }

        // Null distribution of the combination statistic, with the same lookup as at runtime.
        // Cross-fitted: each run is scored against tables built from the other half of the
        // runs, since runs look less extreme against tables that include themselves.
        let halves: [Vec<[f64; LEVELS]>; 2] = std::array::from_fn(|h| {
            (0..15)
                .map(|t| {
                    let mut zs: Vec<f64> =
                        runs.iter().skip(h).step_by(2).map(|r| r[t].0).filter(|z| !z.is_nan()).collect();
                    if zs.is_empty() { [f64::NAN; LEVELS] } else { quantiles(&mut zs) }
                })
                .collect()
        });
        let mut ts: Vec<f64> = runs
            .iter()
            .enumerate()
            .map(|(rep, r)| {
                let other = &halves[1 - rep % 2];
                let ps: Vec<f64> = (0..15)
                    .filter(|&t| !r[t].0.is_nan() && !z_rows[s][t][0].is_nan() && !other[t][0].is_nan())
                    .map(|t| calibration::tail_p(&other[t], r[t].0, TWO_SIDED[t]))
                    .collect();
                acat(&ps)
            })
            .filter(|t| !t.is_nan())
            .collect();
        t_rows[s] = quantiles(&mut ts);
    }

    report(&raw_fail);
    for note in &notes {
        println!("{}", note);
    }
    write_table(&args.out, &source, reps, &z_rows, &t_rows)
}

/// Prints how often each raw p-value falls below 0.01 on random data (should be 1%).
fn report(raw_fail: &[[f64; 15]]) {
    println!("\nRaw p < 0.01 on random data, in % (1.0 if the raw p-value were right):");
    print!("{:<14}", "bytes");
    for s in 0..SIZES {
        print!("{:>7}", format!("2^{}", MIN_LOG2 as usize + s));
    }
    println!();
    for t in 0..15 {
        print!("{:<14}", TESTS[t].0);
        for row in raw_fail {
            if row[t].is_nan() { print!("{:>7}", "-") } else { print!("{:>7.1}", 100.0 * row[t]) }
        }
        println!();
    }
    println!();
}

fn write_table(path: &str, source: &str, reps: usize, z_rows: &[[[f64; LEVELS]; 15]], t_rows: &[[f64; LEVELS]]) -> io::Result<()> {
    let row = |q: &[f64; LEVELS]| -> String {
        if q[0].is_nan() {
            return format!("[f64::NAN; {}]", LEVELS);
        }
        let values: Vec<String> = q.iter().map(|v| format!("{:?}", v)).collect();
        format!("[{}]", values.join(", "))
    };

    let mut out = String::new();
    out.push_str("// Generated by `cargo run --release --example calibrate`. Do not edit.\n");
    out.push_str(&format!("// Null source: {}, {} samples per size, 2^{}..2^{} bytes.\n\n", source, reps, MIN_LOG2, MAX_LOG2));
    out.push_str("use super::{LEVELS, SIZES};\n\n");
    out.push_str(&format!("pub const REPS: usize = {};\n", reps));
    out.push_str(&format!("pub const SOURCE: &str = {:?};\n\n", source));
    out.push_str("/// Quantiles of each test's z-score on random data, per size and test.\n");
    out.push_str("#[rustfmt::skip]\npub static Z: [[[f64; LEVELS]; 15]; SIZES] = [\n");
    for (s, tests) in z_rows.iter().enumerate() {
        out.push_str(&format!("    // 2^{} bytes\n    [\n", MIN_LOG2 as usize + s));
        for (t, q) in tests.iter().enumerate() {
            out.push_str(&format!("        // {}\n        {},\n", TESTS[t].0, row(q)));
        }
        out.push_str("    ],\n");
    }
    out.push_str("];\n\n/// Quantiles of the Cauchy combination statistic on random data, per size.\n");
    out.push_str("#[rustfmt::skip]\npub static T: [[f64; LEVELS]; SIZES] = [\n");
    for (s, q) in t_rows.iter().enumerate() {
        out.push_str(&format!("    // 2^{} bytes\n    {},\n", MIN_LOG2 as usize + s, row(q)));
    }
    out.push_str("];\n");

    File::create(path)?.write_all(out.as_bytes())?;
    println!("Wrote {}", path);
    Ok(())
}

fn check(reps: usize, args: &Args, urandom: &Option<Mutex<File>>) {
    if calibration::REPS == 0 {
        eprintln!("No calibration table compiled in. Generate one first.");
        std::process::exit(1);
    }
    // Calibrated sizes plus some in between and beyond, to exercise interpolation and extrapolation.
    let default_sizes = vec![256, 1000, 1024, 4096, 10_000, 65_536, 100_000, 259_200, 1 << 20, 1_500_000, 4_000_000];
    let sizes = args.sizes.clone().unwrap_or(default_sizes);
    const ALPHAS: [f64; 4] = [0.05, 0.01, 1e-3, 1e-4];
    const BATCH: usize = 100_000;

    println!("Calibrated p below 0.05 / 0.01 / 0.001 / 0.0001 on fresh random data, in % of samples.");
    for (i, &n) in sizes.iter().enumerate() {
        // Counts per test (and combined, last) per alpha, plus the number of non-skipped runs.
        let mut below = [[0usize; 4]; 16];
        let mut total = [0usize; 16];
        for start in (0..reps).step_by(BATCH) {
            let batch = BATCH.min(reps - start);
            // Seeds offset from generation, so these samples were not used to build the table.
            let scores = parallel(batch, n, |rep, buf| {
                sample(buf, urandom, args.seed.wrapping_add(0x5eed), 100 + i, start + rep);
                let score = Onod::score(buf, false);
                let mut ps = [f64::NAN; 16];
                for (p, r) in ps.iter_mut().zip(&score.tests) {
                    *p = r.p_cal.unwrap_or(f64::NAN);
                }
                ps[15] = score.p_combined;
                ps
            });
            for ps in &scores {
                for k in 0..16 {
                    if ps[k].is_nan() {
                        continue;
                    }
                    total[k] += 1;
                    for (a, &alpha) in ALPHAS.iter().enumerate() {
                        below[k][a] += (ps[k] < alpha) as usize;
                    }
                }
            }
        }

        println!("\n{} bytes, {} samples (expected {:.2} / {:.2} / {:.3} / {:.4}):", n, reps, 5.0, 1.0, 0.1, 0.01);
        for k in 0..16 {
            let name = if k < 15 { TESTS[k].0 } else { "Combined" };
            if total[k] == 0 {
                println!("  {:<14} -", name);
                continue;
            }
            let rate = |a: usize| 100.0 * below[k][a] as f64 / total[k] as f64;
            println!("  {:<14} {:>6.2} {:>6.2} {:>7.3} {:>8.4}", name, rate(0), rate(1), rate(2), rate(3));
        }
    }
}

fn main() -> io::Result<()> {
    let args = parse_args();
    let urandom = if args.urandom { Some(Mutex::new(File::open("/dev/urandom")?)) } else { None };
    match args.check {
        Some(reps) => check(reps, &args, &urandom),
        None => generate(&args, &urandom)?,
    }
    Ok(())
}