use std::fs;
use std::path::PathBuf;

const K_MAX: usize = 32;
const WS_MAX: usize = 6;

fn build_n_log_n() -> [f64; K_MAX + 1] {
    let mut t = [0.0f64; K_MAX + 1];
    for n in 1..=K_MAX {
        t[n] = (n as f64) * (n as f64).ln();
    }
    t
}

/// Max achievable entropy over `4^ws` raw sub-words given only `nwords`
/// observations (most-uniform integer partition), per
/// `docmd/theory/entropy.md`.
fn build_emax() -> [[f64; WS_MAX + 1]; K_MAX + 1] {
    let mut t = [[0.0f64; WS_MAX + 1]; K_MAX + 1];
    for k in 2..=K_MAX {
        for ws in 1..=WS_MAX.min(k - 1) {
            let n_raw = 1usize << (ws * 2);
            let nwords = k - ws + 1;
            let c = nwords / n_raw;
            let r = nwords % n_raw;
            let nf = nwords as f64;
            let t1 = if c == 0 || n_raw == r {
                0.0
            } else {
                let f1 = c as f64 / nf;
                (n_raw - r) as f64 * f1 * f1.ln()
            };
            let t2 = if r == 0 {
                0.0
            } else {
                let f2 = (c + 1) as f64 / nf;
                r as f64 * f2 * f2.ln()
            };
            t[k][ws] = -(t1 + t2);
        }
    }
    t
}

fn build_log_nwords() -> [[f64; WS_MAX + 1]; K_MAX + 1] {
    let mut t = [[0.0f64; WS_MAX + 1]; K_MAX + 1];
    for k in 2..=K_MAX {
        for ws in 1..=WS_MAX.min(k - 1) {
            t[k][ws] = ((k - ws + 1) as f64).ln();
        }
    }
    t
}

fn emit_f64_1d(out: &mut String, name: &str, n: usize, values: &[f64]) {
    out.push_str(&format!("pub(crate) const {name}: [f64; {n}] = [\n"));
    for v in values {
        out.push_str(&format!("    {v:?},\n"));
    }
    out.push_str("];\n");
}

fn emit_f64_2d(out: &mut String, name: &str, rows: usize, cols: usize, values: &[[f64; WS_MAX + 1]]) {
    out.push_str(&format!("pub(crate) const {name}: [[f64; {cols}]; {rows}] = [\n"));
    for row in values {
        out.push_str("    [");
        for (i, v) in row.iter().enumerate() {
            if i > 0 { out.push_str(", "); }
            out.push_str(&format!("{v:?}"));
        }
        out.push_str("],\n");
    }
    out.push_str("];\n");
}

fn main() {
    let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap());
    let mut out = String::new();

    let n_log_n = build_n_log_n();
    emit_f64_1d(&mut out, "N_LOG_N", K_MAX + 1, &n_log_n);

    let emax = build_emax();
    emit_f64_2d(&mut out, "EMAX", K_MAX + 1, WS_MAX + 1, &emax);

    let log_nwords = build_log_nwords();
    emit_f64_2d(&mut out, "LOG_NWORDS", K_MAX + 1, WS_MAX + 1, &log_nwords);

    fs::write(out_dir.join("entropy_tables.rs"), out).unwrap();
}