(define-module (graphgif))
(use-modules (cairo)
	     (d-)
	     (srfi srfi-1)
	     (srfi srfi-9 gnu)
	     (srfi srfi-11)
	     (srfi srfi-43))

;;;;;;;;;;;;;;;;;;;
;; Basic Drawing ;;
;;;;;;;;;;;;;;;;;;;

(define-public pi  3.14159) ;; Good enough
(define-public tau (* 2 pi))

(define-public black (cairo-pattern-create-rgb 0 0 0))
(define-public white (cairo-pattern-create-rgb 1 1 1))
(define-public red   (cairo-pattern-create-rgb 1 0 0))

(define-immutable-record-type <node>
  (node coords edges color)
  node?
  (coords node-coords set-node-coords)
  (edges node-edges set-node-edges)
  (color node-color set-node-color))
(export node node?
	node-coords set-node-coords
	node-edges set-node-edges
	node-color set-node-color)

(define (edge-painter cr graph)
  "Creates a closure that draws the edges of graph using cr"
  (lambda (node)
    (cairo-set-source cr black)
    (let-values ([(x y) (car+cdr (node-coords node))]
		 [(edges) (node-edges node)])
      (for-each
       (lambda (edge)
	 (let*-values ([(other) (assq-ref graph edge)]
		       [(ox oy) (car+cdr (node-coords other))])
	   (cairo-move-to cr x y)
	   (cairo-line-to cr ox oy)
	   (cairo-stroke cr)))
       edges))))

(define (node-painter cr)
  "Creates a closure that draws the nodes of graph using cr"
  (lambda (node)
    (let-values ([(x y) (car+cdr (node-coords node))]
		 [(color) (or (node-color node) white)])
      (cairo-arc cr x y 4. 0. tau)
      (cairo-set-source cr color)
      (cairo-fill-preserve cr)
      (cairo-set-source cr black)
      (cairo-stroke cr))))

