use combinations::Combinations;
use nom::{
    bytes::complete::{is_not, tag, take_until},
    multi::separated_list1,
    IResult, Parser,
};
use nom_locate::LocatedSpan;

type Span<'a> = LocatedSpan<&'a str>;
type SpanPoint<'a> = LocatedSpan<&'a str, Point>;

#[derive(Debug, Clone, Ord, Eq, PartialEq, PartialOrd)]
struct Galaxy {
    point: Point,
}

#[derive(Debug, Clone, Ord, Eq, PartialEq, PartialOrd)]
struct Point {
    x: usize,
    y: usize,
}

fn with_xy(span: Span) -> SpanPoint {
    span.map_extra(|_| Point {
        x: span.get_column() as usize - 1,
        y: span.location_line() as usize - 1,
    })
}

fn parse_grid(input: Span) -> IResult<Span, Vec<Galaxy>> {
    let (input, _) = take_until("#")(input)?;
    let (input, galaxies) = separated_list1(
        is_not("#"),
        tag("#").map(with_xy).map(|span| Galaxy {
            point: Point {
                x: span.extra.x,
                y: span.extra.y,
            },
        }),
    )(input)?;
    Ok((input, galaxies))
}

fn expand_universe(galaxies: Vec<Galaxy>, rows: usize, columns: usize, amount: usize) -> Vec<Galaxy> {
    let galaxy_xs = galaxies
        .iter()
        .map(|galaxy| galaxy.point.x)
        .collect::<Vec<_>>();
    let empty_columns = (0..columns).filter(|x| !galaxy_xs.contains(&x));
    let expanded_galaxies = galaxies
        .iter()
        .map(|galaxy| {
            let number_empty_columns = empty_columns
                .clone()
                .filter(|x| x < &galaxy.point.x)
                .count();
            Galaxy {
                point: Point {
                    x: galaxy.point.x + (amount - 1) * number_empty_columns,
                    y: galaxy.point.y,
                },
            }
        })
        .collect::<Vec<Galaxy>>();
    let galaxy_ys = expanded_galaxies
        .iter()
        .map(|galaxy| galaxy.point.y)
        .collect::<Vec<_>>();
    let empty_rows = (0..rows).filter(|y| !galaxy_ys.contains(&y));
    expanded_galaxies
        .iter()
        .map(|galaxy| {
            let number_empty_rows = empty_rows.clone().filter(|y| y < &galaxy.point.y).count();
            Galaxy {
                point: Point {
                    x: galaxy.point.x,
                    y: galaxy.point.y + (amount - 1) * number_empty_rows,
                },
            }
        })
        .collect()
}

pub fn part1(input: &str) -> String {
    let columns = input.lines().count();
    let rows = input.lines().next().unwrap().chars().count();
    let (_, galaxies) = parse_grid(Span::new(input)).unwrap();
    let expanded_galaxies = expand_universe(galaxies, rows, columns, 2);
    Combinations::new(expanded_galaxies, 2).map(|pair| {
        (pair[0].point.x as i64 - pair[1].point.x as i64).abs() + (pair[0].point.y as i64 - pair[1].point.y as i64).abs()
    }).sum::<i64>().to_string()
}

pub fn part2_test(input: &str) -> String {
    let columns = input.lines().count();
    let rows = input.lines().next().unwrap().chars().count();
    let (_, galaxies) = parse_grid(Span::new(input)).unwrap();
    let expanded_galaxies = expand_universe(galaxies, rows, columns, 10);
    Combinations::new(expanded_galaxies, 2).map(|pair| {
        (pair[0].point.x as i64 - pair[1].point.x as i64).abs() + (pair[0].point.y as i64 - pair[1].point.y as i64).abs()
    }).sum::<i64>().to_string()
}

pub fn part2(input: &str) -> String {
    let columns = input.lines().count();
    let rows = input.lines().next().unwrap().chars().count();
    let (_, galaxies) = parse_grid(Span::new(input)).unwrap();
    let expanded_galaxies = expand_universe(galaxies, rows, columns, 1000000);
    Combinations::new(expanded_galaxies, 2).map(|pair| {
        (pair[0].point.x as i64 - pair[1].point.x as i64).abs() + (pair[0].point.y as i64 - pair[1].point.y as i64).abs()
    }).sum::<i64>().to_string()
}

pub mod prelude {
    pub use super::part1;
    pub use super::part2;
    pub use super::part2_test;
}