mirage-engine 0.1.1

//! The sky a frame draws where nothing else was drawn, and the light every
//! surface takes from it.

use core::f32::consts::{PI, TAU};
use core::fmt;
use core::fmt::Debug;
use core::hash::Hash;

use crate::light::rgb;
use crate::math::{UVec2, Vec3, Vec4};
use crate::{Assets, Catalog, Color, TextureData};

/// The color of the sky a frame that sets none is drawn and lit by: one
/// neutral dark grey, the whole way around.
const DEFAULT: Color = Color::rgb(0.1, 0.1, 0.1);

/// The equirect a gradient is drawn into at startup. A gradient changes
/// smoothly, so one this small holds it.
const RASTERIZED: UVec2 = UVec2::new(64, 32);

/// The widest a mip may be for the nine coefficients to be taken from it,
/// in texels.
const COEFFICIENT_WIDTH: u32 = 256;

/// How far under the horizon a ground color reaches its whole share: the
/// `y` of a direction three degrees under level, so the horizon reads as a
/// line and not a step.
const GROUND_BLEND: f32 = 0.052_336;

/// What each coefficient's shape is scaled by, taken twice because a
/// fragment applies none of it; in the order a fragment reads them back.
const SHAPES: [f32; 9] = [
    0.282_095, 0.488_603, 0.488_603, 0.488_603, 1.092_548, 1.092_548, 0.315_392, 1.092_548,
    0.546_274,
];

/// The share of the light from every direction each coefficient takes,
/// over `pi`, so a sky of one color lands exactly that color on a surface
/// facing anywhere.
const TAKEN: [f32; 9] = [
    1.0,
    2.0 / 3.0,
    2.0 / 3.0,
    2.0 / 3.0,
    0.25,
    0.25,
    0.25,
    0.25,
    0.25,
];

/// Turns a game's value into a skybox.
///
/// Values are cache keys: equal values must build the same sky. [`Catalog`]
/// proves every named asset loads before the first frame, and a startup
/// error names a skybox by its [`Debug`].
pub trait Skyboxes: Catalog + Clone + Debug + Eq + Hash {
    /// Builds this value's sky, the first time a frame is drawn by it.
    ///
    /// Must not read a file or the network; loaded data is in `assets`.
    fn build(&self, assets: &Assets) -> SkyboxData;
}

/// The vocabulary of a game with no sky of its own to name.
///
/// No value of it exists, so such a game draws and is lit by the default
/// sky.
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub enum NoSkyboxes {}

impl Catalog for NoSkyboxes {
    fn catalog() -> Vec<Self> {
        Vec::new()
    }
}

impl Skyboxes for NoSkyboxes {
    fn build(&self, _assets: &Assets) -> SkyboxData {
        match *self {}
    }
}

/// The sky one value of a game's vocabulary is drawn and lit from: one of
/// two kinds, an image of the whole sky or a gradient between three colors,
/// the fraction of its own light it lands on a surface, and the ground
/// color under its horizon where a game states one.
#[derive(Clone, Debug)]
pub struct SkyboxData {
    kind: Kind,
    light: f32,
    ground: Option<Color>,
}

/// The closed set of skies; [`SkyboxData`] has one constructor for each.
#[derive(Clone, Debug)]
enum Kind {
    Equirect(TextureData),
    Gradient(Gradient),
}

impl SkyboxData {
    /// The sky `image` holds: the whole way around across it, and zenith to
    /// nadir down it, so it is twice as wide as it is tall.
    ///
    /// Required if you want a sky an image holds. An image not twice as wide
    /// as it is tall stops startup naming the skybox.
    pub fn equirect(image: TextureData) -> Self {
        Self::of(Kind::Equirect(image))
    }

