gwseq-io 0.2.1

//! Does the section-encoding heuristic make smaller files than the obvious
//! rules?
//!
//! The heuristic in `bbi/section.rs` decides, every time a value breaks the
//! shape of the section it is being added to, whether to widen the encoding or
//! close the section and open a narrower one. It does that with a static cost
//! model: bytes of extra item width against a flat 64-byte split cost, with the
//! widening discounted to 25% for how well a near-constant coordinate column
//! deflates.
//!
//! Those constants were assumed rather than derived, and nothing about the code
//! can say whether they *pick well* — only that they pick consistently. This
//! writes the same data through every policy and reports what came out, which
//! is the only thing that can rank them: every policy produces a valid file, so
//! no correctness test has an opinion.
//!
//!     cargo run --release -p gwseq-io --example section_policy
//!     cargo run --release -p gwseq-io --example section_policy -- --sweep
//!     cargo run --release -p gwseq-io --example section_policy -- --bedgraph FILE
//!
//! The metric is the whole file, not the data blocks: a policy that splits more
//! writes a bigger R-tree, and paying for the index is part of the choice.

use std::collections::BTreeMap;

use gwseq_io::bbi::{
    BbiKind, BbiWriter, BbiWriterOptions, CostModel, SectionCounts, SectionPolicy,
};
use gwseq_io::genomic::ChrMap;

const CHR: &str = "chr1";

/// One value as the writer takes it: where it starts, how wide it is, what it
/// holds.
type Item = (i64, i64, f32);

/// A named body of data, and how a caller hands it over.
///
/// `runs` matters as much as the values: the heuristic's second term costs a
/// widening against what the *current call* still has to give it, so the same
/// values offered one at a time and offered in runs make different files. Each
/// case says which it is.
struct Case {
    name: &'static str,
    what: &'static str,
    items: Vec<Item>,
    /// Item counts per `write_value`/`write_values` call. `None` means one call
    /// per item, which is the pessimal case for the heuristic.
    runs: Option<Vec<usize>>,
}

fn rng(seed: u64) -> impl FnMut() -> u64 {
    let mut state = seed | 1;
    move || {
        state ^= state << 13;
        state ^= state >> 7;
        state ^= state << 17;
        state
    }
}

/// A smooth signal with structure, so deflate has something to find and the
/// value column is not incompressible noise.
fn signal(i: usize) -> f32 {
    ((i as f32) * 0.017).sin() * 40.0 + (i % 31) as f32 * 0.5
}

fn cases() -> Vec<Case> {
    let mut out = Vec::new();

    // The easy case: nothing ever breaks, so every policy must agree.
    out.push(Case {
        name: "uniform",
        what: "20k contiguous 10 bp values — fixedStep throughout",
        items: (0..20_000)
            .map(|i| (i as i64 * 10, i as i64 * 10 + 10, signal(i)))
            .collect(),
        runs: Some(vec![20_000]),
    });

    // Uniform with occasional gaps: the starts turn irregular, the spans do
    // not, so varStep is what is needed and only now and then.
    let mut next = rng(11);
    let mut items = Vec::new();
    let mut at = 0i64;
    for i in 0..20_000 {
        if next() % 400 == 0 {
            at += 10 * (1 + (next() % 5) as i64);
        }
        items.push((at, at + 10, signal(i)));
        at += 10;
    }
    out.push(Case {
        name: "gappy",
        what: "20k 10 bp values with a gap every ~400 — varStep now and then",
        items,
        runs: None,
    });

    // Spans that change occasionally: bedGraph is what is needed, rarely.
    let mut next = rng(12);
    let mut items = Vec::new();
    let mut at = 0i64;
    for i in 0..20_000 {
        let span = if next() % 500 == 0 { 25 } else { 10 };
        items.push((at, at + span, signal(i)));
        at += span;
    }
    out.push(Case {
        name: "rare-wide",
        what: "20k values, one in ~500 a different width — bedGraph, rarely",
        items,
        runs: None,
    });

    // Alternating regimes: long uniform runs interrupted by short irregular
    // bursts. What the heuristic's documentation claims to handle well — "a
    // one-off irregularity is absorbed where the start of a differently shaped
    // run splits".
    let mut next = rng(13);
    let mut items = Vec::new();
    let mut runs = Vec::new();
    let mut at = 0i64;
    let mut i = 0usize;
    while items.len() < 20_000 {
        let uniform = 200 + (next() % 800) as usize;
        for _ in 0..uniform {
            items.push((at, at + 10, signal(i)));
            at += 10;
            i += 1;
        }
        runs.push(uniform);
        let burst = 3 + (next() % 20) as usize;
        for _ in 0..burst {
            let span = 3 + (next() % 40) as i64;
            items.push((at, at + span, signal(i)));
            at += span + (next() % 7) as i64;
            i += 1;
        }
        runs.push(burst);
    }
    out.push(Case {
        name: "regimes",
        what: "uniform runs of 200-1000 broken by bursts of 3-22 irregular values",
        items,
        runs: Some(runs),
    });

    // Irregular throughout: bedGraph from the second item, so the only question
    // is whether the first few values cost a split.
    let mut next = rng(14);
    let mut items = Vec::new();
    let mut at = 0i64;
    for i in 0..20_000 {
        let span = 5 + (next() % 200) as i64;
        items.push((at, at + span, signal(i)));
        at += span + (next() % 50) as i64;
    }
    out.push(Case {
        name: "irregular",
        what: "20k intervals of random width and spacing — bedGraph throughout",
        items,
        runs: None,
    });

    // The adversarial one: uniform, but every other value breaks the shape.
    // Whichever way the heuristic decides it decides constantly.
    let mut items = Vec::new();
    let mut at = 0i64;
    for i in 0..20_000 {
        let span = if i % 2 == 0 { 10 } else { 11 };
        items.push((at, at + span, signal(i)));
        at += span;
    }
    out.push(Case {
        name: "alternating",
        what: "20k values alternating 10 bp and 11 bp — a break every other item",
        items,
        runs: None,
    });

    out
}

