use bevy::{platform::collections::HashMap, prelude::*};
use std::f32::consts::PI;

fn main() {
    App::new()
        .insert_resource(ClearColor(Color::srgb_u8(26, 26, 46)))
        .insert_resource(Time::<Fixed>::from_hz(60.0))
        .add_plugins(DefaultPlugins.set(WindowPlugin {
            primary_window: Some(Window {
                fit_canvas_to_parent: true,
                prevent_default_event_handling: false,
                canvas: Some("#bevy".into()),
                ..default()
            }),
            ..default()
        }))
        .add_systems(Startup, setup)
        .add_systems(FixedUpdate, update_boids)
        .run();
}

const BOID_MAX_VELOCITY: f32 = 200.0;

#[derive(Component)]
struct Boid {
    velocity: Vec2,
}

fn setup(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<ColorMaterial>>,
) {
    const NUM_BOIDS: usize = 50;

    commands.spawn(Camera2d);

    // Triangle pointing right
    let triangle_mesh = Mesh2d(meshes.add(Triangle2d::new(
        Vec2::new(10.0, 0.0),
        Vec2::new(-3.0, 3.0),
        Vec2::new(-3.0, -3.0),
    )));
    let triangle_material = MeshMaterial2d(materials.add(Color::WHITE));
    commands.spawn_batch((0..=NUM_BOIDS).map(move |n| {
        let angle = 2.0 * PI / (NUM_BOIDS as f32) * n as f32;
        (
            Boid {
                velocity: Vec2::from_angle(angle) * BOID_MAX_VELOCITY,
            },
            triangle_mesh.clone(),
            triangle_material.clone(),
            Transform::default().rotate(Quat::from_rotation_z(angle)),
        )
    }));
}

struct Model {
    separation: Vec2,
    alignment: Vec2,
    cohesion: Vec2,
}

fn update_boids(
    window: Single<&Window, With<bevy::window::PrimaryWindow>>,
    mut query: Query<(Entity, &mut Boid, &mut Transform)>,
    time: Res<Time<Fixed>>,
) {
    let dt = time.delta_secs();
    let x_bound = window.width() * 0.7;
    let y_bound = window.height() * 0.7;

    const PROTECTED_RANGE: f32 = 30.0;
    const VISIBLE_RANGE: f32 = 100.0;

    let mut model_updates = HashMap::new();
    for (entity, boid, transform) in query.iter() {
        let mut nb_neighbors: usize = 0;
        let mut separation = Vec2::new(0.0, 0.0);
        let mut avg_speed = Vec2::new(0.0, 0.0);
        let mut avg_pos = Vec2::new(0.0, 0.0);
        for (other_entity, other_boid, other_transform) in query.iter() {
            if entity == other_entity {
                continue;
            }
            let distance = transform
                .translation
                .xy()
                .distance(other_transform.translation.xy());

            if distance > 1.0 && distance < PROTECTED_RANGE {
                separation +=
                    (transform.translation.xy() - other_transform.translation.xy()) / distance;
            } else if distance < VISIBLE_RANGE {
                avg_speed += other_boid.velocity;
                avg_pos += other_transform.translation.xy();
                nb_neighbors += 1;
            }
        }
        if nb_neighbors > 0 {
            avg_speed /= nb_neighbors as f32;
            avg_pos /= nb_neighbors as f32;
        }
        model_updates.insert(
            entity,
            Model {
                separation,
                alignment: avg_speed - boid.velocity,
                cohesion: avg_pos - transform.translation.xy(),
            },
        );
    }

    for (entity, mut boid, mut transform) in query.iter_mut() {
        let mut acceleration = Vec2::new(0.0, 0.0);
        if transform.translation.x < -x_bound / 2.0 {
            acceleration.x = 1.0;
        } else if transform.translation.x > x_bound / 2.0 {
            acceleration.x = -1.0;
        }

        if transform.translation.y < -y_bound / 2.0 {
            acceleration.y = 1.0;
        } else if transform.translation.y > y_bound / 2.0 {
            acceleration.y = -1.0;
        }

        let out_of_bounds_acceleration = acceleration.normalize_or_zero();

        let model_update = model_updates.get(&entity).unwrap();

        // Update velocity and apply forces to position
        boid.velocity += (out_of_bounds_acceleration
            + model_update.separation.normalize_or_zero() * 5.
            + model_update.alignment.normalize_or_zero()
            + model_update.cohesion.normalize_or_zero())
            * BOID_MAX_VELOCITY
            * dt;
        boid.velocity = boid.velocity.clamp_length_max(BOID_MAX_VELOCITY);

        transform.translation.x += boid.velocity.x * dt;
        transform.translation.y += boid.velocity.y * dt;
        transform.rotation = Quat::from_rotation_z(boid.velocity.y.atan2(boid.velocity.x));
    }
}