mlt-core 0.19.2

use num_traits::{AsPrimitive, WrappingAdd, WrappingSub};
use zigzag::ZigZag;

use crate::MltError::InvalidPairStreamSize;
use crate::{Decoder, MltResult};

/// ZigZag-encode `data` into `target`.
///
/// `target` is treated as a scratch buffer: cleared before writing.
pub fn encode_zigzag<'a, T: ZigZag>(data: &[T], target: &'a mut Vec<T::UInt>) -> &'a [T::UInt] {
    target.clear();
    target.extend(data.iter().map(|&v| T::encode(v)));
    target
}

/// Delta-then-ZigZag-encode `data` into `target` in a single pass.
///
/// `target` is treated as a scratch buffer: cleared before writing.
/// Fuses the delta and zigzag steps to avoid an intermediate allocation.
pub fn encode_zigzag_delta<'a, T: Copy + ZigZag + WrappingSub<Output = T>>(
    data: &[T],
    target: &'a mut Vec<T::UInt>,
) -> &'a [T::UInt] {
    target.clear();
    target.reserve(data.len());
    let mut prev = T::zero();
    for &v in data {
        target.push(T::encode(v.wrapping_sub(&prev)));
        prev = v;
    }
    target
}

/// Encode signed integer vec2 values using componentwise delta + zigzag into `target`.
///
/// Input: `[x0, y0, x1, y1, ...]`
/// Output: `[zigzag(x0-0), zigzag(y0-0), zigzag(x1-x0), zigzag(y1-y0), ...]`
///
/// `target` is treated as a scratch buffer: cleared before writing.
/// This is the inverse of `decode_componentwise_delta_vec2s`.
pub fn encode_componentwise_delta_vec2s<'a, T>(
    data: &[T],
    target: &'a mut Vec<T::UInt>,
) -> &'a [T::UInt]
where
    T: ZigZag + WrappingSub,
{
    target.clear();
    target.reserve(data.len());
    let mut prev_x = T::zero();
    let mut prev_y = T::zero();
    for &[x, y] in data.as_chunks::<2>().0 {
        target.push(T::encode(x.wrapping_sub(&prev_x)));
        target.push(T::encode(y.wrapping_sub(&prev_y)));
        (prev_x, prev_y) = (x, y);
    }
    target
}

/// Encode signed integer vec3 values using componentwise delta + zigzag into `target`.
///
/// Input: `[x0, y0, z0, x1, y1, z1, ...]`
/// `target` is treated as a scratch buffer: cleared before writing.
/// This is the inverse of `decode_componentwise_delta_vec3s`.
#[cfg(feature = "unstable-v2")]
pub fn encode_componentwise_delta_vec3s<'a, T>(
    data: &[T],
    target: &'a mut Vec<T::UInt>,
) -> &'a [T::UInt]
where
    T: ZigZag + WrappingSub,
{
    target.clear();
    target.reserve(data.len());
    let mut prev = [T::zero(); 3];
    for triple in data.as_chunks::<3>().0 {
        for (&v, p) in triple.iter().zip(&mut prev) {
            target.push(T::encode(v.wrapping_sub(p)));
            *p = v;
        }
    }
    target
}

/// Zigzag of the deltas of the deltas, both running from zero, so `[a, b, c]` becomes `[a, b - 2a, c - 2b + a]`.
#[cfg(feature = "unstable-v2")]
pub fn encode_zigzag_delta2<'a, T: Copy + ZigZag + WrappingSub<Output = T>>(
    data: &[T],
    target: &'a mut Vec<T::UInt>,
) -> &'a [T::UInt] {
    target.clear();
    target.reserve(data.len());
    let (mut prev, mut prev_delta) = (T::zero(), T::zero());
    for &v in data {
        let delta = v.wrapping_sub(&prev);
        target.push(T::encode(delta.wrapping_sub(&prev_delta)));
        (prev, prev_delta) = (v, delta);
    }
    target
}

