use itertools::Itertools;
use std::collections::{HashMap, HashSet};
use std::fmt;
use Tile::*;

#[allow(dead_code)]
pub fn run() {
    let input = "..#.#
#####
.#?..
...#.
##...
";
    let _einput = "....#
#..#.
#..##
..#..
#....";
    task1(Life::from(input));
    task2(input);
}

fn task1(mut life: Life) {
    let mut hashes = HashMap::<Life, usize>::new();
    hashes.insert(life.clone(), 0);
    for round in 1.. {
        life.round();
        if hashes.contains_key(&life) {
            break;
        } else {
            hashes.insert(life.clone(), round);
        }
    }
    println!(
        "Task 1: bio diversity of first reoccuring state is {}",
        life.bio_diversity()
    );
}

fn task2(input: &str) {
    let mut state: HashMap<Field, Tile> = input
        .lines()
        .enumerate()
        .flat_map(|(y, line)| {
            line.chars()
                .enumerate()
                .filter_map(|(x, c)| {
                    if c == '#' || c == '.' {
                        Some((
                            Field(0, x as i32, y as i32),
                            if c == '#' { Bug } else { Empty },
                        ))
                    } else {
                        None
                    }
                })
                .collect_vec()
        })
        .collect();
    for _round in 1..=200 {
        let new: HashMap<Field, Tile> = state
            .keys()
            .map(|field| {
                let v = field
                    .neighbors()
                    .iter()
                    .chain(std::iter::once(field))
                    .map(|f| *f)
                    .collect_vec();
                v
            })
            .flatten()
            .map(|f| (f, new_state_of(f, &state)))
            .collect();
        state = new;
    }
    let bug_count = state.iter().filter(|(_k, tile)| **tile == Bug).count();
    println!("Task 2: we have {} bugs overall", bug_count);
}

#[allow(dead_code)]
fn print_levels(state: &HashMap<Field, Tile>, mark_neighbors: Field) {
    let mut levels = state
        .iter()
        .map(|(Field(level, _, _), _)| *level)
        .collect::<HashSet<i32>>()
        .into_iter()
        .collect_vec();
    levels.sort();

    let neighbors = mark_neighbors.neighbors();

    for level in levels {
        println!("Depth {}:", level);
        for y in 0..5 {
            for x in 0..5 {
                let c = if x == 2 && y == 2 {
                    '?'
                } else if Field(level, x, y) == mark_neighbors {
                    'O'
                } else if neighbors.contains(&Field(level, x, y)) {
                    'x'
                } else if let Some(tile) = state.get(&Field(level, x, y)) {
                    match tile {
                        Bug => '#',
                        Empty => '.',
                    }
                } else {
                    '*'
                };
                print!("{}", c);
            }
            println!("");
        }
    }
}

fn new_state_of(field: Field, state: &HashMap<Field, Tile>) -> Tile {
    let at = |f| *state.get(&f).unwrap_or(&Empty);
    let old = at(field);
    if at(field) == Bug {
        if field
            .neighbors()
            .into_iter()
            .filter(|p| at(*p) == Bug)
            .count()
            != 1
        {
            Empty
        } else {
            old
        }
    } else {
        let count = field
            .neighbors()
            .into_iter()
            .filter(|p| at(*p) == Bug)
            .count();
        if count == 1 || count == 2 {
            Bug
        } else {
            old
        }
    }
}

#[derive(Eq, PartialEq, Hash, Clone, Copy, Debug)]
struct Field(i32, i32, i32);
impl Field {
    fn neighbors(&self) -> Vec<Self> {
        let mut res = vec![
            Self(self.0, self.1 - 1, self.2),
            Self(self.0, self.1 + 1, self.2),
            Self(self.0, self.1, self.2 - 1),
            Self(self.0, self.1, self.2 + 1),
        ];
        res.retain(|p| !(p.1 == 2 && p.2 == 2) && ((0..5).contains(&p.1) && (0..5).contains(&p.2)));

        // link to higher level
        if self.1 == 0 {
            res.push(Self(self.0 - 1, 1, 2));
        }
        if self.1 == 4 {
            res.push(Self(self.0 - 1, 3, 2));
        }
        if self.2 == 0 {
            res.push(Self(self.0 - 1, 2, 1));
        }
        if self.2 == 4 {
            res.push(Self(self.0 - 1, 2, 3));
        }

        // link to lower level
        match self {
            Field(lvl, 2, 1) => (0..5).for_each(|x| res.push(Self(lvl + 1, x, 0))),
            Field(lvl, 2, 3) => (0..5).for_each(|x| res.push(Self(lvl + 1, x, 4))),
            Field(lvl, 1, 2) => (0..5).for_each(|y| res.push(Self(lvl + 1, 0, y))),
            Field(lvl, 3, 2) => (0..5).for_each(|y| res.push(Self(lvl + 1, 4, y))),
            _ => (),
        }

        res
    }
}

