//! returns: 0
//! 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.

constant MAX_NODES: u32 = 16;
constant 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];
    set g.heap[i] = g.heap[j];
    set 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);
            set 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 {
            set smallest = left;
        }
        if right < g.heapSize and g.heap[right].dist < g.heap[smallest].dist {
            set smallest = right;
        }
        if smallest == i {
            return;
        }
        heapSwap(g, i, smallest);
        set i = smallest;
    }
}

fn heapPush(g: *mut Graph, dist: u32, node: u32) {
    set g.heap[g.heapSize] = HeapEntry { dist, node };
    set 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];
    set g.heapSize -= 1;
    set 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) {
    set g.adj[from][to] = weight;
}

fn addBidiEdge(g: *mut Graph, from: u32, to: u32, weight: u32) {
    set g.adj[from][to] = weight;
    set g.adj[to][from] = weight;
}

fn resetGraph(g: *mut Graph, n: u32) {
    set g.numNodes = n;
    let mut i: u32 = 0;
    while i < MAX_NODES {
        let mut j: u32 = 0;
        while j < MAX_NODES {
            set g.adj[i][j] = INF;
            set j += 1;
        }
        set g.dist[i] = INF;
        set g.prev[i] = -1;
        set g.visited[i] = false;
        set i += 1;
    }
    set 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) {
    set g.dist[source] = 0;
    heapPush(g, 0, source);

    while let entry = heapPop(g) {
        let u: u32 = entry.node;

        if g.visited[u] {
            continue;
        }
        set 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] {
                    set g.dist[v] = newDist;
                    set g.prev[v] = u as i32;
                    heapPush(g, newDist, v);
                }
            }
            set 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;

    set path[len] = target;
    set len += 1;

    let mut cur: ?u32 = getPrev(g, target);
    while let node = cur {
        set path[len] = node;
        set len += 1;
        set 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];
        set path[a] = path[b];
        set path[b] = tmp;
        set a += 1;
        set 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 {
                    set diff = i - j;
                }
                addEdge(g, i, j, diff * 3 + 1);
            }
            set j += 1;
        }
        set 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;
        }
        set 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;
}