use nom::{
    bytes::complete::tag,
    character::complete::{
        self, space0, space1,
    },
    multi::separated_list1,
    sequence::{preceded, tuple},
    IResult,
};
use std::iter::zip;

fn parse_races(input: &str) -> IResult<&str, Vec<(u64, u64)>> {
    let (input, times) = preceded(tuple((tag("Time:"),space1)), separated_list1(space1, complete::u64))(input)?;
    let (input, distances) = preceded(tuple((tag("\nDistance:"),space1)), separated_list1(space1, complete::u64))(input)?;
    Ok((input, zip(times, distances).collect()))
}

fn parse_races_alt(input: &str) -> IResult<&str, (u64, u64)> {
    let (input, times) = preceded(tuple((tag("Time:"),space0)), complete::u64)(input)?;
    let (input, distances) = preceded(tuple((tag("\nDistance:"),space0)), complete::u64)(input)?;
    Ok((input, (times, distances)))
}

pub fn part1(input: &str) -> String {
    let (_, races) = parse_races(input).unwrap();
    races.into_iter().map(|(time, distance)| {
        let upper_bound_full = (time as f64 + ((time as f64).powi(2) - 4.0 * distance as f64).sqrt()) / 2.0;
        let upper_bound = if upper_bound_full - upper_bound_full.floor() == 0.0 {
            upper_bound_full.floor() as u64 - 1
        } else {
            upper_bound_full.floor() as u64
        };
        let lower_bound_full = (time as f64 - ((time as f64).powi(2) - 4.0 * distance as f64).sqrt()) / 2.0;
        let lower_bound = if lower_bound_full - lower_bound_full.ceil() == 0.0 {
            lower_bound_full.ceil() as u64 + 1
        } else {
            lower_bound_full.ceil() as u64
        };
        upper_bound - lower_bound + 1
    }).product::<u64>().to_string()
}

pub fn part2(input: &str) -> String {
    let (_, (time, distance)) = parse_races_alt(input.replace(" ", "").as_str()).unwrap();
    let upper_bound_full = (time as f64 + ((time as f64).powi(2) - 4.0 * distance as f64).sqrt()) / 2.0;
    let upper_bound = if upper_bound_full - upper_bound_full.floor() == 0.0 {
        upper_bound_full.floor() as u64 - 1
    } else {
        upper_bound_full.floor() as u64
    };
    let lower_bound_full = (time as f64 - ((time as f64).powi(2) - 4.0 * distance as f64).sqrt()) / 2.0;
    let lower_bound = if lower_bound_full - lower_bound_full.ceil() == 0.0 {
        lower_bound_full.ceil() as u64 + 1
    } else {
        lower_bound_full.ceil() as u64
    };
    (upper_bound - lower_bound + 1).to_string()
}

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