package main

import (
	"fmt"

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

func main() {
	fmt.Println("# Day 17")
	fmt.Println()
	inp := h.StdinToStringSlice()
	fmt.Println("## Part 1:\nAnswer:", run(inp, 3))
	fmt.Println("## Part 2:\nAnswer:", run(inp, 4))
}

func run(inp []string, dim int) int {
	cs := NewCubeMapFromLines(inp)
	for i := 0; i < 6; i++ {
		next := NewCubeMap()
		wrk := cs.Copy()
		for _, c := range cs.ToSlice() {
			wrk.AddMany(c.Neighbors(dim))
		}
		for _, c := range wrk.ToSlice() {
			active := cs.Contains(c)
			cnt := 0
			for _, n := range c.Neighbors(dim) {
				if cs.Contains(n) {
					cnt++
				}
			}
			if cnt == 3 || (active && cnt == 2) {
				next.Add(c)
			}
		}
		cs = next
	}
	return cs.Size()
}

type Cube struct {
	X, Y, Z, W int
}

func (c Cube) Neighbors(d int) []Cube {
	ret := []Cube{}
	for dx := -1; dx <= 1; dx++ {
		for dy := -1; dy <= 1; dy++ {
			for dz := -1; dz <= 1; dz++ {
				for dw := -(d - 3); dw <= d-3; dw++ {
					if dx == 0 && dy == 0 && dz == 0 && dw == 0 {
						continue
					}
					ret = append(ret, Cube{c.X + dx, c.Y + dy, c.Z + dz, c.W + dw})
				}
			}
		}
	}
	return ret
}

type CubeMap struct {
	field map[Cube]bool
}

func NewCubeMap() *CubeMap {
	return &CubeMap{map[Cube]bool{}}
}
func NewCubeMapFromLines(inp []string) *CubeMap {
	cs := NewCubeMap()
	for y, line := range inp {
		for x, b := range line {
			if b == '#' {
				cs.Add(Cube{x, y, 0, 0})
			}
		}
	}
	return cs
}

func (cs *CubeMap) Add(c Cube) {
	cs.field[c] = true
}
func (cs *CubeMap) AddMany(cubes []Cube) {
	for _, c := range cubes {
		cs.Add(c)
	}
}
func (cs *CubeMap) Remove(c Cube) {
	delete(cs.field, c)
}
func (cs *CubeMap) Contains(c Cube) bool {
	_, ok := cs.field[c]
	return ok
}
func (cs *CubeMap) Copy() *CubeMap {
	field := map[Cube]bool{}
	for c := range cs.field {
		field[c] = true
	}
	return &CubeMap{field}
}
func (cs *CubeMap) Size() int {
	return len(cs.field)
}
func (cs *CubeMap) ToSlice() []Cube {
	ret := []Cube{}
	for c := range cs.field {
		ret = append(ret, Cube{c.X, c.Y, c.Z, c.W})
	}
	return ret
}