ph-color-bake 0.1.0

//! Host golden-vector emission against independent host reference arithmetic.
//!
//! LUT knots and RGB→RGB coefficients come from the `f64` baker. Expected
//! apply outputs are computed here from the documented numeric rules and must
//! not call `ph-color` constructors. Replay tests in `ph-color` are the only
//! callers of the target path.

use std::fs::{self, File};
use std::io::{self, Write};
use std::path::Path;

use crate::Primaries;
use crate::lut::{bake_lut, identity, lerp_u16, lookup};
use crate::matrix::{COEF_ONE, rgb_to_rgb};

const COEF_FRAC_BITS: u32 = 28;

/// Write every freeze-time operation into `dir` as CSV.
pub fn write_all(dir: &Path) -> io::Result<()> {
    fs::create_dir_all(dir)?;
    write_quantize(dir)?;
    write_expand(dir)?;
    write_lerp(dir)?;
    write_gain(dir)?;
    write_matrix(dir)?;
    write_lut(dir)?;
    write_gradient(dir)?;
    Ok(())
}

fn create(dir: &Path, name: &str) -> io::Result<File> {
    File::create(dir.join(name))
}

fn sample_u16() -> Vec<u16> {
    let mut out = vec![
        0, 1, 2, 127, 128, 255, 256, 2047, 2048, 2049, 4096, 16384, 32767, 32768, 32769, 49152,
        65534, 65535,
    ];
    let mut x = 0u32;
    while x <= u32::from(u16::MAX) {
        out.push(x as u16);
        x = x.saturating_add(8192);
        if x == 0 {
            break;
        }
    }
    out.sort_unstable();
    out.dedup();
    out
}

fn sample_t() -> Vec<u16> {
    let mut out = sample_u16();
    out.extend([0, 1, 32767, 32768, 65534, 65535]);
    out.sort_unstable();
    out.dedup();
    out
}

fn sample_colors() -> Vec<[u16; 3]> {
    let mut out = Vec::new();
    for c in sample_u16() {
        out.push([c, c, c]);
        out.push([c, 0, 65535]);
    }
    for a in [0u16, 1, 65535] {
        for b in [0u16, 1, 65535] {
            for c in [0u16, 1, 65535] {
                out.push([a, b, c]);
            }
        }
    }
    out.sort_unstable();
    out.dedup();
    out
}

fn expand_bits(bits: u32, q: u16) -> u16 {
    debug_assert!((1..=16).contains(&bits));
    let shift = 16u32.saturating_sub(bits);
    let mask = match 1u32.checked_shl(bits) {
        Some(width) => width.saturating_sub(1) as u16,
        None => u16::MAX,
    };
    let code = q & mask;
    code.checked_shl(shift).unwrap_or(0)
}

fn quantize_bits(bits: u32, v: u16) -> (u16, i32) {
    debug_assert!((1..=16).contains(&bits));
    let shift = 16u32.saturating_sub(bits);
    let q = v.checked_shr(shift).unwrap_or(0);
    let expanded = expand_bits(bits, q);
    let residual = i32::from(v).saturating_sub(i32::from(expanded));
    (q, residual)
}

fn shift_round_sat_u16(acc: i64) -> u16 {
    let shifted = shift_round_away_from_zero(acc, COEF_FRAC_BITS);
    if shifted <= 0 {
        0
    } else if shifted >= i64::from(u16::MAX) {
        u16::MAX
    } else {
        shifted as u16
    }
}

fn shift_round_away_from_zero(value: i64, shift: u32) -> i64 {
    if shift == 0 {
        return value;
    }
    let Some(divisor) = 1i64.checked_shl(shift) else {
        return if value < 0 { i64::MIN } else { i64::MAX };
    };
    let Some(quot) = value.checked_div(divisor) else {
        return value;
    };
    let Some(rem) = value.checked_rem(divisor) else {
        return quot;
    };
    let Some(half) = divisor.checked_shr(1) else {
        return quot;
    };
    if rem.saturating_abs() >= half {
        if value >= 0 {
            quot.saturating_add(1)
        } else {
            quot.saturating_sub(1)
        }
    } else {
        quot
    }
}

fn apply_gain(g: i32, c: u16) -> u16 {
    shift_round_sat_u16(i64::from(g).saturating_mul(i64::from(c)))
}

fn apply_row(row: [i32; 3], ch: [u16; 3]) -> u16 {
    let acc = i64::from(row[0]).saturating_mul(i64::from(ch[0]));
    let acc = acc.saturating_add(i64::from(row[1]).saturating_mul(i64::from(ch[1])));
    let acc = acc.saturating_add(i64::from(row[2]).saturating_mul(i64::from(ch[2])));
    shift_round_sat_u16(acc)
}

fn lerp_color(a: [u16; 3], b: [u16; 3], t: u16) -> [u16; 3] {
    [
        lerp_u16(a[0], b[0], t),
        lerp_u16(a[1], b[1], t),
        lerp_u16(a[2], b[2], t),
    ]
}

