#![cfg_attr(target_arch = "spirv", no_std)]

use glam::{Vec4, Vec4Swizzles, vec2, vec4};
use rand::Rng;
use rand::distributions::Standard;
use rand::prelude::Distribution;
use rand_xoshiro::Xoshiro128Plus;
use spirv_std::spirv;
use spirv_std::num_traits::Float;

#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
#[repr(C)]
pub struct ImageConstants {
    accum_width: i32,
    accum_height: i32,
    viewport_width: i32,
    viewport_height: i32,
    background_color: Vec4,
}

impl ImageConstants {
    pub fn new(
        accum_width: i32,
        accum_height: i32,
        viewport_width: i32,
        viewport_height: i32,
        background_color: Vec4,
    ) -> Self {
        Self {
            accum_width,
            accum_height,
            viewport_width,
            viewport_height,
            background_color,
        }
    }

    pub fn viewport_dimensions(&self) -> glam::IVec2 {
        glam::ivec2(self.viewport_width, self.viewport_height)
    }

    pub fn with_accumulate(&mut self, width: i32, height: i32) {
        self.accum_width = width;
        self.accum_height = height;
    }

    pub fn with_viewport(&mut self, width: i32, height: i32) {
        self.viewport_width = width;
        self.viewport_height = height;
    }
}

#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
#[repr(C)]
pub struct CameraConstants {
    scale: f32,
    zoom: f32,
    rotate: f32,
    offset_x: f32,
    offset_y: f32,
}

impl CameraConstants {
    pub fn camera(&self, image_width: i32, image_height: i32) -> glam::Affine2 {
        let zoom_rotate_offset = glam::Affine2::from_scale_angle_translation(
            glam::Vec2::splat(2f32.powf(self.zoom)),
            self.rotate.to_radians(),
            -vec2(self.offset_x, self.offset_y),
        );
        let ifs_to_pixel = glam::Affine2::from_scale_angle_translation(
            glam::Vec2::splat(self.scale),
            0.0,
            vec2(image_width as f32 / 2.0, image_height as f32 / 2.0),
        );

        ifs_to_pixel * zoom_rotate_offset
    }
}

pub trait Color {
    type Element;
    const BLACK: Self;

    const WHITE: Self;
}

impl Color for Vec4 {
    type Element = f32;

    const BLACK: Self = vec4(0., 0., 0., 1.);
    const WHITE: Self = vec4(1., 1., 1., 1.);
}

pub(crate) fn image_index(pixel_x: i32, pixel_y: i32, image_width: i32) -> i32 {
    pixel_x + pixel_y * image_width
}

pub struct BiUnit;

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

#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
#[repr(C)]
pub struct ThreadState {
    rng_state: [u32; 4],
    point: glam::Vec2,
}

impl ThreadState {
    fn xoshiro_to_state(rng: &Xoshiro128Plus) -> [u32; 4] {
        unsafe { core::mem::transmute::<_, [u32; 4]>(rng.clone()) }
    }

    fn state_to_xoshiro(rng_state: &[u32; 4]) -> Xoshiro128Plus {
        unsafe { core::mem::transmute::<_, Xoshiro128Plus>(*rng_state) }
    }

    pub fn new(rng: &mut Xoshiro128Plus) -> Self {
        let point = vec2(rng.sample(BiUnit), rng.sample(BiUnit));
        let rng_state = Self::xoshiro_to_state(rng);
        rng.jump();
        Self { rng_state, point }
    }

    pub fn get_rng(&self) -> Xoshiro128Plus {
        Self::state_to_xoshiro(&self.rng_state)
    }

    pub fn set_rng(&mut self, rng: &Xoshiro128Plus) {
        self.rng_state = Self::xoshiro_to_state(rng)
    }
}

#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
#[repr(C)]
pub struct Coefs {
    a: f32,
    b: f32,
    c: f32,
    d: f32,
    e: f32,
    f: f32,
}

impl Coefs {
    pub fn new(a: f32, b: f32, c: f32, d: f32, e: f32, f: f32) -> Self {
        Self { a, b, c, d, e, f }
    }

    pub fn apply(&self, point: glam::Vec2) -> glam::Vec2 {
        vec2(
            self.a * point.x + self.b * point.y + self.c,
            self.d * point.x + self.e * point.y + self.f,
        )
    }
}

#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
#[repr(C)]
pub struct Transform {
    coefs: Coefs,
    variation_offset: u32,
    variation_count: u32,
    weight: f32,
}

impl Transform {
    pub fn new(coefs: Coefs, variation_offset: u32, variation_count: u32, weight: f32) -> Self {
        Self {
            coefs,
            variation_offset,
            variation_count,
            weight,
        }
    }

    pub fn apply(&self, variations: &[Variation], point: glam::Vec2) -> glam::Vec2 {
        let point = self.coefs.apply(point);
        let mut variation_point = glam::Vec2::ZERO;
        let variation_begin = self.variation_offset as usize;
        let variation_end = (self.variation_offset + self.variation_count) as usize;
        for i in variation_begin..variation_end {
            variation_point += variations[i].apply(point);
        }
        variation_point
    }
}

#[derive(Copy, Clone, Debug)]
#[repr(u32)]
pub enum VariationKind {
    Linear = 0,
}

// UNSAFE: Sound because enum has guaranteed layout (u32) and defined zero-value
unsafe impl bytemuck::Zeroable for VariationKind {}
// UNSAFE: Sound because enum has guaranteed layout (u32) and defined zero-value
unsafe impl bytemuck::Pod for VariationKind {}