/// Write one case with one policy and report the file size.
fn write(case: &Case, policy: SectionPolicy, cost: CostModel, level: u32) -> (u64, SectionCounts) {
    let dir = std::env::temp_dir().join("gwseq_section_policy");
    std::fs::create_dir_all(&dir).expect("a writable temp directory");
    let path = dir.join(format!("{}-{policy:?}.bigwig", case.name));
    let path = path.to_str().expect("utf-8 path");

    let end = case.items.last().map(|(_, e, _)| *e).unwrap_or(1);
    let mut w = BbiWriter::create(
        path,
        BbiWriterOptions {
            kind: BbiKind::BigWig,
            chr_sizes: Some(ChrMap::from_entries([(CHR.to_string(), end + 1000)])),
            parallel: 1,
            compression_level: level,
            section_policy: policy,
            cost_model: cost,
            ..Default::default()
        },
    )
    .expect("a writable file");

    match &case.runs {
        // Handed over in runs, which is what tells the heuristic how much is
        // still coming. A contiguous run of one span goes in as one extend.
        Some(runs) => {
            let mut at = 0usize;
            for len in runs {
                let chunk = &case.items[at..at + len];
                let contiguous = chunk.windows(2).all(|p| p[0].1 == p[1].0)
                    && chunk
                        .iter()
                        .all(|(s, e, _)| e - s == chunk[0].1 - chunk[0].0);
                if contiguous && chunk.len() > 1 {
                    let values: Vec<f32> = chunk.iter().map(|(_, _, v)| *v).collect();
                    w.write_values(CHR, chunk[0].0, chunk[0].1 - chunk[0].0, &values)
                        .expect("valid values");
                } else {
                    for (s, e, v) in chunk {
                        w.write_value(CHR, *s, *e, *v).expect("valid values");
                    }
                }
                at += len;
            }
        }
        None => {
            for (s, e, v) in &case.items {
                w.write_value(CHR, *s, *e, *v).expect("valid values");
            }
        }
    }
    w.close().expect("a finished file");
    let counts = w.section_counts();
    let size = std::fs::metadata(path)
        .expect("the file just written")
        .len();
    let _ = std::fs::remove_file(path);
    (size, counts)
}

