Day 19 AoC: The factory must grow

Part 1 takes 13 seconds, part 2 50 seconds. Pretty bad but cbf to find more optimizations methods. Also feels like some parts of the code are somewhat redundant(e.g. all the buy_xxx_robot functions, that could most likely be handled much more elegantly in just one function.

#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <format>
#include <stack>
#include <set>
#include <queue>
#include <array>
#include <chrono>
#include <scn/scn.h>

struct Blueprint {
	int id;
	int ore_robot;
	int clay_robot;
	int obsidian_robot_ore;
	int obsidian_robot_clay;
	int geode_robot_ore;
	int geode_robot_obsidian;

	int max_ore() {
		return std::max({ ore_robot, clay_robot, obsidian_robot_ore, geode_robot_ore });
	}
	int max_clay() {
		return obsidian_robot_clay;
	}
	int max_obsidian() {
		return geode_robot_obsidian;
	}
};

struct State {
	int time;
	int ore;
	int clay;
	int obsidian;
	int geode;
	int ore_robot;
	int clay_robot;
	int obsidian_robot;
	int geode_robot;

	std::array<bool, 4> prev_skipped = { false, false, false, false };

	auto operator<=>(const State& s) const = default;
};

std::ostream& operator<<(std::ostream& os, Blueprint& bp) {
	os << std::format("ID: {}\nOre Robot: {}\nClay Robot:{}\nObsidian Robot:\n\t{}\n\t{}\nGeode Robot:\n\t{}\n\t{}",
		bp.id, bp.ore_robot, bp.clay_robot, bp.obsidian_robot_ore, bp.obsidian_robot_clay, bp.geode_robot_ore, bp.geode_robot_obsidian);
	return os;
}

std::ostream& operator<<(std::ostream& os, State& s) {
	os << std::format("Minute {}\n{} Ore robot, total ore {}\n{} Clay Robot Total Clay {}\n{} Obsidian Robot, Total Obsidian {}\n{} Geode Robot, Total Geode {}",
		s.time, s.ore_robot, s.ore, s.clay_robot, s.clay, s.obsidian_robot, s.obsidian, s.geode_robot, s.geode);
	return os;
}

void add_resources(State& current) {
	current.ore += current.ore_robot;
	current.clay += current.clay_robot;
	current.obsidian += current.obsidian_robot;
	current.geode += current.geode_robot;
}

std::pair<bool, State> buy_ore_robot(State current, Blueprint& bp) {
	if (current.prev_skipped[0])
		return { false, current };
	bool bought = false;
	if (current.ore >= bp.ore_robot && !current.prev_skipped[0] && current.ore_robot < bp.max_ore()) {
		add_resources(current);
		current.ore -= bp.ore_robot;
		current.ore_robot++;
		bought = true;
	}

	return { bought, current };
}

std::pair<bool, State> buy_clay_robot(State current, Blueprint& bp) {
	bool bought = false;
	if (current.ore >= bp.clay_robot && !current.prev_skipped[1] && current.clay_robot < bp.max_clay()) {
		add_resources(current);
		current.ore -= bp.clay_robot;
		current.clay_robot++;
		bought = true;
	}

	return { bought, current };
}

std::pair<bool, State> buy_obsidian_robot(State current, Blueprint& bp) {
	bool bought = false;
	if (current.ore >= bp.obsidian_robot_ore && current.clay >= bp.obsidian_robot_clay && !current.prev_skipped[2] && current.obsidian_robot < bp.max_obsidian()) {
		add_resources(current);
		current.ore -= bp.obsidian_robot_ore;
		current.clay -= bp.obsidian_robot_clay;
		current.obsidian_robot++;
		bought = true;
	}

	return { bought, current };
}

std::pair<bool, State> buy_geode_robot(State current, Blueprint& bp) {
	bool bought = false;
	if (current.ore >= bp.geode_robot_ore && current.obsidian >= bp.geode_robot_obsidian) {
		add_resources(current);
		current.ore -= bp.geode_robot_ore;
		current.obsidian -= bp.geode_robot_obsidian;
		current.geode_robot++;
		bought = true;
	}

	return { bought, current };
}

int64_t run(std::vector<Blueprint>& blueprints, bool partTwo = false)
{
	int64_t total = 0;
	int64_t partTwo_res = 1;
	int minutes = partTwo ? 31 : 23;
	int run = 0;
	for (auto bp : blueprints) {
		if (partTwo && run > 2) break;
		run++;

		State max_geode{ 1,0,0,0,0,1,0,0,0 };

		State start_state = { 0,0,0,0,0,1,0,0,0 };
		std::stack<State> state_stack;
		std::set<State> state_set;
		state_stack.push(start_state);

		while (!state_stack.empty()) {
			State current_state = state_stack.top();
			state_stack.pop();
			if (current_state.time > minutes) continue;
			current_state.time++;

			if (state_set.contains(current_state)) continue;
			else state_set.insert(current_state);

			int n = (minutes - current_state.time) + 2;
			int total_possible = (n * (n - 1)) / 2;
			if ((current_state.geode + current_state.geode_robot * n + total_possible) < max_geode.geode) continue;

			auto geode = buy_geode_robot(current_state, bp);
			if (geode.first) state_stack.push(geode.second);
			else {
				auto obsidian = buy_obsidian_robot(current_state, bp);
				if (obsidian.first) state_stack.push(obsidian.second);

				auto clay = buy_clay_robot(current_state, bp);
				if (clay.first) state_stack.push(clay.second);

				auto ore = buy_ore_robot(current_state, bp);
				if (ore.first) state_stack.push(ore.second);
				current_state.prev_skipped[0] = ore.first;
				current_state.prev_skipped[1] = clay.first;
				current_state.prev_skipped[2] = obsidian.first;
			}

			add_resources(current_state);

			if (max_geode.geode < current_state.geode)
				max_geode = current_state;

			// if it could buy all, but didn't don't check the didn't buy anything path
			if(!(geode.first && current_state.prev_skipped[0] && current_state.prev_skipped[1] && current_state.prev_skipped[2]))
				state_stack.push(current_state);

		}
		std::cout << std::format("Blueprint {}: {} geodes opened\n", bp.id, max_geode.geode);
		total += max_geode.geode * bp.id;
		partTwo_res *= max_geode.geode;
	}

	return partTwo ? partTwo_res : total;
}


int main()
{
	std::fstream file("input.txt");
	std::string buf;

	std::vector<Blueprint> blueprints;

	while (std::getline(file, buf)) {
		Blueprint bp;
		scn::scan(buf, "Blueprint {}: Each ore robot costs {} ore. Each clay robot costs {} ore. Each obsidian robot costs {} ore and {} clay. Each geode robot costs {} ore and {} obsidian.:",
			bp.id, bp.ore_robot, bp.clay_robot, bp.obsidian_robot_ore, bp.obsidian_robot_clay, bp.geode_robot_ore, bp.geode_robot_obsidian);
		blueprints.emplace_back(bp);
	}

	auto a = std::chrono::high_resolution_clock::now();
	int64_t part_one = run(blueprints);

	auto b = std::chrono::high_resolution_clock::now();
	std::cout << "Part One: " << part_one << '\n';
	int64_t part_two = run(blueprints, true);
	auto c = std::chrono::high_resolution_clock::now();

	auto ba = std::chrono::duration_cast<std::chrono::milliseconds>(b - a);
	auto cb = std::chrono::duration_cast<std::chrono::milliseconds>(c - b);
	std::cout << "Part Two: " << part_two << '\n';
	std::cout << std::format("Time p1: {} p2: {}", ba, cb) << '\n';

}