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

#[derive(Reflect, Resource)]
#[reflect(Resource)]
struct ModelConfig {
    num_boids: usize,
    x_bound: f32,
    y_bound: f32,
    max_velocity: f32,
    protected_range: f32,
    visible_range: f32,
    alignment_factor: f32,
    cohesion_factor: f32,
    separation_factor: f32,
}

impl Default for ModelConfig {
    fn default() -> Self {
        Self {
            num_boids: 50,
            x_bound: 0.7,
            y_bound: 0.7,
            max_velocity: 200.0,
            protected_range: 50.0,
            visible_range: 150.0,
            alignment_factor: 0.5,
            cohesion_factor: 0.5,
            separation_factor: 200.0,
        }
    }
}

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()
        }))
        .init_resource::<ModelConfig>()
        .register_type::<ModelConfig>()
        .add_systems(Startup, setup)
        .add_systems(FixedUpdate, update_boids)
        .run();
}

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

fn setup(
    model_config: Res<ModelConfig>,
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<ColorMaterial>>,
) {
    let ModelConfig {
        num_boids,
        max_velocity,
        ..
    } = *model_config;

    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) * 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(
    model_config: Res<ModelConfig>,
    window: Single<&Window, With<bevy::window::PrimaryWindow>>,
    time: Res<Time<Fixed>>,
    mut query: Query<(Entity, &mut Boid, &mut Transform)>,
) {
    let dt = time.delta_secs();
    let x_bound = window.width() * model_config.x_bound;
    let y_bound = window.height() * model_config.y_bound;

    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 delta = transform.translation.xy() - other_transform.translation.xy();
            let distance = delta.length();
            if distance < model_config.protected_range {
                if distance > f32::EPSILON {
                    let direction = delta / distance;
                    let strength = 1.0 - distance / model_config.protected_range;
                    separation += direction * strength;
                }
            } else if distance < model_config.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 out_of_bounds_acceleration = Vec2::new(0.0, 0.0);
        if transform.translation.x < -x_bound / 2.0 {
            out_of_bounds_acceleration.x = 1.0;
        } else if transform.translation.x > x_bound / 2.0 {
            out_of_bounds_acceleration.x = -1.0;
        }

        if transform.translation.y < -y_bound / 2.0 {
            out_of_bounds_acceleration.y = 1.0;
        } else if transform.translation.y > y_bound / 2.0 {
            out_of_bounds_acceleration.y = -1.0;
        }
        let out_of_bounds_acceleration = out_of_bounds_acceleration.normalize_or_zero();

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

        let boundary_force = out_of_bounds_acceleration * model_config.max_velocity;
        let separation_force = model_update.separation * model_config.separation_factor;
        let alignment_force = model_update.alignment * model_config.alignment_factor;
        let cohesion_force = model_update.cohesion * model_config.cohesion_factor;

        let acceleration = boundary_force + separation_force + alignment_force + cohesion_force;
        boid.velocity += acceleration * dt;
        boid.velocity = boid.velocity.clamp_length_max(model_config.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));
    }
}