day 22

advent_of_code::solution!(22);
use std::cmp::min;
use std::collections::HashSet;

#[derive(Debug, PartialEq, Eq, Clone)]
struct Brick {
    start: (u32, u32, u32),
    end: (u32, u32, u32),
    cubes: HashSet<(u32, u32)>,
}

impl Brick {
    pub fn parse(c: &str) -> Self {
        let p = |c: &str| -> (u32, u32, u32) {
            let mut w = c.split(",").map(|b| b.parse::<u32>().unwrap());
            (w.next().unwrap(), w.next().unwrap(), w.next().unwrap())
        };

        let (l, r) = c.split_once("~").unwrap();
        let (p1, p2) = (p(l), p(r));

        Brick::new(p1, p2)
    }

    pub fn new(start: (u32, u32, u32), end: (u32, u32, u32)) -> Brick {
        assert!(end.0 >= start.0 && end.1 >= start.1 && end.2 >= start.2);

        Brick {
            start: start,
            end: end,
            cubes: (start.0..(end.0 + 1))
                .map(|i| (start.1..(end.1 + 1)).map(move |j| (i, j)))
                .flatten()
                .collect(),
        }
    }
    // distance from ground
    pub fn bottom(&self) -> u32 {
        min(self.start.2, self.end.2)
    }

    fn top(&self) -> u32 {
        self.end.2
    }

    fn intersects(&self, other: &Brick) -> bool {
        self.cubes.intersection(&(other.cubes)).count() > 0
    }

    pub fn drop(&self, on_ground: &Vec<Brick>) -> Brick {
        let c = self.clone();
        let new_floor = on_ground
            .iter()
            .map(|b| if self.intersects(b) { b.top() + 1 } else { 1 })
            .max()
            .unwrap();

        let mut new_start = self.start.clone();
        let mut new_end = self.end.clone();

        if new_floor < c.bottom() {
            new_end.2 = new_floor + (new_end.2 - new_start.2);
            new_start.2 = new_floor;
        }

        Brick::new(new_start, new_end)
    }

    pub fn supports(&self, other: &Brick) -> bool {
        other.bottom() == (self.top() + 1) && self.intersects(other)
    }
}

fn process(input: &str) -> (Vec<Brick>, Vec<usize>, Vec<Vec<bool>>) {
    let bricks: Vec<Brick> = input.lines().map(|l| Brick::parse(l)).collect();

    let max_h = bricks.iter().map(|b| b.bottom()).max().unwrap();

    let dropped = (2..max_h + 1).fold(
        bricks
            .clone()
            .into_iter()
            .filter(|b| b.bottom() == 1)
            .collect::<Vec<Brick>>(),
        |mut on_ground, h| {
            // look at bricks that start at this level and move them down
            // as far as they can travel
            let dropped = bricks
                .iter()
                .filter(|b| b.bottom() == h)
                .map(|b| b.drop(&on_ground))
                .collect::<Vec<Brick>>();

            dropped.into_iter().for_each(|b| {
                on_ground.push(b);
            });

            on_ground
        },
    );

    // find which blocks support which others: i supports j
    let supports: Vec<Vec<bool>> = (0..dropped.len())
        .map(|i| {
            (0..dropped.len())
                .map(|j| {
                    if i == j {
                        false
                    } else {
                        dropped[i].supports(&dropped[j])
                    }
                })
                .collect::<Vec<bool>>()
        })
        .collect();

    let r: Vec<usize> = (0..dropped.len())
        .filter(|i| {
            // block can be removed if it supports nothing
            // or the things it supports are also supported by something else
            supports[*i]
                .iter()
                .enumerate()
                .filter(|(_, s)| **s)
                .filter(|(j, _)| (0..dropped.len()).filter(|k| supports[*k][*j]).count() == 1)
                .count()
                == 0
        })
        .collect();

    (dropped, r, supports)
}

pub fn part_one(input: &str) -> Option<usize> {
    let (_, removable, _) = process(input);

    Some(removable.iter().count())
}

pub fn part_two(input: &str) -> Option<usize> {
    let (dropped, removable, supports) = process(input);

    let supported_by: Vec<Vec<usize>> = (0..dropped.len())
        .map(|i| (0..dropped.len()).filter(|j| supports[*j][i]).collect())
        .collect();

    Some(
        (0..dropped.len())
            .filter(|i| !removable.contains(i))
            .map(|i| {
                // if i is removed, find those that fall, and continue the process
                ((i + 1)..dropped.len())
                    .fold(vec![], |mut fall, j| {
                        // find the number of bricks supporting j once i has been removed
                        if dropped[j].bottom() > 1
                            && supported_by[j]
                                .iter()
                                // count the number of remaining blocks that support j
                                .filter(|&k| *k != i && !fall.contains(k))
                                .count()
                                == 0
                        {
                            // nothing supporting j - it will collapse
                            fall.push(j);
                        }
                        fall
                    })
                    .len()
            })
            .sum::<usize>(),
    )
}