el_diablo/src/level_generator.rs

use petgraph::algo::min_spanning_tree;
use petgraph::data::*;
use petgraph::graph::Graph;
use petgraph::graph::UnGraph;

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

use crate::artifacts::Artifact;
use crate::artifacts::Chest;
use crate::artifacts::Potion;
use crate::constants::get_monsters_per_level;
use crate::constants::ROOM_HEIGHT;
use crate::constants::ROOM_WIDTH;
use crate::monster::create_monster_by_type;
use crate::monster::Monster;
use crate::position::Position;
use crate::room::Connection;
use crate::{
    constants::{
        get_room_type_per_level, LEVEL_HEIGHT, LEVEL_WIDTH, ROOMS_HORIZONTAL, ROOMS_VERTICAL,
    },
    level::{Level, StructureElement},
    room::{Room, RoomType},
};

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

impl LevelGenerator {
    pub fn generate_rooms_to_place(
        rng: &mut ThreadRng,
        level: usize,
        first: bool,
        last: bool,
    ) -> Vec<Room> {
        let mut rooms_to_place: Vec<Room> = Vec::with_capacity(ROOMS_VERTICAL * ROOMS_HORIZONTAL);

        let mut start_room = Room::new(rng);
        if first {
            start_room.kind = RoomType::Start;
        } else {
            start_room.kind = RoomType::StairUp;
        }
        rooms_to_place.push(start_room);
        for _ in 2..ROOMS_HORIZONTAL * ROOMS_VERTICAL {
            let mut room = Room::new(rng);
            room.kind = LevelGenerator::select_room_type(level, rng);
            if room.kind != RoomType::EmptyRoom {
                rooms_to_place.push(room);
            }
        }
        let mut end_room = Room::new(rng);
        if last {
            end_room.kind = RoomType::End;
        } else {
            end_room.kind = RoomType::StairDown;
        }
        rooms_to_place.push(end_room);
        rooms_to_place
    }

    pub fn place_rooms(
        rng: &mut ThreadRng,
        rooms_to_place: &mut Vec<Room>,
    ) -> [[Room; ROOMS_VERTICAL]; ROOMS_HORIZONTAL] {
        let mut rooms: [[Room; 7]; 8] = [[Room::new(rng); ROOMS_VERTICAL]; ROOMS_HORIZONTAL];
        let mut room_row = rng.gen_range(0..ROOMS_VERTICAL);
        let mut room_col = rng.gen_range(0..ROOMS_HORIZONTAL);
        rooms[room_col][room_row] = rooms_to_place.pop().unwrap();
        while let Some(room) = rooms_to_place.pop() {
            let mut placed = false;
            // randomize going horizontal or vertical
            let mut directions_to_try = vec![Direction::Horizontal, Direction::Vertical];
            directions_to_try.shuffle(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_HORIZONTAL {
                            if rooms[col][room_row].kind == RoomType::EmptyRoom {
                                free_cols.push(col);
                            }
                        }
                        if free_cols.is_empty() {
                            continue;
                        }
                        free_cols.shuffle(rng);
                        room_col = *free_cols.first().unwrap();
                        rooms[room_col][room_row] = room;
                        placed = true;
                        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(rng);
                        room_row = *free_rows.first().unwrap();
                        rooms[room_col][room_row] = room;
                        placed = true;
                        break;
                    }
                }
            }

