(in-package "CL-USER")

;; data is any data associated with the node
;;   (make-node x) creates a node
;;   - edges is a list of (<node> <flow> <cur-flow>) (last is for min-cut)
;;   - tag is used when finding paths
(defstruct (node
            (:constructor make-node (data))
            (:print-function (lambda (node s k)
                               (declare (ignore k))
                               (format s "<~W>" (node-data node)))))
  data (edges nil) (tag nil))

;; connect `from' to `to' with the given `capacity' (creates a zero edge back
;; if none was added so far)
(defun connect-node (from to capacity)
  (when (eq from to)
    (error "connect-node: cannot connect node ~s to itself" to))
  (let ((link (assoc to (node-edges from))))
    (when link
      (if (zerop (second link))
        (setf (second link) capacity)
        (error "connect-node: target ~s is already connected from source ~s"
               to from)))
    (push (list to capacity 0) (node-edges from))
    (unless (assoc from (node-edges to))
      (push (list from 0 0) (node-edges to)))))

;; same, but can connect multiple nodes
(defun connect-nodes (from to capacity &rest more)
  (connect-node from to capacity)
  (when more (apply #'connect-nodes more)))

;; expects a source name, a target name, and a list of node-node-capacity
;; triplets (nodes are identified by `equal' names), returns source and target
;; nodes in a list of two elements
(defun list->nodes (src-data tgt-data connections)
  (let ((table (make-hash-table :test #'equal)))
    (flet ((node (data)
             (if data
               (or (gethash data table)
                   (setf (gethash data table) (make-node data)))
               (error "list->nodes: got NIL for a data"))))
      (dolist (c connections)
        (connect-node (node (first c)) (node (second c)) (third c)))
      (list (node src-data) (node tgt-data)))))

;; connect `from' to `to' *and back* with the given `capacity' (undirected
;; edge)
(defun connect-node* (from to capacity)
  (connect-node from to capacity)
  (connect-node to from capacity))

;; same, but can connect multiple nodes
(defun connect-nodes* (from to capacity &rest more)
  (connect-node* from to capacity)
  (when more (apply #'connect-nodes* more)))

;; same as `list->nodes' but for undirected graphs, also `:verify-flat-graph?'
;; will make it verify that all nodes are connected directly to the source and
;; target nodes
(defun list->nodes* (src-data tgt-data connections &key verify-flat-graph?)
  (let ((table (make-hash-table :test #'equal)))
    (flet ((node (data)
             (if data
               (or (gethash data table)
                   (setf (gethash data table) (make-node data)))
               (error "list->nodes*: got NIL for a data"))))
      (dolist (c connections)
        (connect-node* (node (first c)) (node (second c)) (third c)))
      (let ((src (node src-data))
            (tgt (node tgt-data)))
        (when verify-flat-graph?
          (maphash
           #'(lambda (data node)
               (declare (ignore data))
               (unless (or (eq node src) (eq node tgt)
                           (and (assoc node (node-edges src))
                                (assoc node (node-edges tgt))))
                 (error "list->nodes*: node ~s is not connected to ~s & ~s"
                        node src tgt)))
           table))
        (list src tgt)))))

;; find a path between the two nodes, returns a list that begins with `from' and
;; ends with `to'
(defvar *find-path-tag*) ; used for avoiding circles
(defun find-path (from to)
  (setf (node-tag from) (incf *find-path-tag*))
  (let ((paths (list (list from))))
    (loop until (eq to (first (first paths))) do
      (let ((curpath (pop paths)))
        (unless curpath (return-from find-path nil)) ; no path!
        (setf paths
              (nconc (delete-if
                      #'null
                      (mapcar #'(lambda (e)
                                  (and (plusp (third e))
                                       (not (eq *find-path-tag*
                                                (node-tag (first e))))
                                       (progn (setf (node-tag (first e))
                                                    *find-path-tag*)
                                              (cons (first e) curpath))))
                              (node-edges (first curpath))))
                     paths))))
    (nreverse (car paths))))

;; returns the maxflow of the path (the minimum of the available flows)
(defun path-maxflow (path)
  (let* ((node (pop path))
         (maxflow (third (assoc (first path) (node-edges node)))))
    (loop until (null (rest path))
          do (let ((node (pop path)))
               (setf maxflow
                     (min (third (assoc (first path) (node-edges node)))
                          maxflow))))
    maxflow))

;; subtracts a flow along a given path (add to reverse edges)
(defun subtract-flow (path flow)
  (loop until (null (rest path))
        do (let ((node (pop path)))
             (decf (third (assoc (first path) (node-edges node))) flow)
             (incf (third (assoc node (node-edges (first path)))) flow))))

;; returns a list of all reachable nodes from `node'
(defun reachable-nodes (node &optional use-flow?)
  (let ((todo   (list node))
        (result '()))
    (loop until (null todo) do
      (let ((node (pop todo)))
        (unless (member node result)
          (push node result)
          (let ((edges (node-edges node)))
            (setf todo (append (mapcar #'first
                                       (if use-flow?
                                         (remove-if-not #'(lambda (e)
                                                            (plusp (third e)))
                                                        edges)
                                         edges))
                               todo))))))
    result))

;; perform a min-cut using the given source and target
(defun min-cut (src tgt)
  (let ((*find-path-tag* 0)
        (all-nodes (reachable-nodes src)))
    (unless (member tgt all-nodes)
      (error "min-cut: target node ~s is not reachable from source node ~s"
             tgt src))
    (dolist (node all-nodes) ; initialize flows
      (dolist (edge (node-edges node))
        (setf (third edge) (second edge))))
    (loop do (let ((path (find-path src tgt)))
               (if path
                 (subtract-flow path (path-maxflow path))
                 (let* ((src-nodes (reachable-nodes src t))
                        (tgt-nodes (set-difference all-nodes src-nodes)))
                   (unless (member tgt tgt-nodes)
                     (error "min-cut: bad results (tgt in src-nodes)"))
                   (return-from min-cut (list src-nodes tgt-nodes))))))))

(defun tests ()
  (format t "tests from ~a~%"
          "http://www-b2.is.tokushima-u.ac.jp/~ikeda/suuri/maxflow/")
  (let ((src (make-node 's))
        (tgt (make-node 't))
        (v1  (make-node 1))
        (v2  (make-node 2))
        (v3  (make-node 3))
        (v4  (make-node 4))
        (v5  (make-node 5))
        (v6  (make-node 6))
        (v7  (make-node 7))
        (v8  (make-node 8))
        (v9  (make-node 9)))
    (connect-nodes src v1  5
                   src v2  5
                   src v3  5
                   v1  v2  2
                   v2  v3  2
                   v1  v4  3
                   v2  v5  5
                   v3  v6  3
                   v1  v5  2
                   v2  v6  2
                   v5  v4  1
                   v6  v5  1
                   v4  v7  5
                   v5  v8  3
                   v6  v9  3
                   v5  v7  2
                   v5  v9  2
                   v7  v8  2
                   v8  v9  2
                   v7  tgt 5
                   v8  tgt 4
                   v9  tgt 4
                   )
    (format t ">>> test1: should be (<S> <3>)~%      ~s~%"
            (min-cut src tgt)))
  (let ((s* (make-node 's))
        (t* (make-node 't))
        (u* (make-node 'u))
        (v* (make-node 'v))
        (w* (make-node 'w))
        (x* (make-node 'x))
        (y* (make-node 'y))
        (z* (make-node 'z)))
    (connect-nodes s* u* 7 u* y* 3 y* w* 3 w* s* 3
                   v* x* 1 x* t* 5 v* z* 3 z* t* 4
                   s* v* 2 u* x* 3 y* t* 3 w* z* 4)
    (format t ">>> test2: should be (<S> <U>)~%      ~s~%"
            (min-cut s* t*)))
  (let ((s* (make-node 's))
        (t* (make-node 't))
        (u* (make-node 'u))
        (v* (make-node 'v))
        (w* (make-node 'w)))
    (connect-nodes
     s* u* 5 s* v* 7 s* w* 3 w* v* 4 v* u* 2 u* t* 6 v* t* 5 w* t* 8)
    (format t ">>> test3: should be (<S> <V> <U>)~%      ~s~%"
            (min-cut s* t*)))
  (format t "test from lillian's paper~%")
  (let ((s* (make-node 's))
        (t* (make-node 't))
        (y* (make-node 'y))
        (n* (make-node 'n))
        (m* (make-node 'm)))
    (connect-nodes* s* y* .8 s* m* .5 s* n* .1
                    y* t* .2 m* t* .5 n* t* .9
                    y* m* 1. m* n* .2 y* n* .1)
    (format t ">>> test4: should be (<S> <Y> <M>)~%      ~s~%"
            (min-cut s* t*)))
  (format t "test from lillian's paper, using `list->nodes*'~%")
  (let ((src-tgt (list->nodes* 's 't
                               '((s y .8) (s m .5) (s n .1)
                                 (y t .2) (m t .5) (n t .9)
                                 (y m 1.) (m n .2) (y n .1))
                               :verify-flat-graph? t)))
    (format t ">>> test4*: should be (<S> <Y> <M>)~%      ~s~%"
            (apply #'min-cut src-tgt))))