// Package graph is an undirected graph management library.
// Do not support multithreaded updates.
package graph

import (
	"fmt"
	"math"
	"math/rand"
	"slices"
	"strings"

	"gitea.paas.celticinfo.fr/oabrivard/abrolgo/algo"
	"gitea.paas.celticinfo.fr/oabrivard/abrolgo/container"
	"gitea.paas.celticinfo.fr/oabrivard/abrolgo/maths"
	"golang.org/x/exp/constraints"
)

type Edge[T comparable, V constraints.Integer] struct {
	ID     int
	src    T
	dest   T
	weight V
}

type Graph[T comparable, V constraints.Integer] struct {
	edgeID        int
	v, e          int // number of edges and vertices
	adjacencyList map[T][]Edge[T, V]
	edges         []Edge[T, V]
}

// NewGraph creates a new instance of the undirected Graph data structure.
// It initializes the graph with zero vertices and zero edges.
// The graph uses a map to store the adjacency list and a slice to store the edges.
// The type parameters T and V represent the types of the vertices and edge values, respectively.
// The vertices must be comparable, and the edge values must be integers.
// Returns a pointer to the newly created Graph.
func NewGraph[T comparable, V constraints.Integer]() *Graph[T, V] {
	return &Graph[T, V]{
		edgeID:        0,
		v:             0,
		e:             0,
		adjacencyList: map[T][]Edge[T, V]{},
		edges:         []Edge[T, V]{},
	}
}

// AddVertex adds a new vertex to the graph.
// This function initializes an empty adjacency list for the new vertex.
func (g *Graph[T, V]) AddVertex(vertex T) {
	g.adjacencyList[vertex] = []Edge[T, V]{}
}

// AddEdge adds an edge to the graph between the source node (src) and the destination node (dest),
// with the given weight (weight).
// It adds the edge in both directions (src -> dest and dest -> src).
// If vertices src and dest do not exist in the graph, they are added to the graph.
func (g *Graph[T, V]) AddEdge(src, dest T, weight V) {
	g.edgeID++
	srcToDest := Edge[T, V]{g.edgeID, src, dest, weight}
	destToSrc := Edge[T, V]{g.edgeID, dest, src, weight}
	g.adjacencyList[src] = append(g.adjacencyList[src], srcToDest)
	g.adjacencyList[dest] = append(g.adjacencyList[dest], destToSrc)
	g.edges = append(g.edges, srcToDest)
	g.e++
	g.v = len(g.adjacencyList)
}

// RemoveEdge removes the first edge between the source node (src) and the destination node (dest).
// It updates the adjacency list, the list of edges, and the number of edges in the graph.
// If the edge does not exist, it does nothing.
func (g *Graph[T, V]) RemoveEdge(src, dest T) {
	var ID int
	edgeFound := false

	// Check if source vertex exists
	if _, exists := g.adjacencyList[src]; !exists {
		// Optionally handle the case where the source vertex does not exist
		return
	}

	// Remove the edge from the adjacency list
	for i, edge := range g.adjacencyList[src] {
		if edge.dest == dest {
			g.adjacencyList[src] = append(g.adjacencyList[src][:i], g.adjacencyList[src][i+1:]...)
			ID = edge.ID
			edgeFound = true
			break
		}
	}

	// Remove the reciprocal edge
	for i, edge := range g.adjacencyList[dest] {
		if edge.ID == ID {
			g.adjacencyList[dest] = append(g.adjacencyList[dest][:i], g.adjacencyList[dest][i+1:]...)
			break
		}
	}

	// If no edge was removed from the adjacency list, return to avoid the uneceesary loop
	if !edgeFound {
		return
	}

	// Remove the edge from the list of edges
	for i, edge := range g.edges {
		if edge.ID == ID {
			g.edges = append(g.edges[:i], g.edges[i+1:]...)
			break
		}
	}

	g.e--
}

// RemoveEdge removes the first edge between the source node (src) and the destination node (dest).
// It updates the adjacency list, the list of edges, and the number of edges in the graph.
// If the edge does not exist, it does nothing.
func (g *Graph[T, V]) RemoveEdgeWithID(ID int) {
	edgeFound := false
	var edgeToRemove Edge[T, V]

	// Remove the edge from the list of edges
	for i, edge := range g.edges {
		if edge.ID == ID {
			edgeFound = true
			edgeToRemove = edge
			g.edges = append(g.edges[:i], g.edges[i+1:]...)
			break
		}
	}

	// Return if no edge found
	if !edgeFound {
		return
	}

	src := edgeToRemove.src
	dest := edgeToRemove.dest

	// Remove the edge from the adjacency list
	for i, edge := range g.adjacencyList[src] {
		if edge.ID == ID {
			g.adjacencyList[src] = append(g.adjacencyList[src][:i], g.adjacencyList[src][i+1:]...)
			break
		}
	}

	// Remove the reciprocal edge
	for i, edge := range g.adjacencyList[dest] {
		if edge.ID == ID {
			g.adjacencyList[dest] = append(g.adjacencyList[dest][:i], g.adjacencyList[dest][i+1:]...)
			break
		}
	}

	g.e--
}