    /// The sky `zenith` straight up, `horizon` level, and `nadir` straight
    /// down, each texel between them the share its own elevation takes.
    ///
    /// Required if you want a sky with no image at all: one color in all
    /// three lands that color on every surface, from every direction.
    pub fn gradient(zenith: Color, horizon: Color, nadir: Color) -> Self {
        Self::of(Kind::Gradient(Gradient::new(zenith, horizon, nadir)))
    }

    /// The same sky landing `fraction` of its own light on every surface,
    /// and a surface that reflects it reflecting that share. The sky the
    /// frame draws is unchanged.
    ///
    /// Required if you want the frame's own lights to show under a bright
    /// sky. `1.0` where never called, held at zero and above; past `1.0`
    /// lands more light than the sky holds.
    pub fn lit_by(mut self, fraction: f32) -> Self {
        self.light = fraction.max(0.0);
        self
    }

    /// The same sky reading as `color` in every direction under the horizon,
    /// blended over the three degrees below it: what a metal reflects there
    /// and what the sky lands on a surface facing down take that color, and
    /// so does the sky the frame draws where nothing covers it.
    ///
    /// Required if you want a metal to reflect the floor it stands on rather
    /// than the sky's own underside. No ground where never called.
    pub fn with_ground(mut self, color: Color) -> Self {
        self.ground = Some(color);
        self
    }

    fn of(kind: Kind) -> Self {
        Self {
            kind,
            light: 1.0,
            ground: None,
        }
    }

    /// This sky as the engine keeps it for the run, or the error its image
    /// is no sky for.
    pub(crate) fn resident(&self) -> Result<Resident, SkyboxError> {
        let largest = match &self.kind {
            Kind::Equirect(image) => Mip::of(image)?,
            Kind::Gradient(gradient) => Mip::between(*gradient),
        };
        let largest = match self.ground {
            Some(ground) => largest.over(rgb(ground)),
            None => largest,
        };

        Ok(Resident::of(largest).lit_by(self.light))
    }
}

impl Default for SkyboxData {
    /// The sky a name no source holds returns: black the whole way around.
    fn default() -> Self {
        Self::equirect(TextureData::rgba8(UVec2::new(2, 1), vec![0; 8]))
    }
}

/// The three colors a gradient sky is drawn between.
#[derive(Clone, Copy, Debug, PartialEq)]
pub(crate) struct Gradient {
    /// Straight up.
    zenith: Color,
    /// Level with the horizon.
    horizon: Color,
    /// Straight down.
    nadir: Color,
}

impl Gradient {
    /// The sky `zenith` straight up, `horizon` level, and `nadir` straight
    /// down.
    pub(crate) const fn new(zenith: Color, horizon: Color, nadir: Color) -> Self {
        Self {
            zenith,
            horizon,
            nadir,
        }
    }
}

impl Default for Gradient {
    /// The sky a frame that sets none is drawn and lit by: one neutral dark
    /// grey, the whole way around.
    fn default() -> Self {
        Self::new(DEFAULT, DEFAULT, DEFAULT)
    }
}

/// Why an image is no sky at all, which startup reports under the name of
/// the skybox value it was built for.
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub(crate) enum SkyboxError {
    /// It is not twice as wide as it is tall.
    Sides { width: u32, height: u32 },
    /// It holds another count of bytes than the four its every texel needs.
    Pixels {
        width: u32,
        height: u32,
        bytes: usize,
    },
}

impl fmt::Display for SkyboxError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Sides { width, height } => write!(
                f,
                "is {width}x{height}; a skybox image is twice as wide as it is tall"
            ),
            Self::Pixels {
                width,
                height,
                bytes,
            } => write!(
                f,
                "is {width}x{height} over {bytes} bytes, where every texel needs four"
            ),
        }
    }
}

/// One sky as the engine keeps it: the mips it is drawn and reflected from,
/// largest first, and the light it lands on a surface.
pub(crate) struct Resident {
    largest: Mip,
    /// Every halving of `largest`, down to the one that is a single texel.
    halved: Vec<Mip>,
    irradiance: Irradiance,
    /// What a surface's reflection of this sky is scaled by.
    light: f32,
}

