#lang typed/racket

;; dijkstra's shortest-path algorithm

;; we'll represent a directed, weighted graph as a
;; hash mapping node names to
;; lists of weighted edges, where a weighted edge contains
;; a (target) name and a weight.

;; here's an example graph. There's a path of two edges
;; to node d, but the shortest-weight one goes through f

(define-type Edge-List (Listof (List Symbol Nonnegative-Real)))
(define-type Graph (Immutable-HashTable Symbol Edge-List))

(define example-graph
  (make-immutable-hash
   '((a . ((b 2) (c 1)))
     (b . ((d 10) (e 3)))
     (c . ((e 5)))
     (d . ((f 4)))
     (e . ((f 2)))
     (f . ((d 1))))))

;; a dist-hash is a hash table mapping each node to its best known
;; shortest-path, or 'inf if no path is yet known. initially,
;; all nodes will be at distance inf.

(define-type Dist (U 'inf Nonnegative-Real))
(define-type Dist-Hash (Immutable-HashTable Symbol Dist))

;; in the "reached" nodes, 'inf is not allowed:

(define-type Num-Dist-Hash (Immutable-HashTable Symbol Nonnegative-Real))

;; the algorithm will maintain two of these, one for finished
;; nodes and one for unfinished nodes. The second of these could
;; certainly be implemented as a priority queue, if desired.

(define (make-init-hash [g : Graph]) : Dist-Hash
  (for/hash : Dist-Hash ([name : Symbol (in-list (hash-keys g))])
     (values name 'inf)))

;; given an up-to-date dist-hash,
;; return the node with the shortest known path
(define (best [dist-hash : Dist-Hash]) : (Pairof Symbol Nonnegative-Real)
  (define all-dists (hash->list dist-hash))
  (define non-inf-dists (remove-infs all-dists))
  (when (empty? non-inf-dists)
    (raise-argument-error 'best
                          "dist-hash with some non-inf distances"
                          0 dist-hash))
  ;; return the one with the smallest distance:
  (argmin (inst cdr Any Real) non-inf-dists))

;; is this pair a number pair?
(: num-pair? (-> (Pairof Symbol Dist) Boolean : (Pairof Symbol Nonnegative-Real)));
(define (num-pair? d)
  (number? (cdr d)))

;; given a list of pairs, remove those whose cdr is 'inf
(define (remove-infs [pairs : (Listof (Pairof Symbol Dist))])
  : (Listof (Pairof Symbol Nonnegative-Real))
  (filter num-pair? pairs))

;; given a finished dist-hash and a to-be-reached dist-hash
;; and a newly-added name and the full graph, update the
;; to-be-reached dist-hash by updating nodes reachable from
;; the newly-added node
(define (update-dists [reached : Num-Dist-Hash] [unreached : Dist-Hash] [new : Symbol] [g : Graph])
  : Dist-Hash
  (define new-dist (hash-ref reached new))
  (define new-edges (hash-ref g new))
  (for/fold ([unreached : Dist-Hash unreached])
            ([edge (in-list new-edges)]
             #:when (not (hash-has-key? reached (first edge))))
    (define tgt-name (first edge))
    (define edge-weight (second edge))
    (define old-dist (hash-ref unreached tgt-name))
    (cond [(or (equal? old-dist 'inf)
               (< (+ new-dist edge-weight) old-dist))
           (hash-set unreached tgt-name (+ new-dist edge-weight))]
          [else unreached])))

;; given a 'reached' hash-dist and an 'unreached' hash-dist and the
;; graph g, perform
;; a single iteration of dijkstra's algorithm, returning the new
;; 'reached' and 'unreached' hashes
(define (dijkstra-step [reached : Num-Dist-Hash] [unreached : Dist-Hash] [g : Graph])
  : (List Num-Dist-Hash Dist-Hash)
  (match-define (cons new-node best-dist) (best unreached))
  (define new-reached (hash-set reached new-node
                                best-dist))
  (define new-unreached (hash-remove unreached new-node))
  (define new2-unreached (update-dists new-reached
                                       new-unreached new-node g))
  (list new-reached new2-unreached))

;; given a reached and unreached dist-hash, iterate until all nodes are
;; reached
(define (dijkstra-loop [reached : Num-Dist-Hash] [unreached : Dist-Hash] [g : Graph])
  : Dist-Hash
  (cond [(= 0 (hash-count unreached))
         reached]
        [else (match-define (list new-reached new-unreached)
                (dijkstra-step reached unreached g))
              (dijkstra-loop new-reached new-unreached g)]))

;; given a graph and a starting node, perform dijkstra's algorithm
;; and return the resulting 'reached' dist-hash
(define (dijkstra [g : Graph] [starting-node : Symbol])
  (dijkstra-loop (hash) (hash-set (make-init-hash g) starting-node 0) g))

(require typed/rackunit)

(check-equal? (update-dists (hash 'a 0 'c 1)
                            (hash 'b 2 'd 'inf 'e 'inf 'f 'inf)
                            'c
                            example-graph)
              (hash 'b 2 'd 'inf 'e 6 'f 'inf))

(check-equal? (dijkstra-step (hash 'a 0 'c 1)
                             (hash 'b 2 'd 'inf 'e 6 'f 'inf)
                             example-graph)
              (list
               (hash 'a 0 'c 1 'b 2)
               (hash 'd 12 'e 5 'f 'inf)))
(check-equal? (best (hash 'a 42
                          'b 'inf
                          'c 9))
              (cons 'c 9))

(check-equal? (dijkstra example-graph 'a)
              (hash 'a 0 'b 2 'c 1 'd 8 'e 5 'f 7))