use aoc_runner_derive::{aoc, aoc_generator};
use grid::{Coord2d, Grid, MirrorAxis};
use itertools::Itertools;
use std::collections::HashSet;
use std::fmt::Display;
use std::iter::repeat_n;

type CellType = bool;

#[derive(Debug, Clone)]
struct PresentShape {
    variants: Vec<Grid<CellType>>,
}

impl From<&str> for PresentShape {
    fn from(value: &str) -> Self {
        let (_idx, rest) = value.split_once(":\n").unwrap();

        let shape_raw: Grid<u8> = rest.parse().unwrap();
        let mut shape: Grid<CellType> = shape_raw.same_shape(false);
        for pos in shape_raw.find_all(&b'#') {
            shape.set(&pos, true);
        }

        Self {
            variants: (0..4)
                .map(|rot| shape.rotated(rot))
                .chain([MirrorAxis::X, MirrorAxis::Y].map(|axis| shape.mirrored(axis)))
                .unique()
                .collect(),
        }
    }
}

impl Display for PresentShape {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("{}", self.variants[0]))
    }
}

impl PresentShape {
    fn popcount(&self) -> u64 {
        self.variants[0].count(&true) as u64
    }
}

#[derive(Debug, Clone)]
struct Present {
    shape: usize,
}

#[derive(Debug, Clone)]
struct ChristmasTree {
    area: Grid<CellType>,
    presents: Vec<Present>,
}

#[derive(Debug, Clone)]
struct PresentsProblem {
    trees: Vec<ChristmasTree>,
    shapes: Vec<PresentShape>,
}

#[aoc_generator(day12)]
fn parse(input: &str) -> PresentsProblem {
    let mut parts = input.split("\n\n").collect_vec();
    let trees_s = parts.pop().unwrap();

    let shapes: Vec<PresentShape> = parts.into_iter().map(|p| p.into()).collect();
    let mut trees = Vec::new();

    for l in trees_s.lines() {
        let (d, el) = l.split_once(": ").unwrap();
        let (w, h) = d.split_once('x').unwrap();
        let present_counts: Vec<usize> = el
            .split_ascii_whitespace()
            .map(|count| count.parse().unwrap())
            .collect_vec();
        let presents = present_counts
            .iter()
            .enumerate()
            .flat_map(|(i, count)| repeat_n(Present { shape: i }, *count))
            .collect();
        trees.push(ChristmasTree {
            area: Grid::with_shape(w.parse().unwrap(), h.parse().unwrap(), false),
            presents,
        });
    }

    PresentsProblem { trees, shapes }
}

// place b on a and test if any # overlap any non-.
// fn overlaps(a: &Grid<CellType>, b: &Grid<CellType>, p: &Coord2d) -> bool {
//     b.find_all(&true)
//         .any(|c| a.get(&(c + p)).is_some_and(|c| *c))
// }

impl PresentsProblem {
    fn place_shape(
        &self,
        mut t: ChristmasTree,
        idx: usize,
        pos: Coord2d,
        variant: usize,
    ) -> Option<ChristmasTree> {
        let variant = &self.shapes[t.presents[idx].shape].variants[variant];

        for xy in variant.find_all(&true) {
            if let Some(was) = t.area.set(&(pos + &(xy)), true)
                && was
            {
                // overlapped
                return None;
            }
        }
        Some(t)
    }

    fn place_next(
        &self,
        t: ChristmasTree,
        cur: usize,
        cache: &mut HashSet<(usize, Grid<CellType>)>, // impossible shapes
    ) -> Option<ChristmasTree> {
        if cur == t.presents.len() {
            // done
            return Some(t);
        }
        if cache.contains(&(cur, t.area.clone())) {
            // doomed
            return None;
        }
        let variants = &self.shapes[t.presents[cur].shape].variants;
        for (i, _v) in variants.iter().enumerate() {
            for y in 0..&t.area.height() - 2 {
                for x in 0..&t.area.width() - 2 {
                    if let Some(new_t) = self.place_shape(
                        t.clone(),
                        cur,
                        Coord2d {
                            x: x as i64,
                            y: y as i64,
                        },
                        i,
                    ) {
                        if let Some(solution) = self.place_next(new_t, cur + 1, cache) {
                            return Some(solution);
                        }
                    }
                }
            }
        }
        cache.insert((cur, t.area.clone()));
        None
    }
}

#[aoc(day12, part1)]
fn part1(input: &PresentsProblem) -> u64 {
    let input = input.clone();
    let mut count = 0;

    for t in input.trees.iter() {
        let area = (t.area.width() * t.area.height()) as u64;
        let (occupied_area, present_count) = t
            .presents
            .iter()
            .map(|p| (input.shapes[p.shape].popcount(), 1))
            .reduce(|(o_a, pc_a), (o_b, pc_b)| (o_a + o_b, pc_a + pc_b))
            .unwrap();
        if occupied_area > area {
            // definitely impossible
            continue;
        }
        let available_3x3s = (t.area.width() / 3) * (t.area.height() / 3);
        if available_3x3s >= present_count {
            // definitely_possible
            count += 1;
            continue;
        }
        if input
            .place_next(t.clone(), 0, &mut HashSet::new())
            .is_some()
        {
            count += 1;
            continue;
        }
    }

    count
}

#[aoc(day12, part2)]
fn part2(_input: &PresentsProblem) -> u64 {
    0
}

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

    const EXAMPLE: &str = "0:
###
##.
##.

1:
###
##.
.##

2:
.##
###
##.

3:
##.
###
##.

4:
###
#..
###

5:
###
.#.
###

4x4: 0 0 0 0 2 0
12x5: 1 0 1 0 2 2
12x5: 1 0 1 0 3 2";

    #[rstest]
    #[case(EXAMPLE, 2)]
    fn part1_example(#[case] input: &str, #[case] expected: u64) {
        assert_eq!(part1(&parse(input)), expected);
    }

    #[rstest]
    #[case(EXAMPLE, 0)]
    fn part2_example(#[case] input: &str, #[case] expected: u64) {
        assert_eq!(part2(&parse(input)), expected);
    }
}