//! # Chaos Game
//!
//! Fractal flames are a class of
//! [iterated function systems](https://en.wikipedia.org/wiki/Iterated_function_system)
//! that generate images following a simple algorithm:
//!
//! - Pick a starting point `(x, y)`
//! - Iterate:
//!     - Pick a [`Transform`] from the set of available transforms
//!     - Apply the current point to the chosen transform, generating a new point `(x, y)`
//!     - Plot the new point `(x, y)`
//!
//! This algorithm is also known as the ["chaos game"](https://en.wikipedia.org/wiki/Chaos_game),
//! and it forms the basic system for producing images.

use crate::transform::Transform;
use crate::variation::Variation;
use rand::distr::{Distribution, StandardUniform};
use rand::{Rng, RngExt};
use spirv_std::glam::{Vec2, vec2};

struct BiUnit;
impl Distribution<f32> for BiUnit {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f32 {
        rng.sample::<f32, _>(StandardUniform) * 2.0 - 1.0
    }
}

/// Iterate one step in the chaos game; choose the next transform, apply it,
/// and return the resulting point. Also returns the transform index so that
/// path-dependent weights (the "Xaos" table in Apophysis) can be chosen
/// for the next iteration step.
///
/// # Arguments
///
/// * `weights` - Weights are assumed to be normalized; adding all elements together should return the value 1
pub fn step_chaos_game<R: Rng>(
    point: Vec2,
    rng: &mut R,
    transforms: &[Transform],
    weights: &[f32],
    variations: &[Variation],
) -> (Vec2, u32) {
    let mut choice_weight = rng.sample::<f32, _>(StandardUniform);
    let mut transform_index: u32 = 0;

    for i in 0..weights.len() {
        choice_weight -= weights[i];
        if choice_weight <= 0.0 {
            break;
        }

        transform_index += 1;
    }

    (
        transforms[transform_index as usize].transform_point(rng, variations, point),
        transform_index,
    )
}

/// Iterator for chaos game state. Holds the current point and references to all other data
/// necessary to generate fractal flame images.
///
/// New points in the chaos game are produced by iterating on the chaos game.
pub struct ChaosGame<'a, R: Rng> {
    current_point: Vec2,
    rng: &'a mut R,
    transforms: &'a [Transform],
    weights: &'a [f32],
    variations: &'a [Variation],
}

impl<'a, R: Rng> ChaosGame<'a, R> {
    /// Create a new chaos game iterator
    pub fn new(
        rng: &'a mut R,
        transforms: &'a [Transform],
        weights: &'a [f32],
        variations: &'a [Variation],
    ) -> Self {
        let current_point = vec2(rng.sample(BiUnit), rng.sample(BiUnit));
        ChaosGame {
            current_point,
            rng,
            transforms,
            weights,
            variations,
        }
    }
}

impl<'a, R: Rng> Iterator for ChaosGame<'a, R> {
    type Item = Vec2;

    fn next(&mut self) -> Option<Self::Item> {
        let (next_point, _) = step_chaos_game(
            self.current_point,
            self.rng,
            self.transforms,
            self.weights,
            self.variations,
        );
        self.current_point = next_point;

        Some(next_point)
    }
}

/// Shader entry point for running the chaos game to produce new IFS coordinates
pub mod entry {
    use crate::chaos_game::ChaosGame;
    use crate::rng::xoshiro_from_state;
    use crate::transform::Transform;
    use crate::variation::Variation;
    use glam::Vec2;
    use spirv_std::spirv;

    /// Given a set of fractal flame parameters, generate new IFS coordinates
    /// and store them in the output array.
    #[spirv(compute(entry_point_name = "main_chaos_game", threads(1)))]
    pub fn main_chaos_game(
        #[spirv(spec_constant(id = 1, default = 20))] iteration_discard: u32,
        #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] _rng_seed: &[u8],
        #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] transforms: &[Transform],
        #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] weights: &[f32],
        #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] variations: &[Variation],
        #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] output: &mut [Vec2],
    ) {
        let rng_seed_actual = [0u8; 32];

        let mut rng = xoshiro_from_state(rng_seed_actual);
        let mut chaos_game = ChaosGame::new(&mut rng, transforms, weights, variations);

        for _ in 0..iteration_discard {
            chaos_game.next().unwrap();
        }

        for i in 0..output.len() {
            output[i] = chaos_game.next().unwrap();
        }
    }
}