package main

import (
	"fmt"
	"os"
	"sort"
	"strconv"
)

type SetEntry struct {
	parent, size int
}

type DisjointSet struct {
	set []SetEntry
}

func NewDisjointSet(size int) *DisjointSet {
	set := make([]SetEntry, size)

	for i := range len(set) {
		set[i].parent = i
		set[i].size = 1
	}

	instance := new(DisjointSet)
	instance.set = set

	return instance
}

func (d *DisjointSet) Find(item int) int {
	for d.set[item].parent != item {
		d.set[item].parent = d.set[d.set[item].parent].parent
		item = d.set[item].parent
	}

	return item
}

func (d *DisjointSet) Union(x, y int) bool {
	xp := d.Find(x)
	yp := d.Find(y)

	if xp == yp {
		return false
	}

	if yp < xp {
		xp, yp = yp, xp
	}

	d.set[xp].size += d.set[yp].size
	d.set[yp].parent = xp

	return true
}

func (d *DisjointSet) Size(item int) int {
	if d.set[item].parent == item {
		return d.set[item].size
	} else {
		return 0
	}
}

func read_input(filename string) [][3]int {
	var points [][3]int
	var point [3]int

	file, err := os.Open(filename)

	if err != nil {
		panic(err)
	}

	for {
		parsed, err := fmt.Fscanf(file, "%d,%d,%d\n", &point[0], &point[1], &point[2])
		if err != nil || parsed != 3 {
			break
		}
		points = append(points, point)
	}

	file.Close()

	return points
}

func usage() {
	fmt.Printf("Usage: %v INPUT_FILE [connections]\n", os.Args[0])
	os.Exit(1)
}

func compute_group_sizes(groups *DisjointSet, len int) []int {
	var sizes []int
	for i := range len {
		size := groups.Size(i)
		if size > 0 {
			sizes = append(sizes, size)
		}
	}

	sort.Ints(sizes)

	return sizes
}

func main() {
	if len(os.Args) < 2 {
		usage()
	}

	connections := 1000
	if len(os.Args) >= 3 {
		parsed, err := strconv.Atoi(os.Args[2])
		if err != nil {
			usage()
		}
		connections = parsed
	}

	points := read_input(os.Args[1])

	var distances [][3]int

	for i, a := range points {
		for j := i + 1; j < len(points); j += 1 {
			b := points[j]

			square_dist := (a[0]-b[0])*(a[0]-b[0]) + (a[1]-b[1])*(a[1]-b[1]) + (a[2]-b[2])*(a[2]-b[2])

			distances = append(distances, [3]int{square_dist, i, j})
		}
	}

	sort.Slice(distances, func(i, j int) bool { return distances[i][0] < distances[j][0] })

	groups := NewDisjointSet(len(points))

	for _, d := range distances[:connections] {
		groups.Union(d[1], d[2])
	}

	sizes := compute_group_sizes(groups, len(points))

	product := 1
	for i := len(sizes) - 3; i < len(sizes); i += 1 {
		product *= sizes[i]
	}

	fmt.Printf("Part 1: %v\n", product)

	to_merge := len(sizes) - 1
	for _, d := range distances[connections:] {
		if groups.Union(d[1], d[2]) {
			to_merge -= 1
			if to_merge == 0 {
				fmt.Printf("Part 2: %v\n", points[d[1]][0]*points[d[2]][0])
				return
			}
		}
	}

	println("I should not get here")
}