use itertools::Itertools;

#[derive(Debug)]
pub struct HeightMap {
    pub map: Vec<Vec<usize>>,
    pub basin: Vec<Vec<Option<usize>>>,
}

impl HeightMap {
    pub fn new(map: Vec<Vec<usize>>) -> Self {
        let mut basin = Vec::new();
        for line in &map {
            let mut basin_line = Vec::new();
            for _point in line {
                basin_line.push(None)
            }
            basin.push(basin_line);
        }
        let mut new_self = Self { map, basin };
        let mut basin_counter = 0;
        for y in 0..new_self.map.len() {
            for x in 0..new_self.map[y].len() {
                if new_self.map[y][x] < 9 && new_self.basin[y][x].is_none() {
                    new_self.flood_basin((x, y), basin_counter);
                    basin_counter += 1;
                }
            }
        }
        new_self
    }

    pub fn get_low_points(&self) -> Vec<(isize, isize)> {
        let mut low_points = Vec::new();
        for y in 0..self.map.len() {
            for x in 0..self.map[y].len() {
                if self.is_low_point((x as isize, y as isize)) {
                    low_points.push((x as isize, y as isize))
                }
            }
        }
        low_points
    }

    pub fn is_low_point(&self, position: (isize, isize)) -> bool {
        let val = match self.get(position) {
            Some(pos) => pos,
            None => return false,
        };
        let (x, y) = position;
        for (x_off, y_off) in [(1isize, 0isize), (0, 1), (-1, 0), (0, -1)].iter() {
            if let Some(neighbour) = self.get((x + x_off, y + y_off)) {
                if neighbour <= val {
                    return false;
                }
            }
        }
        true
    }

    pub fn flood_basin(&mut self, (x, y): (usize, usize), number: usize) {
        if let Some(basin) = self.basin[y][x] {
            assert_eq!(basin, number);
            return;
        } else {
            self.basin[y][x] = Some(number)
        }
        for (x_off, y_off) in [(1isize, 0isize), (0, 1), (-1, 0), (0, -1)].iter() {
            if let Some(neighbour) = self.get((x as isize + x_off, y as isize + y_off)) {
                if neighbour < 9 {
                    self.flood_basin(
                        ((x as isize + x_off) as usize, (y as isize + y_off) as usize),
                        number,
                    )
                }
            }
        }
    }

    pub fn get(&self, (x, y): (isize, isize)) -> Option<usize> {
        let (x, y): (usize, usize) = match (x.try_into(), y.try_into()) {
            (Ok(x), Ok(y)) => (x, y),
            _ => return None,
        };
        self.map.get(y).map(|val| val.get(x)).flatten().copied()
    }

    pub fn get_risk_level(&self, position: (isize, isize)) -> Option<usize> {
        self.get(position).map(|x| x + 1)
    }

    pub fn get_basin_sizes(&self) -> Vec<usize> {
        self.basin
            .iter()
            .flat_map(|line| line.iter())
            .filter_map(|val| *val)
            .counts()
            .into_values()
            .collect()
    }
}