//! The range decoder of RFC 6716 §4.1, as the DRED payload uses it
//! (draft-ietf-mlcodec-opus-dred-07 §3 and §3.1), and the matching encoder of RFC 6716 §5.1,
//! which only the corpus generator uses.
//!
//! The payload uses three of the decoder's calls: a symbol by its frequency (`ec_decode` then
//! `ec_dec_update`), a symbol by a 15-bit inverse cumulative table (`ec_dec_icdf16`), and a
//! uniform integer below at most 256 (`ec_dec_uint`, which for such a range is one symbol by
//! frequency and reads no raw bits). `tell` is `ec_tell`, the bits the decoder has used, which the
//! payload's end rule reads. Every value here is a 32-bit or 64-bit unsigned integer formed as the
//! RFC forms it; nothing wraps except the shifts the RFC defines as dropping high bits.

const SYM_BITS: u32 = 8;
const CODE_TOP: u32 = 1 << 31;
const CODE_BOT: u32 = CODE_TOP >> SYM_BITS;
const CODE_SHIFT: u32 = 32 - SYM_BITS - 1;
const CODE_EXTRA: u32 = (32 - 2) % SYM_BITS + 1;
const SYM_MAX: u32 = (1 << SYM_BITS) - 1;

fn ilog(x: u32) -> i64 {
    (32 - x.leading_zeros()) as i64
}

/// The range decoder (RFC 6716 §4.1.1–§4.1.6).
pub struct Dec<'a> {
    buf: &'a [u8],
    offs: usize,
    rng: u32,
    val: u32,
    rem: u32,
    ext: u32,
    nbits_total: i64,
}

impl<'a> Dec<'a> {
    /// `ec_dec_init` over `buf`. Past its end the decoder reads zero bytes (§4.1.1).
    pub fn new(buf: &'a [u8]) -> Dec<'a> {
        let mut d = Dec { buf, offs: 0, rng: 1 << CODE_EXTRA, val: 0, rem: 0, ext: 0, nbits_total: 32 + 1 - ((32 - CODE_EXTRA as i64) / 8) * 8 };
        d.rem = d.byte();
        d.val = d.rng - 1 - (d.rem >> (SYM_BITS - CODE_EXTRA));
        d.normalize();
        d
    }

    fn byte(&mut self) -> u32 {
        match self.buf.get(self.offs) {
            Some(&b) => {
                self.offs += 1;
                b as u32
            }
            None => 0,
        }
    }

    fn normalize(&mut self) {
        while self.rng <= CODE_BOT {
            self.nbits_total += SYM_BITS as i64;
            self.rng <<= SYM_BITS;
            let sym = self.rem;
            self.rem = self.byte();
            let sym = ((sym << SYM_BITS) | self.rem) >> (SYM_BITS - CODE_EXTRA);
            self.val = (self.val.wrapping_shl(SYM_BITS).wrapping_add(SYM_MAX & !sym)) & (CODE_TOP - 1);
        }
    }

    /// `ec_decode(ft)`: the cumulative frequency the next symbol lies at, `ft ≥ 1`.
    pub fn decode(&mut self, ft: u32) -> u32 {
        self.ext = self.rng / ft;
        let s = self.val / self.ext;
        ft - (s + 1).min(ft)
    }

    /// `ec_dec_update(fl, fh, ft)`, after [`Self::decode`].
    pub fn update(&mut self, fl: u32, fh: u32, ft: u32) {
        let s = self.ext * (ft - fh);
        self.val -= s;
        self.rng = if fl > 0 { self.ext * (fh - fl) } else { self.rng - s };
        self.normalize();
    }

    /// `ec_dec_uint(ft)` for `2 ≤ ft ≤ 256`: one symbol of `ft` equally likely values.
    pub fn uint(&mut self, ft: u32) -> u32 {
        debug_assert!((2..=256).contains(&ft));
        let s = self.decode(ft);
        self.update(s, s + 1, ft);
        s
    }

    /// `ec_dec_icdf16(icdf, ftb)`: the symbol whose interval the inverse cumulative table
    /// `icdf` (in units of `2^-ftb`, ending with 0) gives.
    pub fn icdf16(&mut self, icdf: &[u16], ftb: u32) -> usize {
        let d = self.val;
        let r = self.rng >> ftb;
        let mut s = self.rng;
        let mut t;
        let mut k = 0usize;
        loop {
            t = s;
            s = r * icdf[k] as u32;
            if d < s {
                k += 1;
            } else {
                break;
            }
        }
        self.val = d - s;
        self.rng = t - s;
        self.normalize();
        k
    }

    /// `ec_tell`: the number of bits the decoder has used so far (RFC 6716 §4.1.6.1).
    pub fn tell(&self) -> i64 {
        self.nbits_total - ilog(self.rng)
    }
}

/// The range encoder (RFC 6716 §5.1), for the corpus generator. Its output, ended as the
/// draft's reference encoder ends a DRED payload (§3: `⌈tell/8⌉` bytes, the flush clipped to
/// them), decodes to the same symbols through [`Dec`].
pub struct Enc {
    buf: Vec<u8>,
    storage: usize,
    rng: u32,
    val: u32,
    rem: i32,
    ext: u32,
    nbits_total: i64,
}

impl Default for Enc {
    fn default() -> Self {
        Self::new()
    }
}

impl Enc {
    pub fn new() -> Enc {
        Enc { buf: Vec::new(), storage: usize::MAX, rng: CODE_TOP, val: 0, rem: -1, ext: 0, nbits_total: 33 }
    }

