color_grading.rs

//! Demonstrates color grading with an interactive adjustment UI.

use std::{
    f32::consts::PI,
    fmt::{self, Formatter},
};

use bevy::{
    camera::Hdr,
    light::CascadeShadowConfigBuilder,
    prelude::*,
    render::view::{ColorGrading, ColorGradingGlobal, ColorGradingSection},
};
use std::fmt::Display;

static FONT_PATH: &str = "fonts/FiraMono-Medium.ttf";

/// How quickly the value changes per frame.
const OPTION_ADJUSTMENT_SPEED: f32 = 0.003;

/// The color grading section that the user has selected: highlights, midtones,
/// or shadows.
#[derive(Clone, Copy, PartialEq)]
enum SelectedColorGradingSection {
    Highlights,
    Midtones,
    Shadows,
}

/// The global option that the user has selected.
///
/// See the documentation of [`ColorGradingGlobal`] for more information about
/// each field here.
#[derive(Clone, Copy, PartialEq, Default)]
enum SelectedGlobalColorGradingOption {
    #[default]
    Exposure,
    Temperature,
    Tint,
    Hue,
}

/// The section-specific option that the user has selected.
///
/// See the documentation of [`ColorGradingSection`] for more information about
/// each field here.
#[derive(Clone, Copy, PartialEq)]
enum SelectedSectionColorGradingOption {
    Saturation,
    Contrast,
    Gamma,
    Gain,
    Lift,
}

/// The color grading option that the user has selected.
#[derive(Clone, Copy, PartialEq, Resource)]
enum SelectedColorGradingOption {
    /// The user has selected a global color grading option: one that applies to
    /// the whole image as opposed to specifically to highlights, midtones, or
    /// shadows.
    Global(SelectedGlobalColorGradingOption),

    /// The user has selected a color grading option that applies only to
    /// highlights, midtones, or shadows.
    Section(
        SelectedColorGradingSection,
        SelectedSectionColorGradingOption,
    ),
}

impl Default for SelectedColorGradingOption {
    fn default() -> Self {
        Self::Global(default())
    }
}

/// Buttons consist of three parts: the button itself, a label child, and a
/// value child. This specifies one of the three entities.
#[derive(Clone, Copy, PartialEq, Component)]
enum ColorGradingOptionWidgetType {
    /// The parent button.
    Button,
    /// The label of the button.
    Label,
    /// The numerical value that the button displays.
    Value,
}

#[derive(Clone, Copy, Component)]
struct ColorGradingOptionWidget {
    widget_type: ColorGradingOptionWidgetType,
    option: SelectedColorGradingOption,
}

/// A marker component for the help text at the top left of the screen.
#[derive(Clone, Copy, Component)]
struct HelpText;

fn main() {
    App::new()
        .add_plugins(DefaultPlugins)
        .init_resource::<SelectedColorGradingOption>()
        .add_systems(Startup, setup)
        .add_systems(
            Update,
            (
                handle_button_presses,
                adjust_color_grading_option,
                update_ui_state,
            )
                .chain(),
        )
        .run();
}

fn setup(
    mut commands: Commands,
    currently_selected_option: Res<SelectedColorGradingOption>,
    asset_server: Res<AssetServer>,
) {
    // Create the scene.
    add_basic_scene(&mut commands, &asset_server);

    // Create the root UI element.
    let font = asset_server.load(FONT_PATH);
    let color_grading = ColorGrading::default();
    add_buttons(&mut commands, &font, &color_grading);

    // Spawn help text.
    add_help_text(&mut commands, &font, &currently_selected_option);

    // Spawn the camera.
    add_camera(&mut commands, &asset_server, color_grading);
}

/// Adds all the buttons on the bottom of the scene.
fn add_buttons(commands: &mut Commands, font: &Handle<Font>, color_grading: &ColorGrading) {
    commands.spawn((
        // Spawn the parent node that contains all the buttons.
        Node {
            flex_direction: FlexDirection::Column,
            position_type: PositionType::Absolute,
            row_gap: px(6),
            left: px(12),
            bottom: px(12),
            ..default()
        },
        children![
            // Create the first row, which contains the global controls.
            buttons_for_global_controls(color_grading, font),
            // Create the rows for individual controls.
            buttons_for_section(SelectedColorGradingSection::Highlights, color_grading, font),
            buttons_for_section(SelectedColorGradingSection::Midtones, color_grading, font),
            buttons_for_section(SelectedColorGradingSection::Shadows, color_grading, font),
        ],
    ));
}

