(ql:quickload 'lparallel)
(setf lparallel:*kernel* (lparallel:make-kernel 4))

(defpackage particle-swarm-optimisation
  (:nicknames :pso)
  (:use #:cl #:lparallel))

(in-package :pso)

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

(defparameter *fiction-coefficient* 1)		; a
(defparameter *global-stiffness-coefficient* 1)	; b_glob
(defparameter *local-stiffness-coefficient* 1)	; b_loc

(defstruct search-space
  (left nil :type (simple-array real *))
  (right nil :type (simple-array real *)))

(defstruct group
  (attractor #() :type (simple-array real *))
  (particles nil :type list))

(defstruct particle
  (local-attractor nil :type (simple-array real *))
  (position nil :type (simple-array real *))
  (velocity nil :type (simple-array real *))
  (group nil :type (or (eql nil) group)))

(defun get-random (from to)
  (+ (random (float (- to from)))
     from))


;;; search-space functions
(defun in-space-p (search-space position)
  (declare (type (simple-array real *) position)
	   (type search-space search-space))
  (with-slots (left right search-space) search-space
    (declare (type (simple-array real *) left right))
    (every #'<= left position right)))

(defun space-dimensions (search-space)
  (declare (type search-space search-space))
  (with-slots (left) search-space
    (declare (type (simple-array real *) left))
    (length left)))

(defun random-position (search-space)
  "Return random position in space"
  (declare (type search-space search-space))
  (with-slots (left right) search-space
    (declare (type (simple-array real *) left right))
    (assert (= (length left) (length right)))
    (assert (every #'<= left right))
    (map 'vector #'get-random left right)))

(defun null-vector (search-space)
  (declare (type search-space search-space))
  (coerce (loop repeat (space-dimensions search-space)
		collect 0)
	  'vector))


;;; vector operations
(defun scale (vector)
  (map 'vector (lambda (v) (* (get-random 0 1) v)) vector))

(defun v+ (&rest vectors)
  (apply #'map 'vector #'+ vectors))

(defun v- (v1 v2)
  (declare (type (simple-array real *) v1 v2))
  (map 'vector #'- v1 v2))

(defun v* (vector scalar)
  (declare (type (simple-array real *) vector)
	   (type real scalar))
  (if (= scalar 1)
      vector
      (map 'vector (lambda (c) (* c scalar)) vector)))


;;; grouping functionsn
(defun ring-grouper (particles)
  (loop for a in particles
	for b in (append (last particles) particles)
	for g = (make-group :particles (list a b))
	do (setf (particle-group a) g
		 (particle-group b) g)
	collect g))

(defun complete-grouper (particles)
  (loop for particle in particles
	collect (setf (particle-group particle)
		      (make-group :particles particles))))


;;; out-of-bounds handlers
(defun infinity (particle next space)
  (declare (ignore particle))
  (assert (not (in-space-p space next)))
  next)

(defun nearest (particle next space)
  (declare (ignore particle))
  (assert (not (in-space-p space next)))
  (map 'vector
       (lambda (coord left right)
	 (cond ((< coord left) left)
	       ((< right coord) right)
	       (coord)))
       next
       (search-space-left space)
       (search-space-right space)))

(defun absorb (particle next space)
  (declare (ignore particle))
  (assert (not (in-space-p space next)))
  (map 'vector
       (lambda (coord)
	 0)				;TODO
       next))

(defun periodic (particle next space)
  (declare (ignore particle))
  (assert (not (in-space-p space next)))
  (map 'vector
       (lambda (coord left right)
	 (+ (mod (+ left coord) (- right left))
	    left))
       next
       (search-space-left space)
       (search-space-right space)))

(defun mirror (particle next space)
  (declare (ignore particle))
 (assert (not (in-space-p space next)))
  (map 'vector
       (lambda (coord)
	 0)				;TODO
       next))

(defun random-restart (particle next space)
  (declare (ignore particle))
  (assert (not (in-space-p space next)))
  (random-position space))


;;; PSO
(defun pso (f n m space &key (grouper #'complete-grouper) (oob-handler #'infinity))
  (declare (type (function ((simple-array real *)) real) f)
	   (type (integer 1 *) n)
	   (type search-space space))
  (labels ((evaluate (position)
	     (if (in-space-p space position)
		 (funcall f position)
		 most-positive-double-float))
	   (better (p1 p2)
	     (< (evaluate p1) (evaluate p2)))
	   (best (particles)
	     (reduce (lambda (p1 p2) (if (better p1 p2) p1 p2))
		     (mapcar #'particle-local-attractor particles)))
	   (update-group (group)
	     (with-slots (particles attractor) group
	       (setf attractor (best particles)))))
    (let* ((particles (loop repeat n
			    for pos = (random-position space)
			    collect (make-particle :velocity (null-vector space)
						   :local-attractor pos
						   :position pos)))
	   (groups (delete-duplicates (funcall grouper particles))))
      (loop for i from 1 to m do
	(pmapc #'update-group groups)
	(pmapc
	 (lambda (particle)
	   (with-slots (position velocity local-attractor group) particle
	     (setf velocity (v+ (v* velocity *fiction-coefficient*)
				(v* (scale (v- local-attractor position))
				    *local-stiffness-coefficient*)
				(v* (scale (v- (group-attractor group)
					       position))
				    *global-stiffness-coefficient*)))
	     (let ((next-position (v+ position velocity)))
	       (setf position
		     (if (in-space-p space next-position)
			 position
			 (funcall oob-handler particle next-position space))))
	     (when (better position local-attractor)
	       (setf local-attractor position))))
	 particles))
      (best particles))))


;;; testing code
(defun test (fn
	     &key
	       (fiction-coefficient 0.72984)
	       (global-stiffness-coefficient 1.496172)
	       (local-stiffness-coefficient 1.496172)
	       (iterations 100)
	       (particle-count 10)
	       dimensions
	       (size 100)
	       (optimum 0)
	       (group #'complete-grouper)
	       (oob #'infinity)
	     &aux
	       (*fiction-coefficient* fiction-coefficient)
	       (*global-stiffness-coefficient* global-stiffness-coefficient)
	       (*local-stiffness-coefficient* local-stiffness-coefficient)
	       (*print-pretty* nil))
  (assert (numberp dimensions))
  (eval-when (:execute)
    (format t "~2&START (iter = ~D, n = ~D, dim = ~D)"
	    iterations particle-count dimensions)
    (let ((result (time (pso  fn particle-count iterations
			      (make-search-space
			       :left (make-array dimensions :initial-element (- size))
			       :right (make-array dimensions :initial-element size))
			      :grouper group
			      :oob-handler oob)))
	  (goal (if (numberp optimum)
		    (make-array dimensions :initial-element optimum)
		    optimum)))
      (format t "~%Final Position: ~{~F ~}" (coerce result 'list))
      (format t "~%Evaluation: ~F" (funcall fn result))
      (format t "~%Error: ~F" (- (funcall fn result)
				 (funcall fn goal)))
      (format t "~%Distance: ~F" (loop for a across result
				       for b across goal
				       sum (* (- a b) (- a b)))))))

(test (lambda (x)			;sphere
	(loop for a across x sum (* a a)))
      :dimensions 2
      :particle-count 100
      :iterations 1000
      :size 10
      :optimum 0)

(test (lambda (x)			;rosenbrock
	    (+ (loop for a across x
		     sum (- (* a a) (* (cos (* 2 pi a)) -10)))
	       (* 10 (length x))))
      :dimensions 2
      :particle-count 100
      :iterations 10
      :size 30
      :optimum 1)

(test (lambda (x)			;rastrigin
	(+ (loop for a across x
		 sum (- (* a a)
			(* (cos (* 2 pi a))
			   -10)))
	   (* 10 (length x))))
      :dimensions 2
      :particle-count 100
      :iterations 10
      :size 5.12
      :optimum 0)

(test (lambda (x)			;schwefel
	(loop for a across x
	      sum (* (sin (sqrt (abs a)))
		     (- a))))
      :dimensions 2
      :particle-count 100
      :iterations 10
      :size 500
      :optimum 420.6987)