;;; -*- mode: Lisp; Syntax: Common-Lisp; -*-
;;;
;;; Copyright (c) 2009 by the authors.
;;;
;;; See LICENCE for details.

(in-package :hu.dwim.genetic-programming)

;;;;;;
;;; Individual

(defstruct individual
  (fitness nil :type (or null number))
  (age 0 :type integer)
  (erroneous nil :type boolean)
  (offspring-count 0 :type integer)
  (lisp-form nil :type list)
  (walked-form nil :type walked-form)
  (compiled-form nil :type (or null function)))

(def print-object (individual :identity nil)
    (format t ":FITNESS ~A :AGE ~A :LISP-FORM ~A" (individual-fitness -self-) (individual-age -self-) (individual-lisp-form -self-)))

(def (function e) make-individual-with-form (&key lisp-form walked-form)
  (bind ((walked-form (or walked-form (walk-form lisp-form)))
         (lisp-form (or lisp-form (second (unwalk-form walked-form)))))
    (make-individual :lisp-form lisp-form :walked-form walked-form)))

(def (function e) ensure-individual-compiled-form (individual)
  (or (individual-compiled-form individual)
      (setf (individual-compiled-form individual)
            (hu.dwim.util:with-muffled-boring-compiler-warnings
              (compile nil (individual-lisp-form individual))))))

;;;;;;
;;; Population

(defstruct population
  (generation 0 :type integer)
  (individuals nil :type sequence))