/// Adds the buttons for the global controls (those that control the scene as a
/// whole as opposed to shadows, midtones, or highlights).
fn buttons_for_global_controls(color_grading: &ColorGrading, font: &Handle<Font>) -> impl Bundle {
    let make_button = |option: SelectedGlobalColorGradingOption| {
        button_for_value(
            SelectedColorGradingOption::Global(option),
            color_grading,
            font,
        )
    };

    // Add the parent node for the row.
    (
        Node::default(),
        children![
            Node {
                width: px(125),
                ..default()
            },
            make_button(SelectedGlobalColorGradingOption::Exposure),
            make_button(SelectedGlobalColorGradingOption::Temperature),
            make_button(SelectedGlobalColorGradingOption::Tint),
            make_button(SelectedGlobalColorGradingOption::Hue),
        ],
    )
}

/// Adds the buttons that control color grading for individual sections
/// (highlights, midtones, shadows).
fn buttons_for_section(
    section: SelectedColorGradingSection,
    color_grading: &ColorGrading,
    font: &Handle<Font>,
) -> impl Bundle {
    let make_button = |option| {
        button_for_value(
            SelectedColorGradingOption::Section(section, option),
            color_grading,
            font,
        )
    };

    // Spawn the row container.
    (
        Node {
            align_items: AlignItems::Center,
            ..default()
        },
        children![
            // Spawn the label ("Highlights", etc.)
            (
                text(&section.to_string(), font, Color::WHITE),
                Node {
                    width: px(125),
                    ..default()
                }
            ),
            // Spawn the buttons.
            make_button(SelectedSectionColorGradingOption::Saturation),
            make_button(SelectedSectionColorGradingOption::Contrast),
            make_button(SelectedSectionColorGradingOption::Gamma),
            make_button(SelectedSectionColorGradingOption::Gain),
            make_button(SelectedSectionColorGradingOption::Lift),
        ],
    )
}

/// Adds a button that controls one of the color grading values.
fn button_for_value(
    option: SelectedColorGradingOption,
    color_grading: &ColorGrading,
    font: &Handle<Font>,
) -> impl Bundle {
    let label = match option {
        SelectedColorGradingOption::Global(option) => option.to_string(),
        SelectedColorGradingOption::Section(_, option) => option.to_string(),
    };

    // Add the button node.
    (
        Button,
        Node {
            border: UiRect::all(px(1)),
            width: px(200),
            justify_content: JustifyContent::Center,
            align_items: AlignItems::Center,
            padding: UiRect::axes(px(12), px(6)),
            margin: UiRect::right(px(12)),
            border_radius: BorderRadius::MAX,
            ..default()
        },
        BorderColor::all(Color::WHITE),
        BackgroundColor(Color::BLACK),
        ColorGradingOptionWidget {
            widget_type: ColorGradingOptionWidgetType::Button,
            option,
        },
        children![
            // Add the button label.
            (
                text(&label, font, Color::WHITE),
                ColorGradingOptionWidget {
                    widget_type: ColorGradingOptionWidgetType::Label,
                    option,
                },
            ),
            // Add a spacer.
            Node {
                flex_grow: 1.0,
                ..default()
            },
            // Add the value text.
            (
                text(
                    &format!("{:.3}", option.get(color_grading)),
                    font,
                    Color::WHITE,
                ),
                ColorGradingOptionWidget {
                    widget_type: ColorGradingOptionWidgetType::Value,
                    option,
                },
            ),
        ],
    )
}

/// Creates the help text at the top of the screen.
fn add_help_text(
    commands: &mut Commands,
    font: &Handle<Font>,
    currently_selected_option: &SelectedColorGradingOption,
) {
    commands.spawn((
        Text::new(create_help_text(currently_selected_option)),
        TextFont {
            font: FontSource::from(font),
            ..default()
        },
        Node {
            position_type: PositionType::Absolute,
            left: px(12),
            top: px(12),
            ..default()
        },
        HelpText,
    ));
}

/// Adds some text to the scene.
fn text(label: &str, font: &Handle<Font>, color: Color) -> impl Bundle + use<> {
    (
        Text::new(label),
        TextFont {
            font: font.into(),
            font_size: FontSize::Px(15.0),
            ..default()
        },
        TextColor(color),
    )
}

fn add_camera(commands: &mut Commands, asset_server: &AssetServer, color_grading: ColorGrading) {
    commands.spawn((
        Camera3d::default(),
        Hdr,
        Transform::from_xyz(0.7, 0.7, 1.0).looking_at(Vec3::new(0.0, 0.3, 0.0), Vec3::Y),
        color_grading,
        DistanceFog {
            color: Color::srgb_u8(43, 44, 47),
            falloff: FogFalloff::Linear {
                start: 1.0,
                end: 8.0,
            },
            ..default()
        },
        EnvironmentMapLight {
            diffuse_map: asset_server.load("environment_maps/pisa_diffuse_rgb9e5_zstd.ktx2"),
            specular_map: asset_server.load("environment_maps/pisa_specular_rgb9e5_zstd.ktx2"),
            intensity: 2000.0,
            ..default()
        },
    ));
}

