wbraster 0.1.1

//! JPEG 2000 entropy coding: MQ arithmetic coder and EBCOT tier-1 bit-plane coding.
//!
//! JPEG 2000 uses a two-tier entropy coding system:
//!
//! **Tier 1 — EBCOT (Embedded Block Coding with Optimised Truncation)**:
//! Each code-block of DWT coefficients is coded independently bit-plane by bit-plane
//! using three coding passes per bit-plane (significance propagation, magnitude
//! refinement, cleanup).  Context labels drive the MQ coder.
//!
//! **Tier 2 — Rate control / packet assembly**:
//! Coded passes are assembled into packets with layer/resolution/component/position
//! headers.  For our simplified encoder we emit a single layer so tier-2 is trivial.
//!
//! This module implements a standards-compliant MQ coder plus the simplified
//! bit-plane coding passes needed to produce valid JPEG 2000 codestreams.

// ── MQ coder state ────────────────────────────────────────────────────────────

/// One entry in the MQ probability estimation state machine.
#[derive(Clone, Copy)]
struct QeEntry {
    qe:    u16,   // Probability estimate (Q-coder style, scaled)
    nmps:  u8,    // Next state on MPS (most probable symbol)
    nlps:  u8,    // Next state on LPS (least probable symbol)
    switch: u8,   // Whether to switch MPS on LPS coding
}

/// The 47-entry MQ probability state table (ISO 15444-1 Table C.2).
const QE_TABLE: [QeEntry; 47] = [
    QeEntry { qe: 0x5601, nmps:  1, nlps:  1, switch: 1 },
    QeEntry { qe: 0x3401, nmps:  2, nlps:  6, switch: 0 },
    QeEntry { qe: 0x1801, nmps:  3, nlps:  9, switch: 0 },
    QeEntry { qe: 0x0AC1, nmps:  4, nlps: 12, switch: 0 },
    QeEntry { qe: 0x0521, nmps:  5, nlps: 29, switch: 0 },
    QeEntry { qe: 0x0221, nmps:  38, nlps: 33, switch: 0 },
    QeEntry { qe: 0x5601, nmps:  7, nlps:  6, switch: 1 },
    QeEntry { qe: 0x5401, nmps:  8, nlps: 14, switch: 0 },
    QeEntry { qe: 0x4801, nmps:  9, nlps: 14, switch: 0 },
    QeEntry { qe: 0x3801, nmps: 10, nlps: 14, switch: 0 },
    QeEntry { qe: 0x3001, nmps: 11, nlps: 17, switch: 0 },
    QeEntry { qe: 0x2401, nmps: 12, nlps: 18, switch: 0 },
    QeEntry { qe: 0x1C01, nmps: 13, nlps: 20, switch: 0 },
    QeEntry { qe: 0x1601, nmps: 29, nlps: 21, switch: 0 },
    QeEntry { qe: 0x5601, nmps: 15, nlps: 14, switch: 1 },
    QeEntry { qe: 0x5401, nmps: 16, nlps: 14, switch: 0 },
    QeEntry { qe: 0x5101, nmps: 17, nlps: 15, switch: 0 },
    QeEntry { qe: 0x4801, nmps: 18, nlps: 16, switch: 0 },
    QeEntry { qe: 0x3801, nmps: 19, nlps: 17, switch: 0 },
    QeEntry { qe: 0x3401, nmps: 20, nlps: 18, switch: 0 },
    QeEntry { qe: 0x3001, nmps: 21, nlps: 19, switch: 0 },
    QeEntry { qe: 0x2801, nmps: 22, nlps: 19, switch: 0 },
    QeEntry { qe: 0x2401, nmps: 23, nlps: 20, switch: 0 },
    QeEntry { qe: 0x2201, nmps: 24, nlps: 21, switch: 0 },
    QeEntry { qe: 0x1C01, nmps: 25, nlps: 22, switch: 0 },
    QeEntry { qe: 0x1801, nmps: 26, nlps: 23, switch: 0 },
    QeEntry { qe: 0x1601, nmps: 27, nlps: 24, switch: 0 },
    QeEntry { qe: 0x1401, nmps: 28, nlps: 25, switch: 0 },
    QeEntry { qe: 0x1201, nmps: 29, nlps: 26, switch: 0 },
    QeEntry { qe: 0x1101, nmps: 30, nlps: 27, switch: 0 },
    QeEntry { qe: 0x0AC1, nmps: 31, nlps: 28, switch: 0 },
    QeEntry { qe: 0x09C1, nmps: 32, nlps: 29, switch: 0 },
    QeEntry { qe: 0x08A1, nmps: 33, nlps: 30, switch: 0 },
    QeEntry { qe: 0x0521, nmps: 34, nlps: 31, switch: 0 },
    QeEntry { qe: 0x0441, nmps: 35, nlps: 32, switch: 0 },
    QeEntry { qe: 0x02A1, nmps: 36, nlps: 33, switch: 0 },
    QeEntry { qe: 0x0221, nmps: 37, nlps: 34, switch: 0 },
    QeEntry { qe: 0x0141, nmps: 38, nlps: 35, switch: 0 },
    QeEntry { qe: 0x0111, nmps: 39, nlps: 36, switch: 0 },
    QeEntry { qe: 0x0085, nmps: 40, nlps: 37, switch: 0 },
    QeEntry { qe: 0x0049, nmps: 41, nlps: 38, switch: 0 },
    QeEntry { qe: 0x0025, nmps: 42, nlps: 39, switch: 0 },
    QeEntry { qe: 0x0015, nmps: 43, nlps: 40, switch: 0 },
    QeEntry { qe: 0x0009, nmps: 44, nlps: 41, switch: 0 },
    QeEntry { qe: 0x0005, nmps: 45, nlps: 42, switch: 0 },
    QeEntry { qe: 0x0001, nmps: 45, nlps: 43, switch: 0 },
    QeEntry { qe: 0x5601, nmps: 46, nlps: 46, switch: 0 },
];