    fn write(&mut self, b: u32) {
        if self.buf.len() < self.storage {
            self.buf.push(b as u8);
        }
    }

    fn carry_out(&mut self, c: u32) {
        if c != SYM_MAX {
            let carry = c >> SYM_BITS;
            if self.rem >= 0 {
                self.write(self.rem as u32 + carry);
            }
            if self.ext > 0 {
                let sym = (SYM_MAX + carry) & SYM_MAX;
                while self.ext > 0 {
                    self.write(sym);
                    self.ext -= 1;
                }
            }
            self.rem = (c & SYM_MAX) as i32;
        } else {
            self.ext += 1;
        }
    }

    fn normalize(&mut self) {
        while self.rng <= CODE_BOT {
            self.carry_out(self.val >> CODE_SHIFT);
            self.val = (self.val << SYM_BITS) & (CODE_TOP - 1);
            self.rng <<= SYM_BITS;
            self.nbits_total += SYM_BITS as i64;
        }
    }

    /// `ec_encode(fl, fh, ft)`.
    pub fn encode(&mut self, fl: u32, fh: u32, ft: u32) {
        let r = self.rng / ft;
        if fl > 0 {
            self.val += self.rng - r * (ft - fl);
            self.rng = r * (fh - fl);
        } else {
            self.rng -= r * (ft - fh);
        }
        self.normalize();
    }

    /// `ec_enc_uint(v, ft)` for `2 ≤ ft ≤ 256`.
    pub fn uint(&mut self, v: u32, ft: u32) {
        self.encode(v, v + 1, ft);
    }

    /// `ec_enc_icdf16(s, icdf, ftb)`.
    pub fn icdf16(&mut self, s: usize, icdf: &[u16], ftb: u32) {
        let r = self.rng >> ftb;
        if s > 0 {
            self.val += self.rng - r * icdf[s - 1] as u32;
            self.rng = r * (icdf[s - 1] as u32 - icdf[s] as u32);
        } else {
            self.rng -= r * icdf[s] as u32;
        }
        self.normalize();
    }

    pub fn tell(&self) -> i64 {
        self.nbits_total - ilog(self.rng)
    }

    /// End the payload as the draft's reference encoder does: `⌈tell/8⌉` bytes, the final
    /// flush of RFC 6716 §5.1.5 clipped to them, and zeros after.
    pub fn finish(mut self) -> Vec<u8> {
        let fill = ((self.tell() + 7) / 8) as usize;
        debug_assert!(self.buf.len() <= fill);
        self.storage = fill;
        let mut l = 32 - ilog(self.rng) as i32;
        let mut msk = (CODE_TOP - 1) >> l;
        let mut end = (self.val.wrapping_add(msk)) & !msk;
        if (end | msk) >= self.val.wrapping_add(self.rng) {
            l += 1;
            msk >>= 1;
            end = (self.val.wrapping_add(msk)) & !msk;
        }
        while l > 0 {
            self.carry_out(end >> CODE_SHIFT);
            end = (end << SYM_BITS) & (CODE_TOP - 1);
            l -= SYM_BITS as i32;
        }
        if self.rem >= 0 || self.ext > 0 {
            self.carry_out(0);
        }
        self.buf.resize(fill, 0);
        self.buf
    }
}

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

    /// Symbols of every kind the payload uses, encoded then decoded: the same symbols, and the
    /// decoder's `tell` after each equals the encoder's.
    #[test]
    fn encoder_and_decoder_agree() {
        let mut seed = 0x1234_5678u32;
        let mut next = || {
            seed ^= seed << 13;
            seed ^= seed >> 17;
            seed ^= seed << 5;
            seed
        };
        let icdf: [u16; 8] = [30000, 20000, 9000, 4000, 1500, 400, 60, 0];
        let mut syms = Vec::new();
        let mut enc = Enc::new();
        let mut tells = Vec::new();
        for i in 0..4000 {
            match i % 3 {
                0 => {
                    let ft = 2 + next() % 255;
                    let v = next() % ft;
                    enc.uint(v, ft);
                    syms.push((0, v, ft));
                }
                1 => {
                    let s = next() % 8;
                    enc.icdf16(s as usize, &icdf, 15);
                    syms.push((1, s, 0));
                }
                _ => {
                    let ft = 2 + next() % 40;
                    let fl = next() % ft;
                    enc.encode(fl, fl + 1, ft);
                    syms.push((2, fl, ft));
                }
            }
            tells.push(enc.tell());
        }
        let bytes = enc.finish();
        let mut dec = Dec::new(&bytes);
        for (k, &(kind, v, ft)) in syms.iter().enumerate() {
            let got = match kind {
                0 => dec.uint(ft),
                1 => dec.icdf16(&icdf, 15) as u32,
                _ => {
                    let s = dec.decode(ft);
                    dec.update(s, s + 1, ft);
                    s
                }
            };
            assert_eq!(got, v, "symbol {k}");
            assert_eq!(dec.tell(), tells[k], "tell after symbol {k}");
        }
        assert!(8 * bytes.len() as i64 - dec.tell() <= 7);
    }

    /// A decoder past the end of its bytes reads zeros, and keeps answering.
    #[test]
    fn decoder_reads_zeros_past_the_end() {
        let mut d = Dec::new(&[]);
        for _ in 0..100 {
            let _ = d.uint(256);
        }
        assert!(d.tell() > 0);
    }
}
