use pyo3::prelude::*;
use rand::Rng;
use rayon::prelude::*;

#[pyfunction]
fn montecarlo(samples: usize) -> PyResult<f64> {
    let cpus = num_cpus::get();
    let num = samples / cpus;

    let hits: u128 = (0..cpus)
        .into_par_iter()
        .map(|_| {
            let mut rng = rand::thread_rng();
            let mut count = 0;
            for _ in 0..num {
                let x: f64 = rng.gen_range(0.0..1.0);
                let y: f64 = rng.gen_range(0.0..1.0);

                let distance = x.mul_add(x, y * y);
                if distance <= 1.0 {
                    count += 1;
                }
            }
            count
        })
        .sum();

    Ok(hits as f64 / samples as f64 * 4.0)
}

/// Python module providing a parallel `montecarlo` implementation to calculate pi.
#[pymodule]
fn rs_montecarlo(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(montecarlo, m)?)?;
    Ok(())
}