agb_eb_ext 0.25.3

use agb::input::{Button, ButtonController};
use agb::rng::RandomNumberGenerator;
use core::ops::{Range, RangeInclusive};

/// RNG seed generator
///
/// Uses input and maths to generate seeds
///
/// # Usage
///
/// Call [`update`](SeedGen::update) up to once per loop/frame
/// then call [`create_rng`](SeedGen::create_rng) to create a [`RandomNumberGenerator`] with that seed
#[derive(Debug, Clone)]
pub struct SeedGen {
    pub seed: [u32; 4],
}

impl Default for SeedGen {
    fn default() -> SeedGen {
        Self::new([0x15f1c1; 4])
    }
}

impl SeedGen {
    /// Create with an initial seed
    pub const fn new(seed: [u32; 4]) -> Self {
        Self { seed }
    }

    /// Create a generator from the current seed
    pub fn create_rng(&self) -> RandomNumberGenerator {
        RandomNumberGenerator::new_with_seed(self.seed)
    }

    /// Mix the current button state and a frame counter into the seed
    pub fn update(&mut self, button_controller: &ButtonController) {
        self.seed[1] = self.seed[1].rotate_left(1);
        self.seed[1] |= button_controller.is_pressed(Button::Left) as u32;
        self.seed[1] |= (button_controller.is_pressed(Button::Right) as u32) << 1;
        self.seed[1] |= (button_controller.is_pressed(Button::Up) as u32) << 2;
        self.seed[1] |= (button_controller.is_pressed(Button::Down) as u32) << 3;
        self.seed[1] |= (button_controller.is_pressed(Button::A) as u32) << 4;
        self.seed[1] |= (button_controller.is_pressed(Button::B) as u32) << 5;
        self.seed[1] |= (button_controller.is_pressed(Button::L) as u32) << 6;
        self.seed[1] |= (button_controller.is_pressed(Button::R) as u32) << 7;
        self.seed[0] = self.seed[0].wrapping_add(1);
        self.seed[2] ^= self.seed[0].wrapping_mul(0x9e3779b9);
        self.seed[3] =
            self.seed[3].wrapping_add(button_controller.just_pressed_vector::<i32>().x as u32);
        self.seed[3] = self.seed[3].rotate_right(1);
    }
}

/// Random index in `0..max` without division (multiply-shift)
///
/// `max` must be greater than 0 (returns 0 otherwise, and panics in debug builds)
#[inline]
pub fn next_in(rng: &mut RandomNumberGenerator, max: usize) -> usize {
    debug_assert!(max > 0, "max must be greater than 0");
    mul_shift(rng.next_i32() as u32, max as u32) as usize
}

/// Random `bool`
#[inline]
pub fn next_bool(rng: &mut RandomNumberGenerator) -> bool {
    rng.next_i32() < 0
}

/// `true` with probability `1 / 2^bits`
///
/// `bits` must be at most 32
/// - `0` = 100% (always true)
/// - `1` = 50%
/// - `2` = 25%
/// - `3` = 12.5%
/// - `4` = 6.25%
/// - `5` = 3.125%
///
/// Negate to do the reverse, so `!one_in_pow2(rng, 3)` is 87.5%
#[inline]
pub fn one_in_pow2(rng: &mut RandomNumberGenerator, bits: u8) -> bool {
    debug_assert!(bits <= 32, "bits must be at most 32");
    if bits == 0 {
        return true;
    }
    (rng.next_i32() as u32) >> (32 - bits as u32) == 0
}

/// `(random * span) >> 32`
#[inline]
fn mul_shift(random_val: u32, span: u32) -> u32 {
    ((random_val as u64 * span as u64) >> 32) as u32
}

/// Random offset in `0..span` where `span` may be up to `1 << 32`
#[inline]
fn random_offset(rng: &mut RandomNumberGenerator, span: u64) -> u32 {
    let random_val = rng.next_i32() as u32;
    if span > u32::MAX as u64 {
        // full range: every u32 is valid, avoid the 64x64 multiply
        random_val
    } else {
        mul_shift(random_val, span as u32)
    }
}

/// Ranges that can produce a random value, see [`next_range`]
pub trait RandomRange {
    /// Element type of the range
    type Output;

