el_diablo/src/level_generator.rs

use std::cmp::{max, min};
use std::ops::Range;

use petgraph::algo::min_spanning_tree;
use petgraph::data::*;
use petgraph::graph::Graph;
use petgraph::graph::UnGraph;
use rand::prelude::SliceRandom;
use rand::Rng;
use rand::rngs::ThreadRng;

use crate::artifacts::{Artifact, Chest, Potion};
use crate::level::{Level, StructureElement};
use crate::monster::{Monster, Orc, Rat};
use crate::position::Position;

const ROOMS_VERTICAL: usize = 7;
const ROOMS_HORIZONTAL: usize = 4;

const ROOM_WIDTH: usize = 7;
const ROOM_HEIGHT: usize = 6;

#[derive(PartialEq, Copy, Clone)]
enum RoomType {
    Start,
    End,
    StairUp,
    StairDown,
    BasicRoom,
    TreasureRoom,
    MonsterRoom,
    EmptyRoom,
}

#[derive(Copy, Clone)]
struct ConnectionInfo {
    offset: usize,
    distance: usize,
}

#[derive(Copy, Clone)]
struct Room {
    pub kind: RoomType,
    pub offset_x: usize,
    pub offset_y: usize,
    pub width: usize,
    pub height: usize,
    pub connection_down: Option<ConnectionInfo>,
    pub connection_right: Option<ConnectionInfo>,
}

impl Room {
    fn new() -> Self {
        Self {
            kind: RoomType::EmptyRoom,
            offset_x: 0,
            offset_y: 0,
            width: 0,
            height: 0,
            connection_down: None,
            connection_right: None,
        }
    }
    /// change the size and position of a room randomly within its bounds
    fn random(&mut self, rng: &mut ThreadRng) {
        let width = rng.gen_range(3..6);
        let height = rng.gen_range(3..5);
        self.width = width;
        self.height = height;
        self.offset_x = rng.gen_range(0..(ROOM_WIDTH - width));
        self.offset_y = rng.gen_range(0..(ROOM_HEIGHT - height));
    }
    fn get_x_range(&self) -> Range<usize> {
        self.offset_x..self.offset_x + self.width
    }
    fn get_y_range(&self) -> Range<usize> {
        self.offset_y..self.offset_y + self.height
    }
}

pub struct LevelGenerator {
    rooms: [[Room; ROOMS_HORIZONTAL]; ROOMS_VERTICAL],
    level: usize,
}

