traaains/sim/train/

systems.rs

use bevy::{
    ecs::{
        entity::Entity,
        query::{With, Without},
        system::{Commands, Query, Res},
    },
    log::tracing,
    math::Quat,
    time::Time,
    transform::components::Transform,
};
use bezier::Bezier;

use crate::{
    sim::{
        Acceleration, Controls, GRAVITY, STEEL_SLIP, Velocity,
        track::{
            node::{ConnectionLocation, TrackNode},
            pos::TrackPos,
            segment::TrackSegment,
        },
        train::{BRAKING_COEFF, Bogie, Carriage, Train, commands::RemoveTrain},
    },
    utils::Dir,
};

/// Updates trains acceleration based on its current velocity, power and mass
pub fn update_trains_acceleration(
    q_trains: Query<(&mut Acceleration, &Velocity, &Train, &Controls)>,
) {
    const ZERO_TOLERANCE: f32 = 1e-6;

    for (mut acc, vel, train, controls) in q_trains {
        let scaled_power = train.total_power * controls.throttle;
        let mut power_force = if vel.0.abs() < ZERO_TOLERANCE {
            if controls.throttle == 0.0 {
                0.0
            } else {
                f32::MAX
            }
        } else {
            scaled_power / vel.0
        };
        power_force = power_force.clamp_magnitude(GRAVITY * train.total_mass * STEEL_SLIP);

        acc.0 = (power_force) / train.total_mass;
    }
}

/// System to apply brakes to trains. Little bit of a hack, since brakes would in reality slow down the wheels which would in turn
/// deliver less force/acceleration to the train.
pub fn apply_brakes(q_trains: Query<(&mut Velocity, &Controls)>) {
    for (mut vel, controls) in q_trains {
        vel.0 += (-vel.0.signum() * controls.brake * BRAKING_COEFF).clamp_magnitude(vel.0);
    }
}

fn move_track_pos_by_meters(i: f32, m: f32, track_pos: &mut TrackPos, segment: &TrackSegment) {
    let delta = segment.shape.displace_by(i, m);
    track_pos.idx = match track_pos.dir.as_ref().unwrap() {
        Dir::Positive => delta,
        Dir::Negative => -delta,
    };
}

/// Returns true if the bogie was successfully handled (not derailed)
fn handle_bogie(
    q_bogies: &mut Query<&mut TrackPos>,
    q_segments: Query<&TrackSegment>,
    meters: f32,
    bogie: Entity,
) -> bool {
    if let Ok(mut bogie) = q_bogies.get_mut(bogie)
        && let Ok(segment) = q_segments.get(bogie.segment)
    {
        move_track_pos_by_meters(bogie.idx, meters, &mut bogie, segment);
        true
    } else {
        false
    }
}

/// System that moves carriage's bogies with respect to the carriage velocity
pub fn move_trains_with_velocity(
    mut cmd: Commands,
    q_trains: Query<(Entity, &Train, &Velocity)>,
    q_carriages: Query<&Carriage>,
    mut q_bogies: Query<&mut TrackPos>,
    q_segments: Query<&TrackSegment>,
    time: Res<Time>,
) {
    for (train_entity, train, vel) in q_trains {
        let meters = vel.0 * time.delta_secs();

        let mut derailed = false;
        for carriage_entity in train.carriages.iter() {
            let carriage = q_carriages.get(*carriage_entity).unwrap();

            // if any of these calls to q_bogies fails, then they derailed and we want to:
            // - skip the rest of the bogies
            // - despawn this train
            match &carriage.bogies {
                super::CarriageBogies::Single(bogie) => {
                    if !handle_bogie(&mut q_bogies, q_segments, meters, *bogie) {
                        derailed = true;
                        break;
                    }
                }
                super::CarriageBogies::Double { front, back } => {
                    if !handle_bogie(&mut q_bogies, q_segments, meters, *front) {
                        derailed = true;
                        break;
                    }

                    if !handle_bogie(&mut q_bogies, q_segments, meters, *back) {
                        derailed = true;
                        break;
                    }
                }
            }
        }

        if derailed {
            cmd.queue(RemoveTrain::new(train_entity));
        }
    }
}

/// System that moves bogies to different track segments if they exceed the segment's domain
pub fn move_on_rail_to_diff_segment(
    mut cmd: Commands,
    q_onrail: Query<(Entity, &mut TrackPos)>,
    q_segments: Query<&TrackSegment>,
    q_nodes: Query<&TrackNode>,
) {
    for (on_rail_entity, mut on_rail_pos) in q_onrail {
        // if guard for only trains that have left their respective segments
        if (0.0..=1.0).contains(&on_rail_pos.idx) {
            continue;
        }

        tracing::debug!("moving {}", on_rail_entity);

        let Ok(segment) = q_segments.get(on_rail_pos.segment) else {
            continue;
        };

        let final_node = q_nodes
            .get(match on_rail_pos.dir {
                Some(Dir::Positive) => segment.to_node,
                Some(Dir::Negative) => segment.from_node,
                None => panic!("On rail entities must always have direction set!"),
            })
            .unwrap();

        let Some(other_connection) = final_node.connections.get_active_for(on_rail_pos.segment)
        else {
            // curr segment isnt active in the node's conn, derail
            tracing::info!("Current segment is not a part of an active connection, derailing...");
            cmd.entity(on_rail_entity).remove::<TrackPos>();
            continue;
        };

        tracing::debug!("  to active conn: {:?}", other_connection);

        // account for the fact that different curves have different lengths and thus ratios will have different meaning
        let starting_idx = other_connection.location.starting_index();
        let starting_dir = other_connection.location.starting_dir();

        let add_offset = if on_rail_pos.idx > 1.0 {
            on_rail_pos.idx - 1.0
        } else {
            -on_rail_pos.idx
        };

        // we will cheat here and say that this is good enough for the length calculcation since the add_offset should be really small anyway
        let add_meters = segment.shape.len() * add_offset;
        let add_ratio = q_segments
            .get(other_connection.segment)
            .unwrap()
            .shape
            .displace_by(starting_idx, add_meters);

        tracing::debug!("additional meters: {}", add_meters);
        tracing::debug!("additional ratio: {}", add_ratio);

        on_rail_pos.segment = other_connection.segment;
        on_rail_pos.idx = starting_idx
            + match other_connection.location {
                ConnectionLocation::Start => add_ratio,
                ConnectionLocation::End => -add_ratio,
            };
        on_rail_pos.dir = Some(starting_dir);
    }
}

/// System that moves carriages to their bogies, aligning them properly
pub fn move_carriages_to_bogies(
    q_carriages: Query<(&Carriage, &mut Transform), Without<Bogie>>,
    q_bogies: Query<&Transform, With<Bogie>>,
) {
    for (carriage, mut transform) in q_carriages {
        match &carriage.bogies {
            super::CarriageBogies::Single(bogie) => {
                *transform = *q_bogies.get(*bogie).unwrap();
            }
            super::CarriageBogies::Double { front, back } => {
                let front_trans = *q_bogies.get(*front).unwrap();
                let back_trans = *q_bogies.get(*back).unwrap();

                transform.translation = (front_trans.translation + back_trans.translation) / 2.0;
                let dir = front_trans.translation.truncate() - back_trans.translation.truncate();
                transform.rotation = Quat::from_rotation_z(f32::atan2(dir.y, dir.x));
            }
        }
    }
}