// RemoveVertex removes the specified vertex from the graph.
// It updates the adjacency list, the list of edges, and the number of edges and vertices in the graph.
// If the vertex does not exist, it does nothing.
func (g *Graph[T, V]) RemoveVertex(vertex T) {
	// Remove the vertex from the adjacency list
	delete(g.adjacencyList, vertex)

	// Remove all edges in the adjacency list that have 'vertex' as their destination
	for src, edges := range g.adjacencyList {
		for i := 0; i < len(edges); {
			if edges[i].dest == vertex {
				// Remove edge from the slice
				edges = append(edges[:i], edges[i+1:]...)
				g.e-- // Decrement the edge count
			} else {
				i++
			}
		}
		g.adjacencyList[src] = edges
	}

	// Remove all edges from the edges slice that involve 'vertex'
	for i := 0; i < len(g.edges); {
		if g.edges[i].src == vertex || g.edges[i].dest == vertex {
			// Remove edge from the slice
			g.edges = append(g.edges[:i], g.edges[i+1:]...)
		} else {
			i++
		}
	}

	g.v = len(g.adjacencyList)
}

// GetVertices returns the list of vertices in the graph.
func (g *Graph[T, V]) GetVertices() []T {
	vertices := make([]T, len(g.adjacencyList))
	i := 0
	for vertex := range g.adjacencyList {
		vertices[i] = vertex
		i++
	}
	return vertices
}

// GetEdges returns the list of edges in the graph.
func (g *Graph[T, V]) GetEdges() []Edge[T, V] {
	return g.edges
}

// GetNeighbors returns the list of edges connected to the specified vertex in the graph.
func (g *Graph[T, V]) GetNeighbors(vertex T) []Edge[T, V] {
	return g.adjacencyList[vertex]
}

// HasNeighbor checks if a given source node has a neighbor node in the graph.
func (g *Graph[T, V]) HasNeighbor(src, neighbor T) bool {
	return slices.ContainsFunc(g.adjacencyList[src], func(edge Edge[T, V]) bool {
		return edge.dest == neighbor
	})
}

// GetEdge returns the first edge between the source node (src) and the destination node (dest).
// If the edge does not exist, it returns an empty Edge.
func (g *Graph[T, V]) GetEdge(src, dest T) Edge[T, V] {
	for _, edge := range g.adjacencyList[src] {
		if edge.dest == dest {
			return edge
		}
	}
	return Edge[T, V]{}
}

// String returns a string representation of the graph.
func (g *Graph[T, V]) String() string {
	var sb strings.Builder
	sb.Grow(512) // Provide an initial capacity estimate

	for vertex, edges := range g.adjacencyList {
		sb.WriteString(fmt.Sprint(vertex))
		sb.WriteString(": ")

		edgeStrings := make([]string, len(edges))
		for i, e := range edges {
			edgeStrings[i] = fmt.Sprintf("(%v, %v, %v, %v)", e.ID, e.src, e.dest, e.weight)
		}

		sb.WriteString(strings.Join(edgeStrings, " "))
		sb.WriteString("\n")
	}

	return sb.String()
}

// CountComponents returns the number of connected components in the graph.
func (g *Graph[T, V]) CountComponents() int {
	visited := map[T]bool{}
	count := 0

	// We perform a breadth-first search starting from each unvisited vertex
	// and increment the count for each new component found.
	for vertex := range g.adjacencyList {
		if !visited[vertex] {
			count++
			queue := container.NewQueue[T]()
			queue.Enqueue(vertex)

			for queue.HasElement() {
				curVert := queue.Dequeue()
				visited[curVert] = true

				for _, linkedNode := range g.adjacencyList[curVert] {
					if !visited[linkedNode.dest] {
						queue.Enqueue(linkedNode.dest)
					}
				}
			}
		}
	}
	return count
}

// IsConnected returns true if the graph is connected, false otherwise.
func (g *Graph[T, V]) IsConnected() bool {
	return g.CountComponents() == 1
}

// IsCyclic returns true if the graph contains a cycle, false otherwise.
func (g *Graph[T, V]) IsCyclic() bool {
	visited := map[T]bool{}
	recStack := map[T]bool{}
	edgeIDs := map[int]bool{}

	for vertex := range g.adjacencyList {
		if g.isCyclicUtil(vertex, visited, recStack, edgeIDs) {
			return true
		}
	}
	return false
}

