//! The context-adaptive binary arithmetic coding of an NNC integer tensor
//! (`extensions/prior-carriage.md` §12.4): the arithmetic decoder and encoder, the
//! two-rate probability estimate of each context, the context selection, and the
//! binarisation of an integer (significance, sign, greater-than flags, an exponential
//! Golomb remainder), for the case §12 reads: dependent quantisation off (`dq_flag` 0) and
//! the row-wise scan (`scan_order` 0).
//!
//! Everything is integer arithmetic on fixed widths, so a decoder that follows it gives the
//! same integers on every machine. The tables are the coding's own; §12.4 says where they
//! were read from.

use super::{refuse, Refusal};

/// The range of the least probable symbol, by `(|state| >> 7) + (range & 0xe0)`.
#[rustfmt::skip]
pub const RLPS: [u8; 256] = [
    128, 112, 97, 84, 74, 65, 57, 50, 45, 39, 34, 30, 27, 23, 20, 18, 15, 14, 12, 11, 10, 9, 7, 7, 5, 5, 4, 4, 3, 3, 2, 2,
    142, 125, 108, 93, 82, 72, 63, 56, 50, 43, 38, 33, 30, 26, 22, 20, 17, 16, 13, 12, 11, 10, 8, 8, 6, 6, 5, 5, 3, 3, 2, 2,
    156, 137, 119, 103, 90, 79, 70, 61, 55, 48, 42, 37, 33, 28, 24, 22, 19, 17, 15, 13, 12, 11, 9, 9, 6, 6, 5, 5, 4, 4, 2, 2,
    171, 150, 130, 112, 99, 87, 76, 67, 60, 52, 46, 40, 36, 31, 27, 24, 21, 19, 16, 15, 13, 12, 10, 10, 7, 7, 6, 6, 4, 4, 3, 3,
    185, 162, 141, 121, 107, 94, 82, 73, 65, 56, 50, 43, 39, 34, 29, 26, 22, 21, 17, 16, 14, 13, 11, 11, 8, 8, 6, 6, 4, 4, 3, 3,
    199, 175, 152, 131, 115, 101, 89, 78, 70, 61, 54, 47, 42, 36, 31, 28, 24, 22, 19, 17, 15, 14, 12, 12, 8, 8, 7, 7, 5, 5, 3, 3,
    213, 187, 163, 140, 123, 108, 95, 84, 75, 65, 58, 50, 45, 39, 33, 30, 26, 24, 20, 18, 16, 15, 13, 13, 9, 9, 7, 7, 5, 5, 3, 3,
    228, 200, 174, 150, 132, 116, 102, 90, 80, 70, 62, 54, 48, 42, 36, 32, 28, 26, 22, 20, 18, 16, 14, 14, 10, 10, 8, 8, 6, 6, 4, 4,
];

/// The adaptation curve, indexed by `16 + (sign · state >> k)`.
#[rustfmt::skip]
pub const CURVE: [i32; 32] = [
    2512, 2288, 2064, 1840, 1616, 1392, 1168, 944, 720, 560, 464, 368, 272, 208, 144, 80,
    64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 0,
];

/// The nine initialisations a context may be given: two adaptation shifts (each plus a
/// base) and the starting value of the combined state and of the fast state.
#[rustfmt::skip]
pub const INIT: [[i32; 4]; 9] = [
    [1, 4, 0, 0], [1, 4, -1310, -41], [1, 4, 3039, 95], [0, 5, 0, 0], [2, 6, 962, 30],
    [2, 6, 3039, 95], [2, 6, -673, -21], [3, 5, 0, 0], [3, 5, 962, 30],
];

/// The number of contexts for `num_gtx` greater-than flags: 24 significance (three for
/// each of eight quantiser states), three sign, two per greater-than flag, 32 for the
/// remainder's prefix, and four more (the one at `len − 2` codes the initialisations).
pub fn context_count(num_gtx: u32) -> usize {
    8 * 3 + 3 + num_gtx as usize * 2 + 32 + 4
}

