package year25

import (
	"cmp"
	"fmt"
	"log"
	"math"
	"os"
	"slices"
	"strconv"

	"git.bizdoc.ro/gabi-public/Advent-of-Code.git/aoc"
	"git.bizdoc.ro/private/devkit.git/collections/geometry/v2"
)

type ClusterId int

type ClustersDB[T any] struct { // TODO: move to devkit
	elementIdToCluster []ClusterId
	clusterToSize      []int
}

func (s *ClustersDB[T]) GetClusterId(elementId int) ClusterId {
	return s.elementIdToCluster[elementId]
}

func (s *ClustersDB[T]) Move(elementId int, to ClusterId) {
	elementCluster := s.GetClusterId(elementId)

	if elementCluster != to {
		s.clusterToSize[elementCluster] -= 1
		s.clusterToSize[to] += 1

		if s.clusterToSize[elementCluster] < 0 {
			panic("invalid cluster to size")
		}

		s.elementIdToCluster[elementId] = to
	}
}

func (s *ClustersDB[T]) JoinClustersByElements(a, b int) {
	elementCluster1 := s.GetClusterId(a)
	elementCluster2 := s.GetClusterId(b)

	for el, clusterId := range s.elementIdToCluster {
		if clusterId == elementCluster2 {
			s.Move(el, elementCluster1)
		}
	}
}

func NewClustersDB[T any](data []T) *ClustersDB[T] {
	points := make([]ClusterId, len(data))
	sizes := make([]int, len(data))

	for i := range points {
		points[i] = ClusterId(i)
		sizes[i] = 1
	}

	return &ClustersDB[T]{
		elementIdToCluster: points,
		clusterToSize:      sizes,
	}
}

func Day8Part1(ctx aoc.Context, steps int) (int, error) {
	ctx.Logger = log.New(os.Stdout, "", 0)
	boxes, err := day8ParseInput(ctx)
	if err != nil {
		return 0, fmt.Errorf("day8: failed to parse input %w", err)
	}

	db := NewClustersDB(boxes)
	for _, d := range SortedDistances(boxes)[:steps] {
		if d.Distance != 0 && (db.GetClusterId(d.E1) != db.GetClusterId(d.E2)) {
			db.JoinClustersByElements(d.E1, d.E2)
		}
	}

	ints := slices.Clone(db.clusterToSize)
	slices.Sort(ints)
	return aoc.Prod(ints[len(ints)-3:]), nil
}

func Day8Part2(ctx aoc.Context) (int, error) {
	ctx.Logger = log.New(os.Stdout, "", 0)
	boxes, err := day8ParseInput(ctx)
	if err != nil {
		return 0, fmt.Errorf("day8: failed to parse input %w", err)
	}

	db := NewClustersDB(boxes)
	for _, d := range SortedDistances(boxes) {
		if d.Distance != 0 && (db.GetClusterId(d.E1) != db.GetClusterId(d.E2)) {
			db.JoinClustersByElements(d.E1, d.E2)
			if db.clusterToSize[db.GetClusterId(d.E1)] == len(boxes) {
				return int(boxes[d.E1][0]) * int(boxes[d.E2][0]), nil
			}
		}
	}

	panic("unreachable")
}

func day8ParseInput(c aoc.Context) ([]Vector, error) {
	g, err := geometry.ReaderToIntsGrid(c.Body, ",")
	if err != nil {
		return nil, fmt.Errorf("day8: failed to read input %w", err)
	}

	points := make([]Vector, g.Rows())
	for i := range points {
		points[i] = make(Vector, g.Cols())
		for j := range points[i] {
			v := g.At(geometry.NewPoint(i, j))
			points[i][j] = float64(v)
		}
	}
	return points, nil
}

type Vector []float64 // TODO: move to devkit

func (v Vector) String() string {
	x := make([]string, len(v))
	for i, v := range v {
		x[i] = strconv.FormatFloat(v, 'f', -1, 64)
	}
	return fmt.Sprintf("[%g,%g,%g]", v[0], v[1], v[2])
}

func SquaredDelta(b, a Vector, dim int) float64 {
	x := b[dim] - a[dim]
	return x * x
}
func SquaredEuclideanDistance(a, b Vector) (s float64) {
	for i := range a {
		s += SquaredDelta(a, b, i)
	}
	return
}
func EuclideanDistance(a, b Vector) float64 {
	return math.Sqrt(SquaredEuclideanDistance(a, b))
}

type Pair struct {
	E1, E2   int
	Distance float64
}

func SortedDistances(points []Vector) []Pair {
	n := len(points)
	pairs := make([]Pair, 0, n*(n-1)/2)

	for i := 0; i < n; i++ {
		for j := i + 1; j < n; j++ {
			pairs = append(pairs, Pair{
				E1: i,
				E2: j,
				// no need to spend time on sqrt. order does not change
				Distance: SquaredEuclideanDistance(points[i], points[j]),
			})
		}
	}

	slices.SortFunc(pairs, func(a, b Pair) int {
		return cmp.Compare(a.Distance, b.Distance)
	})

	return pairs
}