fn add_basic_scene(commands: &mut Commands, asset_server: &AssetServer) {
    // Spawn the main scene.
    commands.spawn(SceneRoot(asset_server.load(
        GltfAssetLabel::Scene(0).from_asset("models/TonemappingTest/TonemappingTest.gltf"),
    )));

    // Spawn the flight helmet.
    commands.spawn((
        SceneRoot(
            asset_server
                .load(GltfAssetLabel::Scene(0).from_asset("models/FlightHelmet/FlightHelmet.gltf")),
        ),
        Transform::from_xyz(0.5, 0.0, -0.5).with_rotation(Quat::from_rotation_y(-0.15 * PI)),
    ));

    // Spawn the light.
    commands.spawn((
        DirectionalLight {
            illuminance: 15000.0,
            shadow_maps_enabled: true,
            ..default()
        },
        Transform::from_rotation(Quat::from_euler(EulerRot::ZYX, 0.0, PI * -0.15, PI * -0.15)),
        CascadeShadowConfigBuilder {
            maximum_distance: 3.0,
            first_cascade_far_bound: 0.9,
            ..default()
        }
        .build(),
    ));
}

impl Display for SelectedGlobalColorGradingOption {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let name = match *self {
            SelectedGlobalColorGradingOption::Exposure => "Exposure",
            SelectedGlobalColorGradingOption::Temperature => "Temperature",
            SelectedGlobalColorGradingOption::Tint => "Tint",
            SelectedGlobalColorGradingOption::Hue => "Hue",
        };
        f.write_str(name)
    }
}

impl Display for SelectedColorGradingSection {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let name = match *self {
            SelectedColorGradingSection::Highlights => "Highlights",
            SelectedColorGradingSection::Midtones => "Midtones",
            SelectedColorGradingSection::Shadows => "Shadows",
        };
        f.write_str(name)
    }
}

impl Display for SelectedSectionColorGradingOption {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let name = match *self {
            SelectedSectionColorGradingOption::Saturation => "Saturation",
            SelectedSectionColorGradingOption::Contrast => "Contrast",
            SelectedSectionColorGradingOption::Gamma => "Gamma",
            SelectedSectionColorGradingOption::Gain => "Gain",
            SelectedSectionColorGradingOption::Lift => "Lift",
        };
        f.write_str(name)
    }
}

impl Display for SelectedColorGradingOption {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            SelectedColorGradingOption::Global(option) => write!(f, "\"{option}\""),
            SelectedColorGradingOption::Section(section, option) => {
                write!(f, "\"{option}\" for \"{section}\"")
            }
        }
    }
}

impl SelectedSectionColorGradingOption {
    /// Returns the appropriate value in the given color grading section.
    fn get(&self, section: &ColorGradingSection) -> f32 {
        match *self {
            SelectedSectionColorGradingOption::Saturation => section.saturation,
            SelectedSectionColorGradingOption::Contrast => section.contrast,
            SelectedSectionColorGradingOption::Gamma => section.gamma,
            SelectedSectionColorGradingOption::Gain => section.gain,
            SelectedSectionColorGradingOption::Lift => section.lift,
        }
    }

    fn set(&self, section: &mut ColorGradingSection, value: f32) {
        match *self {
            SelectedSectionColorGradingOption::Saturation => section.saturation = value,
            SelectedSectionColorGradingOption::Contrast => section.contrast = value,
            SelectedSectionColorGradingOption::Gamma => section.gamma = value,
            SelectedSectionColorGradingOption::Gain => section.gain = value,
            SelectedSectionColorGradingOption::Lift => section.lift = value,
        }
    }
}

impl SelectedGlobalColorGradingOption {
    /// Returns the appropriate value in the given set of global color grading
    /// values.
    fn get(&self, global: &ColorGradingGlobal) -> f32 {
        match *self {
            SelectedGlobalColorGradingOption::Exposure => global.exposure,
            SelectedGlobalColorGradingOption::Temperature => global.temperature,
            SelectedGlobalColorGradingOption::Tint => global.tint,
            SelectedGlobalColorGradingOption::Hue => global.hue,
        }
    }

    /// Sets the appropriate value in the given set of global color grading
    /// values.
    fn set(&self, global: &mut ColorGradingGlobal, value: f32) {
        match *self {
            SelectedGlobalColorGradingOption::Exposure => global.exposure = value,
            SelectedGlobalColorGradingOption::Temperature => global.temperature = value,
            SelectedGlobalColorGradingOption::Tint => global.tint = value,
            SelectedGlobalColorGradingOption::Hue => global.hue = value,
        }
    }
}

