(require 'alexandria)

(defpackage #:ant-colony-optimisation
  (:nicknames #:aco)
  (:shadow cl:search)
  (:use #:cl #:alexandria))

(in-package #:aco)

(declaim (optimize (debug 3) (speed 3) (space 3)))

(load "graph.cl")


;;; ACO

(declaim (type (function (t t t) single-float)
	       *default-pheromone*
	       *default-heuristic*)
	 (type single-float
	       *pheromone-amplifier*
	       *heuristic-amplifier*
	       *pheromone-decay*
	       *q3*))

(defparameter *default-pheromone* (constantly 0.0))
(defparameter *default-heuristic* (constantly 1.0))
(defparameter *pheromone-amplifier* 1.0)
(defparameter *heuristic-amplifier* 2.0)
(defparameter *pheromone-decay* 0.0)
(defparameter *q3* 100.0)

(defstruct ant
  (path nil :type list)
  started-at)

(defun search (count graph legal-p eval end-p)
  (declare (type fixnum count)
	   (type (function (list) boolean) legal-p end-p)
	   (type (function (list) single-float) eval))
  (let* ((n (length (the list (nodes graph))))
	 (pheromone (make-array (list n n) :initial-element 0.0
					   :element-type 'single-float))
	 (heuristic (make-array (list n n) :initial-element 0.0
					   :element-type 'single-float))
	 best-solution)
    (dolist (edge (edges graph))
      (setf (aref pheromone (car edge) (cadr edge))
	    (apply *default-pheromone* edge))
      (setf (aref heuristic (car edge) (cadr edge))
	    (apply *default-heuristic* edge)))
    (labels
	((currently-at (ant)
	   (first (ant-path ant)))
	 (neighbour-states (ant)
	   (outgoing graph (currently-at ant)))
	 (normalize (alist)
	   (declare (type list alist))
	   (let ((alpha (coerce (loop for x in alist sum (cdr x))
				'single-float)))
	     (mapcar (if (= alpha 0.0)
			 (lambda (x)
			   (declare (type (cons t single-float) x))
			   (cons (car x) (/ (length alist))))
			 (lambda (x)
			   (declare (type (cons t single-float) x))
			   (cons (car x) (/ (cdr x) alpha))))
		     alist)))
	 (choose (alist)
	   (declare (type list alist))
	   (if (= (length alist) 1)
	       (caar alist)
	       (let ((x (the single-float (random 1.0))) (sum 0.0))
		 (declare (type single-float sum))
		 (dolist (ent (normalize alist))
		   (declare (type (cons t real) ent))
		   (when (> (incf sum (cdr ent)) x)
		     (return-from choose (car ent)))))))
	 (move-chance (ant i j)
	   ;; (assert (> (expt (aref pheromone i j) *pheromone-amplifier*) 0)
	   ;; 	   (i j (aref pheromone i j)))
	   ;; (assert (> (expt (aref heuristic i j) *heuristic-amplifier*) 0)
	   ;; 	   (i j (aref heuristic i j)))
	   (/ (* (expt (aref pheromone i j) *pheromone-amplifier*)
		 (expt (aref heuristic i j) *heuristic-amplifier*))
	      (loop for s in (neighbour-states ant)
		    sum (* (expt (aref pheromone i s) *pheromone-amplifier*)
			   (expt (aref heuristic i s) *heuristic-amplifier*)))))
	 (choose-path (ant)
	   (choose (loop with i = (currently-at ant)
			 for j in (neighbour-states ant)
			 unless (member j (ant-path ant))
			   collect (cons j (coerce (move-chance ant i j)
						   'single-float)))))
	 (new-ant (&aux (start (random-elt (nodes graph))))
	   (make-ant :path (list start) :started-at start)))
      (loop for ants = (loop repeat count collect (new-ant)) do
	(loop repeat n do
	  (dolist (ant ants)
	    (let ((next (choose-path ant)))
	      (when next (push next (ant-path ant)))))
	  (dotimes (i n)
	    (dotimes (j n)
	      (setf (aref pheromone i j)
		    (+ (loop for ant in ants
			     when (cl:search (list i j) (ant-path ant))
			       sum (/ *q3* (length (ant-path ant))))
		       (* *pheromone-decay* (aref pheromone i j))))
	      (when (or (= i j) (> (aref pheromone i j) 0))
		(incf (aref pheromone i j) 0.0001))
	      ;; (assert (or (= i j) (> (aref pheromone i j) 0)) (i j))
	      )))
	(when (and (null best-solution) ants)
	  (setf best-solution (first (remove-if-not legal-p ants :key #'ant-path))))
	(dolist (ant ants)
	  (when (and (funcall legal-p (ant-path ant))
		     (> (funcall eval (ant-path best-solution))
			(funcall eval (ant-path ant))))
	    (setf best-solution ant)))
	(when (funcall end-p (mapcar #'ant-path ants))
	  (return (and best-solution
		       (funcall legal-p (ant-path best-solution))
		       (ant-path best-solution))))))))


;;; TSP

(defun tsp (count graph rounds)
  (declare (type fixnum count rounds))
  (let ((*default-pheromone* (constantly (coerce (/ count 200) 'single-float)))
	(*default-heuristic* (lambda (from to dist)
			       (declare (ignore from to))
			       (coerce (/ dist) 'single-float))))
    (flet ((legal-p (path)
	     (set-equal nodes path))
	   (score (assignment)
	     (declare (type list path))
	     (loop ))
	   (end-p (state)
	     (declare (ignore state))
	     (>= 0 (decf rounds))))
      (let ((res (search count graph #'legal-p #'score #'end-p)))
	(values res (score res))))))


;;; testing



(eval-when (:execute)

  (time (tsp 100 *ulysseus-16* 1000))

  )