    /// Random value within the range
    fn random(self, rng: &mut RandomNumberGenerator) -> Self::Output;
}

macro_rules! impl_random_range {
    ($($t:ty),*) => {$(
        impl RandomRange for Range<$t> {
            type Output = $t;

            #[inline]
            fn random(self, rng: &mut RandomNumberGenerator) -> $t {
                debug_assert!(self.start < self.end, "range must not be empty");
                if self.start >= self.end {
                    return self.start;
                }
                let span = (self.end as i64 - self.start as i64) as u64;
                self.start.wrapping_add(random_offset(rng, span) as $t)
            }
        }

        impl RandomRange for RangeInclusive<$t> {
            type Output = $t;

            #[inline]
            fn random(self, rng: &mut RandomNumberGenerator) -> $t {
                let (start, end) = (*self.start(), *self.end());
                debug_assert!(start <= end, "range must not be empty");
                if start >= end {
                    return start;
                }
                let span = (end as i64 - start as i64) as u64 + 1;
                start.wrapping_add(random_offset(rng, span) as $t)
            }
        }
    )*};
}

impl_random_range!(u8, u16, u32, i8, i16, i32);

/// Returns a random value from a `Range` or `RangeInclusive` of `u8`, `u16`, `u32`, `i8`, `i16` or `i32`
///
/// An empty range returns its start value (and panics in debug builds)
///
/// ```ignore
/// let dmg = next_range(&mut rng, 3..=7);
/// let idx = next_range(&mut rng, 0u8..4);
/// let knockback = next_range(&mut rng, -2i8..=2);
/// ```
#[inline]
pub fn next_range<R: RandomRange>(rng: &mut RandomNumberGenerator, range: R) -> R::Output {
    range.random(rng)
}

/// Shuffles a slice (or array/`Vec` via deref) in place using Fisher–Yates
pub fn shuffle<T>(rng: &mut RandomNumberGenerator, items: &mut [T]) {
    for i in (1..items.len()).rev() {
        let j = next_in(rng, i + 1);
        items.swap(i, j);
    }
}

/// Returns a random element from a slice (or array/`Vec` via deref), or `None` if it is empty
#[inline]
pub fn random_get<'a, T>(rng: &mut RandomNumberGenerator, items: &'a [T]) -> Option<&'a T> {
    if items.is_empty() {
        None
    } else {
        Some(&items[next_in(rng, items.len())])
    }
}

