pub mod alloc;
pub mod complex;
pub mod containers;
pub mod hash;
pub mod interval;
pub mod math;
pub mod noise;
pub mod options;
pub mod ptr;
pub mod quaternion;
pub mod rng;
pub mod sampling;
pub mod sobol;
pub mod splines;
pub mod transform;

pub use options::PBRTOptions;
pub use ptr::Ptr;
pub use transform::{AnimatedTransform, Transform, TransformGeneric};
pub use containers::Array2D;

use crate::Float;
use core::sync::atomic::{AtomicU32, Ordering};

#[inline]
pub fn find_interval<F>(sz: u32, pred: F) -> u32
where
    F: Fn(u32) -> bool,
{
    let mut first = 0;
    let mut len = sz;

    while len > 0 {
        let half = len >> 1;
        let middle = first + half;

        if pred(middle) {
            first = middle + 1;
            len -= half + 1;
        } else {
            len = half;
        }
    }

    let ret = (first as i32 - 1).max(0) as u32;
    ret.min(sz.saturating_sub(2))
}

#[inline]
pub fn partition_slice<T, F>(data: &mut [T], predicate: F) -> usize
where
    F: Fn(&T) -> bool,
{
    let mut i = 0;
    for j in 0..data.len() {
        if predicate(&data[j]) {
            data.swap(i, j);
            i += 1;
        }
    }
    i
}

#[repr(C)]
#[derive(Debug)]
pub struct AtomicFloat {
    bits: AtomicU32,
}

impl Default for AtomicFloat {
    fn default() -> Self {
        Self::new(0.)
    }
}

impl Clone for AtomicFloat {
    fn clone(&self) -> Self {
        Self::new(self.get())
    }
}

impl AtomicFloat {
    pub fn new(val: f32) -> Self {
        Self {
            bits: AtomicU32::new(val.to_bits()),
        }
    }

    pub fn get(&self) -> Float {
        Float::from_bits(self.bits.load(Ordering::Relaxed))
    }

    pub fn set(&self, val: Float) {
        self.bits.store(val.to_bits(), Ordering::Relaxed);
    }

    /// Atomically adds `val` to the current value.
    /// Compare-And-Swap loop.
    pub fn add(&self, val: f32) {
        let mut current_bits = self.bits.load(Ordering::Relaxed);
        loop {
            let current_val = f32::from_bits(current_bits);
            let new_val = current_val + val;
            let new_bits = new_val.to_bits();

            match self.bits.compare_exchange_weak(
                current_bits,
                new_bits,
                Ordering::Relaxed,
                Ordering::Relaxed,
            ) {
                Ok(_) => break,
                Err(x) => current_bits = x,
            }
        }
    }
}

#[inline(always)]
pub fn gpu_array_from_fn<T, const N: usize>(mut f: impl FnMut(usize) -> T) -> [T; N] {
    unsafe {
        let mut arr: core::mem::MaybeUninit<[T; N]> = core::mem::MaybeUninit::uninit();
        let ptr = arr.as_mut_ptr() as *mut T;
        let mut i = 0;
        while i < N {
            ptr.add(i).write(f(i));
            i += 1;
        }
        arr.assume_init()
    }
}

