;;; "Revolver"
;;;
;;; A minimalist genetic programming engine written in Scheme, based on the
;;; book "Genetic Programming" by John Koza.  Includes a couple of
;;; improvements suggested in the appendix on speeding things up: building a
;;; fitness index for fast fitness-weighted selection, and not reevaluating
;;; the fitness of individuals that are copied into the next generation
;;; without change.
;;;
;;; Copyright (c) 2007 Thomas Munro who asserts his right to be identified as
;;; the author and sole copyright holder of this work.  This software is
;;; released under the terms of the GPL licence (version 2).

(library (revolver)
  (export evolve make-individual make-population show-best)
  (import (rnrs)
          (rnrs eval)
          (prefix (random) random:))

  ; Retrieves a node from a tree by index number, with a numbering scheme
  ; as shown:
  ;
  ; (+ x (* y z))   ->    +     ->     0
  ;                      / \          / \
  ;                     x   *        1   2
  ;                        / \          / \
  ;                       y   z        3   4
  (define (tree-ref tree index)
    ; TODO find out how to do this without using a mutable counter to number
    ; the nodes (I have found only a horrible inefficient way)
    (define count 0)
    (call/cc
     (lambda (return)
       (define (parameters node)
         (cond [(null? node) #f]
               [else (test (car node))
                     (parameters (cdr node))]))
       (define (test node)
         (if (= count index) (return node))
         (set! count (+ 1 count))
         (if (pair? node) (parameters (cdr node))))
       (test tree)
       (error "TREE-REF -- index out of range"))))
  
  ; Counts the atoms in a tree of lists and atoms.
  (define (tree-atoms tree)
    (cond [(null? tree) 0]
          [(pair? tree) (+ (tree-atoms (car tree)) (tree-atoms (cdr tree)))]
          [else 1]))
  
  ; Copies a tree, inserting 'guest' into the place identified by 'index'.
  ;
  ;     +                     +
  ;    / \          %        / \
  ;   x  <*>   +   / \   =  x   %
  ;      / \      x  42        / \
  ;     y   x                 x   42
  ;
  ; TODO as above, I need to find out how to do this without the mutable count
  (define (tree-combine host index guest)
    (define count 0)
    (define (parameters node)
      (if (null? node) '()
          (cons (copy (car node))
                (parameters (cdr node)))))
    (define (copy node)
      (cond [(= count index)
             (set! count (+ 1 count))
             guest]
            [else 
             (set! count (+ 1 count))
             (if (pair? node)
                 (cons (car node) (parameters (cdr node)))
                 node)]))
    (copy host))
  
  ; Creates random offspring by combining two program trees.
  (define (tree-random-recombine male female)
    (let* ([male-index    (random:random-integer (tree-atoms male))]
           [female-index  (random:random-integer (tree-atoms female))]
           [male-fragment (tree-ref male male-index)])
      (tree-combine female female-index male-fragment)))
  
  ; Builds a random program from the provided vectors of functions and 
  ; terminals.
  (define (make-random-program functions terminals min-depth max-depth)
    (define (iter depth)
      (if (and (< depth max-depth)
               (or (< depth min-depth) (zero? (random:random-integer 2))))
          ; generate a function application
          (let* ([pair (vector-ref functions
                                   (random:random-integer (vector-length
                                                           functions)))]
                 [function  (car pair)]
                 [num-args  (cdr pair)]
                 [arguments (let loop ([n num-args])
                              (if (zero? n)
                                  '()
                                  (cons (iter (+ 1 depth)) (loop (- n 1)))))])
            (cons function arguments))
          ; generate a terminal node
          (let ([terminal (vector-ref 
                           terminals
                           (random:random-integer (vector-length terminals)))])
            (cond [(eq? terminal 'random-real) 
                   (- (* 10 (random:random-real)) 5.0)]
                  [(eq? terminal 'random-int) 
                   (random:random-integer 10)]
                  [else terminal]))))
    (iter 0))
  
  ; Creates a list of random (program . fitness) pairs from the list of
  ; functions and terminals provided.
  (define (make-population size functions terminals min-depth max-depth)
    (define (iter n output)
      (if (zero? n)
          output
          (iter (- n 1)
                (cons (make-individual (make-random-program
                                        functions terminals
                                        min-depth max-depth)
                                       #f)
                              output))))
    (iter size '()))
  
  ; Computes and sets the fitness of each individual, and sorts into fitness
  ; order (fittest at the beginning).
  (define (evaluate-population population fitness-function)
    (define (iter input output)
      (if (null? input)
          output
          (let ([program (caar input)]
                [fitness (cdar input)])
            ; we only compute the fitness if it was not already known (#f)
            (if (eq? fitness #f)
                (iter (cdr input)
                      (cons (make-individual program (fitness-function program))
                            output))
                (iter (cdr input)
                      (cons (make-individual program fitness)
                            output))))))
    (list-sort (lambda (left right) (> (cdr left) (cdr right)))
               (iter population '())))
  
  ; Builds a vector of cumulative sums.  This is useful for throwing a dart
  ; to pick an individual with fitness probability.
  (define (make-fitness-index population)
    (define (iter input output sum)
      (cond [(null? input) output]
            [else (iter (cdr input)
                        (cons sum output)
                        (+ sum (individual-fitness (car input))))]))
    (list->vector (reverse (iter population '() 0.0))))
  
  ; Finds a record which has a value equal to or immediately
  ; preceding a given value.
  ; get-ref - a function that can get a record by index
  ; size - the size of the set of records
  ; value - the search value
  ; precedes? - a function which can compare keys
  (define (search get-ref size value precedes?)
    (let loop ([start 0]
               [stop (- size 1)])
      (if (< stop start)
          (if (< stop 0) #f stop)
          (let* ([mid-point (quotient (+ start stop) 2)]
                 [key (get-ref mid-point)])
            (cond [(precedes? value key) (loop start (- mid-point 1))]
                  [(precedes? key value) (loop (+ mid-point 1) stop)]
                  [else mid-point])))))
  
  ;(define vec #(0 3 5))
  ;(find (lambda (n) (vector-ref vec n)) 3 6 <)
  
  ; Creates an individual from a program and a fitnes.
  (define (make-individual program fitness)
    (cons program fitness))
  
  ; Returns the program tree of an individual.
  (define (individual-program individual)
    (car individual))
  
  ; Returns the fitness of an individual.
  (define (individual-fitness individual)
    (cdr individual))
  
  ; Breeds a new population.  The new population is created by cloning and
  ; crossing invididuals in the source population.  The population must be
  ; evaluated and sorted (fittest first).
  (define (breed-population population crossover-probability)
    (let* ([fitness-index     (make-fitness-index population)]
           [individuals       (vector-length fitness-index)]
           [population-vector (list->vector population)]
           [crossovers        (floor (* individuals crossover-probability))]
           [copies            (- individuals crossovers 1)]
           [get-sum           (lambda (n) (vector-ref fitness-index n))]
           [total-fitness     (vector-ref fitness-index (- individuals 1))])
      ; Selects an individual with fitness probability.
      (define (select-individual)
        (let* ([dart  (* (random:random-real) total-fitness)]
               [index (search get-sum individuals dart <)])
          (vector-ref population-vector index))) 
      ; Builds a list of 'count' combined individuals.  Their fitness is set
      ; to #f because it will need to be recomputed.
      (define (combine output count)
        (if (zero? count) output
            (let* ([male      (individual-program (select-individual))]
                   [female    (individual-program (select-individual))]
                   [offspring (tree-random-recombine male female)])
              (combine (cons (make-individual offspring #f) output)
                       (- count 1)))))
      ; Builds a list of 'count' copies individuals (preserving their fitness
      ; so it will not be recomputed in the next generation).
      (define (copy output count)
        (if (zero? count) output
            (copy (cons (select-individual) output)
                  (- count 1))))
      ; We want a list of combined and copied individuals, with the fittest
      ; individual thrown in at the front for good measure.
      ;
      (cons (car population)
            (append (copy '() copies)
                    (combine '() crossovers)))))
  
  ; Evolves a population for 'generations' generations, or until the target
  ; fitness is reached, and returns the result.
  (define (evolve population fitness-function generations crossover-probability
           target-fitness)
    (let loop ([generations generations]
               [population population])
        (define evaluated (evaluate-population population fitness-function))
        (show-best evaluated)
        (cond [(zero? generations) evaluated]
              [(<= target-fitness (individual-fitness (car evaluated))) 
               evaluated]
              [else 
               (loop (- generations 1)
                     (breed-population evaluated crossover-probability))])))
  
  (define (show-best population)
    (define best (car population))
    (display "Best individual: ")
    (display (individual-program best))
    (newline)
    (display "Fitness: ")
    (display (individual-fitness best))
    (newline))
)