package main

import (
	"bufio"
	"log"
	"os"
	"strings"
)

type insert map[string]byte

type polimer struct {
	template string
	inserts  insert
}

func main() {
	if err := myMain(); err != nil {
		log.Println(err)
	}
}

func myMain() error {
	p, err := parseInput("input.txt")
	if err != nil {
		return err
	}

	log.Println(solveFirst(p, 10))

	log.Println(solveFirst(p, 40))

	return nil
}

func solveFirst(p polimer, steps int) int {
	var counter = map[string]int{}
	for ch := range p.template {
		if ch == len(p.template)-1 {
			break
		}
		counter[p.template[ch:ch+2]]++
	}

	for i := 0; i < steps; i++ {
		newcounter := map[string]int{}
		for k, v := range counter {
			newcounter[string(k[0])+string(p.inserts[k])] += v
			newcounter[string(p.inserts[k])+string(k[1])] += v
		}
		counter = newcounter
	}

	chars := map[byte]int{}

	// last item in the template skipped on counting, put it back!
	chars[p.template[len(p.template)-1]]++

	for twoChars, v := range counter {
		chars[twoChars[0]] += v
	}

	var min, max int
	for _, v := range chars {
		if min == 0 {
			min = v
		}
		if v > max {
			max = v
		}
		if v < min {
			min = v
		}
	}

	return max - min
}

func parseInput(fileName string) (polimer, error) {
	var ret = polimer{}
	ret.inserts = insert{}
	fd, err := os.Open(fileName)
	if err != nil {
		return ret, err
	}
	defer fd.Close()

	buf := bufio.NewScanner(fd)
	for buf.Scan() {
		if buf.Text() == "" {
			continue
		}
		tmp := strings.Split(buf.Text(), " -> ")
		if len(tmp) == 2 {
			ret.inserts[tmp[0]] = tmp[1][0]
		} else { // polymer template
			ret.template = tmp[0]
		}
	}

	return ret, nil
}