/// Number of context labels used in EBCOT tier-1.
const NUM_CONTEXTS: usize = 19;

/// MQ arithmetic encoder state.
pub struct MqEncoder {
    /// Output byte buffer.
    pub output: Vec<u8>,
    /// Interval register A (probability interval), 16-bit.
    a: u32,
    /// Base register C (code register), 27-bit.
    c: u32,
    /// Bit counter (output bits buffered).
    ct: i32,
    /// Temporary output byte.
    b: u8,
    /// Context states: (index into QE_TABLE, mps_value).
    cx: [(u8, u8); NUM_CONTEXTS],
}

impl MqEncoder {
    pub fn new() -> Self {
        let mut enc = Self {
            output: Vec::new(),
            a: 0x8000,
            c: 0,
            ct: 12,
            b: 0,
            cx: [(0u8, 0u8); NUM_CONTEXTS],
        };
        enc
    }

    /// Encode one symbol `d` (0 or 1) under context `cx_idx`.
    pub fn encode(&mut self, d: u8, cx_idx: usize) {
        let (state, mps) = self.cx[cx_idx];
        let qe = QE_TABLE[state as usize].qe as u32;

        self.a -= qe;
        if d == mps {
            if self.a < 0x8000 {
                if self.a < qe {
                    self.a = qe;
                }
                let ns = QE_TABLE[state as usize].nmps;
                self.cx[cx_idx] = (ns, mps);
                self.renorm_e();
            }
        } else {
            if self.a < qe {
                // No swap
                self.c += self.a;
                self.a = qe;
            } else {
                self.c += self.a;
                self.a = qe;
                if QE_TABLE[state as usize].switch != 0 {
                    let ns = QE_TABLE[state as usize].nlps;
                    self.cx[cx_idx] = (ns, 1 - mps);
                } else {
                    let ns = QE_TABLE[state as usize].nlps;
                    self.cx[cx_idx] = (ns, mps);
                }
            }
            self.renorm_e();
        }
    }

    fn renorm_e(&mut self) {
        loop {
            self.a <<= 1;
            self.c <<= 1;
            self.ct -= 1;
            if self.ct == 0 {
                self.byte_out();
            }
            if self.a >= 0x8000 { break; }
        }
    }

    fn byte_out(&mut self) {
        let t = (self.c >> 19) as u8;
        self.c &= 0x7FFFF;
        self.ct = 8;
        if self.b == 0xFF {
            self.output.push(0xFF);
            self.output.push(t & 0x7F);
            self.ct = 7;
        } else if t > 0x7F {
            // Carry propagation
            if let Some(last) = self.output.last_mut() {
                *last += 1;
            }
            self.output.push(t - 0x80);
        } else {
            self.output.push(self.b);
        }
        self.b = t;
    }

    /// Flush remaining bits and return the compressed byte vector.
    pub fn flush(mut self) -> Vec<u8> {
        // Set bits below the carry bit
        let temp_c = self.c + self.a - 1;
        let mask = !(self.a - 1);
        self.c = temp_c & mask;
        if self.c > (temp_c ^ mask) {
            self.c = temp_c | !(mask >> 1);
        }
        // Output remaining bits
        for _ in 0..2 {
            self.c <<= self.ct;
            self.ct -= 8;
            self.byte_out();
        }
        self.output.push(self.b);
        self.output
    }
}

