#lang racket

(require plot)

(define function-probability 0.8)
(define arg-probability 0.7)
(define min-val -2)
(define max-val  2)
(define mutate-arg 0.1)
(define mutate-num 0.6)
(define mutate-fun 0.1)
(define (mutate-min x) -1)
(define (mutate-max x)  1)

(define fitness-min -7)
(define fitness-max  7)
(define fitness-points 100)

(define population-size 50)
(define individuals-survive 20)
(define exception-weight 1000)

(define gens
  #((+ expr expr)
    (- expr expr)
    (* expr expr)
    (/ expr expr)
    ;(expt expr expr)
    (sin expr)
    (cos expr)
    (tan expr)))

(define args
  #(x))

(define args-list
  (vector->list args))

(define gens-number (vector-length gens))
(define args-number (vector-length args))

(define (target-fun x)
  (+ (* x 1.46327) (/ 0.3423 x)))

(define (fitness-compare f1 f2)
  (for/sum ([x (in-range fitness-min
                         fitness-max
                         (exact->inexact (/ (- fitness-max fitness-min) fitness-points)))])
    (abs (- (f1 x) (f2 x)))))

(define (generate-values fun)
  (for/list ([x (in-range fitness-min
                          fitness-max
                          (exact->inexact (/ (- fitness-max fitness-min) fitness-points)))])
    (fun x)))

(define (random-real-interval from to)
  (+ from (* (random) (- to from))))

(define (gen-fun depth)
  (define prob (random))
  {if [or (< depth 1) (> prob function-probability)]
      [if (< (random) arg-probability)
          (vector-ref args (random args-number))
          (random-real-interval min-val max-val)]
      [let ([expr (vector-ref gens (random gens-number))])
        (if (list? expr)
            (map (lambda (x)
                   (if (eq? x 'expr) (gen-fun (sub1 depth))
                       x))
                 expr)
            expr)]})

(define (mutate fun depth)
  (define mut (random))
  {cond
    [(and (symbol? fun) (< mut mutate-arg))
     (gen-fun depth)]
    [(and (number? fun) (< mut mutate-num))
     (+ fun (random-real-interval (mutate-min fun) (mutate-max fun)))]
    [(and (list? fun) (< mut mutate-fun))
     (gen-fun (sub1 depth))]
    [(list? fun)
     (cons (car fun) (map (λ (x) (mutate x (sub1 depth))) (cdr fun)))]
    [else
     fun]})

(define (test1 d x)
  (define var (gen-fun d))
  (printf "(~A)(~A) = ~A~%" var x (eval `((λ (x) ,var) ,x))))

(define (compare-data data1 data2)
  (for/sum [(d1 data1) (d2 data2)]
    (abs (- (real-part d1) (real-part d2)))))

(define (test2 d iterations fun)
  (define population (for/vector ((x (in-range population-size))) (gen-fun d)))
  (define values (generate-values fun))
  {for ((n (in-range iterations)))
    [vector-sort! population
                  {λ (f1 f2)
                    (with-handlers [(exn:fail:contract:divide-by-zero? (λ (e) exception-weight))]
                      (define d1 (generate-values (eval `(λ (x) ,f1))))
                      (define d2 (generate-values (eval `(λ (x) ,f2))))
                      (< (compare-data d1 values) (compare-data d2 values)))}
                  #:cache-keys? #t]
    [for ((i (in-range individuals-survive)))
      (vector-set! population i (mutate (vector-ref population i) d))]
    [for ((i (in-range individuals-survive population-size)))
      (vector-set! population i (gen-fun d))]
    (when (= 0 (remainder n 100)) (printf "Iteration ~A...~%" n))}
  population)

(define (enforce-compile arg)
  (if (procedure? arg)
      arg
      (eval `(λ ,args-list ,arg))))

(define (plot2 f1 f2)
  (plot (list (function (enforce-compile f1) min-val max-val)
              (function (enforce-compile f2) min-val max-val))))