package main

import (
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"
)

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

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

func puzzle1(input string) int {
	return do(input, 80)
}

func puzzle2(input string) int {
	return do(input, 256)
}

func do(input string, days int) int {
	var fish []int
	total := 0
	// Parse initial state
	for _, s := range strings.Split(input, ",") {
		f, _ := strconv.Atoi(s)
		fish = append(fish, f)
	}
	// Run simulation for each fish
	for _, f := range fish {
		total += countDescendants(f, days)
	}
	return total
}

var cache = map[string]int{}

func countDescendants(f int, days int) int {
	if days == 0 || f-days >= 0 {
		// Won't grow further
		return 1
	}
	// Check cache
	if c, ok := cache[fmt.Sprintf("%d:%d", f, days)]; ok {
		return c
	}
	result := 0
	if f == 0 {
		// Add new fish
		result = countDescendants(6, days-1) + countDescendants(8, days-1)
	} else {
		// Continue counting
		result = countDescendants(f-1, days-1)
	}
	// Cache result
	cache[fmt.Sprintf("%d:%d", f, days)] = result
	return result
}