;;
;;
;; Chicken Scheme bindings for the CMA-ES optimization library.
;;
;; CMA-ES is copyright 1996, 2003, 2007 Nikolaus Hansen.
;; Chicken Scheme code is copyright 2012 Ivan Raikov.
;;
;; This program is free software; you can redistribute it and/or
;; modify it under the terms of the GNU General Public License,
;; version 2, as published by the Free Software Foundation.
;;
;; This program is distributed in the hope that it will be useful, but
;; WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;; General Public License for more details.
;;
;; You should have received a copy of the GNU General Public License
;; along with this program.  If not, see  <http://www.gnu.org/licenses/>.
;;
;;


(module cmaes 

	(init run 
         init-from-file terminated? get-parameter get-value
	 sample-population update-distribution
	 read-signals-from-file read-signals
	 write-to-file
	 terminate free
	 )

	(import scheme chicken foreign)
	(require-extension srfi-4)
	(require-library lolevel posix files srfi-1 srfi-13)
	(import (only lolevel move-memory! make-pointer-vector pointer-vector-ref pointer-vector-set! free)
		(only srfi-1 fold every)
		(only srfi-13 string-concatenate)
		(only data-structures ->string)
                (only posix file-close file-write open-output-file* file-mkstemp)
		(only files make-pathname)
		)

(define tempdir (make-parameter "/tmp"))


(foreign-declare #<<EOF

#include <assert.h>

#include "cmaes.h"

extern double * cmaes_init(cmaes_t *, int dimension , double *xstart, 
		double *stddev, long seed, int lambda, 
		const char *input_parameter_filename);
void cmaes_resume_distribution(cmaes_t *evo_ptr, char *filename);
void cmaes_exit(cmaes_t *);

double * const * cmaes_SamplePopulation(cmaes_t *);
double *         cmaes_UpdateDistribution(cmaes_t *, 
					  const double *rgFitnessValues);
const char *     cmaes_TestForTermination(cmaes_t *);

double * const * cmaes_ReSampleSingle( cmaes_t *t, int index);
double const *   cmaes_ReSampleSingle_old(cmaes_t *, double *rgx); 
double *         cmaes_SampleSingleInto( cmaes_t *t, double *rgx);
void             cmaes_UpdateEigensystem(cmaes_t *, int flgforce);

double         cmaes_Get(cmaes_t *, char const *keyword);
const double * cmaes_GetPtr(cmaes_t *, char const *keyword); /* e.g. "xbestever" */
double *       cmaes_GetNew( cmaes_t *t, char const *keyword); 
double *       cmaes_GetInto( cmaes_t *t, char const *keyword, double *mem); 

void           cmaes_ReadSignals(cmaes_t *, char const *filename);
void           cmaes_WriteToFile(cmaes_t *, const char *szKeyWord,
                                 const char *output_filename); 
char *         cmaes_SayHello(cmaes_t *);

double *       cmaes_NewDouble(int n);  
void           cmaes_FATAL(char const *s1, char const *s2, char const *s3, 
			   char const *s4);


cmaes_t * cmaes_alloc(void)
{
    cmaes_t * h;
    assert ((h = malloc (sizeof(*h))) != NULL);

    return h;
}

double *cmaes_init_from_file (cmaes_t *h, char *filepath)
{
    double *result;

    result = cmaes_init (h, 0, NULL, NULL, 0, 0, filepath); 

    return result;
}

double cmaes_lambda (cmaes_t *h)
{
    return cmaes_Get (h, "lambda");
}

int cmaes_dimension (cmaes_t *h)
{
    return (int)cmaes_Get (h, "N");
}

double *double_array_pointer_deref (double **ptr, int offset)
{
    return *(ptr+offset);
}



EOF
)

(define-foreign-type cmaesobj "cmaes_t")
  
(define c_double_array_pointer_deref
  (foreign-safe-lambda (nonnull-c-pointer double)  "double_array_pointer_deref" (c-pointer (c-pointer double )) int))
  
(define c_cmaes_alloc
  (foreign-safe-lambda (nonnull-c-pointer cmaesobj)  "cmaes_alloc" ))
  
(define c_cmaes_lambda
  (foreign-safe-lambda int   "cmaes_lambda" (nonnull-c-pointer cmaesobj) ))

(define c_cmaes_dimension
  (foreign-safe-lambda int   "cmaes_dimension" (nonnull-c-pointer cmaesobj) ))

(define c_cmaes_init_from_file
  (foreign-safe-lambda (nonnull-c-pointer double)  "cmaes_init_from_file" (nonnull-c-pointer cmaesobj) c-string  ))

(define c_cmaes_read_signals_from_file
  (foreign-safe-lambda void  "cmaes_ReadSignals" (nonnull-c-pointer cmaesobj) c-string  ))

(define c_cmaes_terminationp
  (foreign-safe-lambda c-string "cmaes_TestForTermination" (nonnull-c-pointer cmaesobj) ))

(define c_cmaes_sample_population
  (foreign-safe-lambda (c-pointer (c-pointer double)) "cmaes_SamplePopulation" (nonnull-c-pointer cmaesobj) ))

(define c_cmaes_get_parameter
  (foreign-safe-lambda double "cmaes_Get" (nonnull-c-pointer cmaesobj) c-string ))

(define c_cmaes_update_distribution
  (foreign-safe-lambda (c-pointer double) "cmaes_UpdateDistribution" (nonnull-c-pointer cmaesobj) f64vector ))

(define c_cmaes_get_new 
  (foreign-safe-lambda (c-pointer double) "cmaes_GetNew" (nonnull-c-pointer cmaesobj) c-string ))

(define c_cmaes_write_to_file 
  (foreign-safe-lambda void "cmaes_WriteToFile" (nonnull-c-pointer cmaesobj) c-string c-string ))

(define c_cmaes_exit 
  (foreign-safe-lambda void "cmaes_exit" (nonnull-c-pointer cmaesobj)))


(define (make-init-parameter-string kv)

  (if (not (pair? kv)) 
      (error 'make-init-parameter-string "invalid parameter" kv))

  (let ((k (car kv)) (v (cdr kv)))
    (case k
      ((N)  ; problem dimension
       (if (and (integer? v) (positive? v))
           (sprintf "~A ~A # problem dimension~%" k v)
           (error 'make-init-parameter-string "parameter N (problem dimension) must be a positive integer")))

      ((initialX)  ; initial search point
       (if (every number? v)
           (let ((n (length v)))
             (string-concatenate
              (cons (sprintf "~A ~A: # initial search point~%" k n)
                    (map (lambda (n) (sprintf "  ~A~%" n)) v))
              ))
           (error 'make-init-parameter-string "parameter initialX (initial search point) must be a list of numbers")))

      ((typicalX)  ; typical search point
       (if (every number? v)
           (let ((n (length v)))
             (string-concatenate
              (cons (sprintf "~A ~A: # typical search point (overwritten by initialX)~%" k n)
                    (map (lambda (n) (sprintf "  ~A~%" n)) v))
              ))
           (error 'make-init-parameter-string "parameter typicalX (typical search point) must be a list of numbers")))

      ((initialStandardDeviations)
       (if (every number? v)
           (let ((n (length v)))
             (string-concatenate
              (cons (sprintf "~A ~A: # typical search point (overwritten by initialX)~%" k n)
                    (map (lambda (n) (sprintf "  ~A~%" n)) v))
              ))
           (error 'make-init-parameter-string "parameter initialStandardDeviations (typical search point) must be a list of numbers")))
      
      ((stopMaxFunEvals)
       (if (and (integer? v) (positive? v))
           (sprintf "~A ~A # max number of f-evaluations~%" k v)
           (error 'make-init-parameter-string "parameter stopMaxFunEvals (max number of f-evaluations) must be a positive integer")))


      ((stopMaxIter)
       (if (and (integer? v) (positive? v))
           (sprintf "~A ~A # max number of iterations (generations)~%" k v)
           (error 'make-init-parameter-string "parameter stopMaxIter (max number of iterations) must be a positive integer")))

      ((stopTolFun)
       (if (and (number? v) (positive? v))
           (sprintf "~A ~A # stop if function value differences are smaller than this value ~%" k v)
           (error 'make-init-parameter-string "parameter stopTolFun (stop if value differences are too small) must be a positive number")))

      ((stopTolFunHist)
       (if (and (number? v) (positive? v))
           (sprintf "~A ~A # stop if best function value differences are smaller than this value ~%" k v)
           (error 'make-init-parameter-string "parameter stopTolFun (stop if best value differences are too small) must be a positive number")))

      ((stopTolX)
       (if (and (number? v) (positive? v))
           (sprintf "~A ~A # stop if step sizes/steps in x-space are smaller than this value ~%" k v)
           (error 'make-init-parameter-string "parameter stopTolX (stop if step sizes are too small) must be a positive number")))

      ((stopTolUpXFactor)
       (if (and (number? v) (positive? v))
           (sprintf "~A ~A # stop if std dev increases more than this value ~%" k v)
           (error 'make-init-parameter-string "parameter stopTolUpXFactor (stop if std dev increases are too big) must be a positive number")))

      ((seed)
       (if (number? v)
           (sprintf "~A ~A ~%" k v)
           (error 'make-init-parameter-string "parameter seed must be a number")))

      (else (error 'make-init-parameter-string "unknown parameter" kv))
       
      )))
  

(define (init-from-file filepath)
  (let* ((h (c_cmaes_alloc))
         (fitvals (c_cmaes_init_from_file h filepath))
	 (lam (c_cmaes_lambda h)))
    (let ((v (make-f64vector lam)))
      (move-memory! fitvals v (* 8 lam)) 
      (values h v))
    ))
  
              

(define (read-signals-from-file h filepath)
  (c_cmaes_read_signals_from_file h filepath))


(define (make-signal-parameter-string kv)

  (if (not (pair? kv)) 
      (error 'make-signal-parameter-string "invalid parameter" kv))

  (let ((k (car kv)) (v (cdr kv)))

    (if (pair? v)
	(sprintf "~A ~A ~A~%" k (car v) (cadr v))
	(sprintf "~A ~A~%" k v))
    
    ))
  

(define (read-signals h parameters)

  (let-values (((fd temp-path) (file-mkstemp (make-pathname (tempdir) "cmaes.sig.XXXXXX"))))

    (let ((strs (map make-signal-parameter-string parameters)))
      (for-each (lambda (s) (file-write fd s)) strs))

    (read-signals-from-file h temp-path)
    
    (file-close fd)

    (delete-file temp-path)

    ))

              
(define (terminated? h)
  (c_cmaes_terminationp h))


(define (sample-population h)
  (let ((ptr (c_cmaes_sample_population h ))
	(lam (c_cmaes_lambda h))
	(n (c_cmaes_dimension h)))
    (let ((vptr (make-pointer-vector lam))
	  (v (make-vector lam)))

      (let recur ((i 0))
	(if (< i lam)
	    (begin
	      (pointer-vector-set! vptr i (c_double_array_pointer_deref ptr i))
	      (recur (+ 1 i)))))
      
      (let recur ((i 0))
	(if (< i lam)
	    (let ((vv (make-f64vector n 0.)))
	      (move-memory! (pointer-vector-ref vptr i) vv (* 8 n))
	      (vector-set! v i vv)
	      (recur (+ 1 i)))
	    v)
	))
    ))


(define (get-parameter h s)
  (c_cmaes_get_parameter h (->string s)))


(define (update-distribution h lam)
  (let ((ptr (c_cmaes_update_distribution h lam))
	(n (c_cmaes_dimension h)))
    (let ((v (make-f64vector n)))
      (move-memory! ptr v (* 8 n))
      v)))


(define (get-value h s)
  (let ((ptr (c_cmaes_get_new h (->string s)))
	(n (c_cmaes_dimension h)))
    (let ((v (make-f64vector n)))
      (move-memory! ptr v (* 8 n))
      v)))
    

(define (write-to-file h s fn)
  (c_cmaes_write_to_file h (->string s) fn))


(define (terminate h)
  (c_cmaes_exit h))



(define (init parameters)

  (let-values (((fd temp-path) (file-mkstemp (make-pathname (tempdir) "cmaes.par.XXXXXX"))))
    (let ((strs (map make-init-parameter-string parameters)))
      (for-each (lambda (s) (file-write fd s)) strs))
    
    (let-values (((h funvals) (init-from-file temp-path)))

      (file-close fd)
      
      (delete-file temp-path)
    
      (values h funvals)
      ))
  )

 

(define (run f h funvals signals #!key 
	     (output-file "all.dat") 
	     (result-values '(xbest xmean)))
  
  (read-signals h signals)

  (let-values (((fd temp-path) (file-mkstemp (make-pathname (tempdir) "cmaes.sig.XXXXXX"))))

    (let ((strs (map make-signal-parameter-string signals)))
      (for-each (lambda (s) (file-write fd s)) strs))

    (file-close fd)

    (let recur ((h h))
      
      (let ((stop (terminated? h)))
	
	(if (not stop)
	    
	    (let ((pop (sample-population h))
		  (lam (get-parameter h 'lambda))
		  (n   (get-parameter h 'dim)))
	      
	      (let inner-recur ((i 0))
		(if (< i lam)
		    (begin
		      (f64vector-set! funvals i (f (vector-ref pop i)))
		      (inner-recur (+ 1 i)))))
	      
	      (update-distribution h funvals)
	    
	      (read-signals-from-file h temp-path)
	    
	      (recur h)
	      
	      )
	    
	    (print stop)
	    )))

    (delete-file temp-path)

    )

  (if output-file (write-to-file h 'all output-file))

  (let ((result (map (lambda (x) (cons x (get-value h (->string x)))) result-values)))

    (terminate h)

    (free h)

    result
    )
)


 
)