impl LevelGenerator {
    pub fn generate(level: usize, first: bool, last: bool) -> Self {
        let mut rng = rand::thread_rng();
        let mut rooms = [[Room::new(); ROOMS_HORIZONTAL]; ROOMS_VERTICAL];
        let mut graph = UnGraph::<(usize, usize), u16>::default();

        // trick the room_connectable function into failing on the first iteration
        rooms[0][0].kind = RoomType::BasicRoom;

        while !LevelGenerator::rooms_connectable(&rooms) {
            let mut room_types: Vec<RoomType> = Vec::with_capacity(ROOMS_HORIZONTAL * ROOMS_VERTICAL);
            // level 0 contains a start room, all others contain a stair up
            if first {
                room_types.push(RoomType::Start);
            } else {
                room_types.push(RoomType::StairUp);
            }
            // level 24 (the last) contains an end room, all others a stair down
            if last {
                room_types.push(RoomType::End);
            } else {
                room_types.push(RoomType::StairDown);
            }
            room_types.push(RoomType::MonsterRoom);
            room_types.push(RoomType::TreasureRoom);
            // generate a random set of rooms and shuffle them
            for _ in room_types.len()..ROOMS_HORIZONTAL * ROOMS_VERTICAL {
                match rng.gen_range(1..=100) {
                    // TODO tune room type distribution
                    1..=33 => { room_types.push(RoomType::EmptyRoom) }
                    34..=66 => { room_types.push(RoomType::TreasureRoom) }
                    67..=90 => { room_types.push(RoomType::MonsterRoom) }
                    _ => { room_types.push(RoomType::BasicRoom) }
                }
            }
            room_types.shuffle(&mut rng);

            graph.clear();
            // place the rooms in the array an add nodes to the graph for every non empty room
            for c in 0..ROOMS_VERTICAL {
                for r in 0..ROOMS_HORIZONTAL {
                    rooms[c][r].kind = room_types.pop().unwrap();
                    if rooms[c][r].kind != RoomType::EmptyRoom {
                        rooms[c][r].random(&mut rng);
                        graph.add_node((c, r));
                    }
                }
            }
        }

        // add edges to the graph connecting each room to all of its neighbours (max 4 of them)
        for c in 0..ROOMS_VERTICAL {
            for r in 0..ROOMS_HORIZONTAL {
                if rooms[c][r].kind == RoomType::EmptyRoom {
                    continue;
                }
                let src_index = graph.node_indices().find(|i| graph[*i] == (c, r)).unwrap();
                for r_1 in r + 1..ROOMS_HORIZONTAL {
                    if rooms[c][r_1].kind != RoomType::EmptyRoom {
                        let tgt_index = graph.node_indices().find(|i| graph[*i] == (c, r_1)).unwrap();
                        // todo use random weight for edge
                        graph.add_edge(src_index, tgt_index, 1);
                        break;
                    }
                }
                for c_1 in c + 1..ROOMS_VERTICAL {
                    if rooms[c_1][r].kind != RoomType::EmptyRoom {
                        let tgt_index = graph.node_indices().find(|i| graph[*i] == (c_1, r)).unwrap();
                        // todo use random weight for edge
                        graph.add_edge(src_index, tgt_index, 1);
                        break;
                    }
                }
            }
        }

        // calculate a minimum spanning tree
        let mst: Graph<(usize, usize), u16, petgraph::Undirected> = Graph::from_elements(min_spanning_tree(&graph));
        for edge in mst.raw_edges() {
            let src = mst[edge.source()];
            let tgt = mst[edge.target()];

            let src_room = rooms[src.0][src.1];
            let mut tgt_room = rooms[tgt.0][tgt.1];

            // cols are the same, either up or down
            if src.0 == tgt.0 {
                let range = LevelGenerator::range_overlap(src_room.get_x_range(), tgt_room.get_x_range());
                let position: usize;
                if range.is_empty() {
                    position = range.start;
                } else {
                    position = rng.gen_range(range);
                }
                if src.1 < tgt.1 {
                    // src to tgt
                    rooms[src.0][src.1].connection_down = Some(ConnectionInfo { offset: position, distance: tgt.1 - src.1 });
                } else {
                    // tgt to src
                    tgt_room.connection_down = Some(ConnectionInfo { offset: position, distance: src.1 - tgt.1 });
                }
            }
            // rows are the same, either left or right
            if src.1 == tgt.1 {
                let range = LevelGenerator::range_overlap(src_room.get_y_range(), tgt_room.get_y_range());
                let mut position: usize;
                if range.is_empty() {
                    position = range.start;
                } else {
                    position = rng.gen_range(range);
                }
                if src.1 == 0 && position == 0 {
                    position = 1;
                }
                if src.0 < tgt.0 {
                    // src to tgt
                    rooms[src.0][src.1].connection_right = Some(ConnectionInfo { offset: position, distance: tgt.0 - src.0 });
                } else {
                    // tgt to src
                    tgt_room.connection_right = Some(ConnectionInfo { offset: position, distance: src.1 - tgt.1 });
                }
            }
        }

        LevelGenerator {
            rooms,
            level,
        }
    }

    fn range_overlap(r1: Range<usize>, r2: Range<usize>) -> Range<usize> {
        max(r1.start, r2.start)..min(r1.end, r2.end)
    }

