use processing_glfw::GlfwContext;

use bevy::math::Vec3;

use processing::prelude::*;

use processing_render::render::command::DrawCommand;

const SPIN_SHADER: &str = r#"

@group(0) @binding(0) var<storage, read_write> rotation: array<f32>;

@group(0) @binding(1) var<uniform> params: vec4<f32>; // x = dt

fn quat_mul(a: vec4<f32>, b: vec4<f32>) -> vec4<f32> {

return vec4<f32>(

a.w * b.xyz + b.w * a.xyz + cross(a.xyz, b.xyz),

a.w * b.w - dot(a.xyz, b.xyz),

);

}

@compute @workgroup_size(64)

fn main(@builtin(global_invocation_id) gid: vec3<u32>) {

let i = gid.x;

let count = arrayLength(&rotation) / 4u;

if i >= count {

return;

}

let dt = params.x;

let q = vec4<f32>(

rotation[i * 4u + 0u],

rotation[i * 4u + 1u],

rotation[i * 4u + 2u],

rotation[i * 4u + 3u],

);

let half_angle = dt * 0.5;

let dq = vec4<f32>(0.0, sin(half_angle), 0.0, cos(half_angle));

let q_new = quat_mul(q, dq);

rotation[i * 4u + 0u] = q_new.x;

rotation[i * 4u + 1u] = q_new.y;

rotation[i * 4u + 2u] = q_new.z;

rotation[i * 4u + 3u] = q_new.w;

}

"#;

fn main() {

sketch().unwrap();

exit(0).unwrap();

}

fn sketch() -> error::Result<()> {

let mut glfw_ctx = GlfwContext::new(900, 700)?;

init(Config::default())?;

let surface = glfw_ctx.create_surface(900, 700)?;

let graphics = graphics_create(surface, 900, 700, TextureFormat::Rgba16Float)?;

graphics_mode_3d(graphics)?;

transform_set_position(graphics, Vec3::new(0.0, 6.0, 18.0))?;

transform_look_at(graphics, Vec3::new(0.0, 0.0, 0.0))?;

let _light =

light_create_directional(graphics, bevy::color::Color::srgb(0.9, 0.85, 0.8), 300.0)?;

let cube = geometry_box(0.6, 0.6, 0.6)?;

let capacity: u32 = 125;

let mut positions: Vec<f32> = Vec::with_capacity(capacity as usize * 3);

let mut rotations: Vec<f32> = Vec::with_capacity(capacity as usize * 4);

let mut scales: Vec<f32> = Vec::with_capacity(capacity as usize * 3);

for x in 0..5 {

for y in 0..5 {

for z in 0..5 {

positions.push((x as f32 - 2.0) * 1.6);

positions.push((y as f32 - 2.0) * 1.6);

positions.push((z as f32 - 2.0) * 1.6);

// identity quat

rotations.push(0.0);

rotations.push(0.0);

rotations.push(0.0);

rotations.push(1.0);

let s = 0.5 + ((x + y + z) as f32 * 0.06);

scales.push(s);

scales.push(s);

scales.push(s);

}

}

}

let position_attr = geometry_attribute_position();

let rotation_attr = geometry_attribute_rotation();

let scale_attr = geometry_attribute_scale();

let p = particles_create(capacity, vec![position_attr, rotation_attr, scale_attr])?;

let position_buf =

particles_buffer(p, position_attr)?.ok_or(error::ProcessingError::ParticlesNotFound)?;

let rotation_buf =

particles_buffer(p, rotation_attr)?.ok_or(error::ProcessingError::ParticlesNotFound)?;

let scale_buf =

particles_buffer(p, scale_attr)?.ok_or(error::ProcessingError::ParticlesNotFound)?;

buffer_write(

position_buf,

positions.iter().flat_map(|f| f.to_le_bytes()).collect(),

)?;

buffer_write(

rotation_buf,

rotations.iter().flat_map(|f| f.to_le_bytes()).collect(),

)?;

buffer_write(

scale_buf,

scales.iter().flat_map(|f| f.to_le_bytes()).collect(),

)?;

let pbr = material_create_pbr()?;

material_set(pbr, "roughness", shader_value::ShaderValue::Float(0.4))?;

let spin_shader = shader_create(SPIN_SHADER)?;

let spin = compute_create(spin_shader)?;

while glfw_ctx.poll_events() {

graphics_begin_draw(graphics)?;

graphics_record_command(

graphics,

DrawCommand::BackgroundColor(bevy::color::Color::srgb(0.06, 0.06, 0.08)),

)?;

graphics_record_command(

graphics,

DrawCommand::Fill(bevy::color::Color::srgb(0.9, 0.5, 0.3)),

)?;

graphics_record_command(graphics, DrawCommand::Material(pbr))?;

graphics_record_command(

graphics,

DrawCommand::Particles {

particles: p,

geometry: cube,

},

)?;

graphics_end_draw(graphics)?;

compute_set(

spin,

"params",

shader_value::ShaderValue::Float4([0.015, 0.0, 0.0, 0.0]),

)?;

particles_apply(p, spin)?;

}

Ok(())

}