/// One context: the fast state (a signed byte), the combined state (a signed 16-bit
/// value) and the two shifts.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Ctx {
    s0: i8,
    s01: i16,
    sh0: u32,
    sh1: u32,
}

impl Ctx {
    pub fn new(init: usize) -> Ctx {
        let e = INIT[init];
        Ctx { s0: e[3] as i8, s01: e[2] as i16, sh0: (e[0] + 4) as u32, sh1: e[1] as u32 }
    }
    /// The most probable symbol.
    pub fn mps(&self) -> u32 {
        if self.s01 >= 0 { 1 } else { 0 }
    }
    fn rlps(&self, range: u32) -> Result<u32, Refusal> {
        let i = (i32::from(self.s01) >> 7).unsigned_abs() + (range & 0xe0);
        RLPS.get(i as usize).map(|&v| u32::from(v)).ok_or_else(|| Refusal::new("nnc_malformed", "a context state outside its table"))
    }
    fn update(&mut self, bin: u32) -> Result<(), Refusal> {
        let sign: i32 = if bin == 1 { 1 } else { -1 };
        let (s0, s1) = (i32::from(self.s0), i32::from(self.s01) - i32::from(self.s0) * 16);
        let c = |x: i32, k: u32| -> Result<i32, Refusal> {
            let i = 16 + ((sign * x) >> k);
            CURVE.get(usize::try_from(i).unwrap_or(usize::MAX)).copied().ok_or_else(|| Refusal::new("nnc_malformed", "a context state outside its curve"))
        };
        let s0n = (s0 + sign * (c(s0, 3)? >> self.sh0)) as i8;
        let s1n = s1 + sign * (c(s1, 7)? >> self.sh1);
        self.s0 = s0n;
        self.s01 = (s1n + 16 * i32::from(s0n)) as i16;
        Ok(())
    }
}

/// The arithmetic decoder over the bytes of one entropy-coded payload.
pub struct Decoder<'a> {
    b: &'a [u8],
    p: usize,
    range: u32,
    value: u32,
    bits_needed: i32,
}

