use aoc2024::matrix;
use std::collections::HashSet;

static PLUS_DELTAS: [(isize, isize); 4] = [
    (-1, 0), (0, -1), (0, 1), (1, 0)
];

type Price = u32;
type Matrix = matrix::Matrix<char>;

fn adj_same_region(m: &Matrix, pos: matrix::Pos) -> Vec<matrix::Pos> {
    let label = m.get(pos);
    PLUS_DELTAS
        .iter()
        .filter_map(|&d| m.pos_move(pos, d))
        .filter(|&p| m.get(p) == label)
        .collect()
}

fn regions_price(m: &Matrix, version: u8) -> Price {
    let mut result = 0;
    let mut visited: HashSet<matrix::Pos> = HashSet::new();
    for i in 0..m.limit.0 {
        for j in 0..m.limit.1 {
            let pos = (i, j);
            if ! visited.contains(&pos) {
                let (price, pss) = region(m, pos, version);
                result += price;
                for p in pss {
                    visited.insert(p);
                }
            }
        }
    }
    result
}

fn region_v1(m: &Matrix, pos: matrix::Pos) -> (Price, HashSet<matrix::Pos>) {
    let mut density = 0;
    let mut visited: HashSet<matrix::Pos> = HashSet::new();
    let mut pending: HashSet<matrix::Pos> = [pos].into();
    let mut new_pending: HashSet<matrix::Pos>;
    while pending.len() > 0 {
        new_pending = HashSet::new();
        for p in pending {
            let adj = adj_same_region(m, p);
            density += adj.len();
            for a in adj {
                if !visited.contains(&a) {
                    new_pending.insert(a);
                }
            }
            visited.insert(p);
        }
        pending = new_pending;
    }
    let area = visited.len() as Price;
    let perimeter = area * 4 - density as Price;
    (area * perimeter, visited)
}

fn p1(input: &str) -> String {
    let m: Matrix = matrix::Matrix::new(input, |c| c);
    let result = regions_price(&m, 1);

    result.to_string()
}

fn pos_move(p: matrix::Pos, d: matrix::PosDelta) -> (isize, isize) {
    (p.0 as isize + d.0, p.1 as isize + d.1)
}

#[derive(Debug)]
struct Lines {
    hs: Vec<(bool, matrix::Pos)>,
    vs: Vec<(bool, matrix::Pos)>,
}

impl Lines {
    fn new() -> Self {
        Lines { hs: vec![], vs: vec![] }
    }

    fn from_pos(pos: matrix::Pos, in_region: &HashSet<matrix::Pos>) -> Self {
        // println!("Checking pos {pos:?}");
        let mut hs = vec![];
        let mut vs = vec![];
        for p in PLUS_DELTAS.iter().map(|&d| pos_move(pos, d)) {
            // println!("  Checking p {p:?}");
            let p_u = (p.0 as usize, p.1 as usize);
            // println!("  Checking p_u {p_u:?}");
            if ! in_region.contains(&p_u) {
                if p.0 == pos.0 as isize {
                    // println!("    PUSHING V ({} vs {}!", p.1, pos.1);
                    vs.push(
                        if p.1 < pos.1 as isize {
                            (true, pos)
                        } else {
                            (false, (pos.0, pos.1+1))
                        }
                    )
                } else {
                    // println!("    PUSHING H ({} vs {}!", p.0, pos.0);
                    hs.push(
                        if p.0 < pos.0 as isize {
                            (true, pos)
                        } else {
                            (false, (pos.0+1, pos.1))
                        }
                    )
                }
            }
        }

        // println!("  Result: hs={hs:?} vs={vs:?}");
        Lines { hs, vs }
    }

    fn append(&mut self, other: &mut Lines) {
        self.hs.append(&mut other.hs);
        self.vs.append(&mut other.vs);
    }

    fn count(&mut self) -> u32 {
        self.hs[..].sort();
        self.vs[..].sort_by(|(d1, (x1, y1)), (d2, (x2, y2))| (d1, y1, x1).cmp(&(d2, y2, x2)));
        // println!("Sorted self: {self:?}");
        self.hs.chunk_by(|(d1, (x1, y1)), (d2, (x2, y2))| d1 == d2 && x1 == x2 && *y1 == y2-1).collect::<Vec<_>>().len() as u32 +
        self.vs.chunk_by(|(d1, (x1, y1)), (d2, (x2, y2))| d1 == d2 && y1 == y2 && *x1 == x2-1).collect::<Vec<_>>().len() as u32
    }
}

// fn count_consec_chunks(xs: &[matrix::Pos]) -> u32 {
//     xs.chunk_by(|().len()
// }


fn region(m: &Matrix, pos: matrix::Pos, version: u8) -> (Price, HashSet<matrix::Pos>) {
    match version {
        1 => region_v1(m, pos),
        2 => {
            let region = region_v2(m, pos);
            let mut lines = Lines::new();
            // println!("\n\n\n\nRegion: {region:?}");
            for pos in region.iter() {
                lines.append(&mut Lines::from_pos(*pos, &region));
                // println!("Lines after processing {pos:?}: {lines:?}");
            }
            let nlines = lines.count();
            let price = nlines * region.len() as Price;
            // println!("Price of region {pos:?} = {price} ({:?} x {:?})", nlines, region.len());
            (price, region)
        }
        _ => panic!(),
    }
}

fn region_v2(m: &Matrix, pos: matrix::Pos) -> HashSet<matrix::Pos> {
    let mut visited: HashSet<matrix::Pos> = HashSet::new();
    let mut pending: HashSet<matrix::Pos> = [pos].into();
    while pending.len() > 0 {
        let mut new_pending: HashSet<matrix::Pos> = HashSet::new();
        for pos in pending {
            // println!("Processing pos={pos:?}");
            let region_adj = adj_same_region(m, pos);
            for a in region_adj {
                if !visited.contains(&a) {
                    new_pending.insert(a);
                }
            }
            visited.insert(pos);
        }
        pending = new_pending;
    }
    visited
}

fn p2(input: &str) -> String {
    let m: Matrix = matrix::Matrix::new(input, |c| c);
    let result = regions_price(&m, 2);

    result.to_string()
}

fn main() {
    aoc2024::run_day("12", p1, p2);
}

#[cfg(test)]
mod tests {
    use super::{p1, p2};

    #[test]
    fn day12_p1() {
        assert_eq!(p1(&aoc2024::read_input("12")), "1344578");
    }

    #[test]
    fn day12_p2() {
        assert_eq!(p2(&aoc2024::read_input("12")), "814302");
    }
}