package day25

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

	"gitea.paas.celticinfo.fr/oabrivard/aoc2023/utils"
)

type Graph map[string]map[string]bool
type Edge struct {
	left  string
	right string
}

func buildGraph(lines []string) Graph {
	graph := map[string]map[string]bool{}

	for _, line := range lines {
		parts := strings.Split(line, ":")

		src := parts[0]
		if _, found := graph[src]; !found {
			graph[src] = map[string]bool{}
		}
		links := graph[src]

		dests := strings.Fields(strings.TrimSpace(parts[1]))
		for _, dst := range dests {
			links[dst] = true

			if _, found := graph[dst]; !found {
				graph[dst] = map[string]bool{}
			}
			reverseLink := graph[dst]
			reverseLink[src] = true
		}
	}

	return graph
}

func printGraph(graph Graph) {
	for k1, v1 := range graph {
		fmt.Print(k1, ": ")
		for k2 := range v1 {
			fmt.Print(k2, " ")
		}
		fmt.Println("")
	}
}

func countComponents(graph Graph) int {
	visited := map[string]bool{}
	count := 0

	for node := range graph {
		if !visited[node] {
			count++
			queue := utils.NewQueue[string]()
			queue.Enqueue(node)

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

				for linkedNode := range graph[curNode] {
					if !visited[linkedNode] {
						queue.Enqueue(linkedNode)
					}
				}
			}
		}
	}
	return count
}

func buildVerticesLists(graph Graph) []map[string]bool {
	visited := map[string]bool{}
	result := []map[string]bool{}
	currSubgraph := 0

	for node := range graph {
		if !visited[node] {
			result = append(result, map[string]bool{})
			result[currSubgraph][node] = true

			queue := utils.NewQueue[string]()
			queue.Enqueue(node)

			for queue.HasElement() {
				curNode := queue.Dequeue()
				visited[curNode] = true
				result[currSubgraph][curNode] = true

				for linkedNode := range graph[curNode] {
					if !visited[linkedNode] {
						queue.Enqueue(linkedNode)
					}
				}
			}
			currSubgraph++
		}
	}
	return result
}

func buildEdgeList(graph Graph) []Edge {
	edges := map[Edge]bool{}

	for node, connectedNodes := range graph {
		for connectedNode := range connectedNodes {
			_, found1 := edges[Edge{node, connectedNode}]
			_, found2 := edges[Edge{connectedNode, node}]

			if !found1 && !found2 {
				edges[Edge{node, connectedNode}] = true
			}
		}
	}

	result := []Edge{}
	for e := range edges {
		result = append(result, e)
	}
	return result
}

func findEdgesToDelete(graph Graph, edges []Edge) (bool, *Edge, *Edge, *Edge) {
	for i := 0; i < len(edges); i++ {
		for j := i + 1; j < len(edges); j++ {
			for k := j + 1; k < len(edges); k++ {
				// Delete edges i, j and k from graph
				delete(graph[edges[i].left], edges[i].right)
				delete(graph[edges[i].right], edges[i].left)
				delete(graph[edges[j].left], edges[j].right)
				delete(graph[edges[j].right], edges[j].left)
				delete(graph[edges[k].left], edges[k].right)
				delete(graph[edges[k].right], edges[k].left)

				// count components
				count := countComponents(graph)

				// restore Graph
				graph[edges[i].left][edges[i].right] = true
				graph[edges[i].right][edges[i].left] = true
				graph[edges[j].left][edges[j].right] = true
				graph[edges[j].right][edges[j].left] = true
				graph[edges[k].left][edges[k].right] = true
				graph[edges[k].right][edges[k].left] = true

				if count == 2 {
					return true, &edges[i], &edges[j], &edges[k]
				}
			}
		}
	}

	return false, nil, nil, nil
}

// Too basic. This brute force approach does not work with a large graph
func Part1(lines []string) int {
	graph := buildGraph(lines)
	printGraph(graph)
	edges := buildEdgeList(graph)
	fmt.Println(edges)
	success, edge1, edge2, edge3 := findEdgesToDelete(graph, edges)

	if !success {
		return 0
	}

	fmt.Println(*edge1)
	fmt.Println(*edge2)
	fmt.Println(*edge3)

	delete(graph[edge1.left], edge1.right)
	delete(graph[edge1.right], edge1.left)
	delete(graph[edge2.left], edge2.right)
	delete(graph[edge2.right], edge2.left)
	delete(graph[edge3.left], edge3.right)
	delete(graph[edge3.right], edge3.left)

	result := buildVerticesLists(graph)

	fmt.Println(result[0])
	fmt.Println(result[1])

	return len(result[0]) * len(result[1])
}

// Try to solve it using https://en.wikipedia.org/wiki/Karger%27s_algorithm
func randomStringKey[T any](dic map[string]T) string {
	i := rand.Int() % len(dic)

	for key := range dic {
		i--
		if i < 0 {
			return key
		}
	}

	return ""
}