fn compare(level: u32) {
    let policies = [
        SectionPolicy::Cost,
        SectionPolicy::Room,
        SectionPolicy::Runt,
        SectionPolicy::Adaptive,
        SectionPolicy::Split,
        SectionPolicy::Widen,
    ];
    println!("compression level {level}, items_per_slot 1024 (the default)\n");
    println!(
        "{:<12} {:>10} {:>9} {:>9} {:>9} {:>9} {:>9}",
        "case", "Cost", "Room", "Runt", "Adapt", "Split", "Widen"
    );
    println!("{}", "-".repeat(80));

    let mut totals = [0u64; 6];
    for case in cases() {
        let mut sizes = Vec::new();
        let mut counts = None;
        for policy in policies {
            let (size, c) = write(&case, policy, CostModel::default(), level);
            if policy == SectionPolicy::Cost {
                counts = Some(c);
            }
            sizes.push(size);
        }
        let best = *sizes.iter().min().expect("four policies");
        for (i, s) in sizes.iter().enumerate() {
            totals[i] += s;
        }
        let cell = |s: u64| {
            if s == best {
                format!("{s}*")
            } else {
                format!("{:.2}%", (s as f64 / best as f64 - 1.0) * 100.0)
            }
        };
        let _ = counts.expect("Cost ran");
        println!(
            "{:<12} {:>10} {:>9} {:>9} {:>9} {:>9} {:>9}",
            case.name,
            cell(sizes[0]),
            cell(sizes[1]),
            cell(sizes[2]),
            cell(sizes[3]),
            cell(sizes[4]),
            cell(sizes[5]),
        );
    }
    let best = *totals.iter().min().expect("some policies");
    println!("{}", "-".repeat(80));
    println!(
        "{:<12} {:>10} {:>9} {:>9} {:>9} {:>9} {:>9}",
        "total",
        format!("{:+.2}%", (totals[0] as f64 / best as f64 - 1.0) * 100.0),
        format!("{:+.2}%", (totals[1] as f64 / best as f64 - 1.0) * 100.0),
        format!("{:+.2}%", (totals[2] as f64 / best as f64 - 1.0) * 100.0),
        format!("{:+.2}%", (totals[3] as f64 / best as f64 - 1.0) * 100.0),
        format!("{:+.2}%", (totals[4] as f64 / best as f64 - 1.0) * 100.0),
        format!("{:+.2}%", (totals[5] as f64 / best as f64 - 1.0) * 100.0),
    );
    println!("\n(a `*` marks the smallest file for that case; the rest are how much bigger)");
    for case in cases() {
        println!("  {:<12} {}", case.name, case.what);
    }
}

/// Sweep the two constants the cost model is built from.
///
/// The question is not whether 64 and 25 are defensible — they are — but
/// whether the file size is sensitive to them at all. If a tenfold change in
/// either moves the total by a fraction of a percent, the model is not doing
/// much work and the honest thing is to say so.
fn sweep(level: u32) {
    let cases = cases();
    let mut by_split: BTreeMap<i64, u64> = BTreeMap::new();
    for split_cost in [0i64, 16, 32, 64, 128, 256, 1024, 8192] {
        let mut total = 0;
        for case in &cases {
            let cost = CostModel {
                split_cost,
                ..Default::default()
            };
            total += write(case, SectionPolicy::Cost, cost, level).0;
        }
        by_split.insert(split_cost, total);
    }
    let base = by_split[&64];
    println!("split_cost (compression_percent held at 25), total bytes over all six cases");
    for (k, v) in &by_split {
        println!(
            "  {k:>6} {v:>10}  {:+.3}%{}",
            (*v as f64 / base as f64 - 1.0) * 100.0,
            if *k == 64 { "   <- shipped" } else { "" }
        );
    }

    let mut by_floor: BTreeMap<i64, u64> = BTreeMap::new();
    for min_split_items in [1i64, 4, 8, 16, 32, 64, 128, 512] {
        let mut total = 0;
        for case in &cases {
            let cost = CostModel {
                min_split_items,
                ..Default::default()
            };
            total += write(case, SectionPolicy::Adaptive, cost, level).0;
        }
        by_floor.insert(min_split_items, total);
    }
    let base = by_floor[&32];
    println!("\nmin_split_items under Adaptive, total bytes over all six cases");
    for (k, v) in &by_floor {
        println!(
            "  {k:>6} {v:>10}  {:+.3}%{}",
            (*v as f64 / base as f64 - 1.0) * 100.0,
            if *k == 32 { "   <- shipped" } else { "" }
        );
    }

    let mut by_patience: BTreeMap<u32, u64> = BTreeMap::new();
    for runt_patience in [1u32, 2, 3, 4, 8, 16] {
        let mut total = 0;
        for case in &cases {
            let cost = CostModel {
                runt_patience,
                ..Default::default()
            };
            total += write(case, SectionPolicy::Adaptive, cost, level).0;
        }
        by_patience.insert(runt_patience, total);
    }
    let base = by_patience[&2];
    println!("\nrunt_patience under Adaptive, total bytes over all six cases");
    for (k, v) in &by_patience {
        println!(
            "  {k:>6} {v:>10}  {:+.3}%{}",
            (*v as f64 / base as f64 - 1.0) * 100.0,
            if *k == 2 { "   <- shipped" } else { "" }
        );
    }

    let mut by_pct: BTreeMap<i64, u64> = BTreeMap::new();
    for percent in [1i64, 5, 10, 25, 50, 75, 100, 200] {
        let mut total = 0;
        for case in &cases {
            let cost = CostModel {
                compression_percent: percent,
                ..Default::default()
            };
            total += write(case, SectionPolicy::Cost, cost, level).0;
        }
        by_pct.insert(percent, total);
    }
    let base = by_pct[&25];
    println!("\ncompression_percent (split_cost held at 64), total bytes over all six cases");
    for (k, v) in &by_pct {
        println!(
            "  {k:>6} {v:>10}  {:+.3}%{}",
            (*v as f64 / base as f64 - 1.0) * 100.0,
            if *k == 25 { "   <- shipped" } else { "" }
        );
    }
}