impl SelectedColorGradingOption {
    /// Returns the appropriate value in the given set of color grading values.
    fn get(&self, color_grading: &ColorGrading) -> f32 {
        match self {
            SelectedColorGradingOption::Global(option) => option.get(&color_grading.global),
            SelectedColorGradingOption::Section(
                SelectedColorGradingSection::Highlights,
                option,
            ) => option.get(&color_grading.highlights),
            SelectedColorGradingOption::Section(SelectedColorGradingSection::Midtones, option) => {
                option.get(&color_grading.midtones)
            }
            SelectedColorGradingOption::Section(SelectedColorGradingSection::Shadows, option) => {
                option.get(&color_grading.shadows)
            }
        }
    }

    /// Sets the appropriate value in the given set of color grading values.
    fn set(&self, color_grading: &mut ColorGrading, value: f32) {
        match self {
            SelectedColorGradingOption::Global(option) => {
                option.set(&mut color_grading.global, value);
            }
            SelectedColorGradingOption::Section(
                SelectedColorGradingSection::Highlights,
                option,
            ) => option.set(&mut color_grading.highlights, value),
            SelectedColorGradingOption::Section(SelectedColorGradingSection::Midtones, option) => {
                option.set(&mut color_grading.midtones, value);
            }
            SelectedColorGradingOption::Section(SelectedColorGradingSection::Shadows, option) => {
                option.set(&mut color_grading.shadows, value);
            }
        }
    }
}

/// Handles mouse clicks on the buttons when the user clicks on a new one.
fn handle_button_presses(
    mut interactions: Query<(&Interaction, &ColorGradingOptionWidget), Changed<Interaction>>,
    mut currently_selected_option: ResMut<SelectedColorGradingOption>,
) {
    for (interaction, widget) in interactions.iter_mut() {
        if widget.widget_type == ColorGradingOptionWidgetType::Button
            && *interaction == Interaction::Pressed
        {
            *currently_selected_option = widget.option;
        }
    }
}

/// Updates the state of the UI based on the current state.
fn update_ui_state(
    mut buttons: Query<(
        &mut BackgroundColor,
        &mut BorderColor,
        &ColorGradingOptionWidget,
    )>,
    button_text: Query<(Entity, &ColorGradingOptionWidget), (With<Text>, Without<HelpText>)>,
    help_text: Single<Entity, With<HelpText>>,
    mut writer: TextUiWriter,
    cameras: Single<Ref<ColorGrading>>,
    currently_selected_option: Res<SelectedColorGradingOption>,
) {
    // Exit early if the UI didn't change
    if !currently_selected_option.is_changed() && !cameras.is_changed() {
        return;
    }

    // The currently-selected option is drawn with inverted colors.
    for (mut background, mut border_color, widget) in buttons.iter_mut() {
        if *currently_selected_option == widget.option {
            *background = Color::WHITE.into();
            *border_color = Color::BLACK.into();
        } else {
            *background = Color::BLACK.into();
            *border_color = Color::WHITE.into();
        }
    }

    let value_label = format!("{:.3}", currently_selected_option.get(cameras.as_ref()));

    // Update the buttons.
    for (entity, widget) in button_text.iter() {
        // Set the text color.

        let color = if *currently_selected_option == widget.option {
            Color::BLACK
        } else {
            Color::WHITE
        };

        writer.for_each_color(entity, |mut text_color| {
            text_color.0 = color;
        });

        // Update the displayed value, if this is the currently-selected option.
        if widget.widget_type == ColorGradingOptionWidgetType::Value
            && *currently_selected_option == widget.option
        {
            writer.for_each_text(entity, |mut text| {
                text.clone_from(&value_label);
            });
        }
    }

    // Update the help text.
    *writer.text(*help_text, 0) = create_help_text(&currently_selected_option);
}

/// Creates the help text at the top left of the window.
fn create_help_text(currently_selected_option: &SelectedColorGradingOption) -> String {
    format!("Press Left/Right to adjust {currently_selected_option}")
}

/// Processes keyboard input to change the value of the currently-selected color
/// grading option.
fn adjust_color_grading_option(
    mut color_grading: Single<&mut ColorGrading>,
    input: Res<ButtonInput<KeyCode>>,
    currently_selected_option: Res<SelectedColorGradingOption>,
) {
    let mut delta = 0.0;
    if input.pressed(KeyCode::ArrowLeft) {
        delta -= OPTION_ADJUSTMENT_SPEED;
    }
    if input.pressed(KeyCode::ArrowRight) {
        delta += OPTION_ADJUSTMENT_SPEED;
    }

    if delta != 0.0 {
        let new_value = currently_selected_option.get(color_grading.as_ref()) + delta;
        currently_selected_option.set(&mut color_grading, new_value);
    }
}