package main

import (
	"fmt"
	"log"
	"os"

	"git.annabunch.es/annabunches/adventofcode/2020/lib/util"
)

func parseInputCubes(input []string) map[[3]int]bool {
	state := make(map[[3]int]bool)
	for x, line := range input {
		for y, char := range line {
			coords := [3]int{x, y, 0}
			if char == '#' {
				state[coords] = true
			}
		}
	}

	return state
}

func iterateCubes(prev map[[3]int]bool) map[[3]int]bool {
	counts := make(map[[3]int]int)

	for node, v := range prev {
		if !v {
			log.Panicf("Unexpected false value!")
		}
		x := node[0]
		y := node[1]
		z := node[2]
		neighbors := make([][3]int, 26)
		index := 0
		for i := -1; i < 2; i++ {
			for j := -1; j < 2; j++ {
				for k := -1; k < 2; k++ {
					if i == 0 && j == 0 && k == 0 {
						continue
					}
					neighbors[index] = [3]int{x + i, y + j, z + k}
					index++
				}
			}
		}

		for _, n := range neighbors {
			counts[n]++
		}
	}

	next := make(map[[3]int]bool)
	for node, count := range counts {
		if (prev[node] && count == 2) || count == 3 {
			next[node] = true
		}
	}
	return next
}

func parseInputHypercubes(input []string) map[[4]int]bool {
	state := make(map[[4]int]bool)
	for x, line := range input {
		for y, char := range line {
			coords := [4]int{x, y, 0, 0}
			if char == '#' {
				state[coords] = true
			}
		}
	}

	return state
}

func iterateHypercubes(prev map[[4]int]bool) map[[4]int]bool {
	counts := make(map[[4]int]int)

	for node, v := range prev {
		if !v {
			log.Panicf("Unexpected false value!")
		}
		x := node[0]
		y := node[1]
		z := node[2]
		w := node[3]
		neighbors := make([][4]int, 80)
		index := 0
		for i := -1; i < 2; i++ {
			for j := -1; j < 2; j++ {
				for k := -1; k < 2; k++ {
					for l := -1; l < 2; l++ {
						if i == 0 && j == 0 && k == 0 && l == 0 {
							continue
						}
						neighbors[index] = [4]int{x + i, y + j, z + k, w + l}
						index++
					}
				}
			}
		}

		for _, n := range neighbors {
			counts[n]++
		}
	}

	next := make(map[[4]int]bool)
	for node, count := range counts {
		if (prev[node] && count == 2) || count == 3 {
			next[node] = true
		}
	}
	return next
}

func main() {
	step := os.Args[1]
	values := util.InputParserStrings(os.Args[2])
	switch step {
	case "1":
		state := parseInputCubes(values)
		for i := 0; i < 6; i++ {
			state = iterateCubes(state)
		}
		fmt.Println(len(state))
	case "2":
		state := parseInputHypercubes(values)
		for i := 0; i < 6; i++ {
			state = iterateHypercubes(state)
		}
		fmt.Println(len(state))
	}
}