path: root/dungeon/src/wfc.rs

//! The `wfc` module contains the implementation of the
//! Wave Function Collapse algorithm for procedural generation of the `Floor`.

// TODO: Add pattern size input

// Reference map's map size and tile count constants
use crate::map::TILE_COUNT;

use crate::map::Tile;
use crate::pos::Direction;
use crate::pos::Pos;

use rand::prelude::IndexedRandom;

/// The `State` struct represents each possible state of a tile in the WFC algorithm.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub(crate) struct State {
	/// Position of the tile
	#[expect(dead_code)]
	pos: Pos,
	/// Possible tiles for this state
	possible_tiles: Vec<Tile>,
	/// Entropy of the state
	entropy: usize,
}
impl State {
	/// Creates a new State instance
	pub fn new(pos: Pos, possible_tiles: Vec<Tile>) -> Self {
		let entropy = possible_tiles.len();
		Self {
			pos,
			possible_tiles,
			entropy,
		}
	}

	/// Updates the possible tiles and recalculates entropy
	pub fn set_possible_tiles(&mut self, possible_tiles: Vec<Tile>) {
		self.possible_tiles = possible_tiles;
		self.entropy = self.entropy();
	}

	/// Calculates the entropy of the state
	/// TODO: Implement shannon entropy calculation
	pub fn entropy(&self) -> usize {
		self.possible_tiles.len()
	}

	/// Checks if the state has been collapsed to a single tile
	pub fn is_collapsed(&self) -> bool {
		self.possible_tiles.len() == 1
	}

	/// Returns a reference to the possible tiles
	pub fn possible_tiles(&self) -> &Vec<Tile> {
		&self.possible_tiles
	}

	/// Collapses the state to a single tile
	pub fn collapse(&mut self, tile: Tile) {
		self.possible_tiles = vec![tile];
		self.entropy = 1;
	}
}

/// The `Compatability` struct represents the compatibility rules between tiles.
pub(crate) struct Compatability {
	/// The first tile in the compatibility pair
	tile1: Tile,
	/// The second tile in the compatibility pair
	tile2: Tile,
	/// The direction from tile1 -> tile2
	direction: Direction,
}

impl Compatability {
	/// Creates a new Compatability instance
	pub fn new(tile1: Tile, tile2: Tile, direction: Direction) -> Self {
		Self {
			tile1,
			tile2,
			direction,
		}
	}
}

/// The `CompatabilityChecker` struct manages compatibility rules between tiles.
pub(crate) struct CompatabilityChecker {
	compatabilities: Vec<Compatability>,
}

impl CompatabilityChecker {
	/// Creates a new CompatabilityChecker instance
	pub fn new(input: &[Vec<Tile>]) -> Self {
		let mut compatabilities = Vec::new();
		// Which pairs of tiles can be placed next to each other and in which directions
		for row in input {
			for window in row.windows(2) {
				if let [tile1, tile2] = window {
					compatabilities.push(Compatability::new(
						*tile1,
						*tile2,
						Direction::East,
					));
					compatabilities.push(Compatability::new(
						*tile2,
						*tile1,
						Direction::West,
					));
				}
			}
		}

		for col in 0..input[0].len() {
			for row in 0..input.len() - 1 {
				let tile1 = input[row][col];
				let tile2 = input[row + 1][col];
				compatabilities.push(Compatability::new(tile1, tile2, Direction::North));
				compatabilities.push(Compatability::new(tile2, tile1, Direction::South));
			}
		}

		Self { compatabilities }
	}

	/// Checks if two tiles are compatible in the given direction
	pub fn is_valid(&self, tile1: Tile, tile2: Tile, direction: Direction) -> bool {
		self.compatabilities
			.iter()
			.any(|c| c.tile1 == tile1 && c.tile2 == tile2 && c.direction == direction)
	}
}

/// The `Wfc` struct encapsulates the Wave Function Collapse algorithm.
pub struct Wfc {
	/// The compatibility checker for tile constraints
	compatabilities: CompatabilityChecker,
	/// The random number generator
	rng: rand::rngs::StdRng,
	/// The current state of the WFC algorithm (smart pointer for interior mutability)
	states: Vec<std::cell::RefCell<State>>,
	/// The States that have been collapsed (as positions)
	collapsed_states: Vec<Pos>,
	// Note: uses map's map size and tile count constants
}
impl Wfc {
	/// Creates a new Wfc instance with the given seed and map size
	pub fn new(input: &[Vec<Tile>], seed: u64) -> Self {
		let rng = rand::SeedableRng::seed_from_u64(seed);
		Self {
			compatabilities: CompatabilityChecker::new(input),
			rng,
			states: Vec::with_capacity(TILE_COUNT),
			collapsed_states: Vec::new(),
		}
	}

	/// Initializes the states vector with default states
	fn initialize_states(&mut self) {
		// Store a superposition of all possible tiles for each position
		for pos in Pos::values() {
			let possible_tiles = Tile::values().collect();
			let state = State::new(pos, possible_tiles);
			self.states.push(std::cell::RefCell::new(state));
		}
	}

