/*

Well, well took longer than anticipated due to mixing up index-variables (unsurprisingly).

*/

use std::io::{BufReader, BufRead};
use std::fs::File;

type Level = Vec<Vec<char>>;

fn count_bugs_in_row(level_ind: i64, row: usize, habitat: &Vec<Level>) -> usize {
    if level_ind >= 0 && level_ind < habitat.len() as i64 {
        habitat[level_ind as usize][row].iter().filter(|c| **c == '#').count()
    } else {
        0
    }
}

fn count_bugs_in_col(level_ind: i64, col: usize, habitat: &Vec<Level>) -> usize {
    if level_ind >= 0 && level_ind < habitat.len() as i64 {
        habitat[level_ind as usize].iter().filter(|r| r[col] == '#').count()
    } else {
        0
    }
}

fn transform_cell(level_ind: i64, row: i64, col: i64, habitat: &Vec<Level>) -> char {
    let cell_offsets = [(1, 0), (-1, 0), (0, 1), (0, -1)];
    let mut n_bugs = 0;
    for offset in &cell_offsets {
        let n_row = row + offset.0;
        let n_col = col + offset.1;
        if n_row == 2 && n_col == 2 {
            // Contribution from level that is at the center of this level
            if row == 1 {
                n_bugs += count_bugs_in_row(level_ind + 1, 0, habitat)
            } else if row == 3 {
                n_bugs += count_bugs_in_row(level_ind + 1, 4, habitat)
            } else if col == 1 {
                n_bugs += count_bugs_in_col(level_ind + 1, 0, habitat)
            } else if col == 3 {
                n_bugs += count_bugs_in_col(level_ind + 1, 4, habitat)
            }
        } else if n_row >= 0 && n_row <= 4 && n_col >= 0 && n_col <= 4 && habitat[level_ind as usize][n_row as usize][n_col as usize] == '#' {
            // Contribution from the current level
            n_bugs += 1;
        } else if level_ind > 0 as i64 {
            // Contribution from the level that is outer to this level
            if n_row < 0 && habitat[level_ind as usize - 1][1][2] == '#' {
                n_bugs += 1
            } else if n_row > 4 && habitat[level_ind as usize - 1][3][2] == '#' {
                n_bugs += 1
            } else if n_col < 0 && habitat[level_ind as usize - 1][2][1] == '#' {
                n_bugs += 1
            } else if n_col > 4 && habitat[level_ind as usize - 1][2][3] == '#' {
                n_bugs += 1
            }
        }
    }
    if habitat[level_ind as usize][row as usize][col as usize] == '#' && n_bugs != 1 {
        '.'
    } else if habitat[level_ind as usize][row as usize][col as usize] == '.' && (n_bugs == 1 || n_bugs == 2) {
        '#'
    } else {
        habitat[level_ind as usize][row as usize][col as usize]
    }
}

fn new_empty_level() -> Level {
    let mut level = Level::new();
    level.resize(5, vec!('.', '.', '.', '.', '.'));
    level
}

fn count_bugs_in_level(level: &Level) -> i64 {
    level.iter().flatten().filter(|c| **c == '#').count() as i64
}

fn tick_habitat_clock(habitat: &Vec<Level>) -> Vec<Level> {
    let mut expanded_habitat = habitat.clone();
    expanded_habitat.insert(0, new_empty_level());
    expanded_habitat.push(new_empty_level());
    let mut transformed_habitat = expanded_habitat.clone();
    for ind in 0..transformed_habitat.len() {
        for row in 0..5 {
            for col in 0..5 {
                if row != 2 || col != 2 {
                    transformed_habitat[ind][row][col] = transform_cell(ind as i64, row as i64, col as i64, &expanded_habitat)
                }
            }
        }    
    }
    if count_bugs_in_level(transformed_habitat.last().unwrap()) == 0 {
        transformed_habitat.pop();
    }
    if count_bugs_in_level(transformed_habitat.first().unwrap()) == 0 {
        transformed_habitat.remove(0);
    }
    transformed_habitat
}

fn count_bugs_in_habitat(habitat: &Vec<Level>) -> usize {
    habitat.iter().flatten().flatten().filter(|c| **c == '#').count()
}

fn main() {
    let habitat_file = "habitat".to_string();
    let mut habitat = load_habitat(&habitat_file);
    for _tick in 0..200 {
        habitat = tick_habitat_clock(&habitat);
    }
    let number_of_bugs = count_bugs_in_habitat(&habitat);
    println!("Total number of bugs after 200 minutes={}", number_of_bugs)
}

fn load_habitat(habitat_file: &String) -> Vec<Level> {
    let mut level = Level::new();
    let file = File::open(habitat_file).expect("Could not open habitat file!");
    let reader = BufReader::new(file);
    for line in reader.lines() {
        let line = line.expect("No string there.");
        let chrs: Vec<char> = line.as_str().trim().chars().collect();
        if chrs.len() != 5 {
            panic!("Malformed habitat file: {}", habitat_file)
        }
        level.push(chrs)
    }
    if level.len() != 5 {
        panic!("Malformed habitat file: {}", habitat_file)
    }
    vec!(level)
}

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

#[test]
    #[test]
    fn test_1() {
        let start_file = "test1".to_string();
        let start_habitat = load_habitat(&start_file);

        let mut habitat = start_habitat.clone();
        for _ in 0..10 {
            habitat = tick_habitat_clock(&habitat);
        }
        assert_eq!(11, habitat.len());
        assert_eq!(99, count_bugs_in_habitat(&habitat))
    }
}