fn sample_stops<const N: usize>(stops: &[[u16; 3]; N], t: u16) -> [u16; 3] {
    let segs = N.saturating_sub(1) as u64;
    let den = u64::from(u16::MAX);
    let pos = u64::from(t).saturating_mul(segs);
    let idx = pos.checked_div(den).unwrap_or(0) as usize;
    let last = N.saturating_sub(1);
    if idx >= last {
        return stops[last];
    }
    let rem = match pos.checked_rem(den) {
        Some(r) if r <= u64::from(u16::MAX) => r as u16,
        _ => 0,
    };
    lerp_color(stops[idx], stops[idx.saturating_add(1)], rem)
}

fn write_quantize(dir: &Path) -> io::Result<()> {
    let mut f = create(dir, "quantize.csv")?;
    writeln!(f, "bits,v,code,residual,expanded")?;
    let mut rows: Vec<(u32, u16)> = Vec::new();
    for bits in 1u32..=16 {
        for v in [0u16, 1, 65535] {
            rows.push((bits, v));
        }
        for v in sample_u16() {
            rows.push((bits, v));
        }
    }
    for v in 0..=u16::MAX {
        rows.push((8, v));
    }
    rows.sort_unstable();
    rows.dedup();
    for (bits, v) in rows {
        let (code, residual) = quantize_bits(bits, v);
        let expanded = expand_bits(bits, code);
        writeln!(f, "{bits},{v},{code},{residual},{expanded}")?;
    }
    Ok(())
}

fn write_expand(dir: &Path) -> io::Result<()> {
    let mut f = create(dir, "expand.csv")?;
    writeln!(f, "bits,q,expanded")?;
    for bits in [1u32, 4, 8] {
        let width = 1u32.checked_shl(bits).unwrap_or(0);
        for q in 0..width {
            let q = q as u16;
            writeln!(f, "{bits},{q},{}", expand_bits(bits, q))?;
        }
    }
    for q in [0u16, 1, 32768, 65535] {
        writeln!(f, "16,{q},{}", expand_bits(16, q))?;
    }
    Ok(())
}

fn write_lerp(dir: &Path) -> io::Result<()> {
    let mut f = create(dir, "lerp.csv")?;
    writeln!(f, "a0,a1,a2,b0,b1,b2,t,o0,o1,o2")?;
    let a = [0u16, 0, 0];
    let b = [65535u16, 32768, 1];
    for t in 0..=u16::MAX {
        write_lerp_row(&mut f, a, b, t)?;
    }
    for &a_ch in &[0u16, 1, 65535] {
        for &b_ch in &[0u16, 1, 65535] {
            let a = [a_ch, 0, 65535];
            let b = [b_ch, 65535, 0];
            for t in sample_t() {
                write_lerp_row(&mut f, a, b, t)?;
            }
        }
    }
    Ok(())
}

fn write_lerp_row(f: &mut File, a: [u16; 3], b: [u16; 3], t: u16) -> io::Result<()> {
    let o = lerp_color(a, b, t);
    writeln!(
        f,
        "{},{},{},{},{},{},{t},{},{},{}",
        a[0], a[1], a[2], b[0], b[1], b[2], o[0], o[1], o[2]
    )
}

fn write_gain(dir: &Path) -> io::Result<()> {
    let mut f = create(dir, "gain.csv")?;
    writeln!(f, "g0,g1,g2,c0,c1,c2,o0,o1,o2")?;
    let two = COEF_ONE.saturating_mul(2);
    let gains = [
        [COEF_ONE, COEF_ONE, COEF_ONE],
        [two, two, two],
        [0, 0, 0],
        [COEF_ONE.saturating_neg(), COEF_ONE, two],
        [i32::MAX, i32::MIN, 0],
        [COEF_ONE, 0, COEF_ONE],
    ];
    for g in gains {
        for c in sample_colors() {
            let out = [
                apply_gain(g[0], c[0]),
                apply_gain(g[1], c[1]),
                apply_gain(g[2], c[2]),
            ];
            writeln!(
                f,
                "{},{},{},{},{},{},{},{},{}",
                g[0], g[1], g[2], c[0], c[1], c[2], out[0], out[1], out[2]
            )?;
        }
    }
    Ok(())
}