(define-public (draw-abstract-graph graph)
  "Creates a cairo surface with graph drawn on it"
  (let* ([surface (cairo-image-surface-create 'rgb24 400 400)]
	 [cr (cairo-create surface)])
    (define paint-edges (edge-painter cr graph))
    (define paint-nodes (node-painter cr))
    ;; White background
    (cairo-rectangle cr 0 0 400 400)
    (cairo-set-source cr white)
    (cairo-fill cr)
    (for-each paint-edges (map cdr graph))
    (for-each paint-nodes (map cdr graph))
    (cairo-destroy cr)
    surface))

;;;;;;;;;;;;;;;;;;;;;;
;; Graph Generation ;;
;;;;;;;;;;;;;;;;;;;;;;

(define (idx->x i w)
  "Given i and width w, returns the x value"
  (modulo i w))

(define (idx->y i w)
  "Given i and width w, returns the y value"
  (quotient i w))

(define (xy->idx x y w)
  "Returns the index of x y on width x in a linear buffer"
  (+ (* y w) x))

(define (idx->edges i w)
  "Returns the edges needed to fully connect i on graph width w"
  (filter-map
   (lambda (offset)
     (let* ([x (idx->x i w)]
	    [y (idx->y i w)]
	    [ox (+ x (car offset))]
	    [oy (+ y (cdr offset))])
       (and
	(not (negative? ox))
	(not (negative? oy))
	(< ox w)
	(xy->idx ox oy w))))
   ;; Auto-connect these directions if legal indices
   ;; This combined with forwards and back connections create a
   ;; fully connected graph
   '(( 0 . -1) ;; Above
     (-1 .  0) ;; Right
     (-1 . -1) ;; Right-above
     (+1 . -1) ;; Left above
     )))

(define lset-unionq (partial lset-union eq?))

(define (forward-connect node graph idx)
  (cons idx
	(set-node-edges
	 node
	 (lset-unionq
	  (node-edges node)
	  (fold (lambda (entry set)
		  (if (memq idx (node-edges (cdr entry)))
		      (cons (car entry) set)
		      set))
		'()
		graph)))))

(define (back-connect edges graph idx)
  (map
   (lambda (entry)
     (cons (car entry)
	   (if (memq (car entry) edges)
	       (set-node-edges
		(cdr entry)
		(lset-unionq (node-edges (cdr entry)) `(,idx)))
	       (cdr entry))))
   graph))

(define-public (connect-bidirectionally graph)
  "Takes a directed graph and connects all edges to make it fully bidirectional"
  (if (null? graph) '()
      (let* ([idx (caar graph)]
	     [node (cdar graph)])
	(cons (forward-connect node (cdr graph) idx)
	      (back-connect
	       (node-edges node)
	       (connect-bidirectionally (cdr graph))
	       idx)))))

(define-public (generate-web w h)
  "Creates a fully connected graph of width w and height h"
  (define (make-node i)
    (cons i
	  (node
	   (cons (+ (* 30 (idx->x i w)) 10)
		 (+ (* 30 (idx->y i w)) 10))
	   (idx->edges i w)
	   #f)))
  (let loop ([i 0]
	     [lst '()])
    (if (>= i (* w h))
	(connect-bidirectionally (reverse lst))
	(loop (1+ i) (cons (make-node i) lst)))))

(define-public (remove-rect graph w x1 y1 x2 y2)
  "Removes a rectangular section of a fully connected rectangular graph"
  ;; Assumption is that is a graph generated by generate-web
  ;; and that it has not been re-indexed
  ;; x1 < x2, y1 < y2
  (define edges-to-remove
    (unfold (lambda (s) (> s (xy->idx x2 y2 w)))
	    identity
	    (lambda (s)
	      (if (>= (idx->x s w) x2)
		  (xy->idx x1 (1+ (idx->y s w)) w)
		  (1+ s)))
	    (xy->idx x1 y1 w)))
  (filter-map
   (lambda (idx/node)
     (if (memq (car idx/node) edges-to-remove)
	 #f
	 (cons (car idx/node)
	       (set-node-edges
		(cdr idx/node)
		(remove
		 (lambda (edge) (memq edge edges-to-remove))
		 (node-edges (cdr idx/node)))))))
   graph))

;;;;;;;;;;;;;;;;;;;
;; ALGORITHMS!!! ;;
;;;;;;;;;;;;;;;;;;;

(define-public (heap-insert priority item heap)
  (define (fix-up! i heap)
    (if (zero? i)
	heap
	(let* ([parent (quotient (1- i) 2)]
	       [i-priority (car (vector-ref heap i))]
	       [parent-priority (car (vector-ref heap parent))])
	  (when (< i-priority parent-priority)
	    (vector-swap! heap i parent))
	  (fix-up! parent heap))))
  (fix-up! (vector-length heap)
	   (vector-append heap (vector (cons priority item)))))

(define-public (heap-peek heap)
  (if (vector-empty? heap) #f (vector-ref heap 0)))

(define-public (heap-pop heap)
  (define (fix-down! heap i)
    (let* ([len (vector-length heap)]
	   [left (1+ (* 2 i))]
	   [right (1+ left)])
      (if (< left len)
	  (let ([min (if (< right len)
			 (argmin (compose car (partial vector-ref heap)) < i left right)
			 (argmin (compose car (partial vector-ref heap)) < i left))])
	    (if (= min i)
		heap
		(begin
		  (vector-swap! heap i min)
		  (fix-down! heap min))))
	  heap)))
  (if (or (vector-empty? heap) (= (vector-length heap) 1))
      #()
      (let ([new-heap (vector-copy heap 0 (1- (vector-length heap)))])
	(vector-set! new-heap 0 (vector-ref heap (1- (vector-length heap))))
	(fix-down! new-heap 0))))

(define (make-inf-queue graph)
  "Creates a heap of the graph where all priorities are +inf"
  (list->vector
   (map (lambda (entry)
	  (cons (inf) (car entry)))
	graph)))

(define* (djikstra graph source sink #:optional update)
  (define unvisited (make-inf-queue graph))
  (let iter ([heap (heap-insert 0 source unvisited)]
	     [visited '()])
    (define-values (dist idx) (car+cdr (heap-peek heap)))
    (cond
     [(eqv? sink idx) dist]
     [(memv idx visited) (iter (heap-pop heap) visited)]
     [(inf? dist) #f]
     [else
      (when update (update visited heap))
      (iter (fold
	     (lambda (edge heap)
	       (heap-insert (+ dist 1) edge heap))
	     (heap-pop heap)
	     (node-edges (assv-ref graph idx)))
	    (cons idx visited))])))
(export djikstra)

(define* (a* graph source sink #:optional update)
  (define (chebychev idx) ;; Uses the Chebychev distance as a heuristic
    (let-values ([(x1 y1) (car+cdr (node-coords (assv-ref graph idx)))]
		 [(x2 y2) (car+cdr (node-coords (assv-ref graph sink)))])
      (max (abs (- x1 x2)) (abs (- y1 y2)))))
  (define unvisited (make-inf-queue graph))
  (let iter ([heap (heap-insert (chebychev source) source unvisited)]
	     [visited '()])
    (define-values (dist idx) (car+cdr (heap-peek heap)))
    (cond
     [(eqv? sink idx) dist]
     [(memv idx visited) (iter (heap-pop heap) visited)]
     [(inf? dist) #f]
     [else
      (when update (update visited heap))
      (iter (fold
	     (lambda (edge heap)
p	       (heap-insert (+ dist (chebychev edge) 1) edge heap))
	     (heap-pop heap)
	     (node-edges (assv-ref graph idx)))
	    (cons idx visited))])))
(export a*)

;;;;;;;;;;;;
;; Output ;;
;;;;;;;;;;;;

(define (output-to-file filename surface-gen)
  "Takes a filename and a coroutine to generate surfaces, and creates an animation"
  (define pngdir (mkdtemp "/tmp/graphgif_XXXXXX"))
  (let loop ([surface (surface-gen)]
	     [i 1])
    (when surface
      (cairo-surface-write-to-png
       surface
       (string-append pngdir "/img-" (number->string i) ".png"))
      (loop (surface-gen) (1+ i))))
  (system* "ffmpeg" "-y"
	   "-i" (string-append pngdir "/img-%d.png")
	   "-loop" "0" filename))

(define-public (write-graphs-to-file filename graph-gen)
  "Takes a filename and a couroutine to generate graphs, and creates an animation"
  (define (surface-gen)
    (let ([graph (graph-gen)])
      (and graph (draw-abstract-graph graph))))
  (output-to-file filename surface-gen))