(def (function e) make-random-population (&key (individual-count 1) lambda-form-generator)
  (make-population :individuals (iter (repeat individual-count)
                                      (for individual = (make-individual-with-form :walked-form (funcall lambda-form-generator)))
                                      (collect individual :result-type 'vector))))

(def (function e) population-best-individual (population)
  (first-elt (population-individuals population)))


(def (function e) select-random-individual-weighted-by-fitness (population individual-weight)
  (bind ((total-fitness (iter (for individual :in-sequence (population-individuals population))
                              (awhen (funcall individual-weight individual)
                                (summing it))))
         (selected-fitness (random total-fitness)))
    (iter (for individual :in-sequence (population-individuals population))
          (awhen (funcall individual-weight individual)
            (summing it :into sum))
          (when (> sum selected-fitness)
            (return individual)))))

(def function population-aging (population)
  (incf (population-generation population))
  (iter (for individual :in-sequence (population-individuals population))
        (incf (individual-age individual))))

;;;;;;
;;; Crossover

(def (function e) make-crossovered-individuals (individual-1 individual-2)
  (bind ((walked-form-1 (individual-walked-form individual-1))
         (walked-form-2 (individual-walked-form individual-2))
         (count-1 (count-forms walked-form-1 +rewritable-form-types+))
         (count-2 (count-forms walked-form-2 +rewritable-form-types+))
         (walked-subform-1 (find-nth-form walked-form-1 +rewritable-form-types+ (random count-1)))
         (walked-subform-2 (find-nth-form walked-form-2 +rewritable-form-types+ (random count-2))))
    (incf (individual-offspring-count individual-1))
    (incf (individual-offspring-count individual-2))
    (when (or (and (typep walked-subform-1 'constant-form)
                   (typep walked-subform-2 'constant-form))
              (and (not (typep walked-subform-1 'constant-form))
                   (not (typep walked-subform-2 'constant-form))))
      (list (make-individual-with-form
             :walked-form (rewrite-nth-form (deep-copy-ast walked-form-1) +rewritable-form-types+ (random count-1)
                                            (lambda (form)
                                              ;; TODO: what about return-type, argument-types?
                                              (declare (ignore form))
                                              (deep-copy-ast walked-subform-2))))
            (make-individual-with-form
             :walked-form (rewrite-nth-form (deep-copy-ast walked-form-2) +rewritable-form-types+ (random count-2)
                                            (lambda (form)
                                              ;; TODO: what about return-type, argument-types?
                                              (declare (ignore form))
                                              (deep-copy-ast walked-subform-1))))))))

;;;;;;
;;; Mutation

#|
wrap form, move form up, flat form -> constant, constant -> flat form
|#

(def (function e) make-mutated-individuals (individual)
  (bind ((walked-form (individual-walked-form individual))
         (count (count-forms walked-form +rewritable-form-types+)))
    (incf (individual-offspring-count individual) 2)
    (list (make-individual-with-form
           :walked-form (rewrite-nth-form (deep-copy-ast walked-form) +rewritable-form-types+ (random count)
                                          (lambda (form)
                                            (etypecase form
                                              (constant-form
                                               (apply 'generate-constant-form (attributes-of form)))
                                              (variable-reference-form
                                               (apply 'generate-variable-reference-form (attributes-of form)))
                                              (free-application-form
                                               (apply' generate-free-application-form (attributes-of form)))))))
          (make-individual-with-form
           :walked-form (rewrite-nth-form (deep-copy-ast walked-form) +rewritable-form-types+ (random count)
                                          (lambda (form)
                                            (etypecase form
                                              (constant-form
                                               (apply 'generate-constant-form (attributes-of form)))
                                              (variable-reference-form
                                               (apply 'generate-variable-reference-form (attributes-of form)))
                                              (free-application-form
                                               ;; KLUDGE: argument types, etc.
                                               (bind ((original-function-descriptor (find-function-descriptor (operator-of form)))
                                                      (matching-function-descriptors (collect-matching-function-descriptors (function-descriptor-return-type original-function-descriptor)
                                                                                                                            (function-descriptor-argument-types original-function-descriptor)
                                                                                                                            (form-attribute form :function-descriptors))))
                                                 (if matching-function-descriptors
                                                     (make-instance 'free-application-form
                                                                    :attributes (attributes-of form)
                                                                    :operator (function-descriptor-name (random-elt matching-function-descriptors))
                                                                    :arguments (mapcar 'deep-copy-ast (arguments-of form)))
                                                     (apply' generate-free-application-form (attributes-of form))))))))))))

;;;;;;
;;; Fitness

(def function evaluate-fitness (individuals fitness-calculator)
  (iter (for individual :in-sequence individuals)
        (unless (or (individual-fitness individual)
                    (individual-erroneous individual))
          (setf (individual-fitness individual) (block nil
                                                  (handler-bind ((serious-condition (lambda (condition)
                                                                                      (declare (ignore condition))
                                                                                      (setf (individual-erroneous individual) t)
                                                                                      (return nil))))
                                                    (sb-ext:with-timeout 10
                                                      (funcall fitness-calculator individual))))))))

(def function sort-individuals (individuals fitness-sort-predicate)
  (sort individuals (lambda (a b)
                      (cond ((null a)
                             nil)
                            ((null b)
                             t)
                            (t
                             (funcall fitness-sort-predicate a b))))
        :key #'individual-fitness))

;;;;;;
;;; Evolution

(def (function e) evolve (&key initial-population-generator evolution-termination-predicate
                               fitness-calculator fitness-sort-predicate population-printer
                               survivor-selector reproduction-function)
  "Evolves a population according to the environment described by the following functions:
INITIAL-POPULATION-GENERATOR is a zero argument function the produces a population. The returned population will form the 0th generation.
FITNESS-CALCULATOR is a one argument function that takes an individual and returns a number representing the fitness value.
FITNESS-SORT-PREDICATE is a two argument function that is used as a predicate to sort the individuals within the population after the fitness function has been evaluated.
EVOLUTION-TERMINATION-PREDICATE is a one argument function that takes the population and returns a boolean value. It determines when evolving the population will be finished.
SURVIVOR-SELECTOR is a one argument function that takes the population and returns the individuals that survive to the next generation.
REPRODUCTION-FUNCTION is a one argument function that takes the population and returns new individuals reproduced from the current generation."
  (check-type initial-population-generator (or symbol function))
  (check-type evolution-termination-predicate (or symbol function))
  (check-type fitness-calculator (or symbol function))
  (check-type fitness-sort-predicate (or symbol function))
  (check-type population-printer (or symbol function))
  (check-type survivor-selector (or symbol function))
  (check-type reproduction-function (or symbol function))
  (iter (with population = (funcall initial-population-generator))
        (initially (setf (population-generation population) 0))
        (for parent-individuals = (population-individuals population))
        (evaluate-fitness parent-individuals fitness-calculator)
        (sort-individuals parent-individuals fitness-sort-predicate)
        (until (funcall evolution-termination-predicate population))
        (for offspring-individuals = (funcall reproduction-function population))
        (evaluate-fitness offspring-individuals fitness-calculator)
        (for new-individuals = (concatenate 'vector parent-individuals offspring-individuals))
        (sort-individuals new-individuals fitness-sort-predicate)
        (setf (population-individuals population) new-individuals)
        (setf (population-individuals population) (funcall survivor-selector population))
        (sort-individuals (population-individuals population) fitness-sort-predicate)
        (population-aging population)
        (when population-printer
          (funcall population-printer population))
        (finally (return population))))