fn write_matrix(dir: &Path) -> io::Result<()> {
    let mut f = create(dir, "matrix.csv")?;
    writeln!(f, "m00,m01,m02,m10,m11,m12,m20,m21,m22,c0,c1,c2,o0,o1,o2")?;
    let identity = rgb_to_rgb(Primaries::srgb(), Primaries::srgb())
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, format!("{e:?}")))?;
    let two = COEF_ONE.saturating_mul(2);
    let matrices = [
        identity,
        [[two, 0, 0], [0, two, 0], [0, 0, two]],
        [[COEF_ONE, COEF_ONE, 0], [0, COEF_ONE, 0], [0, 0, COEF_ONE]],
        [[i32::MAX, 0, 0], [0, i32::MIN, 0], [0, 0, COEF_ONE]],
        [
            [1 << 27, -(1 << 27), 0],
            [COEF_ONE, COEF_ONE.saturating_neg(), 0],
            [COEF_ONE, -(1 << 27), 0],
        ],
    ];
    let extra = [[2u16, 1, 0], [32768, 32767, 0], [65535, 32768, 1]];
    for m in matrices {
        for c in sample_colors().into_iter().chain(extra) {
            let out = [apply_row(m[0], c), apply_row(m[1], c), apply_row(m[2], c)];
            writeln!(
                f,
                "{},{},{},{},{},{},{},{},{},{},{},{},{},{},{}",
                m[0][0],
                m[0][1],
                m[0][2],
                m[1][0],
                m[1][1],
                m[1][2],
                m[2][0],
                m[2][1],
                m[2][2],
                c[0],
                c[1],
                c[2],
                out[0],
                out[1],
                out[2]
            )?;
        }
    }
    Ok(())
}

fn write_lut(dir: &Path) -> io::Result<()> {
    let identity17 = bake_lut::<17>(identity);
    let square17 = bake_lut::<17>(|x| x * x);
    write_lut_named(dir, "identity17", &identity17, 1)?;
    write_lut_named(dir, "square17", &square17, 16)?;
    Ok(())
}

fn write_lut_named<const N: usize>(
    dir: &Path,
    name: &str,
    baked: &crate::lut::BakedLut<N>,
    stride: u16,
) -> io::Result<()> {
    let mut knots = create(dir, &format!("lut_{name}_knots.csv"))?;
    writeln!(knots, "index,knot,max_err_lsb")?;
    for (i, knot) in baked.knots.iter().enumerate() {
        writeln!(knots, "{i},{knot},{}", baked.max_err_lsb)?;
    }

    let mut lookup_file = create(dir, &format!("lut_{name}_lookup.csv"))?;
    writeln!(lookup_file, "x,y")?;
    for x in 0..=u16::MAX {
        let keep_edge = matches!(x, 0 | 1 | 65534 | 65535);
        if stride > 1 && x % stride != 0 && !keep_edge {
            continue;
        }
        writeln!(lookup_file, "{x},{}", lookup(&baked.knots, x))?;
    }
    let last = N.saturating_sub(1) as u32;
    if stride > 1 && last > 0 {
        for i in 0..=last {
            let x = match i.saturating_mul(65535).checked_div(last) {
                Some(v) if v <= u32::from(u16::MAX) => v as u16,
                _ => 0,
            };
            writeln!(lookup_file, "{x},{}", lookup(&baked.knots, x))?;
            if i < last {
                let next = match (i.saturating_add(1))
                    .saturating_mul(65535)
                    .checked_div(last)
                {
                    Some(v) if v <= u32::from(u16::MAX) => v as u16,
                    _ => 0,
                };
                let mid = x.saturating_add(next.saturating_sub(x) / 2);
                writeln!(lookup_file, "{mid},{}", lookup(&baked.knots, mid))?;
            }
        }
    }
    Ok(())
}

fn write_gradient(dir: &Path) -> io::Result<()> {
    let mut two = create(dir, "gradient2.csv")?;
    writeln!(two, "a0,a1,a2,b0,b1,b2,t,o0,o1,o2")?;
    let g2 = [[0u16, 0, 0], [65535, 32768, 1]];
    for t in 0..=u16::MAX {
        let o = sample_stops(&g2, t);
        writeln!(two, "0,0,0,65535,32768,1,{t},{},{},{}", o[0], o[1], o[2])?;
    }

    let mut three = create(dir, "gradient3.csv")?;
    writeln!(three, "a0,a1,a2,b0,b1,b2,c0,c1,c2,t,o0,o1,o2")?;
    let g3 = [[0u16, 0, 0], [32768, 1000, 2000], [65535, 65535, 65535]];
    for t in sample_t() {
        let o = sample_stops(&g3, t);
        writeln!(
            three,
            "0,0,0,32768,1000,2000,65535,65535,65535,{t},{},{},{}",
            o[0], o[1], o[2]
        )?;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn writes_expected_headers_and_quantize_identity_row() {
        let dir = std::env::temp_dir().join(format!("ph-color-golden-{}", std::process::id()));
        write_all(&dir).expect("write goldens");
        let q = fs::read_to_string(dir.join("quantize.csv")).expect("quantize.csv");
        assert!(q.starts_with("bits,v,code,residual,expanded\n"));
        assert!(q.contains("\n8,0,0,0,0\n"));
        // Packed 8-bit code of 0x8080 is 0x80 with residual 0x80; expand is 0x8000.
        assert!(q.contains("\n8,32896,128,128,32768\n"));
        assert!(dir.join("lut_identity17_knots.csv").exists());
        assert!(dir.join("matrix.csv").exists());
        let _ = fs::remove_dir_all(&dir);
    }
}