/// Returns a random mutable element from a slice (or array/`Vec` via deref), or `None` if it is empty
#[inline]
pub fn random_get_mut<'a, T>(
    rng: &mut RandomNumberGenerator,
    items: &'a mut [T],
) -> Option<&'a mut T> {
    if items.is_empty() {
        None
    } else {
        let idx = next_in(rng, items.len());
        Some(&mut items[idx])
    }
}

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

    #[test_case]
    fn next_range_stays_in_range(_gba: &mut agb::Gba) {
        let mut rng = RandomNumberGenerator::new_with_seed([25, 26, 27, 28]);
        for _ in 0..1000 {
            let a = next_range(&mut rng, 3u8..7);
            assert!((3..7).contains(&a), "{a} out of range");
            let b = next_range(&mut rng, 3u8..=7);
            assert!((3..=7).contains(&b), "{b} out of range");
            let c = next_range(&mut rng, 1000u16..2000);
            assert!((1000..2000).contains(&c), "{c} out of range");
            let d = next_range(&mut rng, 1000u16..=2000);
            assert!((1000..=2000).contains(&d), "{d} out of range");
            let e = next_range(&mut rng, 100_000u32..200_000);
            assert!((100_000..200_000).contains(&e), "{e} out of range");
            let f = next_range(&mut rng, 100_000u32..=200_000);
            assert!((100_000..=200_000).contains(&f), "{f} out of range");
        }
    }

    #[test_case]
    fn next_range_handles_extremes(_gba: &mut agb::Gba) {
        let mut rng = RandomNumberGenerator::new_with_seed([29, 30, 31, 32]);
        assert_eq!(next_range(&mut rng, 7u8..=7), 7);
        assert_eq!(next_range(&mut rng, 7u16..8), 7);
        for _ in 0..100 {
            next_range(&mut rng, 0u8..=u8::MAX);
            next_range(&mut rng, 0u16..=u16::MAX);
            next_range(&mut rng, 0u32..=u32::MAX);
            assert_eq!(next_range(&mut rng, u32::MAX..=u32::MAX), u32::MAX);
            let v = next_range(&mut rng, u32::MAX - 1..u32::MAX);
            assert_eq!(v, u32::MAX - 1);
        }
    }

    #[test_case]
    fn shuffle_keeps_all_elements(_gba: &mut agb::Gba) {
        let mut rng = RandomNumberGenerator::new_with_seed([9, 10, 11, 12]);
        let mut items = [1u8, 2, 3, 4, 5, 6, 7, 8];
        shuffle(&mut rng, &mut items);
        for expected in 1u8..=8 {
            assert!(
                items.contains(&expected),
                "{expected} missing after shuffle"
            );
        }
    }

    #[test_case]
    fn shuffle_handles_small_slices(_gba: &mut agb::Gba) {
        let mut rng = RandomNumberGenerator::new_with_seed([13, 14, 15, 16]);
        let mut empty: [u8; 0] = [];
        shuffle(&mut rng, &mut empty);
        let mut single = [42u8];
        shuffle(&mut rng, &mut single);
        assert_eq!(single, [42]);
    }

    #[test_case]
    fn choose_returns_contained_element(_gba: &mut agb::Gba) {
        let mut rng = RandomNumberGenerator::new_with_seed([17, 18, 19, 20]);
        let items = [10u8, 20, 30, 40, 50];
        for _ in 0..100 {
            let v = *random_get(&mut rng, &items).unwrap();
            assert!(items.contains(&v), "{v} not in source slice");
        }
    }

    #[test_case]
    fn choose_handles_empty_and_mut(_gba: &mut agb::Gba) {
        let mut rng = RandomNumberGenerator::new_with_seed([21, 22, 23, 24]);
        let empty: [u8; 0] = [];
        assert!(random_get(&mut rng, &empty).is_none());

        let mut items = [1u8, 1, 1];
        if let Some(v) = random_get_mut(&mut rng, &mut items) {
            *v = 9;
        }
        assert_eq!(items.iter().filter(|&&v| v == 9).count(), 1);
    }
    #[test_case]
    fn signed_ranges_stay_in_range(_gba: &mut agb::Gba) {
        let mut rng = RandomNumberGenerator::new_with_seed([33, 34, 35, 36]);
        let mut seen_neg = false;
        for _ in 0..1000 {
            let a = next_range(&mut rng, -2i8..=2);
            assert!((-2..=2).contains(&a), "{a} out of range");
            seen_neg |= a < 0;
            let b = next_range(&mut rng, -100i16..50);
            assert!((-100..50).contains(&b), "{b} out of range");
            let c = next_range(&mut rng, -1_000_000i32..=1_000_000);
            assert!((-1_000_000..=1_000_000).contains(&c), "{c} out of range");
            next_range(&mut rng, i8::MIN..=i8::MAX);
            next_range(&mut rng, i32::MIN..=i32::MAX);
            next_range(&mut rng, i32::MIN..i32::MAX);
        }
        assert!(seen_neg, "never produced a negative value");
        assert_eq!(next_range(&mut rng, -5i32..=-5), -5);
        assert_eq!(next_range(&mut rng, i32::MIN..i32::MIN + 1), i32::MIN);
    }

    #[test_case]
    fn bool_and_pow2(_gba: &mut agb::Gba) {
        let mut rng = RandomNumberGenerator::new_with_seed([37, 38, 39, 40]);
        let (mut t, mut f) = (0, 0);
        for _ in 0..1000 {
            if next_bool(&mut rng) { t += 1 } else { f += 1 }
        }
        assert!(t > 300 && f > 300, "biased: {t} true / {f} false");

        for _ in 0..100 {
            assert!(one_in_pow2(&mut rng, 0));
        }
        let hits = (0..4000).filter(|_| one_in_pow2(&mut rng, 2)).count();
        assert!((700..1300).contains(&hits), "1/4 chance hit {hits} of 4000");
        let rare = (0..1000).filter(|_| one_in_pow2(&mut rng, 32)).count();
        assert!(rare <= 1, "1/2^32 chance hit {rare} times");
    }
}