func swapConverged(graph *Graph, target, converged, vertex string) {
	connected := (*graph)[target]

	val, found := connected[converged]
	if !found {
		log.Fatal("Should never happen")
	}
	delete(connected, converged)
	connected[vertex] = val
}

func kargerMinCut(graph Graph) (*Graph, int) {
	count := 0

	for len(graph) > 2 {
		// randomly select a vertex and one of its connected vertex (called 'converged') to be merged
		vertex := randomStringKey(graph)
		converged := randomStringKey(graph[vertex])

		// Remove link between vertex and converged vertex
		delete(graph[converged], vertex)
		delete(graph[vertex], converged)
		count++

		// Merge converged vertex into vertex
		for k, v := range graph[converged] {
			graph[vertex][k] = v
		}

		// swap converged vertex with vertex in all adjacency list vertices that shared an edge with converged vertex
		for k := range graph[converged] {
			verticesToLink := graph[k]
			verticesToLink[vertex] = true
			delete(verticesToLink, converged)
			//swapConverged(&graph, k, converged, vertex)
		}

		// remove converged vertex from graph (adjency list)
		delete(graph, converged)
		count = len(graph[vertex])
	}

	return &graph, count
}

func duplicateGraph(graph *Graph) *Graph {
	result := Graph{}

	for k1, v1 := range *graph {
		result[k1] = map[string]bool{}
		for k2, v2 := range v1 {
			result[k1][k2] = v2
		}
	}

	return &result
}

// buggy since the adjency list is implemented with a map of map. Shoud be a map of slices to work
func Part1WithKerger(lines []string) int {
	srcGraph := buildGraph(lines)
	printGraph(srcGraph)
	edges := buildEdgeList(srcGraph)
	fmt.Println(edges)

	minCut := 0
	var graphParts *Graph

	for minCut != 3 {
		graph := duplicateGraph(&srcGraph)

		graphParts, minCut = kargerMinCut(*graph)
	}

	if len(*graphParts) != 2 {
		log.Fatal("Should always be 2")
	}

	result := 1
	for _, v := range *graphParts {
		result *= len(v)
	}

	return result
}

type GraphHolder struct {
	v, e  int
	edges []Edge
	graph Graph
}

type Subset struct {
	parent string
	rank   int
}

// see https://ondrej-kvasnovsky-2.gitbook.io/algorithms/finding/union-find-algorithms
// or https://jojozhuang.github.io/algorithm/algorithm-union-find/
func find(subsets map[string]*Subset, e string) string {
	// find root and make root as parent of e
	// (path compression)
	if subsets[e].parent != e {
		subsets[e].parent = find(subsets, subsets[e].parent)
	}
	return subsets[e].parent
}

func Union(subsets map[string]*Subset, x, y string) {
	xroot := find(subsets, x)
	yroot := find(subsets, y)

	// Attach smaller rank tree under root of high
	// rank tree (Union by Rank)
	if subsets[xroot].rank < subsets[yroot].rank {
		subsets[xroot].parent = yroot
	} else {
		if subsets[xroot].rank > subsets[yroot].rank {
			subsets[yroot].parent = xroot
		} else {
			// If ranks are same, then make one as root and
			// increment its rank by one
			subsets[yroot].parent = xroot
			subsets[xroot].rank++
		}
	}
}

// Adapted from https://www.geeksforgeeks.org/introduction-and-implementation-of-kargers-algorithm-for-minimum-cut/?ref=lbp
func kargerMinCutUF(graphHolder GraphHolder) (int, *map[string]*Subset) {
	// Get data of given graph
	V := graphHolder.v
	E := graphHolder.e
	edges := graphHolder.edges

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

	// Create V subsets with single elements
	for k := range graphHolder.graph {
		subsets[k] = &Subset{k, 0}
	}

	// Initially there are V vertices in
	// contracted graph
	vertices := V

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

		// Find vertices (or sets) of two corners
		// of current edge
		subset1 := find(subsets, edges[i].left)
		subset2 := find(subsets, edges[i].right)

		// 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--
			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 < E; i++ {
		subset1 := find(subsets, edges[i].left)
		subset2 := find(subsets, edges[i].right)
		if subset1 != subset2 {
			cutedges++
		}
	}

	return cutedges, &subsets
}

func Part1WithKergerUF(lines []string) int {
	srcGraph := buildGraph(lines)
	printGraph(srcGraph)
	edges := buildEdgeList(srcGraph)
	fmt.Println(edges)

	gh := GraphHolder{len(srcGraph), len(edges), edges, srcGraph}

	minCut := 1000000
	for i := 0; i < 1000; i++ {
		cut, subsets := kargerMinCutUF(gh)

		if cut < minCut {
			minCut = cut
		}

		if minCut == 3 {
			counts := map[string]int{}

			for _, v := range *subsets {
				currCount, found := counts[v.parent]
				if !found {
					counts[v.parent] = 1
				} else {
					counts[v.parent] = currCount + 1
				}
			}

			result := 1
			for _, count := range counts {
				result *= count
			}

			return result

		}
	}

	return minCut
}