            // all fields in the row/column was full so we can place it at any empty position
            if !placed {
                let mut free_pos: Vec<(usize, usize)> = vec![];
                for col in 0..ROOMS_HORIZONTAL {
                    for row in 0..ROOMS_VERTICAL {
                        if rooms[col][row].kind == RoomType::EmptyRoom {
                            free_pos.push((col, row));
                        }
                    }
                }
                let selected_pos = free_pos[rng.gen_range(0..free_pos.len())];
                rooms[selected_pos.0][selected_pos.1] = room;
            }
        }
        rooms
    }

    pub fn create_mst(
        rooms: &[[Room; ROOMS_VERTICAL]; ROOMS_HORIZONTAL],
    ) -> Graph<(usize, usize), u16, petgraph::Undirected> {
        let mut graph = UnGraph::<(usize, usize), u16>::default();
        for col in 0..ROOMS_HORIZONTAL {
            for row in 0..ROOMS_VERTICAL {
                if rooms[col][row].kind != RoomType::EmptyRoom {
                    graph.add_node((col, row));
                }
            }
        }
        for col in 0..ROOMS_HORIZONTAL {
            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_HORIZONTAL {
                        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;
                            }
                        }
                    }
                }
            }
        }
        Graph::from_elements(min_spanning_tree(&graph))
        // graph
    }

    pub fn generate(level: usize, first: bool, last: bool) -> Self {
        let mut rng = rand::thread_rng();

        /*
           Fill grid with unconnected rooms
        */

        let mut rooms_to_place: Vec<Room> =
            LevelGenerator::generate_rooms_to_place(&mut rng, level, first, last);
        let mut rooms: [[Room; 7]; 8] = LevelGenerator::place_rooms(&mut rng, &mut rooms_to_place);

        // debug print a text view of the dungeon
        println!("  0 1 2 3 4 5 6 7");
        for r in 0..ROOMS_VERTICAL {
            print!("{} ", r);
            for c in 0..ROOMS_HORIZONTAL {
                match rooms[c][r].kind {
                    RoomType::Start => print!("S "),
                    RoomType::End => print!("E "),
                    RoomType::StairUp => print!("< "),
                    RoomType::StairDown => print!("> "),
                    RoomType::BasicRoom => print!("_ "),
                    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 mst: Graph<(usize, usize), u16, petgraph::Undirected> =LevelGenerator::create_mst(&rooms);
        
        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, rng }
    }

    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 = get_monsters_per_level();
        if level < t.len() {
            for (mtype, range) in &t[level] {
                if range.contains(&value) {
                    return create_monster_by_type(mtype, position);
                }
            }
        }
        panic!("no monster selectable!");
    }

    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);
        let mut monsters: Vec<Box<dyn Monster>> = Vec::with_capacity(10);
        let mut artifacts: Vec<Box<dyn Artifact>> = Vec::with_capacity(10);
        for col in 0..ROOMS_HORIZONTAL {
            for row in 0..ROOMS_VERTICAL {
                let room = self.rooms[col][row];
                let position = room.render(&mut structure, col, row);
                match room.kind {
                    RoomType::Start => start_pos = position,
                    RoomType::End => end_pos = position,
                    RoomType::StairUp => start_pos = position,
                    RoomType::StairDown => end_pos = position,
                    RoomType::BasicRoom => {}
                    RoomType::ArtifactRoom => {
                        match self.rng.gen_range(1..=100) {
                            1..=50 => {
                                artifacts.push(Box::new(Chest::new(Position::new(
                                    self.level, position.0, position.1,
                                ))));
                            }
                            _ => {
                                artifacts.push(Box::new(Potion::new(Position::new(
                                    self.level, position.0, position.1,
                                ))));
                            }
                        };
                    }
                    RoomType::MonsterRoom => {
                        monsters.push(LevelGenerator::select_monster(
                            Position::new(self.level, position.0, position.1),
                            &mut self.rng,
                        ));
                    }
                    RoomType::EmptyRoom => {}
                }
            }
        }
        for col in 0..ROOMS_HORIZONTAL {
            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,
            artifacts,
            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);
// }

#[cfg(test)]
fn find_room_types(rooms: &Vec<Room>) -> (bool, bool, bool, bool) {
    let mut start_found = false;
    let mut end_found: bool = false;
    let mut down_found: bool = false;
    let mut up_found: bool = false;
    for room in rooms {
        if room.kind == RoomType::Start {
            start_found = true;
        }
        if room.kind == RoomType::End {
            end_found = true;
        }
        if room.kind == RoomType::StairDown {
            down_found = true;
        }
        if room.kind == RoomType::StairUp {
            up_found = true;
        }
    }
    (start_found, up_found, down_found, end_found)
}
#[test]
fn test_rooms_to_place_first_level() {
    let mut rng = rand::thread_rng();
    let res = LevelGenerator::generate_rooms_to_place(&mut rng, 0, true, false);
    assert!(
        res.len() <= ROOMS_HORIZONTAL * ROOMS_VERTICAL,
        "too many rooms created"
    );
    assert!(0 < res.len(), "too many rooms created");

    let (start_found, up_found, down_found, end_found) = find_room_types(&res);
    assert!(start_found);
    assert!(!end_found);
    assert!(down_found);
    assert!(!up_found);
}