impl Resident {
    /// The sky `gradient` holds, drawn into an equirect of its own.
    pub(crate) fn gradient(gradient: Gradient) -> Self {
        Self::of(Mip::between(gradient))
    }

    /// Halves `largest` down to a single texel, and takes the nine
    /// coefficients from the first mip no wider than [`COEFFICIENT_WIDTH`].
    fn of(largest: Mip) -> Self {
        let mut halved: Vec<Mip> = Vec::new();
        while let Some(mip) = halved.last().unwrap_or(&largest).halved() {
            halved.push(mip);
        }
        // What a surface takes changes smoothly over the sky, so a small
        // copy of it holds all nine.
        let narrow = halved.iter().fold(&largest, |narrowest, mip| {
            match narrowest.size.x <= COEFFICIENT_WIDTH {
                true => narrowest,
                false => mip,
            }
        });

        Self {
            irradiance: Irradiance::of(narrow),
            largest,
            halved,
            light: 1.0,
        }
    }

    /// The same sky landing `fraction` of its light: the coefficients are
    /// scaled here, and the frame's own values hold the fraction for the
    /// reflection.
    fn lit_by(mut self, fraction: f32) -> Self {
        self.irradiance = self.irradiance.scaled(fraction);
        self.light = fraction;
        self
    }

    /// What a surface's reflection of this sky is scaled by; the
    /// coefficients already hold it.
    pub(crate) fn share(&self) -> f32 {
        self.light
    }

    /// The size of its largest mip, in texels.
    pub(crate) fn size(&self) -> UVec2 {
        self.largest.size
    }

    /// Every mip of this sky, largest first.
    pub(crate) fn mips(&self) -> impl Iterator<Item = &Mip> {
        core::iter::once(&self.largest).chain(&self.halved)
    }

    /// How many mips that is.
    pub(crate) fn mip_count(&self) -> u32 {
        1 + self.halved.len() as u32
    }

    /// The mip level a fully rough surface reflects this sky from: the
    /// smallest mip's, counting `0.0` at the largest.
    pub(crate) fn top_mip(&self) -> f32 {
        self.halved.len() as f32
    }

    /// The light this sky lands on a surface, by the direction the surface
    /// faces.
    pub(crate) fn irradiance(&self) -> Irradiance {
        self.irradiance
    }
}

/// One mip of a sky: its size, and its linear texels row by row from the
/// top left.
pub(crate) struct Mip {
    size: UVec2,
    texels: Vec<Vec3>,
}

impl Mip {
    /// The light `image` holds, or the error that image is no sky for.
    fn of(image: &TextureData) -> Result<Self, SkyboxError> {
        let size = image.size();
        let (width, height) = (size.x, size.y);
        if height == 0 || width != 2 * height {
            return Err(SkyboxError::Sides { width, height });
        }
        // A texture built by hand holds whatever bytes it was given, and
        // every mip under this one reads one texel per texel of its size.
        let bytes = image.pixels().len();
        if bytes as u64 != 4 * u64::from(width) * u64::from(height) {
            return Err(SkyboxError::Pixels {
                width,
                height,
                bytes,
            });
        }

        Ok(Self {
            size,
            texels: image
                .pixels()
                .chunks_exact(4)
                .map(|texel| rgb(Color::of_srgb([texel[0], texel[1], texel[2]])))
                .collect(),
        })
    }

    /// The sky `gradient` holds: each texel the share its own elevation
    /// takes of the color above it and the color below it.
    fn between(gradient: Gradient) -> Self {
        let zenith = rgb(gradient.zenith);
        let horizon = rgb(gradient.horizon);
        let nadir = rgb(gradient.nadir);

        Self {
            size: RASTERIZED,
            texels: Grid::of(RASTERIZED)
                .directions()
                .map(|direction| match direction.y >= 0.0 {
                    true => horizon.lerp(zenith, direction.y),
                    false => horizon.lerp(nadir, -direction.y),
                })
                .collect(),
        }
    }

