use std::fmt::{Display, Write};
use std::fs::File;
use std::io::{BufRead, BufReader, Lines};
use std::iter::repeat;

// BOILERPLATE
type InputIter = Lines<BufReader<File>>;

fn get_input() -> InputIter {
    let f = File::open("input").unwrap();
    let br = BufReader::new(f);
    br.lines()
}

fn main() {
    println!("Problem 1 solution: {}", problem1(get_input()));
    println!("Problem 2 solution: {}", problem2(get_input()));
}

// PARSE

struct GalaxyMap {
    map: Vec<Vec<char>>,
    galaxy_positions: Vec<(usize, usize)>,
}

impl<T: BufRead> From<Lines<T>> for GalaxyMap {
    fn from(lines: Lines<T>) -> Self {
        let map = GalaxyMap {
            map: lines.map(|line| line.unwrap().chars().collect()).collect(),
            galaxy_positions: Vec::new(),
        };
        map
    }
}

impl GalaxyMap {
    fn universe_expansion(&mut self) {
        let mut columns_to_insert = Vec::new();

        for col in 0..self.map[0].len() {
            if self.map.iter().map(|row| row[col]).all(|c| c != '#') {
                columns_to_insert.push(col);
            }
        }

        for row in &mut self.map {
            let mut offset = 0;
            for col in &columns_to_insert {
                row.insert(*col + offset, '.');
                offset += 1;
            }
        }

        let rows_to_insert: Vec<_> = self
            .map
            .iter()
            .enumerate()
            .filter_map(|(idx, row)| if !row.contains(&'#') { Some(idx) } else { None })
            .collect();
        let gen_row: Vec<_> = repeat('.').take(self.map[0].len()).collect();
        let mut offset = 0;
        for idx in rows_to_insert {
            self.map.insert(idx + offset, gen_row.clone());
            offset += 1;
        }
    }
    fn universe_expansion2(&mut self) {
        let mut columns_to_grow = Vec::new();

        for col in 0..self.map[0].len() {
            if self.map.iter().map(|row| row[col]).all(|c| c != '#') {
                columns_to_grow.push(col);
            }
        }

        for row in &mut self.map {
            for col in &columns_to_grow {
                row[*col] = 'x';
            }
        }

        let rows_to_grow: Vec<_> = self
            .map
            .iter()
            .enumerate()
            .filter_map(|(idx, row)| if !row.contains(&'#') { Some(idx) } else { None })
            .collect();
        let gen_row: Vec<_> = repeat('x').take(self.map[0].len()).collect();
        for idx in rows_to_grow {
            self.map[idx] = gen_row.clone();
        }
    }
    fn find_galaxies(&mut self) {
        for (y, row) in self.map.iter().enumerate() {
            for (x, _c) in row.iter().enumerate().filter(|(_x, c)| **c == '#') {
                self.galaxy_positions.push((x, y));
            }
        }
    }

    fn galaxy_distance(&self, a: usize, b: usize) -> u64 {
        let (a_x, a_y) = self.galaxy_positions[a];
        let (b_x, b_y) = self.galaxy_positions[b];

        let x_dist = b_x.abs_diff(a_x);
        let y_dist = b_y.abs_diff(a_y);

        x_dist as u64 + y_dist as u64
    }
    fn galaxy_distance2(&self, a: usize, b: usize) -> u64 {
        let (a_x, a_y) = self.galaxy_positions[a];
        let (b_x, b_y) = self.galaxy_positions[b];
        let x_dist = b_x.abs_diff(a_x);
        let y_dist = b_y.abs_diff(a_y);

        let mut dist = x_dist as u64 + y_dist as u64;

        for row in a_y.min(b_y)..a_y.max(b_y) {
            if self.map[row].iter().all(|c| *c == 'x') {
                dist += 1000000 - 1;
            }
        }

        for col in a_x.min(b_x)..a_x.max(b_x) {
            if self.map.iter().all(|row| row[col] == 'x') {
                dist += 1000000 - 1;
            }
        }

        dist
    }
}

impl Display for GalaxyMap {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for row in &self.map {
            for c in row {
                f.write_char(*c)?;
            }
            f.write_char('\n')?;
        }
        Ok(())
    }
}

// PROBLEM 1 solution

fn problem1<T: BufRead>(input: Lines<T>) -> u64 {
    let mut map = GalaxyMap::from(input);
    map.universe_expansion();
    map.find_galaxies();
    println!("{}", map);

    let mut galaxies: Vec<_> = (0..map.galaxy_positions.len()).collect();
    let mut sum = 0u64;
    while let Some(target) = galaxies.pop() {
        sum += galaxies
            .iter()
            .map(|source| map.galaxy_distance(*source, target))
            .sum::<u64>();
    }
    sum
}

// PROBLEM 2 solution
fn problem2<T: BufRead>(input: Lines<T>) -> u64 {
    let mut map = GalaxyMap::from(input);
    map.universe_expansion2();
    map.find_galaxies();
    println!("{}", map);

    let mut galaxies: Vec<_> = (0..map.galaxy_positions.len()).collect();
    let mut sum = 0u64;
    while let Some(target) = galaxies.pop() {
        sum += galaxies
            .iter()
            .map(|source| map.galaxy_distance2(*source, target))
            .sum::<u64>();
    }
    sum
}

#[cfg(test)]
mod tests {
    use crate::*;
    use std::io::Cursor;

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

    #[test]
    fn problem1_example() {
        let c = Cursor::new(EXAMPLE);
        assert_eq!(problem1(c.lines()), 374);
    }

    #[test]
    fn problem2_example() {
        let c = Cursor::new(EXAMPLE);
        assert_eq!(problem2(c.lines()), 0);
    }
}