package main

import (
	"fmt"
	"math"
	"strconv"
	"strings"

	h "git.bullercodeworks.com/brian/adventofcode/helpers"
)

func main() {
	inp := h.StdinToString()
	mc := part1(parseStones(inp))
	part2(parseStones(inp), mc)
}

func part1(inp []int) *MemoCache {
	mc := NewMemo()
	fmt.Println("# Part 1")
	fmt.Println("Stones After 25 Blinks:", mc.StonesAfter(inp, 25))
	return mc
}

func part2(inp []int, mc *MemoCache) {
	fmt.Println("# Part 2")
	fmt.Println("Stones After 75 Blinks:", mc.StonesAfter(inp, 75))
}

func blink(stone int) []int {
	var ret []int
	if stone == 0 {
		ret = append(ret, 1)
	} else if evenDigits(stone) {
		ret = append(ret, split(stone)...)
	} else {
		ret = append(ret, stone*2024)
	}
	return ret
}

// Attempt 1, I tracked the actual stones, but that's too much.
// Attempt 2, Switch from tracking the actual stones to just the length
type MemoKey struct {
	stone  int
	blinks int
}
type MemoCache struct {
	cache map[MemoKey]int
}

func NewMemo() *MemoCache {
	return &MemoCache{cache: make(map[MemoKey]int)}
}

func (mc *MemoCache) StonesAfter(stones []int, b int) int {
	var res int
	for _, s := range stones {
		res += mc.SizeAfterBlinks(s, b)
	}
	return res
}

func (mc *MemoCache) SizeAfterBlinks(s, b int) int {
	key := MemoKey{stone: s, blinks: b}
	if v, ok := mc.cache[key]; ok {
		return v
	}
	res := blink(s)
	if b == 1 {
		mc.cache[key] = len(res)
		return len(res)
	}
	n := mc.SizeAfterBlinks(res[0], b-1)
	if len(res) == 2 {
		n += mc.SizeAfterBlinks(res[1], b-1)
	}
	mc.cache[key] = n
	return n
}

func countDigits(i int) int {
	if i == 0 {
		return 0
	}
	return int(math.Floor(math.Log10(float64(i))) + 1)
}

func evenDigits(i int) bool {
	return countDigits(i)%2 == 0
}

func split(i int) []int {
	cnt := countDigits(i)
	splitAt := cnt / 2
	return []int{int(float64(i) / math.Pow10(splitAt)), i % int(math.Pow10(splitAt))}
}

func parseStones(inp string) []int {
	var ret []int
	pts := strings.Split(inp, " ")
	for i := range pts {
		w, _ := strconv.Atoi(pts[i])
		ret = append(ret, w)
	}
	return ret
}