#lang racket
(module+ test (require rackunit))

(provide
 (contract-out
  [graph? contract?]

  [new-graph (->* ((listof (non-empty-listof string?))
                   #:directed? boolean?)
                  graph?)]
  ;; Constructs a new graph. If the `#:directed?` argument
  ;; is #t, creates a directed graph. Otherwise, creates an
  ;; undirected graph. By default, the graphs are undirected.

  ;; The neighbors are described in the first, mandatory
  ;; argument. Each inner list starts with a node and the
  ;; subsequent nodes are corrected by a 1-weight edge
  ;; to the first node

  ;; As an example, this call:
  #;
  (new-graph '(("a" "b")
               ("c")))
  ;; Creates an undirected graph with nodes "a", "b", and "c"
  ;; with an edge between "a" and "b" because "b" is in the
  ;; same list as "a".

  ;; as another example, this call:
  #;
  (new-graph '(("a" "b" "c" "d")
               ("b" "d")))
  ;; Creates an undirected graph with nodes "a", "b", "c",
  ;; and "d" as as nodes where "a" has an edge to "b", "c",
  ;; and "d", and "b" also has an edge to "d"

  ;; When building an undirected graph, then a node
  ;; being listed as a neighbor in either direction
  ;; is sufficient to count as a neighbor.
  ;;
  ;; The weights of all of the given edges are 1.
  ;; Use `add-edge!` to change them.


  [graph-directed? (-> graph? boolean?)]

  [same-graph? (-> graph? graph? boolean?)]
  ;; Determines if the two graphs have the same nodes,
  ;; and same edges (with the same weights).

  [graph-node/c (-> graph? contract?)]
  ;; Given a graph, returns a contract that accepts
  ;; strings that are nodes in the graph.

  [copy-graph (-> graph? graph?)]
  ;; Makes a new copy of the given graph, such that
  ;; modifications to the input graph (eg via
  ;; add-edge!) do not affect the result graph or
  ;; vice-versa.

  [remove-edges! (-> graph? void?)]
  ;; Removes all of the edges from the given graph.

  [neighbors (->i ([g graph?]
                   [n (g) (graph-node/c g)])
                  [res (listof string?)])]
  ;; Returns the neighbors of `n` in `g`.

  [remove-edge! (->i ([g graph?]
                      [n (g) (graph-node/c g)]
                      [m (g) (graph-node/c g)])
                     [res void?])]
  ;; Removes the edge between `n` and `m` from `g`.

  [add-edge! (->i ([g graph?]
                   [n (g) (graph-node/c g)]
                   [m (g) (graph-node/c g)])
                  ([weight real?])
                  [res void?])]
  ;; Adds an edge from `n` to `m` if `g` is a directed graph and
  ;; between the edges `n` and `m` if `g` is an undirected graph.
  ;; Also, sets the weight from `n` to `m` to `weight`
  ;; if the edge already exists, but with a different weight, the
  ;; weight is changed to the given weight.

  [edge-weight (->i ([g graph?]
                     [n (g) (graph-node/c g)]
                     [m (g) (graph-node/c g)])
                    [res (or/c real? #f)])]
  ;; Returns the weight of the edge between `n` and `m` in `g`
  ;; or #f if there is no edge between `n` and `m`.

  [nodes (-> graph? (listof string?))]
  ;; Returns a list of all of the nodes in the given graph.

  [invert (-> graph? graph?)]
  ;; returns a graph with the same nodes as the input but where
  ;; every edges goes the opposite way as it did in the input

  [no-incoming-edges (->i ([g graph?])
                          [res (g) (listof (graph-node/c g))])]
  ;; returns the nodes in `g` that have no incoming edges
  
  [has-edges? (-> graph? boolean?)]
  ;; returns #t if there are any edges in the given graph, #f otherwise
  
  [random-graph (->* (#:directed? boolean?)
                     (#:edge-probability (between/c 0 1))
                     graph?)]
  ;; Makes a random graph where nodes are connected with the
  ;; probability given by #:edge-probability.

  [random-connected-graph
   (->* ()
        (#:edge-probability (between/c 0 1))
        (and/c graph?
               (not/c graph-directed?)))]
  ;; Constructs a random undirected graph with edges having the
  ;; probability #:edge-probability and then adds edges until
  ;; the graph is connected.

  [graph->string (-> graph? string?)]
  ;; renders the entire graph (useful for debugging)
  ))


;; directed graphs are represented as
;; tables that map nodes to their neighbors,
;; where nodes are strings and
;; the neighbors are lists of strings
;; undirected graph are similar, except
;; we add edges in both directions

;; edges : hash[string -o> (listof string?)]
;; directed? : boolean
;; weights : hash[(cons/c string? string?) -o> real]
(struct graph (edges directed? weights)
   #:methods gen:custom-write
  [(define (write-proc g port mode)
     (define-values (line col pos) (port-next-location port))
     (define prefix "#<graph: ")
     (fprintf port prefix)
     (write-graph g port 3 (+ (or col 0) (string-length prefix)))
     (fprintf port ">"))])

(define (graph-node/c g)
  (define (graph-has-this-node? n) (and (member n (nodes g)) #t))
  (and/c string?
         graph-has-this-node?))

(define (new-graph neighbors #:directed? directed? #;[directed? #f])
  (define edges (make-hash))
  (define weights (make-hash))
  (for ([node (in-list (flatten neighbors))])
    (hash-set! edges node '()))
  (define g (graph edges directed? weights))
  (for ([node+neighbors (in-list neighbors)])
    (define node (first node+neighbors))
    (for ([neighbor (in-list (rest node+neighbors))])
      (add-edge! g node neighbor)))
  (validate-edges g)
  g)

;; signals an error if the internal edge data structure is wrong
(define (validate-edges g)
  (for ([i (in-list (nodes g))])
    (define adj (neighbors g i))
    (unless (equal? adj (remove-duplicates adj))
      (error 'validate-edges "~s -> ~s" i adj))))

(define (same-graph? g1 g2)
  (and (equal? (graph-directed? g1)
               (graph-directed? g2))
       (equal? (nodes g1)
               (nodes g2))
       (for*/and ([n (in-list (nodes g1))]
                  [m (in-list (nodes g1))])
         (equal? (edge-weight g1 n m)
                 (edge-weight g2 n m)))))

(module+ test
  (check-true
   (same-graph? (new-graph '(("a" "b" "c")
                             ("b" "c"))
                           #:directed? #t)
                (graph (make-hash '(("a" . ("c" "b"))
                                    ("b" . ("c"))
                                    ("c" . ())))
                       #t
                       (make-hash (list (cons (cons "a" "c") 1)
                                        (cons (cons "a" "b") 1)
                                        (cons (cons "b" "c") 1))))))

  ;; check that undirected graphs don't require
  ;; the neighbors to be listed in both directions
  (check-true (same-graph? (new-graph '(("a" "b")
                                        ("b"))
                                      #:directed? #f)
                           (new-graph '(("a")
                                        ("b" "a"))
                                      #:directed? #f)))

  (let ()
    (define g1 (new-graph '(("a") ("b") ("c")) #:directed? #f))
    (define g2 (new-graph '(("a") ("b") ("c")) #:directed? #f))
    (add-edge! g1 "a" "b" 1)
    (add-edge! g1 "a" "c" 1)
    (add-edge! g2 "a" "c" 1)
    (add-edge! g2 "a" "b" 1)
    (check-true (same-graph? g1 g2)))

  (let ()
    (define g1 (new-graph '(("a") ("b")) #:directed? #f))
    (define g2 (new-graph '(("a") ("b")) #:directed? #f))
    (add-edge! g1 "a" "b" 3)
    (add-edge! g2 "a" "b" 4)
    (check-false (same-graph? g1 g2))))

(define (copy-graph g) (graph (hash-copy (graph-edges g))
                              (graph-directed? g)
                              (hash-copy (graph-weights g))))
(module+ test
  (check-true (same-graph?
               (copy-graph (new-graph '(("a" "b" "c")
                                        ("b" "c"))
                                      #:directed? #f))
               (new-graph '(("a" "b" "c")
                            ("b" "c"))
                          #:directed? #f)))
  (check-true (same-graph?
               (copy-graph (new-graph '(("a" "b" "c")
                                        ("b" "c"))
                                      #:directed? #t))
               (new-graph '(("a" "b" "c")
                            ("b" "c"))
                          #:directed? #t)))
  (let ([g (new-graph '(("a" "b" "c")
                        ("b" "c"))
                      #:directed? #f)])
    (check-false (eq? g (copy-graph g)))))

(define (neighbors g n)
  (sort (hash-ref (graph-edges g) n '()) string<?))
(module+ test
  (check-equal? (neighbors (new-graph '(("a" "b" "c")
                                        ("b" "c"))
                                      #:directed? #f)
                           "a")
                '("b" "c")))

(define (remove-edge! g n m)
  (define edges (graph-edges g))
  (define weights (graph-weights g))
  (hash-set! edges n (remove m (hash-ref edges n '())))
  (hash-remove! weights (cons n m))
  (unless (graph-directed? g)
    (hash-set! edges m (remove n (hash-ref edges m '())))
    (hash-remove! weights (cons m n))))
(module+ test
  (let ()
    (define a-g (new-graph '(("a" "b" "c")
                             ("b" "c"))
                           #:directed? #f))
    (remove-edge! a-g "a" "b")
    (check-equal? (neighbors a-g "a") (list "c")))
  (let ()
    (define a-g (new-graph '(("a" "b" "c")
                             ("b" "c"))
                           #:directed? #f))
    (remove-edge! a-g "a" "b")
    (check-equal? (edge-weight a-g "a" "b") #f))
  (let ()
    (define a-g (new-graph '(("a" "b" "c")
                             ("b" "c"))
                           #:directed? #f))
    (remove-edge! a-g "a" "b")
    (check-equal? (edge-weight a-g "b" "a") #f))
  (let ()
    (define a-g (new-graph '(("a" "b")) #:directed? #f))
    (remove-edge! a-g "a" "b")
    (check-equal? (neighbors a-g "b") '())))

(define (remove-edges! g)
  (define edges (graph-edges g))
  (hash-clear! (graph-weights g))
  (for ([node (in-list (nodes g))])
    (hash-set! edges node '())))
(module+ test
  (let ()
    (define g (new-graph '(("a" "b" "c")) #:directed? #f))
    (remove-edges! g)
    (check-equal? (neighbors g "a") '())
    (check-equal? (neighbors g "b") '())
    (check-equal? (neighbors g "c") '())
    (check-equal? (edge-weight g "a" "b") #f)
    (check-equal? (edge-weight g "a" "c") #f)))

(define (add-edge! g n m [weight 1])
  (validate-edges g)
  (define edges (graph-edges g))
  (define weights (graph-weights g))
  (define n-neighbors (hash-ref edges n '()))
  (unless (member m n-neighbors)
    (hash-set! edges n (cons m n-neighbors))
    (unless (graph-directed? g)
      (unless (equal? n m)
        (define m-neighbors (hash-ref edges m '()))
        (hash-set! edges m (cons n m-neighbors)))))
  (hash-set! weights (cons n m) weight)
  (unless (graph-directed? g)
    (hash-set! weights (cons m n) weight))
  (validate-edges g))
(module+ test
  (let ()
    (define a-g (new-graph '(("a" "b" "c")
                             ("b" "c"))
                           #:directed? #t))
    (add-edge! a-g "b" "a")
    (check-equal? (neighbors a-g "b") (list "a" "c")))
  (let ()
    (define a-g (new-graph '(("a" "b")
                             ("c"))
                            #:directed? #f))
    (check-equal? (neighbors a-g "b") (list "a"))
    (add-edge! a-g "b" "c")
    (check-equal? (neighbors a-g "b") (list "a" "c")))
  (let ()
    (define a-g (new-graph '(("a")) #:directed? #f))
    (add-edge! a-g "a" "a" 1)))

(define (edge-weight g n m)
  (hash-ref (graph-weights g) (cons n m) #f))

(module+ test
  (let ()
    (define a-g (new-graph '(("a") ("b")) #:directed? #f))
    (add-edge! a-g "a" "b" 12)
    (check-equal? (edge-weight a-g "a" "b") 12)
    (check-equal? (edge-weight a-g "b" "a") 12))
  (let ()
    (define a-g (new-graph '(("a") ("b")) #:directed? #f))
    (check-equal? (edge-weight a-g "a" "b") #f)
    (check-equal? (edge-weight a-g "b" "a") #f))
  )

(define (nodes g) (sort (hash-keys (graph-edges g)) string<?))

(module+ test
  (check-equal? (nodes (new-graph '(("a" "b" "c")
                                    ("b" "c"))
                                  #:directed? #f))
                (list "a" "b" "c"))
  (check-equal? (nodes (new-graph '(("a" "b" "c")
                                    ("b" "c" "d"))
                                   #:directed? #f))
                (list "a" "b" "c" "d")))

(module+ test
  (check-equal? (let ([g (new-graph '(("a") ("b") ("c")) #:directed? #f)])
                  (add-edge! g "b" "a" 1)
                  (add-edge! g "c" "a" 1)
                  (for/list ([n (in-list (nodes g))])
                    (cons n (neighbors g n))))
                (list (cons "a" (list "b" "c"))
                      (cons "b" (list "a"))
                      (cons "c" (list "a")))))

(define (random-graph #:edge-probability [p 1/2]
                      #:directed? directed?)
  (define (i->n i) (~a "n" i))
  (define size (let loop () (cond [(zero? (random 10)) 0] [else (+ 1 (loop))])))
  (define g (new-graph (for/list ([i (in-range size)]) (list (i->n i)))
                       #:directed? directed?))
  (for* ([n (in-range size)]
         [m (in-range size)])
    (when (<= (random) p)
      (add-edge! g (i->n n) (i->n m))))
  g)

(define (random-connected-graph #:edge-probability [p 1/2])
  (define g (random-graph #:edge-probability p #:directed? #f))
  (let loop ()
    (define missing (find-a-missing-connection g))
    (when missing
      (add-edge! g (car missing) (cdr missing))
      (loop)))
  g)

;; find-a-missing-connection : graph -> (or/c #f (cons/c string? string?))
(define (find-a-missing-connection g)
  (define the-nodes (shuffle (nodes g)))
  (cond
    [(null? the-nodes) #f]
    [else
     (let/ec escape
       (define first-node (car the-nodes))
       (for ([node (in-list (cdr the-nodes))])
         (unless (has-path? g node first-node)
           (escape (cons node first-node))))
       #f)]))
(module+ test
  (let ([g (new-graph (list (list "a")
                            (list "d"))
                      #:directed? #f)])
    (check-true (set-member?
                 (set (cons "a" "d")
                      (cons "d" "a"))
                 (find-a-missing-connection g))))
  (let ([g (new-graph (list (list "a" "d")
                            (list "d"))
                      #:directed? #f)])
    (check-equal? (find-a-missing-connection g)
                  #f)))

(define (has-path? g n m)
  (define visited (make-hash))
  (let loop ([n n])
    (cond
      [(hash-has-key? visited n) #f]
      [(equal? n m) #t]
      [else
       (hash-set! visited n #t)
       (for/or ([neighbor (in-list (neighbors g n))])
         (loop neighbor))])))

(module+ test
  (let ([g (new-graph (list (list "a" "b")
                            (list "b")
                            (list "c"))
                      #:directed? #f)])
    (check-true (has-path? g "a" "b"))
    (check-true (has-path? g "a" "a"))
    (check-true (has-path? g "b" "a"))
    (check-false (has-path? g "a" "c")))

  (let ([g (new-graph (list (list "a" "b")
                            (list "b" "c")
                            (list "c" "d"))
                       #:directed? #f)])
    (check-true (has-path? g "a" "b"))
    (check-true (has-path? g "a" "c"))
    (check-true (has-path? g "a" "d"))
    (check-true (has-path? g "d" "a")))

  (let ([g (new-graph (list (list "a" "b")
                            (list "b" "c")
                            (list "d" "e")
                            (list "e" "f"))
                      #:directed? #f)])
    (check-true (has-path? g "a" "b"))
    (check-true (has-path? g "a" "c"))
    (check-false (has-path? g "a" "d"))
    (check-false (has-path? g "a" "e"))
    (check-false (has-path? g "f" "a"))
    (check-true (has-path? g "f" "d"))))


(module+ test
  (check-true (graph? (random-graph))))
(module+ test
  (check-false (graph-directed? (random-connected-graph))))

(define (write-graph g port max-width/length indent)
  (define the-nodes (nodes g))
  (define (nl)
    (newline port)
    (for ([i (in-range indent)]) (display " " port)))
  (for ([node (in-list the-nodes)]
        [i (if max-width/length
               (in-range max-width/length)
               (in-naturals))])
    (unless (zero? i) (nl))
    (fprintf port "~a ->" node)
    (for ([neighbor (in-list (neighbors g node))]
          [i (if max-width/length
                 (in-range max-width/length)
                 (in-naturals))])
      (fprintf port " ~a" neighbor)
      (define w (edge-weight g node neighbor))
      (unless (= 1 w) (fprintf port "(~a)" w)))
    (when (and max-width/length
               ((length (neighbors g node)) . > . max-width/length))
      (fprintf port " ...")))
  (when (and max-width/length
             ((length the-nodes) . > . max-width/length))
    (nl)
    (fprintf port "...")))

(define (graph->string g)
  (define op (open-output-string))
  (write-graph g op #f 0)
  (newline op)
  (get-output-string op))

(module+ test
  (check-equal?
   (graph->string (new-graph '(("a" "b" "c")) #:directed? #f))
   (string-append "a -> b c\n"
                  "b -> a\n"
                  "c -> a\n"))
  (check-equal?
   (let ([g (new-graph '(("a" "b" "c")) #:directed? #f)])
     (add-edge! g "a" "b" 2)
     (add-edge! g "a" "c" 33)
     (graph->string g))
   (string-append "a -> b(2) c(33)\n"
                  "b -> a(2)\n"
                  "c -> a(33)\n"))
  (check-equal?
   (let ([op (open-output-string)])
     (write (new-graph '(("a" "b" "c" "d" "e")) #:directed? #f) op)
     (get-output-string op))
   (string-append "#<graph: a -> b c d ...\n"
                  "         b -> a\n"
                  "         c -> a\n"
                  "         ...>")))

(define (invert g)
  (define res (new-graph (map (λ (x) (list x)) (nodes g))
                         #:directed? (graph-directed? g)))
  (for ([(node neighbors) (in-hash (graph-edges g))])
    (for ([neighbor (in-list neighbors)])
      (add-edge! res neighbor node (edge-weight g node neighbor))))
  res)
(module+ test
  (check-true (same-graph?
               (invert
                (new-graph '(("a" "b" "c")
                             ("b" "q")
                             ("d" "b")
                             ("q" "b"))
                           #:directed? #t))
               (new-graph '(("c" "a")
                            ("b" "a" "d" "q")
                            ("q" "b"))
                          #:directed? #t))))

(define (no-incoming-edges g)
  (define ig (invert g))
  (for/list ([n (in-list (nodes g))]
             #:when (empty? (neighbors ig n)))
    n))

(define (has-edges? g)
  (for/or ([(node neighbors) (in-hash (graph-edges g))])
    (pair? neighbors)))