    /// This mip with `ground` under the horizon, blended across
    /// [`GROUND_BLEND`] below it.
    fn over(self, ground: Vec3) -> Self {
        let texels = Grid::of(self.size)
            .directions()
            .zip(self.texels)
            .map(|(direction, texel)| {
                let share = (-direction.y / GROUND_BLEND).clamp(0.0, 1.0);
                texel.lerp(ground, share)
            })
            .collect();

        Self {
            size: self.size,
            texels,
        }
    }

    /// This mip at half its size, each texel the average of the four it
    /// covers; `None` for one that is a single texel.
    fn halved(&self) -> Option<Self> {
        if self.size.max_element() <= 1 {
            return None;
        }
        let size = (self.size / 2).max(UVec2::ONE);

        Some(Self {
            size,
            texels: (0..size.y)
                .flat_map(|down| (0..size.x).map(move |across| UVec2::new(across, down)))
                .map(|at| self.averaged(at))
                .collect(),
        })
    }

    /// This mip's size in texels.
    pub(crate) fn size(&self) -> UVec2 {
        self.size
    }

    /// Its texels, row by row from the top left.
    pub(crate) fn texels(&self) -> &[Vec3] {
        &self.texels
    }

    /// The average of the four texels this mip holds under texel `at` of
    /// the mip half its size.
    fn averaged(&self, at: UVec2) -> Vec3 {
        let corner = at * 2;
        let covered: Vec3 = [UVec2::ZERO, UVec2::X, UVec2::Y, UVec2::ONE]
            .into_iter()
            .map(|step| self.texel(corner + step))
            .sum();

        covered / 4.0
    }

    /// The texel `at`, held within this mip.
    fn texel(&self, at: UVec2) -> Vec3 {
        let held = at.min(self.size - UVec2::ONE);

        self.texels[(held.y * self.size.x + held.x) as usize]
    }
}

/// The texel grid of one equirect image: the direction each texel holds,
/// and how much of the sky it covers.
///
/// `forward.wgsl` reads a direction back through the same mapping.
#[derive(Clone, Copy)]
struct Grid {
    size: UVec2,
}

impl Grid {
    fn of(size: UVec2) -> Self {
        Self { size }
    }

    /// Each texel's direction, row by row from the top left.
    fn directions(self) -> impl Iterator<Item = Vec3> {
        self.rows().flat_map(Row::directions)
    }

    /// Each texel's direction and the fraction of the whole sky it covers,
    /// row by row from the top left.
    fn texels(self) -> impl Iterator<Item = (Vec3, f32)> {
        self.rows().flat_map(|row| {
            row.directions()
                .map(move |direction| (direction, row.covered))
        })
    }

    /// The rows from the top, each at its own angle from the zenith.
    fn rows(self) -> impl Iterator<Item = Row> {
        let size = self.size.as_vec2();
        let texel = (PI / size.y) * (TAU / size.x);

        (0..self.size.y).map(move |down| {
            let latitude = (down as f32 + 0.5) / size.y * PI;
            let (latitude_sin, latitude_cos) = latitude.sin_cos();
            Row {
                latitude_sin,
                latitude_cos,
                columns: self.size.x,
                covered: latitude_sin * texel,
            }
        })
    }
}

/// One row of a [`Grid`]: its angle from the zenith, and how much of the
/// sky each of its texels covers.
#[derive(Clone, Copy)]
struct Row {
    latitude_sin: f32,
    latitude_cos: f32,
    columns: u32,
    covered: f32,
}

impl Row {
    /// Each texel's direction, from the left.
    fn directions(self) -> impl Iterator<Item = Vec3> {
        let columns = self.columns as f32;

        (0..self.columns).map(move |across| {
            let longitude = ((across as f32 + 0.5) / columns - 0.5) * TAU;
            let (longitude_sin, longitude_cos) = longitude.sin_cos();
            Vec3::new(
                self.latitude_sin * longitude_sin,
                self.latitude_cos,
                -self.latitude_sin * longitude_cos,
            )
        })
    }
}

