package main

import (
	"log"
	"math"
	"os"
	"strings"
)

func main() {
	inputBytes, _ := os.ReadFile("input.txt")
	input := strings.Split(string(inputBytes), "\n")

	log.Println("Puzzle 1:", puzzle1(input))
	log.Println("Puzzle 2:", puzzle2(input))
}

func puzzle1(input []string) int {
	return do(input, 10)
}

func puzzle2(input []string) int {
	return do(input, 40)
}

func parseInput(input []string) (template string, rules map[string]string) {
	template = input[0]
	rules = map[string]string{}
	for _, line := range input[2:] {
		parts := strings.Split(line, " -> ")
		rules[parts[0]] = parts[1]
	}
	return
}

func do(input []string, steps int) int {
	// Parse input
	template, rules := parseInput(input)
	// Polymerize
	pairs := map[string]int{}
	for i := 0; i < len(template)-1; i++ {
		pairs[template[i:i+2]]++
	}
	for i := 0; i < steps; i++ {
		newPairs := map[string]int{}
		for pair, count := range pairs {
			rule := rules[pair]
			newPairs[pair[0:1]+rule] += count
			newPairs[rule+pair[1:2]] += count
		}
		pairs = newPairs
	}
	// Sum up counts for each letter
	counts := map[byte]int{}
	for pair, count := range pairs {
		counts[pair[0]] += count
	}
	// Add count for last letter in template
	counts[template[len(template)-1]]++
	// Get minimum and maximum counts
	max, min := math.MinInt, math.MaxInt
	for _, count := range counts {
		if count > max {
			max = count
		}
		if count < min {
			min = count
		}
	}
	// Return difference
	return max - min
}