avian3d-sandbox/src/bike.rs

use avian3d::{
    math::{Scalar, FRAC_PI_2},
    prelude::{
        AngularDamping, Collider, ColliderDensity, CollisionLayers, ExternalForce, ExternalTorque,
        LinearDamping, RayCaster, RigidBody, SleepingDisabled, SpatialQueryFilter,
        TransformInterpolation,
    },
};
use bevy::{
    asset::Handle,
    ecs::{bundle::Bundle, entity::Entity},
    prelude::{
        children, AlphaMode, App, Assets, Capsule3d, Color, Commands, Component, Dir3, Mesh,
        Mesh3d, MeshMaterial3d, Name, Plugin, Reflect, ResMut, SpawnRelated, Sphere,
        StandardMaterial, Startup, Transform, Vec3, Visibility,
    },
};

use crate::physics::CatControllerState;

pub const SPRING_CONSTANT: Scalar = 40.0;
pub const DAMPING_CONSTANT: Scalar = 10.0;
pub const WHEEL_RADIUS: Scalar = 0.4;
pub const REST_DISTANCE: Scalar = 1.0 + WHEEL_RADIUS;
pub const FRICTION_COEFF: Scalar = 0.9;
pub const FRONT_ATTACH: Vec3 = Vec3::new(0.0, 0.0, -1.5);
pub const REAR_ATTACH: Vec3 = Vec3::new(0.0, 0.0, 1.5);

#[derive(Component)]
pub struct CyberBikeBody;

#[derive(Component, Clone, Copy, Debug)]
pub enum CyberWheel {
    Front,
    Rear,
}

#[derive(Component, Clone, Copy, Debug, Reflect, Default)]
pub struct WheelState {
    pub displacement: Scalar,
    pub force_normal: Scalar,
    pub sliding: bool,
    pub contact_point: Option<Vec3>,
    pub contact_normal: Option<Vec3>,
}

impl WheelState {
    pub fn reset(&mut self) {
        *self = Self::default();
    }
}

#[derive(Component, Clone, Copy, Debug, Reflect)]
#[require(WheelState)]
pub struct WheelConfig {
    pub attach: Vec3,
    pub rest_dist: Scalar,
    pub konstant: Scalar,
    pub damping: Scalar,
    pub friction: Scalar,
    pub radius: Scalar,
}

impl WheelConfig {
    fn new(
        attach: Vec3,
        rest_dist: Scalar,
        konstant: Scalar,
        damping: Scalar,
        friction: Scalar,
        radius: Scalar,
    ) -> Self {
        WheelConfig {
            attach,
            rest_dist,
            konstant,
            damping,
            friction,
            radius,
        }
    }
}

fn spawn_bike(
    mut commands: Commands,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<StandardMaterial>>,
) {
    let pos = Vec3::new(0.0, 5.0, 0.0);
    let xform = Transform::from_translation(pos); //.with_rotation(Quat::from_rotation_z(0.0));

    let body_collider =
        Collider::capsule_endpoints(0.5, Vec3::new(0.0, 0.0, -0.65), Vec3::new(0.0, 0.0, 0.8));

    let pbr_bundle = {
        let color = Color::srgb(0.7, 0.05, 0.7);
        let mut xform = Transform::default();
        xform.rotate_x(FRAC_PI_2);
        (
            Mesh3d(meshes.add(Capsule3d::new(0.5, 1.45))),
            xform,
            Visibility::Visible,
            MeshMaterial3d(materials.add(color)),
            TransformInterpolation,
        )
    };

    let mut body = commands.spawn((xform, Visibility::default()));
    body.insert((
        Name::new("body"),
        RigidBody::Dynamic,
        body_collider,
        CollisionLayers::from_bits(1, 1),
        SleepingDisabled,
        CyberBikeBody,
        CatControllerState::default(),
        ColliderDensity(1.6),
        AngularDamping(0.2),
        LinearDamping(0.1),
        ExternalForce::ZERO.with_persistence(false),
        ExternalTorque::ZERO.with_persistence(false),
    ));

    body.insert(spawn_children(body.id(), meshes, materials))
        .with_child(pbr_bundle);
}

fn spawn_children(
    parent: Entity,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<StandardMaterial>>,
) -> impl Bundle {
    let mesh: Mesh = Sphere::new(WHEEL_RADIUS).into();
    let wheel_material = StandardMaterial {
        base_color: Color::srgb(0.01, 0.01, 0.01),
        alpha_mode: AlphaMode::Opaque,
        perceptual_roughness: 0.5,
        ..Default::default()
    };

    let mesh = meshes.add(mesh);
    let material = materials.add(wheel_material);

    let front_rake = Vec3::new(0.0, -1.0, 0.0); // about 30 degrees
    let front_wheel_pos = FRONT_ATTACH + (front_rake * REST_DISTANCE);
    let rear_rake = Vec3::new(0.0, -1.0, 0.0);
    let rear_wheel_pos = REAR_ATTACH + (rear_rake * REST_DISTANCE);

    children![
        (
            RayCaster::new(FRONT_ATTACH, Dir3::new_unchecked(front_rake))
                .with_max_hits(1)
                .with_max_distance(REST_DISTANCE)
                .with_query_filter(SpatialQueryFilter::from_excluded_entities([parent])),
            CyberWheel::Front
        ),
        wheel_bundle(
            mesh.clone(),
            material.clone(),
            front_wheel_pos,
            WheelConfig::new(
                FRONT_ATTACH,
                REST_DISTANCE,
                SPRING_CONSTANT - 25.0,
                DAMPING_CONSTANT,
                FRICTION_COEFF,
                WHEEL_RADIUS,
            ),
            CyberWheel::Front,
        ),
        (
            RayCaster::new(REAR_ATTACH, Dir3::new_unchecked(rear_rake))
                .with_max_hits(1)
                .with_max_distance(REST_DISTANCE)
                .with_query_filter(SpatialQueryFilter::from_excluded_entities([parent])),
            CyberWheel::Rear,
        ),
        wheel_bundle(
            mesh,
            material,
            rear_wheel_pos,
            WheelConfig::new(
                REAR_ATTACH,
                REST_DISTANCE,
                SPRING_CONSTANT + 10.0,
                DAMPING_CONSTANT,
                FRICTION_COEFF,
                WHEEL_RADIUS,
            ),
            CyberWheel::Rear,
        ),
    ]
}

fn wheel_bundle(
    mesh: Handle<Mesh>,
    material: Handle<StandardMaterial>,
    position: Vec3,
    config: WheelConfig,
    wheel: CyberWheel,
) -> impl Bundle {
    let xform = Transform::from_translation(position);

    let name = match wheel {
        CyberWheel::Front => "front tire",
        CyberWheel::Rear => "rear tire",
    };

    (
        Name::new(name),
        config,
        Mesh3d(mesh),
        MeshMaterial3d(material),
        xform,
        Collider::sphere(WHEEL_RADIUS),
        ColliderDensity(0.5),
        CollisionLayers::NONE,
        TransformInterpolation,
        wheel,
    )
}

pub struct CyberBikePlugin;

impl Plugin for CyberBikePlugin {
    fn build(&self, app: &mut App) {
        app.register_type::<WheelConfig>();
        app.register_type::<WheelState>();
        app.add_systems(Startup, spawn_bike);
    }
}