// Copyright ©2014 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package path

import (
	"cmp"
	"container/heap"
	"math"
	"slices"

	"gonum.org/v1/gonum/graph"
	"gonum.org/v1/gonum/graph/simple"
)

// WeightedBuilder is a type that can add nodes and weighted edges.
type WeightedBuilder interface {
	AddNode(graph.Node)
	SetWeightedEdge(graph.WeightedEdge)
}

// Prim generates a minimum spanning tree of g by greedy tree extension, placing
// the result in the destination, dst. If the edge weights of g are distinct
// it will be the unique minimum spanning tree of g. The destination is not cleared
// first. The weight of the minimum spanning tree is returned. If g is not connected,
// a minimum spanning forest will be constructed in dst and the sum of minimum
// spanning tree weights will be returned.
//
// Nodes and Edges from g are used to construct dst, so if the Node and Edge
// types used in g are pointer or reference-like, then the values will be shared
// between the graphs.
//
// If dst has nodes that exist in g, Prim will panic.
func Prim(dst WeightedBuilder, g graph.WeightedUndirected) float64 {
	nodes := graph.NodesOf(g.Nodes())
	if len(nodes) == 0 {
		return 0
	}

	q := &primQueue{
		indexOf: make(map[int64]int, len(nodes)-1),
		nodes:   make([]simple.WeightedEdge, 0, len(nodes)-1),
	}
	dst.AddNode(nodes[0])
	for _, u := range nodes[1:] {
		dst.AddNode(u)
		heap.Push(q, simple.WeightedEdge{F: u, W: math.Inf(1)})
	}

	u := nodes[0]
	uid := u.ID()
	for _, v := range graph.NodesOf(g.From(uid)) {
		w, ok := g.Weight(uid, v.ID())
		if !ok {
			panic("prim: unexpected invalid weight")
		}
		q.update(v, u, w)
	}

	var w float64
	for q.Len() > 0 {
		e := heap.Pop(q).(simple.WeightedEdge)
		if e.To() != nil && g.HasEdgeBetween(e.From().ID(), e.To().ID()) {
			dst.SetWeightedEdge(g.WeightedEdge(e.From().ID(), e.To().ID()))
			w += e.Weight()
		}

		u = e.From()
		uid := u.ID()
		for _, n := range graph.NodesOf(g.From(uid)) {
			if key, ok := q.key(n); ok {
				w, ok := g.Weight(uid, n.ID())
				if !ok {
					panic("prim: unexpected invalid weight")
				}
				if w < key {
					q.update(n, u, w)
				}
			}
		}
	}
	return w
}

// primQueue is a Prim's priority queue. The priority queue is a
// queue of edge From nodes keyed on the minimum edge weight to
// a node in the set of nodes already connected to the minimum
// spanning forest.
type primQueue struct {
	indexOf map[int64]int
	nodes   []simple.WeightedEdge
}

func (q *primQueue) Less(i, j int) bool {
	return q.nodes[i].Weight() < q.nodes[j].Weight()
}

func (q *primQueue) Swap(i, j int) {
	q.indexOf[q.nodes[i].From().ID()] = j
	q.indexOf[q.nodes[j].From().ID()] = i
	q.nodes[i], q.nodes[j] = q.nodes[j], q.nodes[i]
}

func (q *primQueue) Len() int {
	return len(q.nodes)
}

func (q *primQueue) Push(x interface{}) {
	n := x.(simple.WeightedEdge)
	q.indexOf[n.From().ID()] = len(q.nodes)
	q.nodes = append(q.nodes, n)
}

func (q *primQueue) Pop() interface{} {
	n := q.nodes[len(q.nodes)-1]
	q.nodes = q.nodes[:len(q.nodes)-1]
	delete(q.indexOf, n.From().ID())
	return n
}

// key returns the key for the node u and whether the node is
// in the queue. If the node is not in the queue, key is returned
// as +Inf.
func (q *primQueue) key(u graph.Node) (key float64, ok bool) {
	i, ok := q.indexOf[u.ID()]
	if !ok {
		return math.Inf(1), false
	}
	return q.nodes[i].Weight(), ok
}

// update updates u's position in the queue with the new closest
// MST-connected neighbour, v, and the key weight between u and v.
func (q *primQueue) update(u, v graph.Node, key float64) {
	id := u.ID()
	i, ok := q.indexOf[id]
	if !ok {
		return
	}
	q.nodes[i].T = v
	q.nodes[i].W = key
	heap.Fix(q, i)
}

// UndirectedWeightLister is an undirected graph that returns edge weights and
// the set of edges in the graph.
type UndirectedWeightLister interface {
	graph.WeightedUndirected
	WeightedEdges() graph.WeightedEdges
}

// Kruskal generates a minimum spanning tree of g by greedy tree coalescence, placing
// the result in the destination, dst. If the edge weights of g are distinct
// it will be the unique minimum spanning tree of g. The destination is not cleared
// first. The weight of the minimum spanning tree is returned. If g is not connected,
// a minimum spanning forest will be constructed in dst and the sum of minimum
// spanning tree weights will be returned.
//
// Nodes and Edges from g are used to construct dst, so if the Node and Edge
// types used in g are pointer or reference-like, then the values will be shared
// between the graphs.
//
// If dst has nodes that exist in g, Kruskal will panic.
func Kruskal(dst WeightedBuilder, g UndirectedWeightLister) float64 {
	edges := graph.WeightedEdgesOf(g.WeightedEdges())
	slices.SortFunc(edges, func(a, b graph.WeightedEdge) int {
		return cmp.Compare(a.Weight(), b.Weight())
	})

	ds := make(djSet)
	it := g.Nodes()
	for it.Next() {
		n := it.Node()
		dst.AddNode(n)
		ds.add(n.ID())
	}

	var w float64
	for _, e := range edges {
		if s1, s2 := ds.find(e.From().ID()), ds.find(e.To().ID()); s1 != s2 {
			ds.union(s1, s2)
			dst.SetWeightedEdge(g.WeightedEdge(e.From().ID(), e.To().ID()))
			w += e.Weight()
		}
	}
	return w
}