el_diablo/src/level_generator.rs

use petgraph::algo::min_spanning_tree;
use petgraph::data::*;
use petgraph::graph::Graph;
use petgraph::graph::UnGraph;
use std::cmp::max;
use std::cmp::min;
use std::ops::Range;

use rand::Rng;
use rand::{rngs::ThreadRng, seq::SliceRandom};

use crate::constants::ROOM_HEIGHT;
use crate::constants::ROOM_WIDTH;
use crate::room::Connection;
use crate::{
    constants::{
        get_room_type_per_level, LEVEL_HEIGHT, LEVEL_WIDTH, ROOMS_HORIZONAL, ROOMS_VERTICAL,
    },
    level::{Level, StructureElement},
    room::{Room, RoomType},
};

pub struct LevelGenerator {
    level: usize,
    rooms: [[Room; ROOMS_VERTICAL]; ROOMS_HORIZONAL],
}
enum Direction {
    Horizontal,
    Vertical,
}

impl LevelGenerator {
    pub fn generate(level: usize, first: bool, last: bool) -> Self {
        let mut rng = rand::thread_rng();
        let mut rooms = [[Room::new(&mut rng); ROOMS_VERTICAL]; ROOMS_HORIZONAL];

        /*
           Fill grid with unconnected rooms
        */

        let mut rooms_to_place: Vec<Room> = Vec::with_capacity(ROOMS_VERTICAL * ROOMS_HORIZONAL);

        let mut start_room = Room::new(&mut rng);
        if first {
            start_room.kind = RoomType::Start;
        } else {
            start_room.kind = RoomType::StairUp;
        }
        rooms_to_place.push(start_room);
        let mut end_room = Room::new(&mut rng);
        if last {
            end_room.kind = RoomType::End;
        } else {
            end_room.kind = RoomType::StairDown;
        }
        rooms_to_place.push(end_room);
        for _ in 2..ROOMS_HORIZONAL * ROOMS_VERTICAL {
            let mut room = Room::new(&mut rng);
            room.kind = LevelGenerator::select_room_type(level, &mut rng);
            if room.kind != RoomType::EmptyRoom {
                rooms_to_place.push(room);
            }
        }
        let mut room_row = rng.gen_range(1..ROOMS_VERTICAL);
        let mut room_col = rng.gen_range(1..ROOMS_HORIZONAL);
        rooms[room_col][room_row] = rooms_to_place.pop().unwrap();
        while let Some(room) = rooms_to_place.pop() {
            let mut directions_to_try = vec![Direction::Horizontal, Direction::Vertical];
            directions_to_try.shuffle(&mut rng);
            while !directions_to_try.is_empty() {
                match directions_to_try.pop().unwrap() {
                    Direction::Horizontal => {
                        let mut free_cols: Vec<usize> = vec![];
                        for col in 0..ROOMS_HORIZONAL {
                            if rooms[col][room_row].kind == RoomType::EmptyRoom {
                                free_cols.push(col);
                            }
                        }
                        if free_cols.is_empty() {
                            continue;
                        }
                        free_cols.shuffle(&mut rng);
                        room_col = *free_cols.first().unwrap();
                        rooms[room_col][room_row] = room;
                        break;
                    }
                    Direction::Vertical => {
                        let mut free_rows: Vec<usize> = vec![];
                        for row in 0..ROOMS_VERTICAL {
                            if rooms[room_col][row].kind == RoomType::EmptyRoom {
                                free_rows.push(row);
                            }
                        }
                        if free_rows.is_empty() {
                            continue;
                        }
                        free_rows.shuffle(&mut rng);
                        room_row = *free_rows.first().unwrap();
                        rooms[room_col][room_row] = room;
                        break;
                    }
                }
            }
        }


        println!("  0 1 2 3 4 5 6 7");
        for r in 0..ROOMS_VERTICAL {
            print!("{} ", r);
            for c in 0..ROOMS_HORIZONAL {
                match rooms[c][r].kind {
                    RoomType::Start => print!("S "),
                    RoomType::End => print!("E "),
                    RoomType::StairUp => print!("< "),
                    RoomType::StairDown => print!("> "),
                    RoomType::BasicRoom => print!("B "),
                    RoomType::ArtifactRoom => print!("A "),
                    RoomType::MonsterRoom => print!("M "),
                    RoomType::EmptyRoom => print!("  "),
                };
            }
            println!();
        }

        /*
           Construct a graph from the unconnected rooms and make a minum spanning tree of it
        */
        let mut graph = UnGraph::<(usize, usize), u16>::default();
        // collect nodes
        for col in 0..ROOMS_HORIZONAL {
            for row in 0..ROOMS_VERTICAL {
                if rooms[col][row].kind != RoomType::EmptyRoom {
                    graph.add_node((col, row));
                }
            }
        }
        // collect edges between nodes, each right and down till we find the next room
        for col in 0..ROOMS_HORIZONAL {
            for row in 0..ROOMS_VERTICAL {
                if rooms[col][row].kind == RoomType::EmptyRoom {
                    continue;
                }
                if let Some(src_index) = graph.node_indices().find(|i| graph[*i] == (col, row)) {
                    for col_1 in col + 1..ROOMS_HORIZONAL {
                        if rooms[col_1][row].kind != RoomType::EmptyRoom {
                            if let Some(tgt_index) =
                                graph.node_indices().find(|i| graph[*i] == (col_1, row))
                            {
                                graph.add_edge(src_index, tgt_index, 1);
                                break;
                            }
                        }
                    }
                    for row_1 in row + 1..ROOMS_VERTICAL {
                        if rooms[col][row_1].kind != RoomType::EmptyRoom {
                            if let Some(tgt_index) =
                                graph.node_indices().find(|i| graph[*i] == (col, row_1))
                            {
                                graph.add_edge(src_index, tgt_index, 1);
                                break;
                            }
                        }
                    }
                }
            }
        }
        let mst: Graph<(usize, usize), u16, petgraph::Undirected> =
            Graph::from_elements(min_spanning_tree(&graph));
        for edge in mst.raw_edges() {
            // the tuples are (col, row)
            let (src_node_col, src_node_row) = mst[edge.source()];
            let (tgt_node_col, tgt_node_row) = mst[edge.target()];
            // println!("MST Edge from {:?} to {:?}", (src_node_col, src_node_row), (tgt_node_col, tgt_node_row));

            if src_node_col == tgt_node_col {
                // println!("Down");
                let start_col = src_node_col * ROOM_WIDTH+ROOM_WIDTH/2;
                let start_row = src_node_row * ROOM_HEIGHT+ROOM_HEIGHT;
                let end_col = tgt_node_col * ROOM_WIDTH+ROOM_WIDTH/2;
                 let end_row = tgt_node_row * ROOM_HEIGHT;
                rooms[src_node_col][src_node_row].connection_down = Some(Connection {
                    start_pos: (start_col, start_row),
                    end_pos: (end_col, end_row),
                });
            } else {
                // println!("Right");
                let start_col = src_node_col * ROOM_WIDTH+ROOM_WIDTH;
                let start_row = src_node_row * ROOM_HEIGHT+ROOM_HEIGHT/2;
                let end_col = tgt_node_col * ROOM_WIDTH;
                 let end_row = tgt_node_row * ROOM_HEIGHT+ROOM_HEIGHT/2;
                rooms[src_node_col][src_node_row].connection_right = Some(Connection {
                    start_pos: (start_col, start_row),
                    end_pos: (end_col, end_row),
                });
            }
        }

        LevelGenerator { level, rooms }
    }