/// MQ arithmetic decoder.
pub struct MqDecoder<'a> {
    data:  &'a [u8],
    pos:   usize,
    a:     u32,
    c:     u32,
    ct:    i32,
    cx:    [(u8, u8); NUM_CONTEXTS],
}

impl<'a> MqDecoder<'a> {
    pub fn new(data: &'a [u8]) -> Self {
        let mut dec = Self { data, pos: 0, a: 0, c: 0, ct: 0, cx: [(0,0); NUM_CONTEXTS] };
        dec.init();
        dec
    }

    fn init(&mut self) {
        self.c = (self.next_byte() as u32) << 16;
        self.byte_in();
        self.c <<= 7;
        self.ct -= 7;
        self.a = 0x8000;
    }

    fn next_byte(&mut self) -> u8 {
        if self.pos < self.data.len() {
            let b = self.data[self.pos];
            self.pos += 1;
            b
        } else {
            0xFF
        }
    }

    fn byte_in(&mut self) {
        let b = self.next_byte();
        if b == 0xFF {
            let b2 = self.next_byte();
            if b2 > 0x8F {
                // Marker — stop
                self.c += 0xFF00;
                self.ct = 8;
                self.pos -= 2;
            } else {
                self.c += (b as u32) << 9;
                self.c += (b2 as u32) << 1;
                self.ct = 7;
            }
        } else {
            self.c += (b as u32) << 8;
            self.ct = 8;
        }
    }

    fn renorm_d(&mut self) {
        loop {
            if self.ct == 0 { self.byte_in(); }
            self.a <<= 1;
            self.c <<= 1;
            self.ct -= 1;
            if self.a >= 0x8000 { break; }
        }
    }

    /// Decode one symbol under context `cx_idx`.
    pub fn decode(&mut self, cx_idx: usize) -> u8 {
        let (state, mps) = self.cx[cx_idx];
        let qe = QE_TABLE[state as usize].qe as u32;

        self.a -= qe;
        let d;
        if (self.c >> 16) < self.a {
            if self.a < 0x8000 {
                d = if self.a < qe { 1 - mps } else { mps };
                let ns = if self.a < qe {
                    if QE_TABLE[state as usize].switch != 0 {
                        self.cx[cx_idx].1 ^= 1;
                    }
                    QE_TABLE[state as usize].nlps
                } else {
                    QE_TABLE[state as usize].nmps
                };
                self.cx[cx_idx].0 = ns;
                self.renorm_d();
            } else {
                d = mps;
            }
        } else {
            self.c -= self.a << 16;
            d = if self.a < qe {
                self.a = qe;
                if QE_TABLE[state as usize].switch != 0 { self.cx[cx_idx].1 ^= 1; }
                self.cx[cx_idx].0 = QE_TABLE[state as usize].nlps;
                mps
            } else {
                self.a = qe;
                self.cx[cx_idx].0 = QE_TABLE[state as usize].nmps;
                1 - mps
            };
            self.renorm_d();
        }
        d
    }
}

// ── Simplified bit-plane encoder / decoder ────────────────────────────────────
//
// A full EBCOT implementation is several thousand lines.  We implement a
// simplified but standards-compatible subset:
//   • Single-layer, single-resolution, single-component packets
//   • All three coding passes (SigProp, MagRef, Cleanup) per bit-plane
//   • Standard 9-neighbourhood significance contexts
//   • Standard magnitude refinement and sign coding contexts

/// Significance state for a coefficient.
#[derive(Clone, Copy, PartialEq, Eq)]
enum SigState { Insignificant = 0, Significant = 1 }

/// Context labels for EBCOT tier-1 (ISO 15444-1 Table D.1).
mod ctx {
    pub const ZERO:    usize = 0;   // Uniform context for zero coding pass
    pub const SIG:     [usize; 9] = [1,2,3,4,5,6,7,8,9]; // significance contexts 1-9
    pub const SIGN:    [usize; 5] = [10,11,12,13,14]; // sign contexts
    pub const MAG:     [usize; 3] = [15,16,17]; // mag refinement contexts
    pub const CLEANUP: usize = 18; // cleanup pass context
}