#[test]
fn test_rooms_to_place_middle_level() {
    let mut rng = rand::thread_rng();
    let res = LevelGenerator::generate_rooms_to_place(&mut rng, 1, false, false);
    assert!(
        res.len() <= ROOMS_HORIZONTAL * ROOMS_VERTICAL,
        "too many rooms created"
    );
    assert!(0 < res.len(), "too many rooms created");

    let (start_found, up_found, down_found, end_found) = find_room_types(&res);
    assert!(!start_found);
    assert!(!end_found);
    assert!(down_found);
    assert!(up_found);
}

#[test]
fn test_rooms_to_place_last_level() {
    let mut rng = rand::thread_rng();
    let res = LevelGenerator::generate_rooms_to_place(&mut rng, 2, false, true);
    assert!(
        res.len() <= ROOMS_HORIZONTAL * ROOMS_VERTICAL,
        "too many rooms created"
    );
    assert!(0 < res.len(), "too many rooms created");

    let (start_found, up_found, down_found, end_found) = find_room_types(&res);
    assert!(!start_found);
    assert!(end_found);
    assert!(!down_found);
    assert!(up_found);
}

#[cfg(test)]
fn check_valid_placement(rooms: &[[Room; ROOMS_VERTICAL]; ROOMS_HORIZONTAL]) -> bool {
    for col in 0..ROOMS_HORIZONTAL {
        for row in 0..ROOMS_VERTICAL {
            if rooms[col][row].kind != RoomType::EmptyRoom {
                let mut count = 0;
                for test_col in 0..ROOMS_HORIZONTAL {
                    if rooms[test_col][row].kind != RoomType::EmptyRoom {
                        count += 1;
                    }
                }
                for test_row in 0..ROOMS_VERTICAL {
                    if rooms[col][test_row].kind != RoomType::EmptyRoom {
                        count += 1;
                    }
                }
                if count < 3 {
                    return false;
                }
            }
        }
    }
    true
}

#[cfg(test)]
fn count_rooms(rooms: &[[Room; ROOMS_VERTICAL]; ROOMS_HORIZONTAL]) -> usize {
    let mut res = 0;
    for col in 0..ROOMS_HORIZONTAL {
        for row in 0..ROOMS_VERTICAL {
            if rooms[col][row].kind != RoomType::EmptyRoom {
                res += 1;
            }
        }
    }
    res
}

#[test]
fn test_place_rooms() {
    let mut rng = rand::thread_rng();
    for count in 2..ROOMS_HORIZONTAL * ROOMS_VERTICAL {
        let mut rooms: Vec<Room> = vec![Room::new(&mut rng), Room::new(&mut rng)];
        rooms[0].kind = RoomType::Start;
        rooms[1].kind = RoomType::End;
        for t in 2..count {
            rooms.push(Room::new(&mut rng));
            rooms[t].kind = RoomType::BasicRoom;
        }
        let res = LevelGenerator::place_rooms(&mut rng, &mut rooms);

        assert_eq!(count_rooms(&res), count, "counting {}", count);
        assert!(check_valid_placement(&res));
    }
}

#[test]
fn test_create_mst() {
    let mut rng = rand::thread_rng();
    let mut rooms = [[Room::new(&mut rng); ROOMS_VERTICAL]; ROOMS_HORIZONTAL];
    let res = LevelGenerator::create_mst(&rooms);
    assert_eq!(res.node_count(), 0);
    assert_eq!(res.edge_count(), 0);

    rooms[1][1].kind = RoomType::BasicRoom;
    let res = LevelGenerator::create_mst(&rooms);
    assert_eq!(res.node_count(), 1);
    assert_eq!(res.edge_count(), 0);

    rooms[1][3].kind = RoomType::BasicRoom;
    let res = LevelGenerator::create_mst(&rooms);
    assert_eq!(res.node_count(), 2);
    assert_eq!(res.edge_count(), 1);

    rooms[3][1].kind = RoomType::BasicRoom;
    let res = LevelGenerator::create_mst(&rooms);
    assert_eq!(res.node_count(), 3);
    assert_eq!(res.edge_count(), 2);

    rooms[3][3].kind = RoomType::BasicRoom;
    let res = LevelGenerator::create_mst(&rooms);
    assert_eq!(res.node_count(), 4);
    assert_eq!(res.edge_count(), 3);

    rooms[3][5].kind = RoomType::BasicRoom;
    let res = LevelGenerator::create_mst(&rooms);
    assert_eq!(res.node_count(), 5);
    assert_eq!(res.edge_count(), 4);
}

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