/// The light a sky lands on a surface, as nine coefficients a fragment
/// reads through the direction the surface faces.
///
/// A coefficient's shape is its own factor at a direction `n`: `1`, `n.y`,
/// `n.z`, `n.x`, `n.x * n.y`, `n.y * n.z`, `3 * n.z * n.z - 1`, `n.x * n.z`
/// and `n.x * n.x - n.y * n.y`, in that order. The coefficients hold the
/// sky's own light with the share a flat surface takes of each direction
/// already scaled into them, so what a fragment reads is what the surface
/// takes.
#[derive(Clone, Copy, Debug)]
pub(crate) struct Irradiance([Vec3; 9]);

impl Irradiance {
    /// The nine coefficients of `mip`.
    fn of(mip: &Mip) -> Self {
        let mut coefficients = [Vec3::ZERO; 9];
        for (texel, (direction, covered)) in mip.texels().iter().zip(Grid::of(mip.size).texels()) {
            let light = *texel * covered;
            for (coefficient, shape) in coefficients.iter_mut().zip(Self::shapes(direction)) {
                *coefficient += light * shape;
            }
        }

        Self(core::array::from_fn(|at| {
            coefficients[at] * TAKEN[at] * SHAPES[at] * SHAPES[at]
        }))
    }

    /// The same light, `fraction` of it.
    fn scaled(self, fraction: f32) -> Self {
        Self(self.0.map(|coefficient| coefficient * fraction))
    }

    /// The nine coefficients as the frame's own values hold them, each in
    /// the first three lanes of a `Vec4` of its own.
    pub(crate) fn lanes(&self) -> [Vec4; 9] {
        self.0.map(|coefficient| coefficient.extend(0.0))
    }