// isCyclicUtil is a helper function for IsCyclic.
// It performs a depth-first search on the graph starting from the specified vertex.
// It returns true if the graph contains a cycle, false otherwise.
func (g *Graph[T, V]) isCyclicUtil(vertex T, visited, recStack map[T]bool, edgeIDs map[int]bool) bool {
	visited[vertex] = true
	recStack[vertex] = true

	for _, neighbor := range g.adjacencyList[vertex] {
		// If the edge has already been visited via a reciprocal link, skip it
		if edgeIDs[neighbor.ID] {
			continue
		}
		edgeIDs[neighbor.ID] = true

		if !visited[neighbor.dest] && g.isCyclicUtil(neighbor.dest, visited, recStack, edgeIDs) {
			return true
		} else if recStack[neighbor.dest] {
			return true
		}
	}
	recStack[vertex] = false
	return false
}

// IsTree returns true if the graph is a tree, false otherwise.
func (g *Graph[T, V]) IsTree() bool {
	return g.IsConnected() && !g.IsCyclic()
}

// IsBipartite returns true if the graph is bipartite, false otherwise.
func (g *Graph[T, V]) IsBipartite() bool {
	visited := map[T]bool{}
	colors := map[T]bool{}
	edgeIDs := map[int]bool{}

	for vertex := range g.adjacencyList {
		if !visited[vertex] {
			if !g.isBipartiteUtil(vertex, visited, colors, edgeIDs) {
				return false
			}
		}
	}
	return true
}

// isBipartiteUtil is a helper function for IsBipartite.
// It performs a depth-first search on the graph starting from the specified vertex.
// It returns true if the graph is bipartite, false otherwise.
func (g *Graph[T, V]) isBipartiteUtil(vertex T, visited, colors map[T]bool, edgeIDs map[int]bool) bool {
	visited[vertex] = true

	for _, neighbor := range g.adjacencyList[vertex] {
		// If the edge has already been visited via a reciprocal link, skip it
		if edgeIDs[neighbor.ID] {
			continue
		}
		edgeIDs[neighbor.ID] = true

		if !visited[neighbor.dest] {
			colors[neighbor.dest] = !colors[vertex]
			if !g.isBipartiteUtil(neighbor.dest, visited, colors, edgeIDs) {
				return false
			}
		} else if colors[vertex] == colors[neighbor.dest] {
			return false
		}
	}
	return true
}

// kargerMinCutUF performs the Karger's algorithm to find the minimum cut in the graph using the Union-Find data structure.
// It returns the number of cut edges and the subsets of vertices after the contraction.
// Being a monte-carlo algorithm, it can return a wrong result with a very small probability.
// Adapted from https://www.geeksforgeeks.org/introduction-and-implementation-of-kargers-algorithm-for-minimum-cut/?ref=lbp
func (g *Graph[T, V]) kargerMinCutUF(loops int) (int, [][]T) {
	var subsets map[T]*algo.Subset[T]
	minCut := math.MaxInt64

	for loops > 0 {
		loops--

		// Allocate memory for creating V subsets.
		subsets = map[T]*algo.Subset[T]{}

		// Get data of given graph
		edges := g.edges

		// Create V subsets with single elements
		for k := range g.adjacencyList {
			subsets[k] = &algo.Subset[T]{Parent: k, Rank: 0}
		}

		// Initially there are V vertices in
		// contracted graph
		vertices := g.v

		// Keep contracting vertices until there are
		// 2 vertices.
		for vertices > 2 {
			// Pick a random edge
			i := rand.Int() % g.e

			// Find vertices (or sets) of two corners
			// of current edge
			subset1 := algo.Find(subsets, edges[i].src)
			subset2 := algo.Find(subsets, edges[i].dest)

			// If two corners belong to same subset,
			// then no point considering this edge
			if subset1 == subset2 {
				continue
			} else {
				// Else contract the edge (or combine the
				// corners of edge into one vertex)
				vertices--
				algo.Union(subsets, subset1, subset2)
			}
		}

		// Now we have two vertices (or subsets) left in
		// the contracted graph, so count the edges between
		// two components and return the count.
		cutedges := 0
		for i := 0; i < g.e; i++ {
			subset1 := algo.Find(subsets, edges[i].src)
			subset2 := algo.Find(subsets, edges[i].dest)
			if subset1 != subset2 {
				cutedges++
			}
		}

		if cutedges < minCut {
			minCut = cutedges
		}
	}

	// Consolidate the subsets
	consolidatedSubsets := map[T][]T{}
	for k, v := range subsets {
		_, found := consolidatedSubsets[v.Parent]
		if !found {
			consolidatedSubsets[v.Parent] = []T{}
		}
		consolidatedSubsets[v.Parent] = append(consolidatedSubsets[v.Parent], k)
	}

	resultSubsets := [][]T{}
	for _, s := range consolidatedSubsets {
		resultSubsets = append(resultSubsets, s)
	}

	return minCut, resultSubsets
}

