path: root/2019/src/bin/day_14.rs

use rpds::map::red_black_tree_map::RedBlackTreeMap;
use rpds::rbt_map;
use rpds::vector::Vector;
use std::fmt;
use std::io;
use std::io::prelude::*;
use std::iter;
use std::process;
use std::str::FromStr;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "Day 14: Space Stoichiometry")]
/// Finds how much ore you need to produce one fuel.
///
/// Recipes are passed in on stdin, one per line.
///
/// See https://adventofcode.com/2019/day/14 for details.
struct Opt {
    #[structopt(long = "available-ore")]
    available_ore: Option<i64>,
}

fn main() {
    let stdin = io::stdin();
    let opt = Opt::from_args();

    let recipes = stdin
        .lock()
        .lines()
        .map(|x| exit_on_failed_assertion(x, "Error reading input"))
        .map(|x| exit_on_failed_assertion(x.parse::<Recipe>(), "Input was not a valid recipe"))
        .collect::<Vec<_>>();

    match opt.available_ore {
        Some(ore) => println!("{}", max_fuel_production(ore, &recipes)),
        None => println!("{}", Desires::new(1).min_ore_required(&recipes)),
    }
}

fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message: &str) -> A {
    match data {
        Ok(data) => data,
        Err(e) => {
            eprintln!("{}: {}", message, e);
            process::exit(1);
        }
    }
}

fn max_fuel_production(ore_max: i64, recipes: &[Recipe]) -> i64 {
    binary_search_max_fuel_production(
        ore_max / Desires::new(1).min_ore_required(&recipes),
        2 * ore_max / Desires::new(1).min_ore_required(&recipes),
        ore_max,
        recipes,
    )
}

fn binary_search_max_fuel_production(
    fuel_min: i64,
    fuel_max: i64,
    ore_max: i64,
    recipes: &[Recipe],
) -> i64 {
    if fuel_max - fuel_min <= 1 {
        fuel_min
    } else if Desires::new((fuel_min + fuel_max) / 2).min_ore_required(recipes) <= ore_max {
        binary_search_max_fuel_production((fuel_min + fuel_max) / 2, fuel_max, ore_max, recipes)
    } else {
        binary_search_max_fuel_production(fuel_min, (fuel_min + fuel_max) / 2, ore_max, recipes)
    }
}

#[derive(Debug, Clone)]
struct Recipe {
    ingredients: Vector<Chemical>,
    output: Chemical,
}

#[derive(Default, Debug, Clone)]
struct Chemical {
    name: String,
    quantity: i64,
}

impl FromStr for Recipe {
    type Err = ParseErr;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        // 2 XSNKB, 15 ZVMCB, 3 KDFNZ => 2 RFLX
        s.replace(" => ", "=")
            .replace(", ", ",")
            .split(|c| c == ',' || c == '=')
            .map(|chem| chem.parse::<Chemical>())
            .collect::<Result<Vector<Chemical>, ParseErr>>()
            .map(|chemicals| Recipe {
                ingredients: chemicals
                    .drop_last()
                    .expect("Assertion failed: line did not have any chemicals"),
                output: chemicals
                    .last()
                    .cloned()
                    .expect("Assertion failed: line did not have any chemicals"),
            })
    }
}

impl Recipe {
    fn required_scale(&self, desired_quantity: i64) -> i64 {
        (desired_quantity + self.output.quantity - 1) / self.output.quantity
    }
}

impl FromStr for Chemical {
    type Err = ParseErr;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        // 1 FUEL
        match s.split(' ').collect::<Vec<_>>()[..] {
            [quantity, name] => quantity
                .parse::<i64>()
                .map_err(|_| ParseErr)
                .map(|q| Chemical {
                    name: name.to_string(),
                    quantity: q,
                }),
            _ => Err(ParseErr),
        }
    }
}

impl Chemical {
    fn scale(&self, scale: i64) -> Chemical {
        Chemical {
            name: self.name.clone(),
            quantity: self.quantity * scale,
        }
    }
}

#[derive(Debug, Clone, Copy)]
struct ParseErr;

impl fmt::Display for ParseErr {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Error parsing input")
    }
}

impl std::error::Error for ParseErr {}

#[derive(Debug, Clone)]
struct Desires {
    chemicals: RedBlackTreeMap<String, i64>,
}

impl Desires {
    fn new(fuel: i64) -> Desires {
        Desires {
            chemicals: rbt_map!["FUEL".to_string() => fuel, "ORE".to_string() => 0],
        }
    }

    fn min_ore_required(&self, recipes: &[Recipe]) -> i64 {
        iter::successors(Some(self.clone()), |prev| Some(prev.next(recipes)))
            .find(|desires| desires.is_only_ore())
            .unwrap()
            .chemicals
            .get("ORE")
            .cloned()
            .unwrap()
    }

    fn is_only_ore(&self) -> bool {
        !self
            .chemicals
            .iter()
            .any(|(name, quantity)| *quantity > 0 && name != "ORE")
    }

    fn next(&self, recipes: &[Recipe]) -> Desires {
        self.chemicals
            .iter()
            .find(|(name, quantity)| **quantity > 0 && *name != "ORE")
            .map(|(mixing, quantity)| {
                self.with_mixed_recipe(
                    recipes
                        .iter()
                        .find(|recipe| recipe.output.name == *mixing)
                        .expect("Required chemical without a recipe"),
                    *quantity,
                )
            })
            .unwrap_or(self.clone())
    }

    fn with_mixed_recipe(&self, recipe: &Recipe, desired_quantity: i64) -> Desires {
        recipe.ingredients.iter().fold(
            self.with_chemical(
                recipe
                    .output
                    .scale(-1 * recipe.required_scale(desired_quantity)),
            ),
            |acc, next_ingredient| {
                acc.with_chemical(next_ingredient.scale(recipe.required_scale(desired_quantity)))
            },
        )
    }

    fn with_chemical(&self, chemical: Chemical) -> Desires {
        Desires {
            chemicals: match self.chemicals.get(&chemical.name) {
                Some(existing_quantity) => self
                    .chemicals
                    .insert(chemical.name.clone(), existing_quantity + chemical.quantity),
                None => self
                    .chemicals
                    .insert(chemical.name.clone(), chemical.quantity),
            },
        }
    }
}