use nom::{
    bytes::complete::tag, character::complete::u64 as nom_u64, combinator::map,
    multi::separated_list1, IResult,
};
use std::fs;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_7.txt")?;
    let crabs = parse_swarm(&input).unwrap().1;
    dbg!(crabs.linear_min_fuel_sum());
    dbg!(crabs.exponential_min_fuel_sum());
    Ok(())
}

#[derive(
    Default, Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, derive_more::Add, derive_more::Sum,
)]
struct Fuel(u64);

#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct CrabPosition(u64);

impl CrabPosition {
    fn linear_fuel(&self, rhs: &Self) -> Fuel {
        if self > rhs {
            Fuel(self.0 - rhs.0)
        } else {
            Fuel(rhs.0 - self.0)
        }
    }

    fn exponential_fuel(&self, rhs: &Self) -> Fuel {
        let linear_difference = if self > rhs {
            self.0 - rhs.0
        } else {
            rhs.0 - self.0
        };
        Fuel(linear_difference * (linear_difference + 1) / 2)
    }
}

#[derive(Default, Debug)]
struct CrabSwarm {
    crabs: Vec<CrabPosition>,
}

impl CrabSwarm {
    fn new(mut crabs: Vec<CrabPosition>) -> CrabSwarm {
        crabs.sort();
        CrabSwarm { crabs }
    }

    fn linear_min_fuel_sum(&self) -> (CrabPosition, Fuel) {
        (self.crabs[0].0..self.crabs[self.crabs.len() - 1].0)
            .map(CrabPosition)
            .map(|pos| {
                (
                    pos,
                    self.crabs.iter().map(|crab| crab.linear_fuel(&pos)).sum(),
                )
            })
            .min_by_key(|(_pos, fuel)| *fuel)
            .expect("Expected at least one crab")
    }

    fn exponential_min_fuel_sum(&self) -> (CrabPosition, Fuel) {
        (self.crabs[0].0..self.crabs[self.crabs.len() - 1].0)
            .map(CrabPosition)
            .map(|pos| {
                (
                    pos,
                    self.crabs
                        .iter()
                        .map(|crab| crab.exponential_fuel(&pos))
                        .sum(),
                )
            })
            .min_by_key(|(_pos, fuel)| *fuel)
            .expect("Expected at least one crab")
    }
}

fn parse_swarm(input: &str) -> IResult<&str, CrabSwarm> {
    map(
        separated_list1(tag(","), map(nom_u64, CrabPosition)),
        CrabSwarm::new,
    )(input)
}