#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
#[repr(C)]
pub struct Variation {
    kind: VariationKind,
    weight: f32,
}

impl Variation {
    pub fn new(kind: VariationKind, weight: f32) -> Self {
        Self { kind, weight }
    }

    pub fn apply(&self, point: glam::Vec2) -> glam::Vec2 {
        (match self.kind {
            VariationKind::Linear => point,
        }) * self.weight
    }
}

fn next_transform<R: Rng + ?Sized>(
    rng: &mut R,
    total_weight: f32,
    transforms: &[Transform],
) -> Transform {
    let mut weight = rng.sample::<f32, _>(Standard) * total_weight;
    for i in 0..transforms.len() {
        weight -= transforms[i].weight;
        if weight <= 0.0 {
            return transforms[i];
        }
    }
    unreachable!()
}

#[spirv(compute(threads(1)))]
pub fn main_cs(
    #[spirv(push_constant)] image_constants: &ImageConstants,
    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] thread_state: &mut [ThreadState],
    #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] transforms: &[Transform],
    #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] variations: &[Variation],
    #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] accum_image: &mut [Vec4],
) {
    let mut rng = thread_state[0].get_rng();
    let mut point = thread_state[0].point;

    let mut transform_weight = 0f32;
    for i in 0..transforms.len() {
        transform_weight += transforms[i].weight;
    }

    // Fuse 20, should be provided by a uniform in the future
    for _i in 0..20 {
        point = next_transform(&mut rng, transform_weight, transforms).apply(variations, point);
    }

    // ...because `<i32>.max(<i32>)` has compilation issues.
    let max_dimension = if image_constants.accum_width > image_constants.accum_height { image_constants.accum_width } else { image_constants.accum_height };

    // Fixed camera, should be provided by a uniform in the future
    let camera = CameraConstants {
        scale: max_dimension as f32 / 4.0,
        zoom: 0.0,
        rotate: 0.0,
        offset_x: 0.0,
        offset_y: 0.0,
    }
    .camera(image_constants.accum_width, image_constants.accum_height);

    // Iterate 100,000, should be provided by a uniform in the future
    for _i in 0..100_000 {
        point = next_transform(&mut rng, transform_weight, transforms).apply(variations, point);

        let pixel_coordinates = camera.transform_point2(point).as_ivec2();
        if pixel_coordinates.x < 0
            || pixel_coordinates.x >= image_constants.accum_width
            || pixel_coordinates.y < 0
            || pixel_coordinates.y >= image_constants.accum_height
        {
            continue;
        }

        let ii = image_index(
            pixel_coordinates.x,
            pixel_coordinates.y,
            image_constants.accum_height,
        );
        accum_image[ii as usize] = Color::WHITE;
    }
}

#[spirv(vertex)]
pub fn main_vs(
    #[spirv(vertex_index)] vert_id: u32,
    #[spirv(position, invariant)] position: &mut Vec4,
) {
    // Create a "quad" that fills the viewport for the fragment shader.
    // The `draw` call issued by the main application will be for three vertex ID's (0, 1, 2).
    // This code maps them to the points (-1, -1), (3, -1), and (-1, 3) respectively.
    // Because the interior of that triangle covers the entire viewport,
    // the GPU clips to the viewport and invokes the fragment shader for each pixel.
    // https://stackoverflow.com/a/59739538
    // https://www.saschawillems.de/blog/2016/08/13/vulkan-tutorial-on-rendering-a-fullscreen-quad-without-buffers/
    let output_uv = vec2(((vert_id << 1) & 2) as f32, (vert_id & 2) as f32);
    *position = (output_uv * 2.0 - 1.0, 0.0, 1.0).into();
}

#[spirv(fragment)]
pub fn main_fs(
    #[spirv(frag_coord)] frag_coord: Vec4,
    #[spirv(push_constant)] ifs_constants: &ImageConstants,
    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] accum_image: &[Vec4],
    output: &mut Vec4,
) {
    // Bootleg texture sampling; map from viewport image pixel coordinates to accumulator image
    // pixel coordinates
    let viewport_coordinate = frag_coord.xy().as_uvec2();

    let a_width = ifs_constants.accum_width as f32;
    let a_height = ifs_constants.accum_height as f32;
    let v_width = ifs_constants.viewport_width as f32;
    let v_height = ifs_constants.viewport_height as f32;

    // Scale both width and height by the same factor; preserves aspect ratio
    let scale_width = a_width / v_width;
    let scale_height = a_height / v_height;
    let scale = scale_width.max(scale_height);

    // Re-center the image in the viewport after scale
    let offset_x = (v_width * scale - a_width) / 2.0;
    let offset_y = (v_height * scale - a_height) / 2.0;

    let accum_coordinate = viewport_coordinate.as_vec2() * scale - vec2(offset_x, offset_y);

    if accum_coordinate.x < 0.0
        || accum_coordinate.x >= ifs_constants.accum_width as f32
        || accum_coordinate.y < 0.0
        || accum_coordinate.y >= ifs_constants.accum_height as f32
    {
        *output = ifs_constants.background_color;
    } else {
        *output = accum_image[image_index(
            accum_coordinate.x as i32,
            accum_coordinate.y as i32,
            ifs_constants.accum_width,
        ) as usize];
    }
}