package main

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

type Position struct {
	X, Y int
}

func ByteToAntenna(b byte) int {
	if b >= '0' && b <= '9' {
		return int(b - '0')
	} else if b >= 'a' && b <= 'z' {
		return int(b - 'a' + 10)
	} else if b >= 'A' && b <= 'Z' {
		return int(b - 'A' + 36)
	} else {
		return -1
	}
}

func IsInMap(p Position, w, h int) bool {
	if p.X < 0 || p.X >= w {
		return false
	}
	if p.Y < 0 || p.Y >= h {
		return false
	}
	return true
}

func ComputeAntinodes(a, b Position, w, h int, first, repeats int) (antinodes []Position) {
	dx := b.X - a.X
	dy := b.Y - a.Y

	for i := first; i <= repeats; i++ {
		n1, n2 := Position{a.X - dx*i, a.Y - dy*i}, Position{b.X + dx*i, b.Y + dy*i}
		n1InMap := IsInMap(n1, w, h)
		n2InMap := IsInMap(n2, w, h)

		if n1InMap {
			antinodes = append(antinodes, n1)
		}

		if n2InMap {
			antinodes = append(antinodes, n2)
		}

		if !n1InMap && !n2InMap {
			break
		}
	}

	return
}

func main() {
	{
		width, height, antinodes := readData("example.txt")
		fmt.Println(countAntinodes(width, height, antinodes, 1, 1))
		fmt.Println(countAntinodes(width, height, antinodes, 0, 999))
	}
	{
		width, height, antinodes := readData("data.txt")
		fmt.Println(countAntinodes(width, height, antinodes, 1, 1))
		fmt.Println(countAntinodes(width, height, antinodes, 0, 999))
	}
}

func countAntinodes(width, height int, antennas map[int][]Position, first, repeats int) (count int) {
	marked := make([][]bool, height)
	for y := 0; y < height; y++ {
		marked[y] = make([]bool, width)
	}

	for _, positions := range antennas {
		for i := 0; i < len(positions)-1; i++ {
			for j := i + 1; j < len(positions); j++ {
				for _, p := range ComputeAntinodes(positions[i], positions[j], width, height, first, repeats) {
					marked[p.Y][p.X] = true
				}
			}
		}
	}

	for y := 0; y < height; y++ {
		for x := 0; x < len(marked[y]); x++ {
			if marked[y][x] {
				count++
			}
		}
	}

	return
}

func readData(fileName string) (width, height int, antennas map[int][]Position) {
	fp, err := os.Open(fileName)
	if err != nil {
		panic(err)
	}

	scanner := bufio.NewScanner(fp)
	antennas = make(map[int][]Position)

	y := 0
	for ; scanner.Scan(); y++ {
		line := strings.TrimSpace(scanner.Text())
		width = len(line)

		for x := 0; x < width; x++ {
			antenna := ByteToAntenna(line[x])
			if antenna >= 0 {
				antennas[antenna] = append(antennas[antenna], Position{x, y})
			}
		}
	}

	height = y

	return
}