package main

import (
    "fmt"
    "io/ioutil"
    "math"
    "strings"
)

const filePath = "14.txt"

func loadData(file string) (string, map[string]string) {
    dataBytes, err := ioutil.ReadFile(filePath)
    if err != nil {
        panic(err)
    }

    data := string(dataBytes)

    dataString := strings.Split(data, "\n")

    var initialState string
    transformations := map[string]string{}

    for i, v := range dataString {
        if i == 0 {
            initialState = v
        } else if v != "" {
            transformation := strings.Split(v, " -> ")

            transformations[transformation[0]] = transformation[1]
        }
    }

    return initialState, transformations
}

func main() {
    //
}

(...)

func partOne(initialState string, transformations map[string]string) {
    // Iterates over the polymerization process 10 items
    for i := 1; i <= 10; i++ {
        newChars := ""

        // Adds the new characters that result from the pair insertion
        // rules to newChars
        for j := 0; j < len(initialState)-1; j++ {
            newChar := transformations[string(initialState[j:j+2])]
            newChars += newChar
        }

        oldInitialState := initialState
        newInitialState := ""

        // Creates the new state by interlocking the current state
        // and the new chars
        for k := 0; k < len(newChars)+len(oldInitialState); k++ {
            if k%2 == 0 {
                newInitialState += string(oldInitialState[k/2])
            } else {
                newInitialState += string(newChars[(k-1)/2])
            }
        }

        initialState = newInitialState
    }

    // Gets the # of repetitions of most common and least common
    // characters
    maxChar := 0
    minChar := math.MaxInt32
    for m := 'A'; m <= 'Z'; m++ {
        repetitions := strings.Count(initialState, string(m))

        if repetitions > maxChar {
            maxChar = repetitions
        }

        if repetitions < minChar && repetitions != 0 {
            minChar = repetitions
        }
    }

    // Prints answer
    fmt.Println(maxChar - minChar)
}

func main() {
    initialState, transformations := loadData(filePath)
    partOne(initialState, transformations)
}

(...)

func partTwo(initialState string, transformations map[string]string, steps int) {
    // Creates a matrix that stores a slice of length 26
    // for each of the base cases (each pair after 1 step of polymerization)
    memo := map[int]map[string]map[string][]int{}
    memo[1] = map[string]map[string][]int{}

    for k, v := range transformations {
        tmp := make([]int, 26)
        for i := range tmp {
            tmp[i] = 0
        }
        tmp[rune(v[0])-65]++

        memo[1][k] = map[string][]int{}
        // The distinction between "nonAcc" and "acc" variants
        // is necessary so that the letters don't get repeated
        // in the cases that will be built-up
        memo[1][k]["acc"] = tmp
        memo[1][k]["nonAcc"] = tmp
    }

    // Builds the rest of the cases for each pair and number of steps
    // (This is the magic part)
    for i := 2; i <= steps; i++ {
        memo[i] = map[string]map[string][]int{}
        for k, v := range transformations {
            transformFirstPair := string(k[0]) + v
            transformSecondPair := v + string(k[1])

            prevFirstPairNonAccum := memo[i-1][transformFirstPair]["nonAcc"]
            prevSecondPairNonAccum := memo[i-1][transformSecondPair]["nonAcc"]
            prevPairAccum := memo[i-1][k]["acc"]

            memo[i][k] = map[string][]int{}
            memo[i][k]["acc"] = sumUpSlices(prevFirstPairNonAccum, prevSecondPairNonAccum, prevPairAccum)
            memo[i][k]["nonAcc"] = sumUpSlices(prevFirstPairNonAccum, prevSecondPairNonAccum)
        }
    }

    // Finally adding up the occurrences of each letter for
    // the cases with the required number of steps and "acc"
    // variant since that's the one that has the complete amount
    // of "children" for a given pair
    finalCount := make([]int, 26)

    for i := 0; i < len(finalCount); i++ {
        finalCount[i] = 0
    }

    for j := 0; j < len(initialState)-1; j++ {
        currRes := memo[steps][initialState[j:j+2]]["acc"]

        for k := 0; k < len(finalCount); k++ {
            finalCount[k] += currRes[k]
        }

        // Since the matrix stores the "children" of each of the original
        // pairs, it is necessary to add the original pairs to the sum.
        // Also, since they overlap, the first letter only gets added after
        // adding the first pair to the final count
        if j == 0 {
            finalCount[rune(initialState[0])-65]++
        }

        finalCount[rune(initialState[j+1])-65]++
    }

    // fmt.Println(finalCount)

    getDiff(finalCount)
}

// Auxiliary functions to make code more clear
func sumUpSlices(args ...[]int) []int {
    c := make([]int, 26)
    for _, v := range args {
        for j := 0; j < len(v); j++ {
            c[j] += v[j]
        }
    }

    return c
}

func getDiff(count []int) {
    min := math.MaxInt64
    max := 0

    for _, v := range count {
        if v < min && v != 0 {
            min = v
        }

        if v > max {
            max = v
        }
    }

    fmt.Println(max - min)
}