#[derive(PartialEq, Eq, Hash, Clone)]
struct Life {
    area: [[Tile; 5]; 5],
}

impl Life {
    fn from(string: &str) -> Self {
        let vecs = string
            .lines()
            .map(|line| {
                line.chars()
                    .map(|c| if c == '#' { Bug } else { Empty })
                    .collect_vec()
            })
            .collect_vec();
        assert_eq!(5, vecs.len());
        assert_eq!(5, vecs[0].len());

        let mut life = Life {
            area: [[Empty; 5]; 5],
        };

        for x in 0..5 {
            for y in 0..5 {
                life.area[x][y] = vecs[y][x];
            }
        }
        life
    }

    fn at(&self, pos: Pos) -> Tile {
        if (0..5).contains(&pos.0) && (0..5).contains(&pos.1) {
            self.area[pos.0 as usize][pos.1 as usize]
        } else {
            Empty
        }
    }

    fn round(&mut self) {
        let mut new = self.area.clone();
        for y in 0..5 {
            for x in 0..5 {
                if self.at(Pos(x, y)) == Bug {
                    if Pos(x, y)
                        .neighbors()
                        .into_iter()
                        .filter(|p| self.at(**p) == Bug)
                        .count()
                        != 1
                    {
                        new[x as usize][y as usize] = Empty;
                    }
                } else {
                    let count = Pos(x, y)
                        .neighbors()
                        .into_iter()
                        .filter(|p| self.at(**p) == Bug)
                        .count();
                    if count == 1 || count == 2 {
                        new[x as usize][y as usize] = Bug;
                    }
                }
            }
        }
        self.area = new;
    }

    fn bio_diversity(&self) -> usize {
        (0..5)
            .map(|x| {
                (0..5)
                    .map(|y| {
                        if self.at(Pos(x, y)).is_bug() {
                            let v = 1usize.wrapping_shl((5 * y + x) as u32);
                            v
                        } else {
                            0 as usize
                        }
                    })
                    .sum::<usize>()
            })
            .sum()
    }
}

impl fmt::Display for Life {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for y in 0..5 {
            for x in 0..5 {
                let r = write!(f, "{}", self.area[x][y]);
                if let fmt::Result::Err(_) = r {
                    return r;
                }
            }
            let r = writeln!(f, "");
            if let fmt::Result::Err(_) = r {
                return r;
            }
        }
        fmt::Result::Ok(())
    }
}

#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
enum Tile {
    Bug,
    Empty,
}

impl Tile {
    fn is_bug(&self) -> bool {
        *self == Bug
    }
}

impl fmt::Display for Tile {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{}",
            match self {
                Bug => '#',
                Empty => '.',
            }
        )
    }
}

#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
struct Pos(i32, i32);

impl Pos {
    fn neighbors(&self) -> [Pos; 4] {
        [
            Pos(self.0 - 1, self.1),
            Pos(self.0 + 1, self.1),
            Pos(self.0, self.1 - 1),
            Pos(self.0, self.1 + 1),
        ]
    }
}

#[allow(unused)]
mod test {
    use super::Field;
    #[test]
    fn neighborhood_bijective() {
        for x in 0..5 {
            for y in 0..5 {
                if x == 2 && y == 2 {
                    continue;
                }
                let field = Field(0, x, y);
                let neighbors = field.neighbors();
                neighbors.into_iter().for_each(|n| {
                    assert_eq!(
                        n.neighbors().contains(&field),
                        true,
                        "{:?} is not a neighbor of its neighbor {:?}",
                        field,
                        n
                    );
                });
            }
        }
    }
}