//! Dijkstra's shortest path algorithm.

//! Find shortest paths in a weighted directed graph using an adjacency matrix

//! and a min-heap priority queue. Reconstruct paths and verify distances.

const MAX_NODES: u32 = 16;

const INF: u32 = 0xFFFFFFFF;

record HeapEntry {

    dist: u32,

    node: u32,

}

record Graph {

    adj: *mut [[u32; 16]],

    dist: *mut [u32],

    prev: *mut [i32],

    visited: *mut [bool],

    numNodes: u32,

    heap: *mut [HeapEntry],

    heapSize: u32,

}

fn heapSwap(g: *mut Graph, i: u32, j: u32) {

    let tmp: HeapEntry = g.heap[i];

    g.heap[i] = g.heap[j];

    g.heap[j] = tmp;

}

fn siftUp(g: *mut Graph, pos: u32) {

    let mut i: u32 = pos;

    while i > 0 {

        let parent: u32 = (i - 1) / 2;

        if g.heap[i].dist < g.heap[parent].dist {

            heapSwap(g, i, parent);

            i = parent;

        } else {

            return;

        }

    }

}

fn siftDown(g: *mut Graph, pos: u32) {

    let mut i: u32 = pos;

    while true {

        let left: u32 = 2 * i + 1;

        let right: u32 = 2 * i + 2;

        let mut smallest: u32 = i;

        if left < g.heapSize and g.heap[left].dist < g.heap[smallest].dist {

            smallest = left;

        }

        if right < g.heapSize and g.heap[right].dist < g.heap[smallest].dist {

            smallest = right;

        }

        if smallest == i {

            return;

        }

        heapSwap(g, i, smallest);

        i = smallest;

    }

}

fn heapPush(g: *mut Graph, dist: u32, node: u32) {

    g.heap[g.heapSize] = HeapEntry { dist, node };

    g.heapSize += 1;

    siftUp(g, g.heapSize - 1);

}

/// Pop from the heap. Returns nil if the heap is empty.

fn heapPop(g: *mut Graph) -> ?HeapEntry {

    if g.heapSize == 0 {

        return nil;

    }

    let result: HeapEntry = g.heap[0];

    g.heapSize -= 1;

    g.heap[0] = g.heap[g.heapSize];

    if g.heapSize > 0 {

        siftDown(g, 0);

    }

    return result;

}

fn addEdge(g: *mut Graph, from: u32, to: u32, weight: u32) {

    g.adj[from][to] = weight;

}

fn addBidiEdge(g: *mut Graph, from: u32, to: u32, weight: u32) {

    g.adj[from][to] = weight;

    g.adj[to][from] = weight;

}

fn resetGraph(g: *mut Graph, n: u32) {

    g.numNodes = n;

    let mut i: u32 = 0;

    while i < MAX_NODES {

        let mut j: u32 = 0;

        while j < MAX_NODES {

            g.adj[i][j] = INF;

            j += 1;

        }

        g.dist[i] = INF;

        g.prev[i] = -1;

        g.visited[i] = false;

        i += 1;

    }

    g.heapSize = 0;

}

/// Get the predecessor as an optional; -1 means no predecessor.

fn getPrev(g: *Graph, node: u32) -> ?u32 {

    let p = g.prev[node];

    if p < 0 {

        return nil;

    }

    return p as u32;

}

fn dijkstra(g: *mut Graph, source: u32) {

    g.dist[source] = 0;

    heapPush(g, 0, source);

    while let entry = heapPop(g) {

        let u: u32 = entry.node;

        if g.visited[u] {

            continue;

        }

        g.visited[u] = true;

        let mut v: u32 = 0;

        while v < g.numNodes {

            if g.adj[u][v] != INF and not g.visited[v] {

                let newDist: u32 = g.dist[u] + g.adj[u][v];

                if newDist < g.dist[v] {

                    g.dist[v] = newDist;

                    g.prev[v] = u as i32;

                    heapPush(g, newDist, v);

                }

            }

            v += 1;

        }

    }

}

/// Reconstruct the shortest path from source to target.

fn reconstructPath(g: *Graph, target: u32, path: *mut [u32]) -> u32 {

    let mut len: u32 = 0;

    path[len] = target;

    len += 1;

    let mut cur: ?u32 = getPrev(g, target);

    while let node = cur {

        path[len] = node;

        len += 1;

        cur = getPrev(g, node);

    }

    // Reverse the path.

    let mut a: u32 = 0;

    let mut b: u32 = len - 1;

    while a < b {

        let tmp: u32 = path[a];

        path[a] = path[b];

        path[b] = tmp;

        a += 1;

        b -= 1;

    }

    return len;

}

