package fourteen

import (
	"bufio"
	"fmt"
	"os"
)

type polymer struct {
	template string
	rules    map[string]string
}

func (p *polymer) load(filename string) error {
	file, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer file.Close()
	scanner := bufio.NewScanner(file)

	p.rules = map[string]string{}
	scannedTemplate := false
	for scanner.Scan() {
		if !scannedTemplate {
			p.template = scanner.Text()

			scannedTemplate = true
			continue
		}

		if scanner.Text() == "" {
			continue
		}

		var pair, insert string
		if count, err := fmt.Sscanf(scanner.Text(), "%s -> %s", &pair, &insert); count != 2 || err != nil {
			return fmt.Errorf("expected to scan 2 values from %q, found %d: %w", scanner.Text(), count, err)
		}
		p.rules[pair] = insert
	}

	return nil
}

func (p *polymer) elementCount(size int) int {
	result := p.step(size)
	elemCounts := map[rune]int{}

	for _, r := range result {
		elemCounts[r]++
	}

	min := len(result)
	max := 0
	for _, v := range elemCounts {
		if v < min {
			min = v
		}

		if v > max {
			max = v
		}
	}

	return max - min
}

func (p *polymer) elementCountMap(size int) int {
	_, counts := p.stepMap(size)

	min := -1
	max := -1
	for _, v := range counts {
		if v < min || min == -1 {
			min = v
		}

		if v > max {
			max = v
		}
	}

	return max - min
}

func (p *polymer) stepMap(steps int) (map[string]int, map[string]int) {
	pairs := map[string]int{}
	counts := map[string]int{}

	// Build initial pairs and counts from template
	for i := 0; i < len(p.template)-1; i++ {
		// Count first element from pair
		counts[string(p.template[i])]++
		pairs[fmt.Sprintf("%c%c", p.template[i], p.template[i+1])]++
		// If this is the last pair add the second element
		if i == len(p.template)-2 {
			counts[string(p.template[i+1])]++
		}
	}

	for i := 0; i < steps; i++ {
		pairs, counts = p.stepMapOne(pairs, counts)
	}

	return pairs, counts
}

func (p *polymer) stepMapOne(pairs map[string]int, counts map[string]int) (map[string]int, map[string]int) {
	newPairs := map[string]int{}
	for k, v := range pairs {
		result := p.rules[k]
		counts[result] += (1 * v)
		pair1 := fmt.Sprintf("%c%s", k[0], result)
		pair2 := fmt.Sprintf("%s%c", result, k[1])
		newPairs[pair1] += 1 * v
		newPairs[pair2] += 1 * v
	}

	return newPairs, counts
}

func (p *polymer) step(steps int) string {
	seed := p.template
	for i := 0; i < steps; i++ {
		seed = p.stepOne(seed)
	}

	return seed
}

func (p *polymer) stepOne(from string) string {
	result := ""
	for i := 0; i < len(from)-1; i++ {
		// Add first element from pair
		result += string(from[i])
		// Lookup and add element result from pair
		pair := fmt.Sprintf("%c%c", from[i], from[i+1])
		result += p.rules[pair]
		// If this is the last pair add the second element
		if i == len(from)-2 {
			result += string(from[i+1])
		}
	}
	return result
}