/// [`encode_zigzag_delta2`] over each component of interleaved `(x, y)` pairs.
#[cfg(feature = "unstable-v2")]
pub fn encode_componentwise_delta2_vec2s<'a, T>(
    data: &[T],
    target: &'a mut Vec<T::UInt>,
) -> &'a [T::UInt]
where
    T: ZigZag + WrappingSub,
{
    target.clear();
    target.reserve(data.len());
    let (mut prev, mut prev_delta) = ([T::zero(); 2], [T::zero(); 2]);
    for pair in data.as_chunks::<2>().0 {
        for c in 0..2 {
            let delta = pair[c].wrapping_sub(&prev[c]);
            target.push(T::encode(delta.wrapping_sub(&prev_delta[c])));
            (prev[c], prev_delta[c]) = (pair[c], delta);
        }
    }
    target
}

/// ZigZag-decode a slice, charging `dec` for the output allocation.
pub fn decode_zigzag<T: ZigZag>(data: &[T::UInt], dec: &mut Decoder) -> MltResult<Vec<T>> {
    dec.consume_items::<T>(data.len())?;
    Ok(data.iter().map(|&v| T::decode(v)).collect())
}

/// Decode a vector of ZigZag-encoded unsigned deltas, charging `dec` for the output allocation.
pub fn decode_zigzag_delta<T: Copy + ZigZag + WrappingAdd + AsPrimitive<U>, U: 'static + Copy>(
    data: &[T::UInt],
    dec: &mut Decoder,
) -> MltResult<Vec<U>> {
    dec.consume_items::<U>(data.len())?;
    Ok(data
        .iter()
        .scan(T::zero(), |state, &v| {
            *state = state.wrapping_add(&T::decode(v));
            Some((*state).as_())
        })
        .collect())
}

/// Invert the zigzag deltas of the deltas by summing twice.
pub fn decode_zigzag_delta2<T: Copy + ZigZag + WrappingAdd + AsPrimitive<U>, U: 'static + Copy>(
    data: &[T::UInt],
    dec: &mut Decoder,
) -> MltResult<Vec<U>> {
    dec.consume_items::<U>(data.len())?;
    let (mut value, mut delta) = (T::zero(), T::zero());
    Ok(data
        .iter()
        .map(|&v| {
            delta = delta.wrapping_add(&T::decode(v));
            value = value.wrapping_add(&delta);
            value.as_()
        })
        .collect())
}

/// Invert the componentwise zigzag deltas of the deltas.
pub fn decode_componentwise_delta2_vec2s<T: ZigZag + WrappingAdd>(
    data: &[T::UInt],
    dec: &mut Decoder,
) -> MltResult<Vec<T>> {
    if !data.len().is_multiple_of(2) {
        return Err(InvalidPairStreamSize(data.len()));
    }
    let mut result = dec.alloc(data.len())?;
    let (mut value, mut delta) = ([T::zero(); 2], [T::zero(); 2]);
    for pair in data.as_chunks::<2>().0 {
        for c in 0..2 {
            delta[c] = delta[c].wrapping_add(&T::decode(pair[c]));
            value[c] = value[c].wrapping_add(&delta[c]);
            result.push(value[c]);
        }
    }
    Ok(result)
}

/// Decode ([`ZigZag`] + delta) for Vec2s, charging `dec` for the output allocation.
// TODO: The encoded process is (delta + ZigZag) for each component
pub fn decode_componentwise_delta_vec2s<T: ZigZag + WrappingAdd>(
    data: &[T::UInt],
    dec: &mut Decoder,
) -> MltResult<Vec<T>> {
    // Zero pairs is a well-formed pair stream: a layer can have no vertices at all.
    if !data.len().is_multiple_of(2) {
        return Err(InvalidPairStreamSize(data.len()));
    }

    let alloc_size = data.len();
    let mut result = dec.alloc(alloc_size)?;
    let mut last1 = T::zero();
    let mut last2 = T::zero();

    for i in (0..data.len()).step_by(2) {
        last1 = last1.wrapping_add(&T::decode(data[i]));
        last2 = last2.wrapping_add(&T::decode(data[i + 1]));
        result.push(last1);
        result.push(last2);
    }

    dec.adjust_alloc(&result, alloc_size)
        .expect("infallible: two values pushed per pair fill alloc_size exactly");
    Ok(result)
}

