(module benchmarks mzscheme
  
  (require
   (lib "comprehensions.ss" "srfi" "42")
   (lib "vector-lib.ss" "srfi" "43")
   (lib "etc.ss")
   (lib "4.ss" "srfi")
   (lib "foreign.ss")
   (planet "test.ss" ("schematics" "schemeunit.plt" 2))
   (planet "benchmark.ss" ("schematics" "benchmark.plt" 1))
   (prefix table: (planet "table.ss" ("soegaard" "galore.plt" 3)))
   (file "pomdp.ss")
   (file "viterbi.ss")
   (file "array.ss"))
  
  (provide benchmarks)

  (define table-matrix
    (table:alist->equal
     (list-ec
      (:range x 0 20)
      (:range y 0 20)
      (cons (vector x y) (* x y)))))
  
  (define vector-matrix
    (vector-ec
     (:range x 0 20)
     (:range y 0 20)
     (* x y)))

  (define array-matrix
    (array-ec
     #(20 20)
     0.0
     (:range x 0 20)
     (:range y 0 20)
     [(x y) <- (exact->inexact (* x y))]))

  (define f32-matrix
    (let ([f32 (make-f32vector 400 0.0)])
      (do-ec (:range x 0 20) (:range y 0 20)
             (f32vector-set! f32 (+ (* x 20) y) (exact->inexact (* x y))))
      f32))
  
  (define (table-multiply)
    (sum-ec
     (:range row 0 20)
     (:range col 0 20)
     (:range elt 0 20)
     (* (table:lookup (vector row elt) table-matrix)
        (table:lookup (vector col elt) table-matrix))))
  
  (define (vector-multiply)
    (sum-ec
     (:range row 0 20)
     (:range col 0 20)
     (:range elt 0 20)
     (* (vector-ref vector-matrix (+ (* row 20) elt))
        (vector-ref vector-matrix (+ (* elt 20) col)))))

  (define (array-multiply)
    (sum-ec
     (:range row 0 20)
     (:range col 0 20)
     (:range elt 0 20)
     (* (array-ref array-matrix (vector row elt))
        (array-ref array-matrix (vector col elt)))))

  (define (f32-multiply)
    (sum-ec
     (:range row 0 20)
     (:range col 0 20)
     (:range elt 0 20)
     (* (f32vector-ref f32-matrix (+ (* row 20) elt))
        (f32vector-ref f32-matrix (+ (* elt 20) col)))))

  (unsafe!)
  (define (c-multiply)
    ((get-ffi-obj "mmult" (ffi-lib "mmult.so") (_fun _f32vector -> _float))
     f32-matrix))

  (define (c-block-multiply)
    ((get-ffi-obj "block_mmult" (ffi-lib "mmult.so") (_fun _f32vector -> _float))
     f32-matrix))
  
  (define data
    '((0 . 0) (0 . 0) (0 . 0) (0 . 1)
      (1 . 1) (1 . 1) (1 . 0) (1 . 0)
      (1 . 1) (1 . 1) (1 . 0) (1 . 0)
      (1 . 0) (0 . 1) (1 . 1) (1 . 0)
      (1 . 1) (1 . 0) (1 . 1) (1 . 1)
      (1 . 1) (2 . 0) (1 . 1) (2 . 0)
      (1 . 1) (2 . 0) (1 . 1) (2 . 1)
      (2 . 0) (1 . 1) (2 . 0) (1 . 0)
      (1 . 0) (1 . 0)))
  
  (define load-unload
    (make-pomdp
     '(0 1 2 3 4 9 8 7 6 5)
     '(0 1 2)
     '(0 1)
     (alist->array
      #(10)
      '((#(0) . 1.0)))
     (alist->array
      #(10 2 10)
      '((#(0 0 0) . 1.0)
        (#(0 1 1) . 1.0)
        (#(1 0 0) . 1.0)
        (#(1 1 2) . 1.0)
        (#(2 0 1) . 1.0)
        (#(2 1 3) . 1.0)
        (#(3 0 2) . 1.0)
        (#(3 1 4) . 1.0)
        (#(4 0 8) . 1.0)
        (#(4 1 9) . 1.0)
        (#(5 0 0) . 1.0)
        (#(5 1 1) . 1.0)
        (#(6 0 5) . 1.0)
        (#(6 1 7) . 1.0)
        (#(7 0 6) . 1.0)
        (#(7 1 8) . 1.0)
        (#(8 0 7) . 1.0)
        (#(8 1 9) . 1.0)
        (#(9 0 8) . 1.0)
        (#(9 1 9) . 1.0)))
     (alist->array
      #(10 3)
      '((#(0 0) . 1.0)
        (#(5 0) . 1.0)
        (#(1 1) . 1.0)
        (#(6 1) . 1.0)
        (#(2 1) . 1.0)
        (#(7 1) . 1.0)
        (#(3 1) . 1.0)
        (#(8 1) . 1.0)
        (#(4 2) . 1.0)
        (#(9 2) . 1.0)))
     (alist->array
      #(10 2)
      '((#(3 1) . 1.0)
        (#(6 0) . 1.0)))))

  (define benchmarks
    (test-suite
     "benchmarks.ss"
     
     (test-case
      "Ensure vectors are faster than tables"
      ;; We perform a 20x20 matrix multiplication and summation
      (check-faster table-multiply vector-multiply))

     (test-case
      "Ensure arrays faster than tables"
      (check-faster table-multiply array-multiply))

     (test-case
      "Ensure f32vector faster than vector"
      (check-faster f32-multiply vector-multiply))

     (test-case
      "Ensure C multiply faster than vector"
      (check-faster vector-multiply c-multiply))

     (test-case
      "Ensure C block multiply faster than C multiply"
      (check-faster c-multiply c-block-multiply))
     
     (benchmark-case
      "Viterbi algorithm"
      ;; This is to ensure that our implementation of the
      ;; Viterbi algorithm, the main loop of model merging,
      ;; is getting faster
      (viterbi-estimate load-unload data))
     ))
  )