fn testLinear(g: *mut Graph) -> i32 {

    resetGraph(g, 5);

    addEdge(g, 0, 1, 10);

    addEdge(g, 1, 2, 10);

    addEdge(g, 2, 3, 10);

    addEdge(g, 3, 4, 10);

    dijkstra(g, 0);

    let expected: [u32; 5] = [0, 10, 20, 30, 40];

    for exp, i in expected {

        if g.dist[i] != exp {

            return i as i32 + 1;

        }

    }

    let mut path: [u32; 16] = [0; 16];

    let len: u32 = reconstructPath(g, 4, &mut path[..]);

    assert len == 5;

    assert path[0] == 0;

    assert path[4] == 4;

    return 0;

}

fn testShortcut(g: *mut Graph) -> i32 {

    resetGraph(g, 4);

    addEdge(g, 0, 1, 10);

    addEdge(g, 1, 2, 10);

    addEdge(g, 0, 3, 5);

    addEdge(g, 3, 2, 3);

    dijkstra(g, 0);

    assert g.dist[2] == 8;

    let prev2 = getPrev(g, 2) else {

        return 2;

    };

    assert prev2 == 3;

    return 0;

}

fn testBidirectional(g: *mut Graph) -> i32 {

    resetGraph(g, 5);

    addBidiEdge(g, 0, 1, 1);

    addBidiEdge(g, 1, 2, 2);

    addBidiEdge(g, 2, 3, 3);

    addBidiEdge(g, 3, 0, 4);

    addBidiEdge(g, 0, 4, 7);

    addBidiEdge(g, 1, 4, 5);

    addBidiEdge(g, 2, 4, 1);

    dijkstra(g, 0);

    let expected: [u32; 5] = [0, 1, 3, 4, 4];

    for exp, i in expected {

        if g.dist[i] != exp {

            return i as i32 + 1;

        }

    }

    return 0;

}

fn testComplete(g: *mut Graph) -> i32 {

    resetGraph(g, 8);

    let mut i: u32 = 0;

    while i < 8 {

        let mut j: u32 = 0;

        while j < 8 {

            if i != j {

                let mut diff: u32 = j - i;

                if i > j {

                    diff = i - j;

                }

                addEdge(g, i, j, diff * 3 + 1);

            }

            j += 1;

        }

        i += 1;

    }

    dijkstra(g, 0);

    let mut k: u32 = 1;

    while k < 8 {

        let expected: u32 = k * 3 + 1;

        if g.dist[k] != expected {

            return k as i32;

        }

        k += 1;

    }

    return 0;

}

fn testDisconnected(g: *mut Graph) -> i32 {

    resetGraph(g, 6);

    addEdge(g, 0, 1, 5);

    addEdge(g, 1, 2, 3);

    addEdge(g, 3, 4, 2);

    addEdge(g, 4, 5, 1);

    dijkstra(g, 0);

    assert g.dist[0] == 0;

    assert g.dist[1] == 5;

    assert g.dist[2] == 8;

    // Unreachable nodes stay at INF.

    assert g.dist[3] == INF;

    assert g.dist[4] == INF;

    assert g.dist[5] == INF;

    // Predecessors of unreachable nodes should be nil.

    assert getPrev(g, 3) == nil;

    return 0;

}

fn testDiamond(g: *mut Graph) -> i32 {

    resetGraph(g, 5);

    addEdge(g, 0, 1, 5);

    addEdge(g, 0, 2, 5);

    addEdge(g, 1, 3, 5);

    addEdge(g, 2, 3, 5);

    addEdge(g, 3, 4, 2);

    dijkstra(g, 0);

    assert g.dist[3] == 10;

    assert g.dist[4] == 12;

    // Predecessor should be 1 or 2.

    let prev3 = getPrev(g, 3) else {

        return 3;

    };

    assert prev3 == 1 or prev3 == 2;

    return 0;

}

@default fn main() -> i32 {

    let mut adj: [[u32; 16]; 16] = [[0xFFFFFFFF; 16]; 16];

    let mut dist: [u32; 16] = [0xFFFFFFFF; 16];

    let mut prev: [i32; 16] = [-1; 16];

    let mut visited: [bool; 16] = [false; 16];

    let mut heap: [HeapEntry; 256] = [HeapEntry { dist: 0, node: 0 }; 256];

    let mut g: Graph = Graph {

        adj: &mut adj[..],

        dist: &mut dist[..],

        prev: &mut prev[..],

        visited: &mut visited[..],

        numNodes: 0,

        heap: &mut heap[..],

        heapSize: 0,

    };

    let r1: i32 = testLinear(&mut g);

    if r1 != 0 {

        return 10 + r1;

    }

    let r2: i32 = testShortcut(&mut g);

    if r2 != 0 {

        return 20 + r2;

    }

    let r3: i32 = testBidirectional(&mut g);

    if r3 != 0 {

        return 30 + r3;

    }

    let r4: i32 = testComplete(&mut g);

    if r4 != 0 {

        return 40 + r4;

    }

    let r5: i32 = testDisconnected(&mut g);

    if r5 != 0 {

        return 50 + r5;

    }

    let r6: i32 = testDiamond(&mut g);

    if r6 != 0 {

        return 60 + r6;

    }

    return 0;

}