package main

import "core:fmt"
import "core:math"

// the graph starts with a source s and ends with a sink t, hence the 1s.
// then there is a stack of n nodes, then a chain of k nodes, then another stack of m nodes.
// all edges have weight 1, except the chain of k nodes, which all have weight c.
// ```
//      an                  bm
//    /    \    c    c    /    \
//  s - .. - r1 - .. - rk - .. - t
//    \    /              \    /
//      a1                  b1
// ```
// we create and return an adjacency matrix representing the graphs.
create_graph :: proc($n, $m, $k: int, c: int) -> (res: [1 + n + k + m + 1][1 + n + k + m + 1]int) {
	// source to n nodes
	for i in 1 ..= n do res[0][i] = 1

	// n nodes to r1
	r1 := n + 1
	for i in 1 ..= n do res[i][r1] = 1

	// r1 through rk
	rk := r1 + k
	for i in r1 ..< rk - 1 do res[i][i + 1] = c

	// rk to m nodes
	for i in rk ..< rk + m do res[rk - 1][i] = 1

	// m nodes to sink
	for i in rk ..< rk + m do res[i][rk + m] = 1

	res += transpose_graph(res)

	return
}

transpose_graph :: proc(g: [$row][$col]$T) -> (res: [col][row]T) {
	for i in 0 ..< row {
		for j in i ..< col {
			res[i][j], res[j][i] = g[j][i], g[i][j]
		}
	}
	return
}

print_graph :: proc(g: [$row][$col]$T, elem_sep := " ", line_sep := "\n") {
	// column markers
	fmt.print(" ")
	fmt.print(elem_sep)
	for i in 0 ..< row {
		fmt.print(rune('a' + i))
		fmt.print(elem_sep)
	}
	// body
	fmt.print(line_sep)
	for i in 0 ..< row {
		fmt.print(rune('a' + i)) // row markers
		fmt.print(elem_sep)
		for j in 0 ..< col {
			fmt.print(g[i][j] < 0 ? 10 + g[i][j] : g[i][j])
			fmt.print(elem_sep)
		}
		fmt.print(line_sep)
	}
}

find_node_with_excess :: proc(excesses: []int) -> (u: int, ok: bool) {
	for e, idx in excesses do if e > 0 do return idx, true
	return -1, false
}

find_admissible_edge :: proc(u: int, residuals, heights: []int) -> (v: int, ok: bool) {
	assert(len(residuals) == len(heights))
	for r, idx in residuals {
		if r > 0 && heights[u] == heights[idx] + 1 {
			return idx, true
		}
	}
	return -1, false
}

push :: proc(u, v: int, residuals, excesses: []int) -> int {
	assert(len(residuals) == len(excesses))
	return min(excesses[u], residuals[v])
}

push_relabel :: proc(graph: [$n][n]int) -> (flow: [n][n]int) {
	heights: [n]int
	heights[0] = n
	flow[0] += graph[0]
	for _, idx in flow do flow[idx][0] -= graph[idx][0]
	excesses: [n]int
	excesses[0] = -math.sum(flow[0][:])
	excesses += flow[0]
	for {
		u := find_node_with_excess(excesses[1:n - 1]) or_break
		u += 1
		residuals := graph[u] - flow[u]
		if v, ok := find_admissible_edge(u, residuals[:], heights[:]); ok {
			assert(heights[u] == heights[v] + 1)
			delta := push(u, v, residuals[:], excesses[:])
			flow[u][v] += delta
			flow[v][u] -= delta
			excesses[u] -= delta
			excesses[v] += delta
		} else {
			assert(excesses[u] > 0)
			for r, v in residuals do if r > 0 do assert(heights[u] <= heights[v])
			m := n
			for h, v in heights {
				if residuals[v] > 0 do m = h < m ? h : m
			}
			heights[u] = 1 + m
		}
		iterations += 1
	}
	return
}

iterations := 0
n :: 2
m :: 100
k :: 1
c :: 100

main :: proc() {
	graph := create_graph(n, m, k, c)
	// print_graph(graph, elem_sep = " ")
	// fmt.println()
	// print_graph(push_relabel(graph), elem_sep = " ")
	push_relabel(graph)
	fmt.printfln(
		"push-relabel(n=%d, m=%d, k=%d, C=%d) took %d iterations.",
		n,
		m,
		k,
		c,
		iterations,
	)
}