/// The strongest input there is: a real track, re-written.
///
/// Takes a bedGraph — `gwseq export FILE out.bedgraph` makes one from any
/// bigWig — converts it with every policy, and reports the sizes. Synthetic
/// cases are chosen by whoever wrote them; this one is not.
fn real(path: &str, level: u32) {
    use gwseq_io::bbi::convert_to_bigwig;

    let dir = std::env::temp_dir().join("gwseq_section_policy");
    std::fs::create_dir_all(&dir).expect("a writable temp directory");
    println!("{path}\n");
    println!("{:<10} {:>12} {:>10}", "policy", "bytes", "vs best");

    let mut results = Vec::new();
    for policy in [
        SectionPolicy::Cost,
        SectionPolicy::Room,
        SectionPolicy::Runt,
        SectionPolicy::Adaptive,
        SectionPolicy::Split,
        SectionPolicy::Widen,
    ] {
        let out = dir.join(format!("real-{policy:?}.bigwig"));
        let out = out.to_str().expect("utf-8 path");
        convert_to_bigwig(
            std::path::Path::new(path),
            std::path::Path::new(out),
            None,
            BbiWriterOptions {
                kind: BbiKind::BigWig,
                parallel: 1,
                compression_level: level,
                section_policy: policy,
                cost_model: CostModel::default(),
                ..Default::default()
            },
            None,
            None,
        )
        .expect("a convertible bedGraph");
        let size = std::fs::metadata(out).expect("the file just written").len();
        let _ = std::fs::remove_file(out);
        results.push((policy, size));
    }
    let best = results.iter().map(|(_, s)| *s).min().expect("policies");
    for (policy, size) in &results {
        println!(
            "{:<10} {:>12} {:>10}",
            format!("{policy:?}"),
            size,
            if *size == best {
                "best".to_string()
            } else {
                format!("{:+.2}%", (*size as f64 / best as f64 - 1.0) * 100.0)
            },
        );
    }
}

fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();
    let level: u32 = args
        .iter()
        .position(|a| a == "--level")
        .and_then(|i| args.get(i + 1))
        .and_then(|v| v.parse().ok())
        .unwrap_or(6);
    if let Some(i) = args.iter().position(|a| a == "--bedgraph") {
        real(args.get(i + 1).expect("--bedgraph needs a path"), level);
    } else if let Some(i) = args.iter().position(|a| a == "--trace") {
        let case = args.get(i + 1).map(String::as_str).unwrap_or("regimes");
        let policy = match args.get(i + 2).map(String::as_str).unwrap_or("Adaptive") {
            "Cost" => SectionPolicy::Cost,
            "Split" => SectionPolicy::Split,
            "Widen" => SectionPolicy::Widen,
            "Room" => SectionPolicy::Room,
            "Runt" => SectionPolicy::Runt,
            _ => SectionPolicy::Adaptive,
        };
        trace_one(case, policy, level);
    } else if args.iter().any(|a| a == "--counts") {
        debug_counts(level);
    } else if args.iter().any(|a| a == "--sweep") {
        sweep(level);
    } else {
        compare(level);
    }
}

/// Write one case with one policy and dump every section it produced.
/// `GWSEQ_TRACE_RUN=1` in the environment adds the writer's own trace.
fn trace_one(case_name: &str, policy: SectionPolicy, level: u32) {
    let case = cases()
        .into_iter()
        .find(|c| c.name == case_name)
        .expect("no such case");
    let (size, counts) = write(&case, policy, CostModel::default(), level);
    println!(
        "{case_name} under {policy:?}: {size} bytes, bg={} vs={} fs={}",
        counts.bedgraph, counts.varstep, counts.fixedstep
    );
}

#[allow(dead_code)]
fn debug_counts(level: u32) {
    for case in cases() {
        print!("{:<12}", case.name);
        for policy in [
            SectionPolicy::Cost,
            SectionPolicy::Adaptive,
            SectionPolicy::Split,
        ] {
            let (size, c) = write(&case, policy, CostModel::default(), level);
            print!(
                "  {policy:?} {size} bg={} vs={} fs={}",
                c.bedgraph, c.varstep, c.fixedstep
            );
        }
        println!();
    }
}