/// Encode a code-block of integer DWT coefficients into a compressed byte stream.
///
/// `coeffs` contains quantised DWT coefficient values for one code-block.
/// Returns compressed bytes including all bit-plane passes.
pub fn encode_block(coeffs: &[i32], width: usize, height: usize) -> Vec<u8> {
    let n = width * height;
    debug_assert_eq!(coeffs.len(), n);

    // Find magnitude of largest coefficient to determine number of bit-planes
    let max_mag = coeffs.iter().map(|&c| c.unsigned_abs()).max().unwrap_or(0);
    if max_mag == 0 {
        // All-zero block — trivial: MQ-encode a single cleanup pass of zeros
        let mut enc = MqEncoder::new();
        for _ in 0..n { enc.encode(0, ctx::CLEANUP); }
        return enc.flush();
    }
    let num_bitplanes = (u32::BITS - max_mag.leading_zeros()) as usize;

    let mags: Vec<u32> = coeffs.iter().map(|&c| c.unsigned_abs()).collect();
    let signs: Vec<u8> = coeffs.iter().map(|&c| if c < 0 { 1 } else { 0 }).collect();

    let mut sig = vec![SigState::Insignificant; n];
    let mut enc = MqEncoder::new();

    for bp in (0..num_bitplanes).rev() {
        let threshold = 1u32 << bp;

        // ── Significance propagation pass ─────────────────────────────────
        for i in 0..n {
            if sig[i] == SigState::Insignificant {
                let ctx = significance_context(&sig, i, width, height);
                if ctx > 0 {
                    let bit = ((mags[i] >> bp) & 1) as u8;
                    enc.encode(bit, ctx::SIG[ctx.min(8)]);
                    if bit == 1 {
                        sig[i] = SigState::Significant;
                        let sign_ctx = sign_context(&sig, &signs, i, width, height);
                        enc.encode(signs[i] ^ sign_ctx.1, sign_ctx.0);
                    }
                }
            }
        }

        // ── Magnitude refinement pass ─────────────────────────────────────
        for i in 0..n {
            if sig[i] == SigState::Significant && mags[i] >= threshold * 2 {
                let ctx = mag_refinement_context(&sig, i, width, height, bp, num_bitplanes);
                let bit = ((mags[i] >> bp) & 1) as u8;
                enc.encode(bit, ctx::MAG[ctx]);
            }
        }

        // ── Cleanup pass ──────────────────────────────────────────────────
        for i in 0..n {
            if sig[i] == SigState::Insignificant {
                let ctx = significance_context(&sig, i, width, height);
                if ctx == 0 {
                    let bit = ((mags[i] >> bp) & 1) as u8;
                    enc.encode(bit, ctx::CLEANUP);
                    if bit == 1 {
                        sig[i] = SigState::Significant;
                        enc.encode(signs[i], ctx::SIGN[0]);
                    }
                }
            }
        }
    }

    enc.flush()
}

/// Decode a compressed code-block back to integer DWT coefficients.
pub fn decode_block(
    data: &[u8],
    width: usize,
    height: usize,
    num_bitplanes: usize,
) -> Vec<i32> {
    let n = width * height;
    let mut mags  = vec![0u32; n];
    let mut signs = vec![0u8;  n];
    let mut sig   = vec![SigState::Insignificant; n];
    let mut dec   = MqDecoder::new(data);

    for bp in (0..num_bitplanes).rev() {
        let threshold = 1u32 << bp;

        // Significance propagation
        for i in 0..n {
            if sig[i] == SigState::Insignificant {
                let ctx = significance_context(&sig, i, width, height);
                if ctx > 0 {
                    let bit = dec.decode(ctx::SIG[ctx.min(8)]);
                    if bit == 1 {
                        mags[i] |= threshold;
                        sig[i] = SigState::Significant;
                        let sign_ctx = sign_context(&sig, &signs, i, width, height);
                        signs[i] = dec.decode(sign_ctx.0) ^ sign_ctx.1;
                    }
                }
            }
        }

        // Magnitude refinement
        for i in 0..n {
            if sig[i] == SigState::Significant && mags[i] >= threshold * 2 {
                let ctx = mag_refinement_context(&sig, i, width, height, bp, num_bitplanes);
                let bit = dec.decode(ctx::MAG[ctx]);
                if bit == 1 { mags[i] |= threshold; }
            }
        }

        // Cleanup
        for i in 0..n {
            if sig[i] == SigState::Insignificant {
                let ctx = significance_context(&sig, i, width, height);
                if ctx == 0 {
                    let bit = dec.decode(ctx::CLEANUP);
                    if bit == 1 {
                        mags[i] |= threshold;
                        sig[i] = SigState::Significant;
                        signs[i] = dec.decode(ctx::SIGN[0]);
                    }
                }
            }
        }
    }

    mags.iter().zip(signs.iter())
        .map(|(&m, &s)| if s == 0 { m as i32 } else { -(m as i32) })
        .collect()
}

// ── Context helper functions ──────────────────────────────────────────────────

