;;
;; 
;; A collection of utility functions for manipulating SRFI-4 vectors.
;;
;;
;; Copyright 2007-2012 Ivan Raikov and the Okinawa Institute of Science and Technology.
;;
;;
;; This program is free software: you can redistribute it and/or
;; modify it under the terms of the GNU General Public License as
;; published by the Free Software Foundation, either version 3 of the
;; License, or (at your option) any later version.
;;
;; 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.
;;
;; A full copy of the GPL license can be found at
;; <http://www.gnu.org/licenses/>.
;;
;;

(module srfi-4-utils

 (f64vector-fold f32vector-fold s32vector-fold u32vector-fold
		 s16vector-fold u16vector-fold s8vector-fold u8vector-fold
		 f64vector-map f32vector-map s32vector-map u32vector-map
		 s16vector-map u16vector-map s8vector-map u8vector-map
		 f64vector-foldi f32vector-foldi s32vector-foldi u32vector-foldi
		 s16vector-foldi u16vector-foldi s8vector-foldi u8vector-foldi
		 f64vector-mapi f32vector-mapi s32vector-mapi u32vector-mapi
		 s16vector-mapi u16vector-mapi s8vector-mapi u8vector-mapi
		 f64vector-quick-sort! f32vector-quick-sort! s32vector-quick-sort! u32vector-quick-sort!
		 s16vector-quick-sort! u16vector-quick-sort! s8vector-quick-sort! u8vector-quick-sort!
		 f64vector-repmat f64vector-meshgrid)

 (import scheme chicken foreign)
 (require-library srfi-1)
 (import (only srfi-1 fold every))
 (require-extension srfi-4 srfi-42 srfi-4-comprehensions)

(define (make-vector-fold vector-length vector-ref)
  (lambda (f x0 v . rest)
    (let ((n   (vector-length v))
	  (vs  (cons v rest)))
      (fold-ec x0 (:range i 0 n) 
	       (map (lambda (v) (vector-ref v i)) vs)
	       (lambda (x ax) (apply f (append x (list ax))))))))

(define f64vector-fold (make-vector-fold f64vector-length f64vector-ref))
(define f32vector-fold (make-vector-fold f32vector-length f32vector-ref))
(define s32vector-fold (make-vector-fold s32vector-length s32vector-ref))
(define u32vector-fold (make-vector-fold u32vector-length u32vector-ref))
(define s16vector-fold (make-vector-fold s16vector-length s16vector-ref))
(define u16vector-fold (make-vector-fold u16vector-length u16vector-ref))
(define s8vector-fold  (make-vector-fold s8vector-length s8vector-ref))
(define u8vector-fold  (make-vector-fold u8vector-length u8vector-ref))


(define (make-vector-foldi vector-length vector-ref)
  (lambda (f x0 v . rest)
    (let ((n   (vector-length v))
	  (vs  (cons v rest)))
      (fold-ec x0 (:range i 0 n) 
	       (cons i (map (lambda (v) (vector-ref v i)) vs))
	       (lambda (x ax) (apply f (append x (list ax))))))))


(define f64vector-foldi (make-vector-foldi f64vector-length f64vector-ref))
(define f32vector-foldi (make-vector-foldi f32vector-length f32vector-ref))
(define s32vector-foldi (make-vector-foldi s32vector-length s32vector-ref))
(define u32vector-foldi (make-vector-foldi u32vector-length u32vector-ref))
(define s16vector-foldi (make-vector-foldi s16vector-length s16vector-ref))
(define u16vector-foldi (make-vector-foldi u16vector-length u16vector-ref))
(define s8vector-foldi  (make-vector-foldi s8vector-length s8vector-ref))
(define u8vector-foldi  (make-vector-foldi u8vector-length u8vector-ref))


(define (f64vector-map f v . rest)  
  (let ((n (f64vector-length v)))
    (f64vector-of-length-ec n (:range i 0 n) (apply f (map (lambda (v) (f64vector-ref v i)) (cons v rest))))))
(define (f32vector-map f v . rest)  
  (let ((n (f32vector-length v)))
    (f32vector-of-length-ec n (:range i 0 n) (apply f (map (lambda (v) (f32vector-ref v i)) (cons v rest))))))
(define (s32vector-map f v . rest)  
  (let ((n (s32vector-length v)))
    (s32vector-of-length-ec n (:range i 0 n) (apply f (map (lambda (v) (s32vector-ref v i)) (cons v rest))))))
(define (u32vector-map f v . rest)  
  (let ((n (u32vector-length v)))
    (u32vector-of-length-ec n (:range i 0 n) (apply f (map (lambda (v) (u32vector-ref v i)) (cons v rest))))))
(define (s16vector-map f v . rest)  
  (let ((n (s16vector-length v)))
    (s16vector-of-length-ec n (:range i 0 n) (apply f (map (lambda (v) (s16vector-ref v i)) (cons v rest))))))
(define (u16vector-map f v . rest)  
  (let ((n (u16vector-length v)))
    (u16vector-of-length-ec n (:range i 0 n) (apply f (map (lambda (v) (u16vector-ref v i)) (cons v rest))))))
(define (s8vector-map f v . rest)  
  (let ((n (s8vector-length v)))
    (s8vector-of-length-ec n (:range i 0 n) (apply f (map (lambda (v) (s8vector-ref v i)) (cons v rest))))))
(define (u8vector-map f v . rest)  
  (let ((n (u8vector-length v)))
    (u8vector-of-length-ec n (:range i 0 n) (apply f (map (lambda (v) (u8vector-ref v i)) (cons v rest))))))


(define (f64vector-mapi f v . rest)  
  (let ((n (f64vector-length v)))
    (f64vector-of-length-ec n (:range i 0 n) (apply f (cons i  (map (lambda (v) (f64vector-ref v i)) (cons v rest)))))))
(define (f32vector-mapi f v . rest)  
  (let ((n (f32vector-length v)))
    (f32vector-of-length-ec n (:range i 0 n) (apply f (cons i  (map (lambda (v) (f32vector-ref v i)) (cons v rest)))))))
(define (s32vector-mapi f v . rest)  
  (let ((n (s32vector-length v)))
    (s32vector-of-length-ec n (:range i 0 n) (apply f (cons i  (map (lambda (v) (s32vector-ref v i)) (cons v rest)))))))
(define (u32vector-mapi f v . rest)  
  (let ((n (u32vector-length v)))
    (u32vector-of-length-ec n (:range i 0 n) (apply f (cons i  (map (lambda (v) (u32vector-ref v i)) (cons v rest)))))))
(define (s16vector-mapi f v . rest)  
  (let ((n (s16vector-length v)))
    (s16vector-of-length-ec n (:range i 0 n) (apply f (cons i  (map (lambda (v) (s16vector-ref v i)) (cons v rest)))))))
(define (u16vector-mapi f v . rest)  
  (let ((n (u16vector-length v)))
    (u16vector-of-length-ec n (:range i 0 n) (apply f (cons i  (map (lambda (v) (u16vector-ref v i)) (cons v rest)))))))
(define (s8vector-mapi f v . rest)  
  (let ((n (s8vector-length v)))
    (s8vector-of-length-ec n (:range i 0 n) (apply f (cons i  (map (lambda (v) (s8vector-ref v i)) (cons v rest)))))))
(define (u8vector-mapi f v . rest)  
  (let ((n (u8vector-length v)))
    (u8vector-of-length-ec n (:range i 0 n) (apply f (cons i (map (lambda (v) (u8vector-ref v i)) (cons v rest)))))))


;;
;; In-place quick sort from SRFI-32 reference implementation.
;; Modified so that the comparison function uses element indices as
;; well as element values:
;;
;; elt< :: i1 v1 i2 v2 -> boolean
;;
;; Copyright (c) 1998 by Olin Shivers. You may do as you please with
;; this code, as long as you do not delete this notice or hold me
;; responsible for any outcome related to its use.
;;

(define (srfi-4-vector-quick-sort! vector-ref vector-set! vector-length)
  (lambda (v elt< . rest)
   (let-optionals rest ((start 0) (end (vector-length v)))
    (let recur ((l start) (r end))	; Sort the range [l,r).
      (if (fx< 1 (fx- r l))
	  
	  ;; Choose the median of V[l], V[r], and V[middle] for the pivot.
	  (let ((median
		 (lambda (i1 i2 i3)
		   (let ((v1 (vector-ref v i1))
			 (v2 (vector-ref v i2))
			 (v3 (vector-ref v i3)))
		     (receive (ilittle little ibig big)
			      (if (elt< i1 v1 i2 v2) (values i1 v1 i2 v2) (values i2 v2 i1 v1))
			      (if (elt< ibig big i3 v3) 
				  (values ibig big)
				  (if (elt< ilittle little i3 v3) 
				      (values i3 v3) 
				      (values ilittle little))))))))
	    
	    (let-values (((ipivot pivot) (median l (quotient (fx+ l r) 2) (fx- r 1))))
   	      (let loop ((i l) (j (fx- r 1)))
		(let ((i (let scan ((i i)) (if (elt< i (vector-ref v i) ipivot pivot)
					       (scan (fx+ i 1))
					       i)))
		      (j (let scan ((j j)) (if (elt< ipivot pivot j (vector-ref v j))
					       (scan (fx- j 1))
					       j))))
		  (if (fx< i j)
		      (let ((tmp (vector-ref v j)))		
			(vector-set! v j (vector-ref v i))	; Swap V[I]
			(vector-set! v i tmp)		;  and V[J].
			(loop (fx+ i 1) (fx- j 1)))
		      
		      (begin (recur l i) (recur (fx+ j 1) r)))))))
	  v)))))


(define f64vector-quick-sort!
  (srfi-4-vector-quick-sort! f64vector-ref f64vector-set! f64vector-length))
(define f32vector-quick-sort!
  (srfi-4-vector-quick-sort! f32vector-ref f32vector-set! f32vector-length))
(define s32vector-quick-sort!
  (srfi-4-vector-quick-sort! s32vector-ref s32vector-set! s32vector-length))
(define u32vector-quick-sort!
  (srfi-4-vector-quick-sort! u32vector-ref u32vector-set! u32vector-length))
(define s16vector-quick-sort!
  (srfi-4-vector-quick-sort! s16vector-ref s16vector-set! s16vector-length))
(define u16vector-quick-sort!
  (srfi-4-vector-quick-sort! u16vector-ref u16vector-set! u16vector-length))
(define s8vector-quick-sort!
  (srfi-4-vector-quick-sort! s8vector-ref s8vector-set! s8vector-length))
(define u8vector-quick-sort!
  (srfi-4-vector-quick-sort! u8vector-ref u8vector-set! u8vector-length))

#>

/* repeat a block of memory rep times */
void cmemrep(double *dest, size_t chunk, int rep)
{
  if(rep == 1) return;
  memcpy(dest + chunk, dest, sizeof(double)*chunk); 
  if(rep & 1) {
    dest += chunk;
    memcpy(dest + chunk, dest, sizeof(double)*chunk);
  }
  /* now repeat using a block twice as big */
  cmemrep(dest, chunk<<1, rep>>1);
}

void crepmat(double *dest, const double *src, int ndim, 
	     int *destdimsize, int *dimsize, const int *dims, int *rep) 
{
  int i, chunk;
  int d = ndim-1;
  /* copy the first repetition into dest */
  if(d == 0) 
  {
    chunk = dimsize[0];
    memcpy(dest,src,sizeof(double)*chunk);
  }
  else 
  {
    /* recursively repeat each slice of src */
    for(i=0; i<dims[d]; i++) {
      crepmat(dest + i*destdimsize[d-1], src + i*dimsize[d-1], 
 	      ndim-1, destdimsize, dimsize, dims, rep);
    }
    chunk = destdimsize[d-1]*dims[d];
  }
  /* copy the result rep-1 times */
  cmemrep(dest,chunk,rep[d]);
}

<#

(define crepmat (foreign-safe-lambda void "crepmat"  f64vector f64vector int s32vector s32vector s32vector s32vector))

(define cmemrep (foreign-safe-lambda void "cmemrep"  f64vector int int))

(define (f64vector-repmat src dims reps)
  (let* ((ndim (length dims))
	 (dimsize (make-s32vector ndim))
	 (vdims (list->s32vector dims)))
    
    (assert (every positive? reps))
    
    ;; compute dimension sizes
    (s32vector-set! dimsize 0 (car dims))
    (let recur ((i 1) (dims (cdr dims)))
      (if (< i ndim)
	  (begin 
	    (s32vector-set! dimsize i (* (car dims) (s32vector-ref dimsize (- i 1))))
	    (recur (+ 1 i) (cdr dims)))
	  ))

    ;; determine repetition vector 
    (let* ((nrep (length reps))
	   (ndimdest (if (> nrep ndim) nrep ndim))
	   (rep (make-s32vector ndimdest)))
      (let recur ((i 0) (reps reps))
	(if (< i nrep)
	    (let ((repv (car reps)))
	      (s32vector-set! rep i repv)
	      (recur (+ 1 i) (cdr reps)))))

      ;; compute output size
      (let ((destdims (make-s32vector ndimdest 0)))
	(let recur ((i 0))
	  (if (< i ndim)
	      (begin (s32vector-set! destdims i 
				    (* (s32vector-ref vdims i)
				       (s32vector-ref rep i )))
		     (recur (+ 1 i)))))

	(let ((extrarep
	       (let recur ((i ndim) (extrarep 1))
		 (if (< i ndimdest)
		     (let ((ri (s32vector-ref rep i)))
		       (s32vector-set! destdims i ri)
		       (recur (+ 1 i) (* extrarep ri)))
		     extrarep
		     ))))

	  (let ((destdimsize (make-s32vector ndimdest)))
	    (s32vector-set! destdimsize 0 (s32vector-ref destdims 0))
	    (let recur ((i 1))
	      (if (< i ndimdest)
		  (begin
		    (s32vector-set! destdimsize i 
				    (* (s32vector-ref destdimsize (- i 1))
				       (s32vector-ref destdims i)))
		    (recur (+ 1 i)))))

	    ;; return array
	    (let* ((destlen (s32vector-ref destdimsize (- ndimdest 1)))
		   (dest (make-f64vector destlen)))

#|
	      (print "destlen = " destlen)
	      (print "dimsize = " dimsize)
	      (print "destdimsize = " destdimsize)
	      (print "rep = " rep)
|#	      
	      (crepmat dest src ndim destdimsize dimsize vdims rep)
	      (if (> ndimdest ndim)
		  (let ((n (s32vector-ref destdimsize (- ndim 1))))
;		    (print "extrarep = " extrarep)
;		    (print "n = " n)
		    (cmemrep dest n extrarep)))
	      dest
	      ))
	  ))
      )))


	  

(define (f64vector-meshgrid x y z)
  (let ((lenx (f64vector-length x))
	(leny (f64vector-length y))
	(lenz (f64vector-length z)))
    (let ((dimx (list 1 lenx))
	  (dimy (list 1 leny))
	  (dimz (list 1 lenz)))
    (let ((xx (f64vector-repmat
	       (f64vector-repmat x dimx (list leny 1)) 
	       (list leny lenx)
	       (list 1 1 lenz)))
	  (yy (f64vector-repmat
	       (f64vector-repmat y dimy (list 1 lenx))
	       (list lenx leny)
	       (list 1 1 lenz)))
	  (zz (f64vector-repmat z dimz (list (* lenx leny) 1))))
      (list xx yy zz))
    )))
      


)