impl<'a> Decoder<'a> {
    pub fn new(b: &'a [u8]) -> Result<Decoder<'a>, Refusal> {
        if b.len() < 2 {
            return refuse("nnc_truncated", "an entropy-coded payload of fewer than 2 bytes");
        }
        Ok(Decoder { b, p: 2, range: 510, value: 256 * u32::from(b[0]) + u32::from(b[1]), bits_needed: -8 })
    }
    /// The bytes read so far.
    pub fn bytes_read(&self) -> usize {
        self.p
    }
    fn byte(&mut self) -> Result<u32, Refusal> {
        let v = *self.b.get(self.p).ok_or_else(|| Refusal::new("nnc_truncated", "the entropy-coded payload ends inside its unit"))?;
        self.p += 1;
        Ok(u32::from(v))
    }
    pub fn bin(&mut self, ctx: &mut Ctx) -> Result<u32, Refusal> {
        let rlps = ctx.rlps(self.range)?;
        let rmps = self.range - rlps;
        let bin = if self.value < (rmps << 7) {
            self.range = rmps;
            ctx.mps()
        } else {
            self.value -= rmps << 7;
            self.range = rlps;
            1 - ctx.mps()
        };
        let n = self.range.leading_zeros() - 23;
        self.range <<= n;
        self.value <<= n;
        self.bits_needed += n as i32;
        if self.bits_needed >= 0 {
            self.value += self.byte()? << self.bits_needed;
            self.bits_needed -= 8;
        }
        ctx.update(bin)?;
        Ok(bin)
    }
    pub fn bin_ep(&mut self) -> Result<u32, Refusal> {
        self.value += self.value;
        self.bits_needed += 1;
        if self.bits_needed >= 0 {
            self.value += self.byte()?;
            self.bits_needed = -8;
        }
        let sr = self.range << 7;
        if self.value >= sr {
            self.value -= sr;
            Ok(1)
        } else {
            Ok(0)
        }
    }
    pub fn bins_ep(&mut self, n: u32) -> Result<u32, Refusal> {
        let mut v = 0u32;
        for _ in 0..n {
            v = (v << 1) | self.bin_ep()?;
        }
        Ok(v)
    }
    /// The terminating bin, then the stop pattern; returns the bytes the payload took.
    pub fn finish(&mut self) -> Result<usize, Refusal> {
        self.range -= 2;
        if self.value < (self.range << 7) {
            return refuse("nnc_malformed", "the entropy-coded payload does not end with its terminating bin");
        }
        let last = u32::from(self.b[self.p - 1]);
        if ((last << (8 + self.bits_needed)) & 0xff) != 0x80 {
            return refuse("nnc_malformed", "no stop pattern at the end of the entropy-coded payload");
        }
        Ok(self.p)
    }
}

/// The arithmetic encoder.
pub struct Encoder {
    out: Vec<u8>,
    low: u32,
    range: u32,
    bits_left: u32,
    buffered: u8,
    n_buffered: u32,
}

impl Default for Encoder {
    fn default() -> Self {
        Encoder { out: Vec::new(), low: 0, range: 510, bits_left: 23, buffered: 0, n_buffered: 0 }
    }
}

impl Encoder {
    fn write_out(&mut self) {
        let lead = self.low >> (24 - self.bits_left);
        self.bits_left += 8;
        self.low &= 0xffff_ffffu32 >> self.bits_left;
        if lead == 0xff {
            self.n_buffered += 1;
        } else if self.n_buffered > 0 {
            let carry = (lead >> 8) as u8;
            self.out.push(self.buffered.wrapping_add(carry));
            self.buffered = (lead & 0xff) as u8;
            let fill = 0xffu8.wrapping_add(carry);
            while self.n_buffered > 1 {
                self.out.push(fill);
                self.n_buffered -= 1;
            }
        } else {
            self.n_buffered = 1;
            self.buffered = lead as u8;
        }
    }
    pub fn bin(&mut self, bin: u32, ctx: &mut Ctx) -> Result<(), Refusal> {
        let rlps = ctx.rlps(self.range)?;
        self.range -= rlps;
        if bin == ctx.mps() {
            if self.range < 256 {
                self.range += self.range;
                self.low = self.low.wrapping_add(self.low);
                self.bits_left -= 1;
                if self.bits_left < 12 {
                    self.write_out();
                }
            }
        } else {
            let n = rlps.leading_zeros() - 23;
            self.low = self.low.wrapping_add(self.range);
            self.range = rlps << n;
            self.low = self.low.wrapping_shl(n);
            self.bits_left -= n;
            if self.bits_left < 12 {
                self.write_out();
            }
        }
        ctx.update(bin)
    }
    pub fn bin_ep(&mut self, bin: u32) {
        self.low = self.low.wrapping_shl(1);
        if bin == 1 {
            self.low = self.low.wrapping_add(self.range);
        }
        self.bits_left -= 1;
        if self.bits_left < 12 {
            self.write_out();
        }
    }
    pub fn bins_ep(&mut self, v: u32, n: u32) {
        for i in (0..n).rev() {
            self.bin_ep((v >> i) & 1);
        }
    }
    /// The terminating bin and the flush; returns the payload's bytes.
    pub fn finish(mut self) -> Vec<u8> {
        self.range -= 2;
        self.low = self.low.wrapping_add(self.range);
        self.low = self.low.wrapping_shl(7);
        self.range = 2 << 7;
        self.bits_left -= 7;
        if self.bits_left < 12 {
            self.write_out();
        }
        if (self.low >> (32 - self.bits_left)) != 0 {
            self.out.push(self.buffered.wrapping_add(1));
            while self.n_buffered > 1 {
                self.out.push(0);
                self.n_buffered -= 1;
            }
            self.low = self.low.wrapping_sub(1 << (32 - self.bits_left));
        } else {
            if self.n_buffered > 0 {
                self.out.push(self.buffered);
            }
            while self.n_buffered > 1 {
                self.out.push(0xff);
                self.n_buffered -= 1;
            }
        }
        self.low >>= 8;
        self.low <<= 1;
        self.low += 1;
        self.bits_left -= 1;
        self.low = self.low.wrapping_shl(32 - (24 - self.bits_left));
        let mut i = 0;
        while i < 24 - self.bits_left {
            self.out.push(((self.low >> (24 - i)) & 0xff) as u8);
            i += 8;
        }
        self.out
    }
}

/// The context of a significance or sign flag after a neighbour of value `prev`.
fn neighbour(prev: i64) -> usize {
    match prev.signum() {
        0 => 0,
        -1 => 1,
        _ => 2,
    }
}

/// Decode `n` integers coded with `num_gtx` greater-than flags (§12.4), from the start of
/// `b`; returns them and the bytes the payload took.
pub fn decode_ints(b: &[u8], n: u64, num_gtx: u32) -> Result<(Vec<i64>, usize), Refusal> {
    let mut d = Decoder::new(b)?;
    let len = context_count(num_gtx);
    let mut ctx: Vec<Ctx> = (0..len).map(|_| Ctx::new(0)).collect();
    let sel = len - 2;
    for i in 0..len - 4 {
        if (3..24).contains(&i) {
            continue; // the significance contexts of the other quantiser states
        }
        let mut c = ctx[sel];
        let mut init = 0usize;
        if d.bin(&mut c)? == 1 {
            init = 1 + d.bins_ep(3)? as usize;
        }
        ctx[sel] = c;
        ctx[i] = Ctx::new(init);
    }
    let base_rem = 27 + 2 * num_gtx as usize;
    // The count is the caller's (bounded by the reader, §12.2); reserve no more than the
    // bytes could plausibly hold, and grow if a run of zeros codes denser than that.
    let mut out = Vec::with_capacity((n.min(b.len() as u64 * 8)) as usize);
    let mut prev = 0i64;
    for _ in 0..n {
        let mut v: i64 = 0;
        if d.bin(&mut ctx[neighbour(prev)])? == 1 {
            v = 1;
            let neg = d.bin(&mut ctx[24 + neighbour(prev)])? == 1;
            let side = usize::from(neg);
            let mut gr = d.bin(&mut ctx[27 + side])?;
            let mut coded = 1u32;
            while gr == 1 && coded < num_gtx {
                v += 1;
                gr = d.bin(&mut ctx[27 + side + 2 * coded as usize])?;
                coded += 1;
            }
            if gr == 1 {
                v += 1;
                let mut rem: i64 = 0;
                let mut idx = base_rem;
                if d.bin(&mut ctx[idx])? == 1 {
                    rem += 1;
                    idx += 1;
                    let mut log2 = 1u32;
                    loop {
                        if idx >= base_rem + 32 {
                            return refuse("nnc_malformed", "a remainder prefix of more than 32 bins");
                        }
                        if d.bin(&mut ctx[idx])? == 0 {
                            break;
                        }
                        rem += 1i64 << log2;
                        idx += 1;
                        log2 += 1;
                    }
                    rem += i64::from(d.bins_ep(log2)?);
                }
                v += rem;
            }
            if v > i64::from(i32::MAX) {
                return refuse("nnc_malformed", "an integer past the 32-bit range");
            }
            if neg {
                v = -v;
            }
        }
        out.push(v);
        prev = v;
    }
    let used = d.finish()?;
    Ok((out, used))
}

/// Encode `values` (each within ±(2^31 − 1)) with `num_gtx` greater-than flags, every
/// context at its first initialisation (§12.4).
pub fn encode_ints(values: &[i64], num_gtx: u32) -> Result<Vec<u8>, Refusal> {
    let mut e = Encoder::default();
    let len = context_count(num_gtx);
    let mut ctx: Vec<Ctx> = (0..len).map(|_| Ctx::new(0)).collect();
    let sel = len - 2;
    for i in 0..len - 4 {
        if (3..24).contains(&i) {
            continue;
        }
        let mut c = ctx[sel];
        e.bin(0, &mut c)?;
        ctx[sel] = c;
    }
    let base_rem = 27 + 2 * num_gtx as usize;
    let mut prev = 0i64;
    for &v in values {
        if v.unsigned_abs() > i32::MAX as u64 {
            return refuse("value_out_of_range", format!("{v} is past the 32-bit range NNC's integer payload carries"));
        }
        e.bin(u32::from(v != 0), &mut ctx[neighbour(prev)])?;
        if v != 0 {
            let side = usize::from(v < 0);
            e.bin(side as u32, &mut ctx[24 + neighbour(prev)])?;
            let mut rem = v.unsigned_abs() - 1;
            let mut gr = u32::from(rem > 0);
            e.bin(gr, &mut ctx[27 + side])?;
            let mut coded = 1u32;
            while gr == 1 && coded < num_gtx {
                rem -= 1;
                gr = u32::from(rem > 0);
                e.bin(gr, &mut ctx[27 + side + 2 * coded as usize])?;
                coded += 1;
            }
            if gr == 1 {
                rem -= 1;
                let mut idx = base_rem;
                if rem == 0 {
                    e.bin(0, &mut ctx[idx])?;
                } else {
                    e.bin(1, &mut ctx[idx])?;
                    let mut base = 1u64;
                    let mut log2 = 1u32;
                    idx += 1;
                    while rem > base + (1u64 << log2) - 1 {
                        e.bin(1, &mut ctx[idx])?;
                        base += 1u64 << log2;
                        idx += 1;
                        log2 += 1;
                    }
                    e.bin(0, &mut ctx[idx])?;
                    e.bins_ep((rem - base) as u32, log2);
                }
            }
        }
        prev = v;
    }
    Ok(e.finish())
}

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