	/// Find the state with the lowest entropy that has not been collapsed yet
	fn find_lowest_entropy_state(&self) -> Option<State> {
		self.states
			.iter()
			.filter_map(|state| {
				let borrowed_state = state.borrow();
				if !borrowed_state.is_collapsed() {
					Some(borrowed_state.clone())
				} else {
					None
				}
			})
			.min_by_key(State::entropy)
	}

	/// Collapses the given state by selecting one of its possible tiles (using the rng)
	fn collapse_state(&mut self, state: &State) {
		// Select a random tile from the possible tiles
		if let Some(&selected_tile) = state.possible_tiles().choose(&mut self.rng) {
			// Update the state to be collapsed with the selected tile
			if let Some(state_ref) = self.states.get(state.pos.idx()) {
				let mut state_mut = state_ref.borrow_mut();
				state_mut.collapse(selected_tile);
				self.collapsed_states.push(state.pos);
			}
		}
	}

	/// Propagates constraints after a state has been collapsed, starting from the given state
	fn propagate(&mut self, state: &State) {
		// Keep a queue of states to propagate constraints from
		let mut queue = vec![state.pos];

		while let Some(current_pos) = queue.pop() {
			// For each direction, check neighboring states
			for direction in Direction::values() {
				if let Some(neighbor_pos) = current_pos.step(direction)
					&& let Some(neighbor_state_ref) = self.states.get(neighbor_pos.idx())
				{
					let mut neighbor_state = neighbor_state_ref.borrow_mut();
					let valid_tiles: Vec<Tile> = neighbor_state
						.possible_tiles()
						.iter()
						.cloned()
						.filter(|&tile| {
							self.compatabilities.is_valid(
								state.possible_tiles()[0],
								tile,
								direction,
							)
						})
						.collect();

					if valid_tiles.len() < neighbor_state.possible_tiles().len() {
						neighbor_state.set_possible_tiles(valid_tiles);
						queue.push(neighbor_pos);
					}
				}
			}
		}
	}

	/// Runs the Wave Function Collapse algorithm to generate the map
	#[expect(clippy::unused_self)]
	pub(crate) fn run(&mut self) {
		// ----- Step 1: Read in input sample and define constraints -----
		// (This step is currently handled in the CompatabilityChecker::new method)

		// ----- Step 2: Initialize the grid of possible states -----
		self.initialize_states();

		// Keep running until all states are collapsed
		while self.collapsed_states.len() != TILE_COUNT {
			// ----- Step 3: Collapse the superpositions -----
			// Find the shannon entropy of all superpositions
			if let Some(state) = self.find_lowest_entropy_state() {
				// Collapse the tile with the lowest entropy
				// (Assign a seeded random superposition)
				self.collapse_state(&state);

				// ----- Step 4: Propagate -----
				// Update neighboring superpositions based on the collapsed tile
				// (Remove any that violate constraints)
				self.propagate(&state);
			}

			// ----- Step 5: Repeat until all superpositions are collapsed -----
			// If any tiles are left with no possible states, restart
			// TODO: backtrack
		}
	}

	/// Returns the generated tiles as a boxed array
	pub fn to_tiles(&self) -> Box<[Tile; TILE_COUNT]> {
		let mut arr = [Tile::Air; TILE_COUNT];

		for (i, state_ref) in self.states.iter().enumerate() {
			let state = state_ref.borrow();
			if state.is_collapsed() {
				arr[i] = state.possible_tiles()[0];
			}
		}

		Box::new(arr)
	}
}

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

	use crate::map::MAP_SIZE_USIZE;

	#[test]
	fn test_wfc_generation() {
		// 16x16 input
		// 0 is wall, 1 is air
		let input = vec![
			vec![0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
			vec![0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0],
			vec![0, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0],
			vec![0, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0],
			vec![0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
			vec![1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1],
			vec![0, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0],
			vec![0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0],
			vec![0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
			vec![0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
			vec![0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 0, 0],
			vec![1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1, 1],
			vec![0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
			vec![0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 0, 0],
			vec![0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0],
			vec![0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
		];

		// Convert input to Tiles
		let input_tiles: Vec<Vec<Tile>> = input
			.iter()
			.map(|row| {
				row.iter()
					.map(|&cell| if cell == 0 { Tile::Wall } else { Tile::Air })
					.collect()
			})
			.collect();

		let seed = 12345;
		let mut wfc = Wfc::new(&input_tiles, seed);
		wfc.run();
		let tiles = wfc.to_tiles();

		// Print the generated tiles
		for y in 0..MAP_SIZE_USIZE {
			for x in 0..MAP_SIZE_USIZE {
				// Print air as "██", walls as "  " and stairs as ">"
				let tile = tiles[y * MAP_SIZE_USIZE + x];
				let symbol = match tile {
					Tile::Air => "██",
					Tile::Wall => "  ",
					Tile::Stairs => ">",
				};
				print!("{symbol}");
			}
			println!();
		}

		// Check that the tiles array has the correct length
		assert_eq!(tiles.len(), TILE_COUNT);

		// Check that there is at least one Air tile
		assert!(tiles.contains(&Tile::Air));
	}
}