use super::*;
use crate::utils::math::compute_radical_inverse_permutations;
use anyhow::{Result, anyhow};
use shared::core::geometry::Point2i;
use shared::core::options::get_options;
use shared::core::sampler::{HaltonSampler, MAX_HALTON_RESOLUTION, RandomizeStrategy};

pub trait CreateHaltonSampler {
    fn new(
        samples_per_pixel: i32,
        full_res: Point2i,
        randomize: RandomizeStrategy,
        seed: u64,
    ) -> Self;
}

impl CreateHaltonSampler for HaltonSampler {
    fn new(
        samples_per_pixel: i32,
        full_res: Point2i,
        randomize: RandomizeStrategy,
        seed: u64,
    ) -> Self {
        let (_, digit_permutations) = compute_radical_inverse_permutations(seed);
        let mut base_scales = [0u64; 2];
        let mut base_exponents = [0u64; 2];
        let bases = [2, 3];
        let res_coords = [full_res.x(), full_res.y()];

        for i in 0..2 {
            let base = bases[i] as u64;
            let mut scale = 1u64;
            let mut exp = 0u64;

            let limit = std::cmp::min(res_coords[i], MAX_HALTON_RESOLUTION) as u64;

            while scale < limit {
                scale *= base;
                exp += 1;
            }

            base_scales[i] = scale;
            base_exponents[i] = exp;
        }

        let mut mult_inverse = [0u64; 2];

        mult_inverse[0] =
            Self::multiplicative_inverse(base_scales[0] as i64, base_scales[0] as i64);
        mult_inverse[1] =
            Self::multiplicative_inverse(base_scales[1] as i64, base_scales[1] as i64);

        Self {
            samples_per_pixel,
            randomize,
            digit_permutations: digit_permutations.as_ptr().into(),
            base_scales,
            base_exponents,
            mult_inverse,
            halton_index: 0,
            dim: 0,
        }
    }
}

impl CreateSampler for HaltonSampler {
    fn create(
        params: &ParameterDictionary,
        full_res: Point2i,
        loc: &FileLoc,
        _arena: &mut Arena,
    ) -> Result<Sampler> {
        let options = get_options();
        let nsamp = options
            .quick_render
            .then_some(1)
            .or(options.pixel_samples)
            .unwrap_or_else(|| params.get_one_int("pixelsamples", 16));
        let seed = params.get_one_int("seed", options.seed);
        let s = match params
            .get_one_string("randomization", "permutedigits")
            .as_str()
        {
            "none" => RandomizeStrategy::None,
            "permutedigits" => RandomizeStrategy::PermuteDigits,
            "fastowen" => RandomizeStrategy::FastOwen,
            "owen" => RandomizeStrategy::Owen,
            _ => {
                return Err(anyhow!(
                    "{}: Unknown randomization strategy for Halton",
                    loc
                ));
            }
        };

        let sampler = HaltonSampler::new(nsamp, full_res, s, seed as u64);
        // arena.alloc(sampler);
        Ok(Sampler::Halton(sampler))
    }
}