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

(provide
 (contract-out
  [graph/c contract?]
  [new-graph (-> (listof (non-empty-listof string?)) graph/c)]
  [graph-node/c (-> graph/c contract?)]
  [copy-graph (-> graph/c graph/c)]
  [neighbors (->i ([g graph/c]
                   [n (g) (graph-node/c g)])
                  [res (listof string?)])]
  [remove-edge! (->i ([g graph/c]
                      [n (g) (graph-node/c g)]
                      [m (g) (graph-node/c g)])
                     [res void?])]
  [add-edge! (->i ([g graph/c]
                   [n (g) (graph-node/c g)]
                   [m (g) (graph-node/c g)])
                  [res void?])]
  [invert (-> graph/c graph/c)]
  [nodes (-> graph/c (listof string?))]
  [no-incoming-edges (-> graph/c (listof string?))]
  [has-edges? (-> graph/c boolean?)]
  [random-graph (->* (#:edge-probability (between/c 0 1))
                     graph/c)]))

;; directed graphs, represented as
;; tables that map nodes to their neighbors,
;; where nodes are strings and
;; the neighbors are lists of strings
(define graph/c
  (hash/c string? (listof string?)))

;; given a graph, returns a contract that accepts
;; strings that are nodes in the graph
(define (graph-node/c g)
  (define (graph-has-this-node? n) (and (member n (hash-keys g)) #t))
  (and/c string?
         graph-has-this-node?))

(define (new-graph neighbors)
  (define g (make-hash))
  (for ([node (in-list (flatten neighbors))])
    (hash-set! g node '()))
  (for ([node+neighbors (in-list neighbors)])
    (define node (first node+neighbors))
    (for ([neighbor (in-list (rest node+neighbors))])
      (add-edge! g node neighbor)))
  g)
(module+ test
  (check-equal? (new-graph '(("a" "b" "c")
                             ("b" "c")))
                (make-hash  '(("a" . ("c" "b"))
                              ("b" . ("c"))
                              ("c" . ())))))

(define (copy-graph g) (hash-copy g))
(module+ test
  (check-equal? (copy-graph (new-graph '(("a" "b" "c")
                                         ("b" "c"))))
                (new-graph '(("a" "b" "c")
                             ("b" "c"))))
  (let ([g (new-graph '(("a" "b" "c")
                        ("b" "c")))])
    (check-false (eq? g (copy-graph g)))))

(define (neighbors g n)
  (hash-ref g n '()))
(module+ test
  (check-equal? (neighbors (new-graph '(("a" "b" "c")
                                        ("b" "c")))
                           "a")
                '("c" "b")))

(define (remove-edge! g n m)
  (hash-set! g n (remove m (hash-ref g n '()))))
(module+ test
  (let ()
    (define a-g (new-graph '(("a" "b" "c")
                             ("b" "c"))))
    (remove-edge! a-g "a" "b")
    (check-equal? (neighbors a-g "a") (list "c"))))

(define (add-edge! g n m)
  (define n-neighbors (hash-ref g n '()))
  (unless (member m n-neighbors)
    (hash-set! g n (cons m n-neighbors))))
(module+ test
  (let ()
    (define a-g (new-graph '(("a" "b" "c")
                             ("b" "c"))))
    (add-edge! a-g "b" "a")
    (check-equal? (neighbors a-g "b") (list "a" "c"))))

(define (invert g)
  (define res (new-graph (map (λ (x) (list x)) (nodes g))))
  (for ([(node neighbors) (in-hash g)])
    (for ([neighbor (in-list neighbors)])
      (add-edge! res neighbor node)))
  res)

(module+ test
  (check-equal? (invert (new-graph '(("a" "b" "c")
                                     ("b" "c"))))
                (new-graph '(("b" "a")
                             ("c" "a" "b")))))

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

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

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

(module+ test
  (check-equal? (no-incoming-edges (new-graph '(("a" "b" "c")
                                                ("b" "c"))))
                (list "a"))
  (check-equal? (no-incoming-edges (new-graph '(("a" "b" "c")
                                                ("b" "c" "d")
                                                ("f" "q"))))
                (list "a" "f")))


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

(define (random-graph #:edge-probability [p 1/2])
  (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)))))
  (for* ([n (in-range size)]
         [m (in-range size)])
    (when (<= (random) p)
      (add-edge! g (i->n n) (i->n m))))
  g)
(module+ test
  (check-true (hash? (random-graph))))