    /// Verifies that for a given matrix of rooms each room has at least one other room in the
    /// same row or column.
    fn rooms_connectable(rooms: &[[Room; ROOMS_HORIZONTAL]; ROOMS_VERTICAL]) -> bool {
        for c in 0..ROOMS_VERTICAL {
            for r in 0..ROOMS_HORIZONTAL {
                if rooms[c][r].kind != RoomType::EmptyRoom {
                    let mut connected = 0;
                    for c1 in 0..ROOMS_VERTICAL {
                        if rooms[c1][r].kind != RoomType::EmptyRoom {
                            connected += 1;
                        }
                    }
                    for r1 in 0..ROOMS_HORIZONTAL {
                        if rooms[c][r1].kind != RoomType::EmptyRoom {
                            connected += 1;
                        }
                    }
                    if connected <= 2 {
                        return false;
                    }
                }
            }
        }
        return true;
    }
    pub fn render(&self) -> Level {
        let mut rng = rand::thread_rng();
        let mut structure = [[StructureElement::Wall; 1 + ROOMS_HORIZONTAL * ROOM_HEIGHT]; 1 + ROOMS_VERTICAL * ROOM_WIDTH];
        let mut artifacts: Vec<Box<dyn Artifact>> = Vec::with_capacity(10);
        let mut enemies: Vec<Box<dyn Monster>> = Vec::with_capacity(10);
        let mut start_x: usize = 0;
        let mut start_y: usize = 0;
        let mut end_x: usize = 0;
        let mut end_y: usize = 0;
        for c in 0..ROOMS_VERTICAL {
            for r in 0..ROOMS_HORIZONTAL {
                let top = 1 + r * ROOM_HEIGHT;
                let left = 1 + c * ROOM_WIDTH;
                let room = self.rooms[c][r];
                for x in 0..room.width {
                    for y in 0..room.height {
                        structure[left + room.offset_x + x][top + room.offset_y + y] = StructureElement::Floor;
                    }
                }
                if room.kind == RoomType::TreasureRoom {
                    let t_x = left + room.offset_x + rng.gen_range(0..room.width);
                    let t_y = top + room.offset_y + rng.gen_range(0..room.height);
                    // TODO randomize artifacts
                    match rng.gen_range(1..=100) {
                        1..=50 => { artifacts.push(Box::new(Chest::new(Position::new(self.level, t_x, t_y)))); }
                        _ => { artifacts.push(Box::new(Potion::new(Position::new(self.level, t_x, t_y)))); }
                    };
                }
                if room.kind == RoomType::MonsterRoom {
                    let t_x = left + room.offset_x + rng.gen_range(0..room.width);
                    let t_y = top + room.offset_y + rng.gen_range(0..room.height);
                    // TODO randomize enemies here
                    match rng.gen_range(1..=100) {
                        1..=50 => {
                            enemies.push(Box::new(Orc::new_with_position(Position::new(self.level, t_x, t_y))));
                        }
                        _ => {
                            enemies.push(Box::new(Rat::new_with_position(Position::new(self.level, t_x, t_y))));
                        }
                    };
                }

                if room.kind == RoomType::End || room.kind == RoomType::StairDown {
                    end_x = left + room.offset_x + rng.gen_range(0..room.width);
                    end_y = top + room.offset_y + rng.gen_range(0..room.height);
                }
                if room.kind == RoomType::StairDown {
                    structure[end_x][end_y] = StructureElement::StairDown;
                }
                if room.kind == RoomType::End {
                    structure[end_x][end_y] = StructureElement::End;
                }

                if room.kind == RoomType::Start || room.kind == RoomType::StairUp {
                    start_x = left + room.offset_x + rng.gen_range(0..room.width);
                    start_y = top + room.offset_y + rng.gen_range(0..room.height);
                }
                if room.kind == RoomType::StairUp {
                    structure[start_x][start_y] = StructureElement::StairUp;
                }
                if room.kind == RoomType::Start {
                    structure[start_x][start_y] = StructureElement::Start;
                }
            }
        }
        //
        for c in 0..ROOMS_VERTICAL {
            for r in 0..ROOMS_HORIZONTAL {
                let src_room = self.rooms[c][r];
                if let Some(x_conn) = src_room.connection_down {
                    let tgt_room = self.rooms[c][r + x_conn.distance];
                    let top = 1 + r * ROOM_HEIGHT + src_room.offset_y;
                    let left = 1 + c * ROOM_WIDTH + x_conn.offset;
                    let bottom = 1 + (r + x_conn.distance) * ROOM_HEIGHT + tgt_room.offset_y + tgt_room.height;
                    for i in top..bottom {
                        if structure[left][i] == StructureElement::Wall {
                            structure[left][i] = StructureElement::Floor;
                        }
                    }
                }
                if let Some(y_conn) = src_room.connection_right {
                    let tgt_room = self.rooms[c + y_conn.distance][r];

                    let top = 1 + r * ROOM_HEIGHT + src_room.offset_y + y_conn.offset - 1;
                    let left = 1 + c * ROOM_WIDTH + src_room.offset_x;
                    let right = 1 + (c + y_conn.distance) * ROOM_WIDTH + tgt_room.offset_x + tgt_room.width;
                    for i in left..right {
                        if structure[i][top] == StructureElement::Wall {
                            structure[i][top] = StructureElement::Floor;
                        }
                    }
                }
            }
        }
        Level {
            level: self.level,
            structure,
            discovered: [[false; 1 + ROOMS_HORIZONTAL * ROOM_HEIGHT]; 1 + ROOMS_VERTICAL * ROOM_WIDTH],
            monsters: enemies,
            artifacts,
            start: (start_x, start_y),
            end: (end_x, end_y),
        }
    }
}

#[test]
fn test_level_gen() {
    let level = LevelGenerator::generate(23, false, false).render();
    assert_eq!(level.level, 23);
    assert_ne!(level.start, (0, 0));
    assert_ne!(level.end, (0, 0));
    assert_ne!(level.start, level.end);
    assert_ne!(level.monsters.len(), 0);
    assert_ne!(level.artifacts.len(), 0);
}