// DFS performs a depth-first search starting from the specified vertex until a goal vertex is reached.
// It returns a slice of vertices visited during the search.
// If no goal vertex is found, an empty slice is returned.
func (g *Graph[T, V]) DFS(start T, isGoal func(vertex T) bool) []T {
	visited := map[T]bool{}
	result := []T{}
	stack := container.NewStack[T]()

	stack.Push(start)

	for stack.HasElement() {
		curVert := stack.Pop()
		result = append(result, curVert)

		if isGoal(curVert) {
			return result
		}

		if !visited[curVert] {
			visited[curVert] = true

			for _, linkedNode := range g.adjacencyList[curVert] {
				stack.Push(linkedNode.dest)
			}
		}
	}

	return []T{}
}

// BFS performs a breadth-first search starting from the specified vertex until a goal vertex is reached.
// It returns a slice of vertices visited during the search.
// If no goal vertex is found, an empty slice is returned.
func (g *Graph[T, V]) BFS(start T, isGoal func(vertex T) bool) []T {
	visited := map[T]bool{}
	result := []T{}
	queue := container.NewQueue[T]()

	visited[start] = true
	queue.Enqueue(start)

	for queue.HasElement() {
		curVert := queue.Dequeue()
		result = append(result, curVert)

		if isGoal(curVert) {
			return result
		}

		for _, linkedNode := range g.adjacencyList[curVert] {
			if !visited[linkedNode.dest] {
				visited[linkedNode.dest] = true
				queue.Enqueue(linkedNode.dest)
			}
		}
	}

	return []T{}
}

// AStar performs the A* search algorithm to find the shortest path between two vertices in the graph.
// It returns the shortest path as a slice of vertices.
// If no path is found, an empty slice is returned.
func (g *Graph[T, V]) AStar(start, goal T, heuristic func(vertex T) V) []T {
	// Initialize the distance map
	dist := map[T]V{}
	for vertex := range g.adjacencyList {
		dist[vertex] = maths.MaxInteger[V]()
	}
	dist[start] = V(0)

	// Initialize the previous map
	prev := map[T]T{}

	// Initialize the priority queue
	pq := container.NewPriorityQueue[T, V](true)
	pq.Enqueue(start, 0)

	for !pq.IsEmpty() {
		curVert := pq.Dequeue()

		if curVert == goal {
			// Reconstruct the path
			path := []T{goal}
			for prevVertex, ok := prev[goal]; ok; prevVertex, ok = prev[prevVertex] {
				path = append([]T{prevVertex}, path...)
			}
			return path
		}

		for _, linkedNode := range g.adjacencyList[curVert] {
			alt := dist[curVert] + linkedNode.weight
			if alt < dist[linkedNode.dest] {
				dist[linkedNode.dest] = alt
				prev[linkedNode.dest] = curVert
				priority := alt + heuristic(linkedNode.dest)
				pq.Enqueue(linkedNode.dest, priority)
			}
		}
	}

	return []T{}
}

// Dijkstra performs Dijkstra's algorithm to find the shortest path between two vertices in the graph.
// It returns the shortest path as a slice of vertices.
// If no path is found, an empty slice is returned.
func (g *Graph[T, V]) Dijkstra(start, goal T) []T {

	return g.AStar(start, goal, func(vertex T) V {
		return V(0)
	})

	/*
		// Initialize the distance map
		dist := map[T]V{}
		for vertex := range g.adjacencyList {
			dist[vertex] = maths.MaxInteger[V]()
		}
		dist[start] = V(0)

		// Initialize the previous map
		prev := map[T]T{}

		// Initialize the priority queue
		pq := container.NewPriorityQueue[T, V](true)
		pq.Enqueue(start, 0)

		for !pq.IsEmpty() {
			curVert := pq.Dequeue()

			if curVert == goal {
				// Reconstruct the path
				path := []T{goal}
				for prevVertex, ok := prev[goal]; ok; prevVertex, ok = prev[prevVertex] {
					path = append([]T{prevVertex}, path...)
				}
				return path
			}

			for _, linkedNode := range g.adjacencyList[curVert] {
				alt := dist[curVert] + linkedNode.weight
				if alt < dist[linkedNode.dest] {
					dist[linkedNode.dest] = alt
					prev[linkedNode.dest] = curVert
					pq.Enqueue(linkedNode.dest, alt)
				}
			}
		}

		return []T{}
	*/
}