/// Decode ([`ZigZag`] + delta) for Vec3s, charging `dec` for the output allocation.
#[cfg(feature = "unstable-v2")]
pub fn decode_componentwise_delta_vec3s<T: ZigZag + WrappingAdd>(
    data: &[T::UInt],
    dec: &mut Decoder,
) -> MltResult<Vec<T>> {
    if !data.len().is_multiple_of(3) {
        return Err(crate::MltError::InvalidTripleStreamSize(data.len()));
    }
    let alloc_size = data.len();
    let mut result = dec.alloc(alloc_size)?;
    let mut last = [T::zero(); 3];
    for triple in data.as_chunks::<3>().0 {
        for (&v, l) in triple.iter().zip(&mut last) {
            *l = l.wrapping_add(&T::decode(v));
            result.push(*l);
        }
    }
    dec.adjust_alloc(&result, alloc_size)
        .expect("infallible: three values pushed per triple fill alloc_size exactly");
    Ok(result)
}

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

    use super::*;
    use crate::test_helpers::{dec, starved_dec};

    #[cfg(feature = "unstable-v2")]
    #[test]
    fn delta2_codes_the_deltas_of_the_deltas() {
        let mut encoded = Vec::new();
        encode_zigzag_delta2(&[10_i32, 13, 16, 20], &mut encoded);
        let signed: Vec<i32> = encoded.iter().map(|&v| i32::decode(v)).collect();
        assert_eq!(signed, [10, -7, 0, 1]);
    }

    proptest! {
        #[test]
        fn test_zigzag_roundtrip_i64(data: Vec<i64>) {
            let mut encoded = Vec::new();
            let decoded = decode_zigzag::<i64>(encode_zigzag(&data, &mut encoded), &mut dec()).unwrap();
            prop_assert_eq!(data, decoded);
        }

        #[test]
        fn test_delta_roundtrip_i32(data: Vec<i32>) {
            if data.is_empty() { return Ok(()); }
            let mut encoded = Vec::new();
            encode_zigzag_delta(&data, &mut encoded);
            let decoded: Vec<i32> = decode_zigzag_delta::<i32, i32>(&encoded, &mut dec()).unwrap();
            prop_assert_eq!(data, decoded);
        }

        #[cfg(feature = "unstable-v2")]
        #[test]
        fn delta2_round_trips_i32(data: Vec<i32>) {
            let mut encoded = Vec::new();
            encode_zigzag_delta2(&data, &mut encoded);
            prop_assert_eq!(decode_zigzag_delta2::<i32, i32>(&encoded, &mut dec()).unwrap(), data);
        }

        #[cfg(feature = "unstable-v2")]
        #[test]
        fn delta2_round_trips_i64(data: Vec<i64>) {
            let mut encoded = Vec::new();
            encode_zigzag_delta2(&data, &mut encoded);
            prop_assert_eq!(decode_zigzag_delta2::<i64, i64>(&encoded, &mut dec()).unwrap(), data);
        }

        #[cfg(feature = "unstable-v2")]
        #[test]
        fn componentwise_delta2_round_trips(data: Vec<(i32, i32)>) {
            let flat: Vec<i32> = data.iter().flat_map(|&(x, y)| [x, y]).collect();
            let mut encoded = Vec::new();
            encode_componentwise_delta2_vec2s(&flat, &mut encoded);
            prop_assert_eq!(decode_componentwise_delta2_vec2s::<i32>(&encoded, &mut dec()).unwrap(), flat);
        }

        #[test]
        fn test_componentwise_delta_vec2s(data: Vec<i32>) {
            if data.len() <= 1 {
                return Err(TestCaseError::reject("data not valid vertices"))
            }
            let data_slice = if data.len().is_multiple_of(2) {
                &data
            } else {
                &data[..data.len() - 1]
            };
            let mut encoded = Vec::new();
            let data = encode_componentwise_delta_vec2s(data_slice, &mut encoded);
            let decoded = decode_componentwise_delta_vec2s::<i32>(data, &mut dec()).unwrap();
            prop_assert_eq!(data_slice, &decoded);
        }
    }

    #[cfg(feature = "unstable-v2")]
    proptest! {
        #[test]
        fn componentwise_delta_vec3s_round_trip(data: Vec<[i32; 3]>) {
            let flat = data.as_flattened();
            let mut encoded = Vec::new();
            let encoded = encode_componentwise_delta_vec3s(flat, &mut encoded);
            let decoded = decode_componentwise_delta_vec3s::<i32>(encoded, &mut dec()).unwrap();
            prop_assert_eq!(flat, &decoded[..]);
        }
    }

    #[cfg(feature = "unstable-v2")]
    #[test]
    fn a_vec3_stream_of_partial_triples_is_rejected() {
        let err = decode_componentwise_delta_vec3s::<i32>(&[0, 1, 2, 3], &mut dec()).unwrap_err();
        assert!(
            matches!(err, crate::MltError::InvalidTripleStreamSize(4)),
            "{err:?}"
        );
    }

    #[test]
    fn test_encode_zigzag_empty() {
        let mut target = Vec::<u32>::new();
        assert_eq!(encode_zigzag::<i32>(&[], &mut target), [] as [u32; 0]);
    }

    #[test]
    fn test_encode_zigzag_delta_empty() {
        let mut target = Vec::<u32>::new();
        encode_zigzag_delta::<i32>(&[], &mut target);
        assert_eq!(target, [] as [u32; 0]);
    }

    #[test]
    fn test_decode_zigzag_i32() {
        let encoded_u32 = [0u32, 1, 2, 3, 4, 5, u32::MAX];
        let expected_i32 = [0i32, -1, 1, -2, 2, -3, i32::MIN];
        let decoded_i32 = decode_zigzag::<i32>(&encoded_u32, &mut dec()).unwrap();
        assert_eq!(decoded_i32, expected_i32);
    }

    #[test]
    fn test_decode_zigzag_i64() {
        let encoded_u64 = [0u64, 1, 2, 3, 4, 5, u64::MAX];
        let expected_i64 = [0i64, -1, 1, -2, 2, -3, i64::MIN];
        let decoded_i64 = decode_zigzag::<i64>(&encoded_u64, &mut dec()).unwrap();
        assert_eq!(decoded_i64, expected_i64);
    }

    #[test]
    fn test_decode_zigzag_empty() {
        assert_eq!(
            decode_zigzag::<i32>(&[], &mut dec()).unwrap(),
            [] as [i32; 0]
        );
    }

    #[test]
    fn test_decode_zigzag_delta_empty() {
        assert_eq!(
            decode_zigzag_delta::<i32, i32>(&[], &mut dec()).unwrap(),
            [] as [i32; 0]
        );
    }

    #[test]
    fn an_odd_length_pair_stream_is_rejected() {
        let err = decode_componentwise_delta_vec2s::<i32>(&[0, 1, 2], &mut dec()).unwrap_err();
        assert!(matches!(err, InvalidPairStreamSize(3)), "{err:?}");
    }

    #[test]
    fn zigzag_decoding_past_the_memory_budget_is_rejected() {
        let err = decode_zigzag::<i32>(&[0, 1], &mut starved_dec()).unwrap_err();
        assert!(
            matches!(err, crate::MltError::MemoryLimitExceeded { .. }),
            "{err:?}"
        );
    }

    #[test]
    fn delta_decoding_past_the_memory_budget_is_rejected() {
        let err = decode_zigzag_delta::<i32, i32>(&[0, 1], &mut starved_dec()).unwrap_err();
        assert!(
            matches!(err, crate::MltError::MemoryLimitExceeded { .. }),
            "{err:?}"
        );
    }

    #[test]
    fn pair_decoding_past_the_memory_budget_is_rejected() {
        let err = decode_componentwise_delta_vec2s::<i32>(&[0, 1], &mut starved_dec()).unwrap_err();
        assert!(
            matches!(err, crate::MltError::MemoryLimitExceeded { .. }),
            "{err:?}"
        );
    }

    #[test]
    fn test_decode_componentwise_delta_vec2s() {
        let values = &[1_u32, 2, 3, 4];
        let decoded = decode_componentwise_delta_vec2s::<i32>(values, &mut dec()).unwrap();
        assert_eq!(&decoded, &[-1_i32, 1, -3, 3]);
    }
}