el_diablo/src/level_generator.rs

use std::cmp::{max, min};
use std::collections::HashMap;
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::rngs::ThreadRng;
use rand::Rng;

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

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

pub const ROOM_WIDTH: usize = 9;
pub 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, rs) in rooms.iter_mut().enumerate().take(ROOMS_VERTICAL) {
                for (r, room) in rs.iter_mut().enumerate().take(ROOMS_HORIZONTAL) {
                    room.kind = room_types.pop().unwrap();
                    if room.kind != RoomType::EmptyRoom {
                        room.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, rs) in rooms.iter().enumerate().take(ROOMS_VERTICAL).skip(c + 1) {
                    if rs[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() {
                    range.start
                } else {
                    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 room in rooms.iter().take(ROOMS_VERTICAL) {
                        if room[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;
                    }
                }
            }
        }
        true
    }

    fn select_monster(position: Position, rng: &mut ThreadRng) -> Box<dyn Monster> {
        let level = position.get_level();
        let value = rng.gen_range(1..=100);

        let t = [
            // level 0
            HashMap::from([(MonsterTypes::Rat, 1..=100)]),
            // level 1
            HashMap::from([(MonsterTypes::Rat, 1..=50), (MonsterTypes::Snake, 51..=100)]),
            // level 2
            HashMap::from([(MonsterTypes::Orc, 1..=100)]),
            HashMap::from([(MonsterTypes::Skeleton, 1..=100)]),
        ];
        if level < t.len() {
            for (mtype, range) in &t[level] {
                if range.contains(&value) {
                    return create_monster_by_type(mtype, position);
                }
            }
        }
        match rng.gen_range(1..=100) {
            1..=30 => Box::new(Orc::new_with_position(position)),
            31..=60 => Box::new(Snake::new_with_position(position)),
            _ => Box::new(Rat::new_with_position(position)),
        }
    }

    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
                    enemies.push(LevelGenerator::select_monster(
                        Position::new(self.level, t_x, t_y),
                        &mut rng,
                    ));
                }

                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 room in structure.iter_mut().take(right).skip(left) {
                        if room[top] == StructureElement::Wall {
                            room[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),
            rng: rand::thread_rng(),
        }
    }
}

#[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);
}