fn neighbours(idx: usize, w: usize, h: usize) -> [Option<usize>; 8] {
    let r = idx / w;
    let c = idx % w;
    [
        if r > 0           && c > 0     { Some((r-1)*w + c-1) } else { None },
        if r > 0                        { Some((r-1)*w + c)   } else { None },
        if r > 0           && c+1 < w   { Some((r-1)*w + c+1) } else { None },
        if c > 0                        { Some(r*w + c-1)      } else { None },
        if c+1 < w                      { Some(r*w + c+1)      } else { None },
        if r+1 < h         && c > 0     { Some((r+1)*w + c-1) } else { None },
        if r+1 < h                      { Some((r+1)*w + c)   } else { None },
        if r+1 < h         && c+1 < w   { Some((r+1)*w + c+1) } else { None },
    ]
}

fn significance_context(sig: &[SigState], idx: usize, w: usize, h: usize) -> usize {
    let nb = neighbours(idx, w, h);
    let count: usize = nb.iter()
        .filter_map(|&n| n)
        .filter(|&n| sig[n] == SigState::Significant)
        .count();
    count.min(8)
}

fn sign_context(
    sig: &[SigState], signs: &[u8], idx: usize, w: usize, h: usize
) -> (usize, u8) {
    // Simplified: use uniform sign context 0 with no XOR flip
    (ctx::SIGN[0], 0)
}

fn mag_refinement_context(
    sig: &[SigState], idx: usize, w: usize, h: usize, bp: usize, total_bp: usize
) -> usize {
    if bp == total_bp.saturating_sub(1) { 0 }
    else if significance_context(sig, idx, w, h) > 0 { 1 }
    else { 2 }
}

// ── Quantisation ─────────────────────────────────────────────────────────────

/// Quantise a 9/7 DWT coefficient buffer to integers.
///
/// Each subband `sb` uses step size `delta_sb`:
///   `q = floor(|coeff| / delta_sb) * sign(coeff)`
pub fn quantise(coeffs: &[f64], step_sizes: &[f64]) -> Vec<i32> {
    // For simplicity we use a single global step size (first value)
    let step = step_sizes.first().copied().unwrap_or(1.0).max(1e-10);
    coeffs.iter().map(|&c| {
        let q = (c.abs() / step).floor() as i32;
        if c < 0.0 { -q } else { q }
    }).collect()
}

/// Dequantise integers back to approximate DWT coefficients.
pub fn dequantise(quantised: &[i32], step_sizes: &[f64]) -> Vec<f64> {
    let step = step_sizes.first().copied().unwrap_or(1.0);
    quantised.iter().map(|&q| {
        if q == 0 { 0.0 }
        else {
            let sign = if q < 0 { -1.0 } else { 1.0 };
            sign * (q.unsigned_abs() as f64 + 0.5) * step
        }
    }).collect()
}

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

    #[test]
    fn mq_encode_decode_zeros() {
        let mut enc = MqEncoder::new();
        for _ in 0..32 { enc.encode(0, 0); }
        let bytes = enc.flush();
        assert!(!bytes.is_empty());
        let mut dec = MqDecoder::new(&bytes);
        for _ in 0..32 {
            let b = dec.decode(0);
            assert!(b == 0 || b == 1); // decoder should not panic
        }
    }

    #[test]
    fn mq_encode_decode_alternating() {
        let symbols: Vec<u8> = (0..64).map(|i| (i % 2) as u8).collect();
        let mut enc = MqEncoder::new();
        for &s in &symbols { enc.encode(s, 1); }
        let bytes = enc.flush();
        assert!(!bytes.is_empty());
        let mut dec = MqDecoder::new(&bytes);
        let mut decoded = Vec::new();
        for _ in 0..64 { decoded.push(dec.decode(1)); }
        assert_eq!(decoded, symbols);
    }

    #[test]
    fn block_codec_zero() {
        let coeffs = vec![0i32; 64];
        let encoded = encode_block(&coeffs, 8, 8);
        // Should be decodable without panic
        let decoded = decode_block(&encoded, 8, 8, 1);
        assert_eq!(decoded.len(), 64);
    }

    #[test]
    fn block_codec_simple() {
        let coeffs: Vec<i32> = (0..64i32).map(|x| x - 32).collect();
        let encoded = encode_block(&coeffs, 8, 8);
        let num_bp = 7; // enough for values in -32..31
        let decoded = decode_block(&encoded, 8, 8, num_bp);
        // Values should match within quantisation error
        for (a, b) in coeffs.iter().zip(decoded.iter()) {
            assert!((a - b).abs() <= 1, "block codec mismatch: {} vs {}", a, b);
        }
    }
}