    /// The shape of each coefficient at `direction`, in the order a
    /// fragment reads them back.
    fn shapes(direction: Vec3) -> [f32; 9] {
        let Vec3 { x, y, z } = direction;

        [
            1.0,
            y,
            z,
            x,
            x * y,
            y * z,
            3.0 * z * z - 1.0,
            x * z,
            x * x - y * y,
        ]
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// A sky of one color, as a gradient and as an image of that color.
    fn flat(color: Color) -> (Resident, Resident) {
        let encoded = |linear: f32| {
            let encoded = match linear <= 0.003_130_8 {
                true => linear * 12.92,
                false => 1.055 * linear.powf(1.0 / 2.4) - 0.055,
            };
            (encoded * 255.0).round() as u8
        };
        let texel = [
            encoded(color.red),
            encoded(color.green),
            encoded(color.blue),
            u8::MAX,
        ];
        let image = TextureData::rgba8(UVec2::new(16, 8), texel.repeat(16 * 8));

        (
            Resident::gradient(Gradient::new(color, color, color)),
            SkyboxData::equirect(image)
                .resident()
                .expect("that image is a sky"),
        )
    }

    #[test]
    fn a_sky_of_one_color_lands_that_color_on_a_surface_facing_anywhere() {
        let color = Color::rgb(0.25, 0.5, 0.75);
        let (gradient, image) = flat(color);

        for sky in [gradient, image] {
            let coefficients = sky.irradiance().0;

            for (read, held) in [coefficients[0].x, coefficients[0].y, coefficients[0].z]
                .into_iter()
                .zip([color.red, color.green, color.blue])
            {
                assert!(
                    (read - held).abs() < 0.01,
                    "the first coefficient holds the color itself, got {read} against {held}"
                );
            }
            for coefficient in &coefficients[1..] {
                assert!(
                    coefficient.length() < 0.01,
                    "and no other holds anything, got {coefficient}"
                );
            }
        }
    }

    #[test]
    fn a_sky_bright_above_lands_more_light_on_a_surface_facing_up() {
        let coefficients =
            Resident::gradient(Gradient::new(Color::WHITE, Color::BLACK, Color::BLACK))
                .irradiance()
                .0;

        assert!(
            coefficients[1].y > 0.0,
            "the coefficient a normal reads through its own `+Y` rises with the \
             light above, got {}",
            coefficients[1].y
        );
        assert!(
            coefficients[0].y > 0.0,
            "and the sky lands light on a surface facing anywhere"
        );
    }

    #[test]
    fn an_image_not_twice_as_wide_as_it_is_tall_is_no_sky_at_all() {
        let square = TextureData::rgba8(UVec2::splat(4), vec![u8::MAX; 4 * 4 * 4]);

        assert_eq!(
            SkyboxData::equirect(square).resident().err(),
            Some(SkyboxError::Sides {
                width: 4,
                height: 4
            })
        );
        assert_eq!(
            SkyboxData::equirect(TextureData::default())
                .resident()
                .err(),
            Some(SkyboxError::Sides {
                width: 0,
                height: 0
            }),
            "and an image with no texels at all is none either"
        );
    }

    #[test]
    fn a_sky_is_halved_down_to_one_texel() {
        let sky = Resident::gradient(Gradient::default());
        let sizes: Vec<UVec2> = sky.mips().map(Mip::size).collect();

        assert_eq!(sizes.first(), Some(&RASTERIZED));
        assert_eq!(sizes.last(), Some(&UVec2::ONE));
        assert_eq!(sizes.len(), 7, "one mip per halving of the widest side");
        assert_eq!(sky.mip_count(), 7);
        assert_eq!(sky.top_mip(), 6.0, "the smallest of them is the last");
        assert!(
            sky.mips()
                .all(|mip| mip.texels().len() == (mip.size().x * mip.size().y) as usize),
            "each of them holds its own texels"
        );
    }

    #[test]
    fn the_smallest_mip_holds_the_average_of_the_sky() {
        let sky = Resident::gradient(Gradient::new(Color::WHITE, Color::BLACK, Color::BLACK));
        let smallest = sky.mips().last().expect("the mips end at one texel");

        assert_eq!(smallest.size(), UVec2::ONE);
        let average = smallest.texels()[0].x;
        assert!(
            (0.1..0.4).contains(&average),
            "a sky white above and black below averages between them, got {average}"
        );
    }

    #[test]
    fn a_ground_colors_every_texel_under_the_horizon_and_none_above_it() {
        let sky = SkyboxData::gradient(Color::WHITE, Color::WHITE, Color::WHITE)
            .with_ground(Color::rgb(0.5, 0.0, 0.0))
            .resident()
            .expect("a gradient is a sky");
        let largest = sky.mips().next().expect("a sky has a largest mip");
        let size = largest.size();
        let top = largest.texel(UVec2::new(size.x / 2, 0));
        let bottom = largest.texel(UVec2::new(size.x / 2, size.y - 1));

        assert!(
            top.abs_diff_eq(Vec3::ONE, 0.01),
            "the zenith keeps the sky's own color, got {top}"
        );
        assert!(
            bottom.abs_diff_eq(Vec3::new(0.5, 0.0, 0.0), 0.01),
            "and the nadir reads as the ground, got {bottom}"
        );
    }

    #[test]
    fn a_ground_darker_than_the_sky_lands_less_light_on_a_surface_facing_down() {
        let coefficients = SkyboxData::gradient(Color::WHITE, Color::WHITE, Color::WHITE)
            .with_ground(Color::BLACK)
            .resident()
            .expect("a gradient is a sky")
            .irradiance()
            .0;

        assert!(
            coefficients[1].x > 0.0,
            "the coefficient a normal reads through its own `+Y` rises with the \
             light above and not below, got {}",
            coefficients[1].x
        );
    }
}