use std::ops::Range;

const INPUT: &str = include_str!("../../data/day17");

#[derive(Clone, Copy, Default)]
struct Dimension {
    offset: u32,
    size: u32,
}

impl Dimension {
    fn map(&self, pos: i32) -> Option<usize> {
        let p = (self.offset as i32) + pos;
        if p >= 0 && p < (self.size as i32) {
            Some(p as usize)
        } else {
            None
        }
    }

    fn include(&self, pos: i32) -> Self {
        let left = std::cmp::min(pos, -(self.offset as i32));
        let right = std::cmp::max(pos, (self.size - self.offset - 1) as i32);
        Self {
            offset: (-left) as u32,
            size: (right - left + 1) as u32,
        }
    }

    fn extend(&self) -> Self {
        Self {
            offset: self.offset + 1,
            size: self.size + 2,
        }
    }
}

impl IntoIterator for Dimension {
    type IntoIter = Range<i32>;
    type Item = i32;

    fn into_iter(self) -> Self::IntoIter {
        -(self.offset as i32)..((self.size - self.offset) as i32)
    }
}

#[derive(Clone)]
struct Field {
    x_dim: Dimension,
    y_dim: Dimension,
    z_dim: Dimension,
    w_dim: Dimension,
    cells: Vec<bool>,
}

impl Default for Field {
    fn default() -> Self {
        Self {
            x_dim: Dimension { offset: 0, size: 1},
            y_dim: Dimension { offset: 0, size: 1},
            z_dim: Dimension { offset: 0, size: 1},
            w_dim: Dimension { offset: 0, size: 1},
            cells: vec![false],
        }
    }
}

impl Field {
    fn alloc(x_dim: Dimension, y_dim: Dimension, z_dim: Dimension, w_dim: Dimension) -> Self {
        let mut cells = Vec::new();
        let size = (x_dim.size * y_dim.size * z_dim.size * w_dim.size) as usize;
        cells.resize(size, false);
        Self { x_dim, y_dim, z_dim, w_dim, cells }
    }

    fn _addr(&self, x: i32, y: i32, z: i32, w: i32) -> Option<usize> {
        let x = self.x_dim.map(x)?;
        let y = self.y_dim.map(y)?;
        let z = self.z_dim.map(z)?;
        let w = self.w_dim.map(w)?;
        let addr = ((w * (self.z_dim.size as usize) + z) * (self.y_dim.size as usize) + y) * (self.x_dim.size as usize) + x;
        Some(addr)
    }

    fn _get(&self, x: i32, y: i32, z: i32, w: i32) -> Option<bool> {
        let addr = self._addr(x, y, z, w)?;
        Some(self.cells[addr])
    }

    fn get(&self, x: i32, y: i32, z: i32, w: i32) -> bool {
        self._get(x, y, z, w).unwrap_or(false)
    }

    fn set(&mut self, x: i32, y: i32, z: i32, w: i32, value: bool) {
        if let Some(addr) = self._addr(x, y, z, w) {
            self.cells[addr] = value;
        } else if value {
            // don't resize to set "false"
            let x_dim = self.x_dim.include(x);
            let y_dim = self.y_dim.include(y);
            let z_dim = self.z_dim.include(z);
            let w_dim = self.w_dim.include(w);
            let mut new_field = Self::alloc(x_dim, y_dim, z_dim, w_dim);
            for w in self.w_dim {
                for z in self.z_dim {
                    for y in self.y_dim {
                        for x in self.x_dim {
                            let addr = new_field._addr(x, y, z, w).unwrap();
                            new_field.cells[addr] = self.get(x, y, z, w);
                        }
                    }
                }
            }
            let addr = new_field._addr(x, y, z, w).unwrap();
            new_field.cells[addr] = value;
            *self = new_field;
        }
    }

    fn parse(data: &str) -> Self {
        let mut this = Self::default();
        for (y, line) in data.lines().enumerate() {
            for (x, chr) in line.chars().enumerate() {
                this.set(x as i32, y as i32, 0, 0, chr == '#');
            }
        }
        this
    }

    fn count_active_neighbors(&self, x: i32, y: i32, z: i32, w: i32) -> u32 {
        let mut count = 0;
        for dw in -1..=1 {
            for dz in -1..=1 {
                for dy in -1..=1 {
                    for dx in -1..=1 {
                        if dw == 0 && dz == 0 && dy == 0 && dx == 0 {
                            continue;
                        }
                        if self.get(x + dx, y + dy, z + dz, w + dw) {
                            count += 1;
                            if count > 3 { return count; } // more than 3 is always bad
                        }
                    }
                }
            }
        }
        count
    }

    fn cycle_3dim(&self) -> Self {
        let mut result = self.clone();
        let w = 0;
        for z in self.z_dim.extend() {
            for y in self.y_dim.extend() {
                for x in self.x_dim.extend() {
                    let count = self.count_active_neighbors(x, y, z, w);
                    let active = count == 3 || (count == 2 && self.get(x, y, z, w));
                    result.set(x, y, z, w, active);
                }
            }
        }
        result
    }

    fn cycle_4dim(&self) -> Self {
        let mut result = self.clone();
        for w in self.w_dim.extend() {
            for z in self.z_dim.extend() {
                for y in self.y_dim.extend() {
                    for x in self.x_dim.extend() {
                        let count = self.count_active_neighbors(x, y, z, w);
                        let active = count == 3 || (count == 2 && self.get(x, y, z, w));
                        result.set(x, y, z, w, active);
                    }
                }
            }
        }
        result
    }
}

impl std::fmt::Debug for Field {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        use std::fmt::Write;
        for z in self.z_dim {
            for w in self.w_dim {
                writeln!(f, "z = {}, w = {}", z, w)?;
                for y in self.y_dim {
                    for x in self.x_dim {
                        if x == 0 && y == 0 {
                            f.write_char(
                                if self.get(x, y, z, w) { 'X' } else { '+' }
                            )?;
                        } else {
                            f.write_char(
                                if self.get(x, y, z, w) { '#' } else { '.' }
                            )?;
                        }
                    }
                    f.write_char('\n')?;
                }
                f.write_char('\n')?;
            }
        }
        Ok(())
    }
}

fn main() {
    let input = Field::parse(INPUT);
    let mut field = input.clone();
    for _ in 0..6 {
        field = field.cycle_3dim();
    }
    println!("Active after 6 3-dim rounds: {}", field.cells.iter().filter(|&&c| c).count());
    let mut field = input.clone();
    for _ in 0..6 {
        field = field.cycle_4dim();
    }
    println!("Active after 6 4-dim rounds: {}", field.cells.iter().filter(|&&c| c).count());
    println!("{:?}", input);
}