    fn select_room_type(level: usize, rng: &mut ThreadRng) -> RoomType {
        let value = rng.gen_range(1..=100);

        let t = get_room_type_per_level();
        if level < t.len() {
            for (mtype, range) in &t[level] {
                if range.contains(&value) {
                    return *mtype;
                }
            }
        }
        panic!("no room selectable!");
    }

    pub fn render(&mut self) -> Level {
        let mut structure = [[StructureElement::Wall; LEVEL_HEIGHT]; LEVEL_WIDTH];
        let mut start_pos = (0, 0);
        let mut end_pos = (0, 0);
        for col in 0..ROOMS_HORIZONAL {
            for row in 0..ROOMS_VERTICAL {
                let position = self.rooms[col][row].render(&mut structure, col, row);
                if self.rooms[col][row].kind == RoomType::Start
                    || self.rooms[col][row].kind == RoomType::StairUp
                {
                    start_pos = position;
                }
                if self.rooms[col][row].kind == RoomType::End
                    || self.rooms[col][row].kind == RoomType::StairDown
                {
                    end_pos = position;
                }
            }
        }
        for col in 0..ROOMS_HORIZONAL {
            for row in 0..ROOMS_VERTICAL {
                if let Some(connection) = self.rooms[col][row].connection_down {
                    // println!("down");
                    connection.render(&mut structure);
                }
                if let Some(connection) = self.rooms[col][row].connection_right {
                    // println!("right");
                    connection.render(&mut structure);
                }
            }
        }
        Level {
            level: self.level,
            structure,
            discovered: [[false; LEVEL_HEIGHT]; LEVEL_WIDTH],
            monsters: vec![],
            artifacts: vec![],
            start: start_pos,
            end: end_pos,
            rng: rand::thread_rng(),
        }
    }
}

#[test]
fn test_level_gen() {
    for _ in 0..1000 {
        LevelGenerator::generate(0, true, false).render();
    }
}

#[test]
fn test_level_gen_respects_level() {
    let level = LevelGenerator::generate(0, true, false).render();
    assert_eq!(0, level.level);
    let level = LevelGenerator::generate(1, true, false).render();
    assert_eq!(1, level.level);
}