    #[test]
    fn integers_round_trip_through_the_coder() {
        let mut vals: Vec<i64> = vec![0, 0, 1, -1, 2, -2, 11, 12, 13, -13, 100, -1000, 65535, -70000, i32::MAX as i64, -(i32::MAX as i64), 0, 5];
        let mut s = 7u64;
        for _ in 0..2000 {
            s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            let r = (s >> 33) as i64;
            vals.push(match r % 5 {
                0 => 0,
                1 => r % 3 - 1,
                2 => r % 41 - 20,
                3 => r % 70001 - 35000,
                _ => (r << 3) % (i32::MAX as i64),
            });
        }
        for gtx in [1u32, 4, 11] {
            let b = encode_ints(&vals, gtx).unwrap();
            let (back, used) = decode_ints(&b, vals.len() as u64, gtx).unwrap();
            assert_eq!(back, vals, "gtx {gtx}");
            assert_eq!(used, b.len(), "gtx {gtx}: the decoder reads exactly the payload");
        }
        let b = encode_ints(&[], 11).unwrap();
        assert_eq!(decode_ints(&b, 0, 11).unwrap(), (vec![], b.len()));
        assert_eq!(encode_ints(&[1i64 << 31], 11).unwrap_err().code, "value_out_of_range");
    }

    #[test]
    fn a_damaged_payload_is_refused_or_decodes_to_other_integers() {
        let vals: Vec<i64> = (0..300).map(|i| (i * 37 % 23) - 11).collect();
        let b = encode_ints(&vals, 11).unwrap();
        assert!(decode_ints(&b[..b.len() - 1], 300, 11).is_err());
        for i in 0..b.len() {
            let mut m = b.clone();
            m[i] ^= 0x10;
            if let Ok((v, used)) = decode_ints(&m, 300, 11) {
                assert!(v != vals || used != b.len() || m == b, "a flip at byte {i} went unnoticed");
            }
        }
    }
}
