(require :asdf)

(pushnew #P"/usr/local/asdf-install/site-systems/" asdf:*central-registry*)

(handler-bind ((style-warning #'muffle-warning))
  (asdf:operate 'asdf:load-op 'cffi-grovel :verbose nil)
  (asdf:operate 'asdf:load-op 'alexandria :verbose nil)
  (asdf:operate 'asdf:load-op 'cffi :verbose nil)
  (asdf:operate 'asdf:load-op 'babel :verbose nil)
  (asdf:operate 'asdf:load-op 'png :verbose nil))

(defpackage #:approx-ft
  (:use "COMMON-LISP"))
(in-package #:approx-ft)

(defun read-png ( input-pathname )
  (with-open-file (input input-pathname
			 :direction :input
			 :element-type '(unsigned-byte 8))
    (png:decode input)))

(defun write-png ( png output-pathname )
  (with-open-file (output output-pathname
			  :direction :output
			  :if-exists :supersede
			  :if-does-not-exist :create
			  :element-type '(unsigned-byte 8))
    (png:encode png output)))

(defun calculate-entropy ( png )
  (let ((counters (make-array '(256)
			      :initial-element 0
			      :element-type 'fixnum))
	(total (* (png:image-height png) (png:image-width png))))
    (dotimes (jj (png:image-height png))
      (dotimes (ii (png:image-width png))
	(incf (aref counters (aref png jj ii 0)))))
    (flet ((entropy ( value )
	     (if (zerop value)
		 0.0d0
		 (let ((pp (/ value total 1.0)))
		   (- (* pp (log pp)))))))
      (loop :for ii :from 0 :below 256
	 :summing (entropy (aref counters ii))))))

(defun calculate-rms ( png )
  (let ((sum 0))
    (dotimes (jj (png:image-height png))
      (dotimes (ii (png:image-width png))
	(let ((vv (- (mod (+ (aref png jj ii 0) 128) 256) 128)))
	  (incf sum (* vv vv)))))
    (sqrt (/ sum (* (png:image-height png) (png:image-width png))))))

(defun box-muller (&optional (random-state *random-state*))
  (let ((u1 (random 1.0 random-state))
        (u2 (random 1.0 random-state)))
    (let ((rr (sqrt (* -2.0 (log u1))))
          (aa (* 2.0 pi u2)))
      (values (coerce (* rr (cos aa)) 'double-float)
              (coerce (* rr (sin aa)) 'double-float)))))

(defun float-to-fixed ( ff )
  (max (min (round (+ (* ff 1024.0) 32768)) 65535) 0))

(defun fixed-to-float ( ff )
  (/ (- ff 32768) 1024.0))

(defparameter +param-count+ 4)
(defparameter +scale+ 0.02d0)
(defparameter +wave-scale+ 1.5d0)

(defun create-random-gene ( count &optional (random-state *random-state*) )
  (let ((nn (* count +param-count+)))
    (let ((ret (make-array (list nn) :element-type '(unsigned-byte 16))))
      (loop :for ii :from 0 :below nn :by 2
	 :do (multiple-value-bind (aa bb)
		 (box-muller random-state)
	       (setf (aref ret (+ ii 0)) (float-to-fixed aa))
	       (when (< (+ ii 1) nn)
		 (setf (aref ret (+ ii 1)) (float-to-fixed bb)))))
      ret)))

(defun gene-to-params ( gg )
  (loop :for ii :from 0 :below (length gg) :by +param-count+
     :collecting (loop :for jj :from 0 :below +param-count+
		    :collecting (fixed-to-float (aref gg (+ ii jj))))))

(defun param-amplitude ( pp )
  (* (first pp) +scale+))
(defun param-frequency ( pp )
  (* (second pp) +wave-scale+))
(defun param-angle ( pp )
  (* (third pp) pi))
(defun param-offset ( pp )
  (* (fourth pp) +wave-scale+))

(defun subtract-image ( orig buf res )
  (let ((ih (png:image-height orig))
	(iw (png:image-width orig)))
    (assert (= (png:image-height orig) (png:image-height buf)))
    (assert (= (png:image-width orig) (png:image-width buf)))
    (assert (= (png:image-channels orig) (png:image-channels buf)))
    (assert (= (png:image-height orig) (png:image-height res)))
    (assert (= (png:image-width orig) (png:image-width res)))
    (assert (= (png:image-channels orig) (png:image-channels res)))
    (loop :for jj :from 0 :below ih
       :do (loop :for ii :from 0 :below iw
	      :do (loop :for kk :from 0 :below (png:image-channels orig)
		     :do (setf (aref res jj ii kk)
			       (mod (- (aref orig jj ii kk)
				       (aref res jj ii kk)) 256)))))
    res))

(defun make-gene-image ( gg buf )
  (let ((params (gene-to-params gg))
	(ih (png:image-height buf))
	(iw (png:image-width buf)))
    (labels ((clip (vv)
	       (declare (type fixnum vv))
	       (cond
		 ((minusp vv) 0)
		 ((< 255 vv) 255)
		 (t vv)))
	     (calc (ii jj params &optional (sum 0.5d0))
	       (declare (type double-float ii jj sum)
			(optimize (speed 3)))
	       (cond
		 ((null params) (clip (coerce (round (* sum 255.0d0))
					      'fixnum)))
		 (t (let ((pp (first params)))
		      (declare (type list pp))
		      (let ((amp (param-amplitude pp))
			    (freq (param-frequency pp))
			    (ang (param-angle pp))
			    (off (param-offset pp)))
			(declare (type double-float amp freq ang off))
			(let ((cos (coerce (cos ang) 'double-float))
			      (sin (coerce (sin ang) 'double-float)))
			  (declare (type double-float cos sin))
			  (let ((xx (+ (* cos ii) (* sin jj))))
			    (declare (type double-float xx))
			    (let ((raw (coerce (cos (* 2.0d0 pi
						       (+ (* xx freq) off)))
					       'double-float)))
			      (declare (type double-float raw))
			      (calc ii jj
				    (rest params)
				    (+ sum (* amp raw))))))))))))
      (loop :for jj :from 0 :below ih
	 :do (loop :with fj = (- (/ jj ih 0.5d0) 1.0d0)
		:for ii :from 0 :below iw
		:do (loop :with fi = (- (/ ii iw 0.5d0) 1.0d0)
		       :for kk :from 0 :below (png:image-channels buf)
		       :do (setf (aref buf jj ii kk)
				 (calc fi fj params))))))
    buf))

(defun subtract-gene-image (gg orig buffer)
  (make-gene-image gg buffer)
  (subtract-image orig buffer buffer))

(defun evaluate-gene ( gg orig &optional buffer )
  (let ((buf (subtract-gene-image gg orig (or buffer (png:copy-image orig)))))
    (* (calculate-rms buf)
       #+nope (calculate-entropy buf))))

(defun read-random-state ( pathname )
  (or (ignore-errors
	(with-open-file (in pathname
			    :direction :input)
	  (read in nil nil)))
      *random-state*))

(defun write-random-state ( pathname rs )
  (with-open-file (out pathname
		       :direction :output
		       :if-exists :supersede
		       :if-does-not-exist :create)
    (write rs :stream out)))

(defun make-population ( genes count evaluator &optional (rs *random-state*) )
    (let ((scored (mapcar #'(lambda (gg)
			      (cons (funcall evaluator gg) gg))
			  (loop :for ii :from 1 :to genes
			     :collecting (create-random-gene count rs)))))
      (make-array genes :initial-contents scored)))

(defun breed (population aa bb dst evaluator &optional (rs *random-state*))
  (let ((bits (* (length (cdr (aref population aa))) 16)))
    (let ((start (random bits rs))
	  (end (random (1- bits) rs))
	  (primary (cdr (aref population aa)))
	  (secondary (cdr (aref population bb)))
	  (gene (cdr (aref population dst))))
      (if (<= start end)
	  (incf end)
	  (rotatef start end))
      (labels ((make-mask (bb &optional (oo 0))
		 (let* ((bb (min (max bb 0) 16))
		       (oo (min (max oo 0) (- 16 bb))))
		   (ash (1- (ash 1 bb)) oo)))
	       (high-bits (vv bb)
		 (logand vv (make-mask bb (- 16 bb))))
	       (mid-bits (vv ss len)
		 (logand vv (make-mask ss (- len ss))))
	       (low-bits (vv bb)
		 (logand vv (make-mask bb)))
	       (combine-word (aa bb ii start end)
		 (let ((start-bit (* ii 16)))
		   (cond
		     ((or (<= (+ start-bit 16) start)
			  (<= end start-bit)) aa)
		     ((and (<= start start-bit)
			   (<=  (+ start-bit 16) end) bb))
		     (t (let ((high (high-bits aa (- start-bit start)))
			      (mid  (mid-bits  bb (- start-bit start)
					          (- end start)))
			      (low  (low-bits  aa (- end (+ start-bit 16)))))
			  (logior high mid low)))))))
	(loop :for word :from 0 :below (/ bits 16)
	   :do (setf (elt gene word)
		     (combine-word (elt primary word)
				   (elt secondary word)
				   word start end))))
      (setf (car (aref population dst))
	    (funcall evaluator gene)))))

(defun mutate (population src dst evaluator &optional (rs *random-state*))
  (let ((in (cdr (aref population src)))
	(out (cdr (aref population dst))))
    (loop :for word :from 0 :below (length in)
       :do (setf (aref out word) (aref in word)))
    (loop :for mutation :from 0 :below (1+ (round (* 3 (abs (box-muller)))))
       :do (let ((bit (random (* (length in) 16) rs)))
	     (let* ((word (floor bit 16))
		    (digit (- bit (* word 16))))
	       (setf (aref out word)
		     (logxor (aref out word) (ash 1 digit))))))
    (setf (car (aref population dst))
	  (funcall evaluator out))))

(defun replace-duplicates (population count evaluator
			   &optional (rs *random-state*))
  (labels ((genes-equal (aa bb &optional (ii -1))
	     (cond
	       ((minusp ii) (= (length aa) (length bb)))
	       ((= ii (length aa)) t)
	       ((= (aref aa ii) (aref bb ii)) (genes-equal aa bb (1+ ii)))
	       (t nil)))
	   (has-duplicate ( population gg ii )
	     (cond
	       ((minusp ii) t)
	       ((genes-equal (cdr (aref population ii)) gg) nil)
	       (t (has-duplicate population gg (1- ii))))))
    (loop :for ii :from 1 :below (length population)
       :do (let ((cons (aref population ii)))
	     (when (has-duplicate population (cdr cons) (1- ii))
	       (let ((gene (create-random-gene count rs)))
		 (setf (car cons) (funcall evaluator gene)
		       (cdr cons) gene)))))))
	
(defun do-generation (population count evaluator &optional (rs *random-state*))
  (let ((nn (floor (/ (length population) 4))))
    (assert (< 1 nn))
    (loop :for ii :from 0 :below nn
       :do (let ((aa (random (* 2 nn)))
		 (bb (random (1- (* 2 nn))))
		 (cc (random nn)))
	     (when (<= aa bb)
		 (incf bb))
	     (breed population aa bb (+ ii nn nn) evaluator rs)
	     (mutate population cc (+ ii nn nn nn) evaluator rs)))
    (replace-duplicates population count evaluator rs)))

(defun make-evaluator ( orig buf )
  (let ((buf (or buf (png:copy-image orig))))
    (lambda (gene)
      (evaluate-gene gene orig buf))))

(defun evolve-image (input-pathname output-pathname
		     &key (population 30)
		     (count 10)
		     (generations 100)
		     (gene-pathname nil)
		     (gene-image-pathname nil))
  (let ((rs (read-random-state #P"random-seed.lisp")))
    (unwind-protect
	 (let* ((png (read-png input-pathname))
		(buf (png:copy-image png))
		(evaluator (make-evaluator png buf))
		(population (make-population population count evaluator rs)))
	   (stable-sort population #'< :key #'car)
	   (dotimes (ii generations)
	     (do-generation population count evaluator rs)
	     (stable-sort population #'< :key #'car)
	     (when gene-pathname
	       (with-open-file (out gene-pathname
				    :direction :output
				    :if-exists :supersede
				    :if-does-not-exist :create)
		 (write (cdr (aref population 0)) :stream out)))
	     (subtract-gene-image (cdr (aref population 0)) png buf)
	     (write-png buf output-pathname)
	     (when gene-image-pathname
	       (make-gene-image (cdr (aref population 0)) buf)
	       (write-png buf (format nil "out-~4,'0D.png" ii))
	       (write-png buf gene-image-pathname))
	     (format t "~S: ~S~%" ii (car (aref population 0)))
	     (force-output))
	   (car (aref population 0)))
      (write-random-state #P"random-seed.lisp" rs))))

(defun do-it ()
  (evolve-image "in.png" "out.png"
		:population 100 :count 50
		:generations (ash 1 12)
		:gene-pathname "gene.lisp"
		:gene-image-pathname "gene.png"))

(do-it)