guile-ploy/test/test-ploy.scm

; (c) Daniel Llorens - 2012-2013
; Tests for (ploy ploy).

; This library 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.

(import (srfi srfi-1) (srfi srfi-9) (srfi srfi-26) (ploy basic) (ploy ploy)
        (ploy test))

(assert (= 99 (array-cell-ref #0(99))))
(assert (= 99 (array-cell-ref #(99) 0)))
(assert (= 99 (array-cell-ref (array-cell-ref #2((55 77) (77 99)) 1) 1)))

(define a (array-copy #t #3(((0 1) (2 3)) ((4 5) (6 7)))))
(array-cell-set! a #(12 13) 0 1)
(T a #3(((0 1) (12 13)) ((4 5) (6 7))))
(array-cell-set! a #2((10 11) (12 13)) 0)
(T a #3(((10 11) (12 13)) ((4 5) (6 7))))

(define (test-A A)
  (for-each
   (lambda (d)
     (let ((data (if (zero? d) '(()) (apply list-product (map iota (take ($ A) d))))))
       (for-each
        (lambda (i)
          (T (apply array-cell-ref A i) (apply array-cell-ref A i)))
        data)))
   (iota (+ 1 (rank A)))))

(test-A (i. 2 3 2))
(test-A (i. 2 3))
(test-A (i. 2))
(test-A (i.))

; @TODO benchmark this other version of i., maybe C.
(define (i.* . args)
  (let* ((a (apply make-array *unspecified* args))
         (a1 (array-contents a))
         (z (tally a1)))
    (let loop ((i 0))
      (cond ((< i z) (array-set! a1 i i) (loop (1+ i)))
            (else a)))))

; ---------------------------------------------
; reshape (@TODO complete).
; ---------------------------------------------

; corner cases, allow placeholder to be zero.

(T (reshape (i. 0 3) #t 2) #2:0:2())
(T (reshape (i. 3 0) #t 2) #2:0:2())
(T (reshape (i. 0 3) 2 #t) #2(() ()))
(T (reshape (i. 3 0) 2 #t) #2(() ()))

; corner cases, flag error if placeholder can't be computed.

(assert-fail (reshape (i. 3 2) 0 #t) "bad dim deduction with empty shape")
(T-msg "bad size deduction from scalar" #(9) (reshape 9 #t))

; doc examples.

(T (reshape (i. 2 3) 6) #(0 1 2 3 4 5))
(T (reshape (i. 2 3) 5) #(0 1 2 3 4))
(T (reshape (i. 2 3) 7) #(0 1 2 3 4 5 0))
(T (reshape (i. 2 3) 3 2) #2((0 1) (2 3) (4 5)))
(T (reshape (i. 2 3) 2 2 2) #3(((0 1) (2 3)) ((4 5) (0 1))))
(T (reshape (i. 2 3) 4 2) #2((0 1) (2 3) (4 5) (0 1)))
(T (reshape (i. 2 3) #t 2) #2((0 1) (2 3) (4 5)))
(T (reshape (i. 2 3) 2 #t) #2((0 1 2) (3 4 5)))
(T (reshape (i. 2 3) #t) #(0 1 2 3 4 5))
(T (reshape (i. 2 3) 0 3) #2:0:3())
(T (reshape (i. 2 3) 0) #())
(T (reshape (i. 2 3) 0 0 0) #3())
(assert-fail (reshape (i. 2 3) 0 #t))
(assert-fail (reshape (i. 2 3) 4 #t))
(T (reshape (i. 2 3)) 0)

; ---------------------------------------------
; simpler ply. Match frames at each nesting level, no types. Lots of copying,
; testing only. @TODO This approach could work for nested (ragged) vectors.
; ---------------------------------------------

(define (match-frame* A f r)
  (cond ((= (length f) (- (rank A) r))
         A)
        ((not (array? A))
         (match-frame* (make-array A) f r))
        (else
         (apply make-shared-array A
                (lambda i (append (take i (- (rank A) r)) (take-right i r)))
                (append f (take-right ($ A) r))))))

(define (prefix-frame* A r)
  "Frame common to all arrays A with cell ranks r"
  (fold (lambda (A r f)
          (let ((fA (drop-right! ($ A) r)))
            (let loop ((s f) (sA fA))
              (cond ((null? sA) f)
                    ((null? s) fA)
                    ((= (car s) (car sA)) (loop (cdr s) (cdr sA)))
                    (else (error "shape clash" A r))))))
        '() A r))

(define (ply* op . a)
  (let ((op (if (verb? op) op (verb op))))
    (let ((ro (apply verb-actual-ri op (map rank a)))
          (sop (if (verb-final? op)
                 (verb-op op)
                 (lambda a (apply ply* (verb-op op) a)))))
      (let ((f (prefix-frame* a ro)))
        (if (null? f)
          (apply sop a)
          (let ((a (map (cut match-frame* <> f <>) a ro)))
            (collapse-array
             #t
             (let ((o (apply make-array *unspecified* f)))
               (apply array-map/frame! o f sop a)
               o))))))))

; prefix-frame*, match-frame*

(define x 3)
(define A #(1 2 3))
(define B #2((1 2) (1 2) (1 2)))

(define f00 (prefix-frame* (list A B) (list 0 0)))
(define f01 (prefix-frame* (list A B) (list 0 1)))
(define f10 (prefix-frame* (list A B) (list 1 0)))
(define f11 (prefix-frame* (list A B) (list 1 1)))

(define f00_ (prefix-frame* (list B A) (list 0 0)))
(define f01_ (prefix-frame* (list B A) (list 1 0)))
(define f10_ (prefix-frame* (list B A) (list 0 1)))
(define f11_ (prefix-frame* (list B A) (list 1 1)))

(assert (equal? f00 f00_ '(3 2)))
(assert (equal? f01 f01_ '(3)))
(assert (equal? f10 f10_ '(3 2)))
(assert (equal? f11 f11_ '(3)))

(T (match-frame* A f00 0) #2((1 1) (2 2) (3 3)))
(T (match-frame* B f00 0) #2((1 2) (1 2) (1 2)))
(T (match-frame* A f01 0) #(1 2 3))
(T (match-frame* B f01 1) #2((1 2) (1 2) (1 2)))
(T (match-frame* A f10 1) #3(((1 2 3) (1 2 3)) ((1 2 3) (1 2 3)) ((1 2 3) (1 2 3))))
(T (match-frame* B f10 0) #2((1 2) (1 2) (1 2)))
(T (match-frame* A f11 1) #2((1 2 3) (1 2 3) (1 2 3)))
(T (match-frame* B f11 1) #2((1 2) (1 2) (1 2)))

(T '(2) (prefix-frame* '(#2((0 0) (0 0)) #(0 0)) '(1 1)))
(T 9 (match-frame* 9 '() 0))

(define x 3)
(define A #(1 2 3))
(define B #2((1 2 3) (1 2 3) (1 2 3)))

(T (ply - A B) (ply* - A B) #2((0 -1 -2) (1 0 -1) (2 1 0)))
(T (ply - B A) (ply* - B A) #2((0 1 2) (-1 0 1) (-2 -1 0)))
(T (ply - x A) (ply* - x A) #(2 1 0))
(T (ply - x B) (ply* - x B) #2((2 1 0) (2 1 0) (2 1 0)))
(T (ply - A x) (ply* - A x) #(-2 -1 0))
(T (ply - B x) (ply* - B x) #2((-2 -1 0) (-2 -1 0) (-2 -1 0)))

; rank 0 inner op, deduced...
(T (ply + #(10 20) #2((1 2 3) (1 2 3)))
   #2((11 12 13) (21 22 23)))

; a case that requires rank>2 loop.
; (10 20) + (2 2 2 $ 1 + i. 8)
(T (ply + #(10 20) (reshape #(1 2 3 4 5 6 7 8) 2 2 2))
   #3(((11 12) (13 14)) ((25 26) (27 28))))

; rank 0 inner op, explicit.
(T (ply (verb + '() 0 0) #(10 20) #2((1 2 3) (1 2 3)))
   #2((11 12 13) (21 22 23)))
; specifying output rank.
(T (ply (verb + '() 0 0) #(10 20) #2((1 2 3) (1 2 3)))
   #2((11 12 13) (21 22 23)))

; rank 0 with wrapped ranks.
; (10 * 1 + i.2) (+"0 0) (2 3 $ 1 2 3 1 2 3)
(T (ply (w/rank + 0 0) #(10 20) #2((1 2 3) (1 2 3)))
   #2((11 12 13) (21 22 23)))
; (10 * 1 + i.3) (+"1 1) (2 3 $ 1 2 3 1 2 3)
(T (ply (w/rank + 1 1) #(10 20 30) #2((1 2 3) (1 2 3)))
   #2((11 22 33) (11 22 33)))
; (10 * 1 + i.3) (+"1 0) (2 3 $ 1 2 3 5 6 7)
(T (ply (w/rank + 1 0) #(10 20 30) #2((1 2 3) (5 6 7)))
   #3(((11 21 31) (12 22 32) (13 23 33)) ((15 25 35) (16 26 36) (17 27 37))))
; (10 * 1 + i.2) (+"1 2) (2 3 $ 1 2 3 5 6 7)
(T (ply (w/rank + 1 2) #(10 20) #2((1 2 3) (5 6 7)))
   #2((11 12 13) (25 26 27)))

; rank 0 with nested ranks.
; 100 200 300 +"0"0 _ (1 2 3 4) NB. From Rich2006 ch. 6
(T (ply (w/rank (w/rank + 0 0) 0 '_) #(100 200 300) #(1 2 3 4))
   #2((101 102 103 104) (201 202 203 204) (301 302 303 304)))
; 100 200 +"0"_ 0 (1 2 3) NB. From Rich2006 ch. 6
(T (ply (w/rank (w/rank + 0 0) '_ 0) #(100 200) #(1 2 3))
   #2((101 201) (102 202) (103 203)))

(define (_sqr a) (* a a))
(define _sqrm
  (verb (lambda (a)
          (array-fold (lambda (a c) (+ c (_sqr (real-part a)) (_sqr (imag-part a))))
                       0. a))
        #f 1))

(define _sqrmd
  (verb (lambda (a b)
          (array-fold (lambda (a b c)
                         (let ((a (- a b)))
                           (+ c (_sqr (real-part a)) (_sqr (imag-part a)))))
                       0. a b))
        #f 1 1))

; rank 1 inner op.
(T-eps 0. (ply _sqrm #2((1 1 1) (2 2 2))) #(3. 12.))
(T (ply (verb tally '() 1) #2((1 1 1) (2 2 2))) #(3 3))

; specifying output rank, so collapse-rank can be elided.
(T (ply (verb from (lambda (x y) '(3)) -1 0)
        (array-copy 's32 (i. 3 3 3))
        #(2 0 1))
; alternative: output-shape will be called with the shapes of the cells.
   (ply (verb from (lambda (x y) (cdr x)) -1 0)
        (array-copy 's32 (i. 3 3 3))
        #(2 0 1))
   #2s32((6 7 8) (9 10 11) (21 22 23)))

; @TODO since the output is pile'd, any output rank is ok. However, this should somehow be known in advance to preallocate the result array.
(T (ply (verb (cut sort <> <) #f 1) #2((4 2 1) (1 3 2)))
   #2((1 2 4) (1 2 3)))

(define _cross (verb (lambda (a b)
                       (let ((a (cut array-ref a <>))
                             (b (cut array-ref b <>)))
                         (vector (- (* (a 1) (b 2)) (* (a 2) (b 1)))
                                 (- (* (a 2) (b 0)) (* (a 0) (b 2)))
                                 (- (* (a 0) (b 1)) (* (a 1) (b 0))))))
                     #f 1 1))

; rank 1 with rank extension.
(T (ply _cross #2((1 0 0) (0 1 0)) #2((0 1 0) (0 0 1)))
   #2((0 0 1) (1 0 0)))
(T (ply _cross #(1 0 0) #(0 0 1))
   #(0 -1 0))
(T (ply _cross #2((1 0 0) (0 1 0)) #(0 0 1))
   #2((0 -1 0) (1 0 0)))
(T (ply _cross #(0 0 1) #2((1 0 0) (0 1 0)))
   #2((0 1 0) (-1 0 0)))

; rank 1 with wrapped rank.
(T (ply (w/rank _cross 2 1) #2((1 0 0) (0 1 0)) #(0 0 1))
   #2((0 -1 0) (1 0 0)))
(T (ply _cross #3(((1 0 0) (0 1 0)) ((2 0 0) (0 2 0))) #(0 0 1))
   #3(((0 -1 0) (1 0 0)) ((0 -2 0) (2 0 0))))
(T (ply _cross #3(((1 0 0) (0 1 0)) ((2 0 0) (0 2 0))) #2((0 0 1) (0 0 2)))
   #3(((0 -1 0) (1 0 0)) ((0 -4 0) (4 0 0))))
(T (ply (w/rank _cross 2 2) #3(((1 0 0) (0 1 0)) ((2 0 0) (0 2 0))) #2((0 0 1) (0 0 2)))
   #3(((0 -1 0) (2 0 0)) ((0 -2 0) (4 0 0))))
(T (ply (w/rank _cross 3 2) #3(((1 0 0) (0 1 0)) ((2 0 0) (0 2 0))) #2((0 0 1) (0 0 2)))
   #3(((0 -1 0) (1 0 0)) ((0 -4 0) (4 0 0))))

; rank 2.
(define (_invert a)
  (let* ((a (cut array-ref a <> <>))
         (D (- (* (a 0 0) (a 1 1)) (* (a 0 1) (a 1 0)))))
    (reshape `#(,(a 1 1) ,(- (a 0 1)) ,(- (a 1 0)) ,(a 0 0)) 2 2)))

(T-eps 0.0
       (ply (verb _invert values 2) #3(((1 0) (0 1)) ((1 1) (0 1))))
       #3(((1.0 0.0) (0.0 1.0)) ((1.0 -1.0) (0.0 1.0))))

; mixed ranks inner op (in J( { b. 0 -> 1 0 _ )
; (2 2 $ 3 2 0 1) { 1 2 3 4
(T (ply (verb from #f 1 0) #(1 2 3 4) #2((3 2) (0 1)))
   #2((4 3) (1 2)))

; 0 1 1 0 { (2 2 $ 3 2 0 1)
(T (ply (verb from #f 2 0) #2((3 2) (0 1)) #(0 1 1 0))
   #2((3 2) (0 1) (0 1) (3 2)))

; 0 1 1 0 ({"1 2) (2 2 $ 3 2 0 1)
(T (ply (w/rank (verb from #f 2 0) 2 1) #2((3 2) (0 1)) #(0 1 1 0))
   #2((3 2) (0 1) (0 1) (3 2)))

; (2 2 $ 0 1 1 0) { (2 2 $ 3 2 0 1)
(T (ply (verb from #f 2 0) #2((3 2) (0 1)) #2((0 1) (1 0)))
   #3(((3 2) (0 1)) ((0 1) (3 2))))

; (2 2 $ 0 1 1 0) ({"1 2) (2 2 $ 3 2 0 1)
(T (ply (w/rank (verb from #f 2 0) 2 1) #2((3 2) (0 1)) #2((0 1) (1 0)))
   #3(((3 2) (0 1)) ((0 1) (3 2))))

; (2 2 $ 0 1 1 0) ({"2 1) (2 2 $ 3 2 0 1)
(T (ply (w/rank (verb from #f '_ 0) 1 2) #2((3 2) (0 1)) #2((0 1) (1 0)))
   #3(((3 2) (2 3)) ((0 1) (1 0))))

; (i. 3 4) +"1"2 i. 2 3 4
(T (ply (w/rank (w/rank + 1 1) 2 2) (i. 3 4) (i. 2 3 4))
   #3(((0 2 4 6) (8 10 12 14) (16 18 20 22)) ((12 14 16 18) (20 22 24 26) (28 30 32 34))))

; re$* provides a case with output rank > input rank. Also a case where inner op rank is larger than out.
; note that J $ uses the list of items to fill the shape, so s $ a ~ (re$* s (from a 0)).
(define (re$* s a)
  (let ((rs (tally s)))
    (apply make-shared-array
           (if (array? a) a (make-typed-array (array-type* a) a))
           (lambda i (drop i rs))
           (append (vector->list s) ($ a)))))

; (2 3) $ 1 2 $ 1 2
(T (ply (verb re$* #f 1 '_) #(2 3) #(1 2))
   #3(((1 2) (1 2) (1 2)) ((1 2) (1 2) (1 2))))

; (2 3) ($"1 0) 1 2 $ 1 2
(T (ply (w/rank (verb re$* #f 1 '_) 1 0) #(2 3) #(1 2))
   #3(((1 1 1) (1 1 1)) ((2 2 2) (2 2 2))))

; exercise a case with in rank 0, out rank > 0, where there's a work around array-map's output.
; (2 3) ($"1 0) 1 3 $ 1 2 3
(T (ply (verb (cut re$* #(2 3) <>) #f 0) #(1 2 3))
   #3(((1 1 1) (1 1 1)) ((2 2 2) (2 2 2)) ((3 3 3) (3 3 3))))

; This is J $. J considers a scalar to have 1 item, so that's a special case.
; @TODO Don't reuse reshape, we know more here and can avoid work.
; @TODO interesting case for oshape (can do better than #f).
(define (reshape-J s A)
  (if (zero? (rank A))
    (apply make-array A (vector->list s))
    (apply reshape A (append (vector->list s) (cdr ($ A))))))

; (2 3) $ 1 2
(T (ply (verb reshape-J #f 1 '_) #(2 3) #(1 2))
   #2((1 2 1) (2 1 2)))

; (2 3) $ 2 2 $ 1 2 3 4
(T (ply (verb reshape-J #f 1 '_) #(2 3) #2((1 2) (3 4)))
   #3(((1 2) (3 4) (1 2)) ((3 4) (1 2) (3 4))))

; (2 3) ($"1 0) 1 2
(T (ply (w/rank (verb reshape-J #f 1 '_) 1 0) #(2 3) #(1 2))
   #3(((1 1 1) (1 1 1)) ((2 2 2) (2 2 2))))

; empty arrays.
(T (ply + (i. 0 2) 7) #2:0:2())
(T (ply + (i. 2 0) 7) #2(() ()))

; ------------------------
; other tests with higher rank arrays.
; ------------------------

; single ply; marred by array copying and no-op looping.
; see w/rank chain in (0 0 0 1) -> (0 0 1 1) -> (0 1 1 1) out.
(define (_meshgrid . l)
  (let ((n (length l)))
    (apply
     ply/t 'f64
     (let loop ((i (- n 1)))
       (if (zero? i)
           vector
           (apply w/rank (loop (- i 1)) (append (make-list (- n i) 0) (make-list i 1)))))
     l)))

(define (_meshgrid-last . l)
  (let ((n (length l)))
    (apply
     ply/t 'f64
     (let loop ((i (- n 1)))
       (if (zero? i)
           vector
           (apply w/rank (loop (- i 1)) (append (make-list i 1) (make-list (- n i) 0)))))
     l)))

(T (apply _meshgrid (map i. (iota 4 1)))
   (apply ply/t 'f64 (w/rank (w/rank (w/rank vector 0 0 0 1) 0 0 1 1) 0 1 1 1)
          (map i. (iota 4 1)))
   #5f64(((((0 0 0 0) (0 0 0 1) (0 0 0 2) (0 0 0 3))
           ((0 0 1 0) (0 0 1 1) (0 0 1 2) (0 0 1 3))
           ((0 0 2 0) (0 0 2 1) (0 0 2 2) (0 0 2 3)))
          (((0 1 0 0) (0 1 0 1) (0 1 0 2) (0 1 0 3))
           ((0 1 1 0) (0 1 1 1) (0 1 1 2) (0 1 1 3))
           ((0 1 2 0) (0 1 2 1) (0 1 2 2) (0 1 2 3))))))

(T (apply _meshgrid-last (map i. (iota 4 1)))
   (apply ply/t 'f64 (w/rank (w/rank (w/rank vector 1 0 0 0) 1 1 0 0) 1 1 1 0)
          (map i. (iota 4 1)))
   #5f64(((((0 0 0 0)) ((0 1 0 0))) (((0 0 1 0)) ((0 1 1 0))) (((0 0 2 0)) ((0 1 2 0))))
         ((((0 0 0 1)) ((0 1 0 1))) (((0 0 1 1)) ((0 1 1 1))) (((0 0 2 1)) ((0 1 2 1))))
         ((((0 0 0 2)) ((0 1 0 2))) (((0 0 1 2)) ((0 1 1 2))) (((0 0 2 2)) ((0 1 2 2))))
         ((((0 0 0 3)) ((0 1 0 3))) (((0 0 1 3)) ((0 1 1 3))) (((0 0 2 3)) ((0 1 2 3))))))

; out, of which the meshgrid loop above is a case.
(T (out * #(10 20 30) (i. 5))
   #2((0 10 20 30 40) (0 20 40 60 80) (0 30 60 90 120)))

; out with op ranks != 0.
(T-eps 0
       (out _sqrmd (i. 3 2) (i. 4 2))
       #2((0 8 32 72) (8 0 8 32) (32 8 0 8)))

; out with different ranks. @TODO A case with rank '_ or negative.
(define (_cons a b) (list->vector (cons a (vector->list b))))
(T (out (verb _cons #f 0 1) (i. 2) (i. 3 4))
   #3(((0 0 1 2 3) (0 4 5 6 7) (0 8 9 10 11))
      ((1 0 1 2 3) (1 4 5 6 7) (1 8 9 10 11))))
(T (out (verb _cons #f 1 1) (i. 2) (i. 3 4))
   #2((#(0 1) 0 1 2 3) (#(0 1) 4 5 6 7) (#(0 1) 8 9 10 11)))

; with more args, giving oshape. @BUG Mismatched shape/rank/args are not caught.
(T (out (verb vector '(3) 0 0 0) #(1 2) #(10 20) #(100 200))
    #4((((1 10 100) (1 10 200)) ((1 20 100) (1 20 200)))
       (((2 10 100) (2 10 200)) ((2 20 100) (2 20 200)))))

; test against once array-slice-for-each bug.
(T (out (verb list '() 1 1) #2((10 45)(10 0)) #2((a b) (c d)))
   #2(((#1(10 45) #1(a b)) (#1(10 45) #1(c d))) ((#1(10 0) #1(a b)) (#1(10 0) #1(c d)))))

; profile giving oshape or not. @BUG Giving it is actually slower.
; ,profile (out (verb vector '(3) 0 0 0) (i. 100) (i. 100) (i. 100))
; ,profile (out (verb vector #f 0 0 0) (i. 100) (i. 100) (i. 100))

; ------------------------
; ply-n/o
; ------------------------

(define o (vector 9 4))
(ply-n/o (verb (lambda (v) (set! o (ply + v o))) #f 1)
         (i. 10 2))
(T o #(99 104))

; ------------------------
; linspace.
; ------------------------

(T (linspace. 0 10 0) #())
(T (linspace. 0 10 1) #(0))
(T (linspace. 0 10 2) #(0 10))
(T (linspace. 0 10 3) #(0 5 10))
(T (linspace. 0 10 4) #(0 10/3 20/3 10))
(T (linspace-m. 0 10 0) #())
(T (linspace-m. 0 10 1) #())
(T (linspace-m. 0 10 2) #(0))
(T (linspace-m. 0 10 3) #(0 5))
(T (linspace-m. 0 10 4) #(0 10/3 20/3))

; ------------------------
; some tests about the need for optimization.
; ------------------------

; cf uniform-grid-cube-points.
(define (_uniform-grid-cube-points rank n)
  (reshape (apply _meshgrid-last (make-list rank (linspace. 0 1 n)))
           (expt n rank) rank))

(define (test-each rank n ext)
  (transpose-array
   (ply + (transpose-array (ply * ext (_uniform-grid-cube-points rank n)) 1 0)
        (reshape (* -.5 ext) rank))
   1 0))

; (ext * ugc) ("+1 1) (rank $ -.5 * ext)
(define (test-each* rank n ext)
  (ply (w/rank + 1 1)
       (ply * ext (_uniform-grid-cube-points rank n))
       (reshape (* -.5 ext) rank)))

(define (test-each** rank n ext)
  (ply* (w/rank + 1 1)
        (ply* * ext (_uniform-grid-cube-points rank n))
        (reshape (* -.5 ext) rank)))

(T (test-each 2 100 1) (test-each* 2 100 1) (test-each** 2 100 1))
(T (test-each 3 30 1) (test-each* 3 30 1) (test-each** 3 30 1))

; maybe
; (+ (each (* ext (uniform-grid-cube-points rank n)))
;    (reshape (* -.5 ext) rank))

; ,profile (test-each 2 100 1)
; ,profile (test-each* 2 100 1)
; ,profile (test-each 3 30 1)
; ,profile (test-each* 3 30 1)

; -----------------------------------------------
; with precomputed nested frames & ranks, or not.
; -----------------------------------------------

(define a (i. 3 2))
(define b (i. 3))
(define c (i. 2))

(T (ply (w/rank + 1 0) a b) (ply* (w/rank + 1 0) a b) #2((0 1) (3 4) (6 7)))
(T (ply (w/rank + 1 1) a c) (ply* (w/rank + 1 1) a c) #2((0 2) (2 4) (4 6)))

; -----------------------------------------------
; developing (from)
; -----------------------------------------------

; How to do cartesian selection with arbitrary indices in a single ply.
; @TODO this could be used e.g. in (from), after the scalars and J-selectors have been 'beaten'.
; Simple enough along a single dimension,

(ply (w/rank (verb array-cell-ref #f '_ 0) 3 1) (i. 10 10 10) #(1 2))
(ply (w/rank (verb array-cell-ref #f '_ 0) 2 1) (i. 10 10 10) #(1 2))
(ply (w/rank (verb array-cell-ref #f '_ 0) 1 1) (i. 10 10 10) #(1 2))

; Let's say we have index vectors a, b, c. Remember what prefix-frame does!
; we start with ->  (w/rank '_ 0 1 1) -> (w/rank '_ 0 0 1) -> (w/rank '_ 0 0 0)
; [---]             [A] | [---]          [AB] | [---]         [ABC] | [---]
; [a]               [a] |                [aB] |               [aBC] |
; [b]               [A] | [b]            [Ab] |               [AbC] |
; [c]               [A] | [c]            [AB] | [c]           [ABc] |
; The last w/rank is redundant, so:

(T (ply (w/rank (w/rank (verb array-cell-ref #f '_ 0 0 0) '_ 0 0 1) '_ 0 1 1)
        (i. 10 10 10)
        #(1 2)
        #(3 4 5)
        #(6 7 8 9))
   (ply (w/rank (verb array-cell-ref #f '_ 0) 1 1)
        (ply (w/rank (verb array-cell-ref #f '_ 0) 2 1)
             (ply (w/rank (verb array-cell-ref #f '_ 0) 3 1)
                  (i. 10 10 10)
                  #(1 2))
             #(3 4 5))
        #(6 7 8 9))
   #3(((136 137 138 139) (146 147 148 149) (156 157 158 159))
      ((236 237 238 239) (246 247 248 249) (256 257 258 259))))

; This gives -> (w/rank '_ 0 2 1) -> (w/rank '_ 0 0 1)  -> etc.
; [---]         [A] | [---]          [ABC] | [---]
; [a]           [a] |                [aBC] |
; [bc]          [A] | [bc]           [Abc] |
; [d]           [A] | [d]            [ABC] | [d]
(T (ply (w/rank (w/rank (verb array-cell-ref #f '_ 0 0 0) '_ 0 0 1) '_ 0 2 1)
        (i. 10 10 10)
        #(0 1)
        #2((2 3 4) (5 6 7))
        #(8 9))
   #4(((( 28  29) ( 38  39) ( 48  49)) (( 58  59) ( 68  69) ( 78  79)))
      (((128 129) (138 139) (148 149)) ((158 159) (168 169) (178 179)))))

; benchmarks
; 1. single ply saves much array-copy.
; 2. fixing args of inf rank (relative to the loop) should be done by ply.
(define i10 (i. 10 10 10))
(define i100 (i. 100 100 100))

(define (test0 iN)
  (repeat 1000 (ply (w/rank (w/rank (verb (cut array-cell-ref iN <> <> <>) '() 0 0 0) 0 0 1) 0 1 1)
                    #(1 2 3) #(4 5 6 7) #(6 7 8 9))))
(define (test1 iN)
  (repeat 1000 (ply (w/rank (w/rank (verb (cut array-cell-ref iN <> <> <>) #f 0 0 0) 0 0 1) 0 1 1)
                    #(1 2 3) #(4 5 6 7) #(6 7 8 9))))
(define (test2 iN)
  (repeat 1000 (ply (w/rank (w/rank (verb array-cell-ref #f '_ 0 0 0) '_ 0 0 1) '_ 0 1 1)
                    iN #(1 2 3) #(4 5 6 7) #(6 7 8 9))))
(define (test3 iN)
  (repeat 1000 (ply (w/rank (verb array-cell-ref #f '_ 0) 1 1)
                    (ply (w/rank (verb array-cell-ref #f '_ 0) 2 1)
                         (ply (w/rank (verb array-cell-ref #f '_ 0) 3 1)
                              iN #(1 2 3))
                         #(4 5 6 7))
                    #(6 7 8 9))))

(define (crude-median . a)
  (list-ref (sort a <) (ceiling (/ (- (length a) 1) 2))))
(define (median-time n proc)
  (apply crude-median (map-in-order (lambda (i) (time (proc))) (iota n))))

(let ((t0 (time (test0 i10)))
      (t1 (time (test1 i10)))
      (t2 (time (test2 i10)))
      (t3 (time (test3 i10))))
  (format! "\n~:{test~a i10 ~a\n~}" (zip (iota 4) (list t0 t1 t2 t3)))
; t2 should be < t3, but too variable too enforce.
  (assert (and (< t1 t2) (< t1 t3))))

(let ((t0 (time (test0 i100)))
      (t1 (time (test1 i100)))
      (t2 (time (test2 i100)))
      (t3 (time (test3 i100))))
  (format! "~:{test~a i100 ~a\n~}" (zip (iota 4) (list t0 t1 t2 t3)))
  (assert (< t1 t2 t3)))

; from

(T (from #2f64((1 2 3 4) (5 6 7 8)) (- 2 1) (J 2 0 2))
   #f64(5.0 7.0))
(T (from #2f64((1 2 3 4) (5 6 7 8)) (- 2 1) (J 2 (- 4 2)))
   #f64(7.0 8.0))
(T (from #2f64((1 2 3 4) (5 6 7 8)) (- 2 1) (J 2 (- 4 4)))
   #f64(5.0 6.0))
(T (from #2f64((1 2 3 4) (5 6 7 8)) (- 2 1) (J 2 (- 4 4) 2))
   #f64(5.0 7.0))

(T (from (i. 2 2) #(0 1) #(0 1)) #2((0 1) (2 3)))
(T (from (i. 2 2) #(1 0) #(0 1)) #2((2 3) (0 1)))
(T (from (i. 2 2) #(0 1) #(1 0)) #2((1 0) (3 2)))
(T (from (i. 2 2) #(1 0) #(1 0)) #2((3 2) (1 0)))

(T (from (i. 2 2) #t #(0 1)) #2((0 1) (2 3)))
(T (from (i. 2 2) #t #(1 0)) #2((1 0) (3 2)))
(T (from (i. 2 2) #(0 1) #t) #2((0 1) (2 3)))
(T (from (i. 2 2) #(1 0) #t) #2((2 3) (0 1)))
(T (from (i. 2 2) #(0 1)) #2((0 1) (2 3)))
(T (from (i. 2 2) #(1 0)) #2((2 3) (0 1)))

(T (from (i. 2 2 2) #(0 1) #(0 1) #(0 1)) #3(((0 1) (2 3)) ((4 5) (6 7))))
(T (from (i. 2 2 2) #(0 1) #(0 1) #(1 0)) #3(((1 0) (3 2)) ((5 4) (7 6))))
(T (from (i. 2 2 2) #(0 1) #(1 0) #(0 1)) #3(((2 3) (0 1)) ((6 7) (4 5))))
(T (from (i. 2 2 2) #(0 1) #(1 0) #(1 0)) #3(((3 2) (1 0)) ((7 6) (5 4))))
(T (from (i. 2 2 2) #(1 0) #(0 1) #(0 1)) #3(((4 5) (6 7)) ((0 1) (2 3))))
(T (from (i. 2 2 2) #(1 0) #(0 1) #(1 0)) #3(((5 4) (7 6)) ((1 0) (3 2))))
(T (from (i. 2 2 2) #(1 0) #(1 0) #(0 1)) #3(((6 7) (4 5)) ((2 3) (0 1))))
(T (from (i. 2 2 2) #(1 0) #(1 0) #(1 0)) #3(((7 6) (5 4)) ((3 2) (1 0))))

(T (from (i. 2 2 2) #t #t #t) #3(((0 1) (2 3)) ((4 5) (6 7))))
(T (from (i. 2 2 2) #t #t #(1 0)) #3(((1 0) (3 2)) ((5 4) (7 6))))
(T (from (i. 2 2 2) #t #(1 0) #t) #3(((2 3) (0 1)) ((6 7) (4 5))))
(T (from (i. 2 2 2) #t #(1 0) #(1 0)) #3(((3 2) (1 0)) ((7 6) (5 4))))
(T (from (i. 2 2 2) #(1 0) #t #t) #3(((4 5) (6 7)) ((0 1) (2 3))))
(T (from (i. 2 2 2) #(1 0) #t #(1 0)) #3(((5 4) (7 6)) ((1 0) (3 2))))
(T (from (i. 2 2 2) #(1 0) #(1 0) #t) #3(((6 7) (4 5)) ((2 3) (0 1))))
(T (from (i. 2 2 2) #(1 0) #(1 0) #(1 0)) #3(((7 6) (5 4)) ((3 2) (1 0))))

; TODONOW
(T (ply (w/rank (w/rank (verb array-cell-ref #f '_ 0 0) 2 0 1) 2 1 1) (i. 2 2 2) #(1 0) #(1 0))
   #3(((3 2) (1 0)) ((7 6) (5 4))))

; with higher rank indices.

(T (from (i. 2 2) #(0 1) #2((0 1) (1 0)))
   (from (i. 2 2) #t #2((0 1) (1 0)))
   #3(((0 1) (1 0)) ((2 3) (3 2))))
(T (from (i. 2 2) #(1 0) #2((0 1) (1 0)))
   #3(((2 3) (3 2)) ((0 1) (1 0))))
(T (from (i. 2 2) #2((0 1) (1 0)))
   (from (i. 2 2) #2((0 1) (1 0)) #t)
   (from (i. 2 2) #2((0 1) (1 0)) #(0 1))
   #3(((0 1) (2 3)) ((2 3) (0 1))))
(T (from (i. 2 2) #2((0 1) (1 0)) #(1 0))
   #3(((1 0) (3 2)) ((3 2) (1 0))))

; with higher rank indices and empty axis in the middle.

(T (from (i. 2 2 2) #2((0 1) (1 0)) #(0 1) #(0 1))
   (from (i. 2 2 2) #2((0 1) (1 0)) #t #(0 1))
   #4((((0 1) (2 3)) ((4 5) (6 7))) (((4 5) (6 7)) ((0 1) (2 3)))))
(T (from (i. 2 2 2) #2((0 1) (1 0)) #(1 0) #(0 1))
   #4((((2 3) (0 1)) ((6 7) (4 5))) (((6 7) (4 5)) ((2 3) (0 1)))))

; -------------------------
; more indexing examples (cf numpy)
; -------------------------

(define z (i. 5 5))
(define i (i. 2 2))
; z[i]
(T (from z i)
   #3(((0 1 2 3 4) (5 6 7 8 9)) ((10 11 12 13 14) (15 16 17 18 19))))
(T (from z i i)                   ;
   #4((((0 1) (2 3)) ((5 6) (7 8))) (((10 11) (12 13)) ((15 16) (17 18)))))
; z[i, i] in numpy
(T (ply (cut from z <> <>) i i)
   #2((0 6) (12 18)))
; idem
(T (ply (verb from #f '_ 0 0) z i i)
   #2((0 6) (12 18)))
; p
(T (from z i #2((0 1 0 1) (2 3 2 3))) ; (cat 1 i i)
   #4((((0 1 0 1) (2 3 2 3)) ((5 6 5 6) (7 8 7 8)))
      (((10 11 10 11) (12 13 12 13)) ((15 16 15 16) (17 18 17 18)))))

; from the doc.

(T (from (i. 2 3) 0) #(0 1 2))
(T (from (i. 2 3) 0 #t) #(0 1 2))
(T (from (i. 2 3) #t 0) #(0 3))
(T (from #(#(0 1) #(2)) 1) #(2))
(assert-fail (from #(#(0 1) #(2)) 1 0))
(T (from (i. 10 2) (J 2 2) #t) #2((4 5) (6 7)))
(T (from #(1 2 3) #2((0 1) (1 2) (2 0))) #2((1 2) (2 3) (3 1)))
(T (from (i. 10 10) #(3 4) #(7 9 2)) #2((37 39 32) (47 49 42)))

; --------------------------------
; reshape, raveling version.
; --------------------------------

; used in one of the reshape cases.
(T (index-rect-lsd-first '(2 3 5) '(0 0 0)) 0)
(T (index-rect-lsd-first '(2 3 5) '(1 0 0)) 1)
(T (index-rect-lsd-first '(2 3 5) '(0 1 0)) 2)
(T (index-rect-lsd-first '(2 3 5) '(0 0 1)) 6)
(T (index-rect-lsd-first '(2 3 5) '(1 1 1)) 9)

(T (index-rect '() '()) 0)
(T (index-rect '(2 3 5) '(0 0 0)) 0)
(T (index-rect '(2 3 5) '(1 0 0)) 15)
(T (index-rect '(2 3 5) '(0 1 0)) 5)
(T (index-rect '(2 3 5) '(0 0 1)) 1)
(T (index-rect '(2 3 5) '(1 1 1)) 21)

; select ravel.
(T (reshape #2((1 2) (3 4)) 2 3) #2((1 2 3) (4 1 2)))
(T (reshape #2((1 2) (3 4)) 2 4 2) #3(((1 2) (3 4) (1 2) (3 4)) ((1 2) (3 4) (1 2) (3 4))))
(T (reshape #(1 2 3 4 5 6 7 8) 2 2 2) #3(((1 2) (3 4)) ((5 6) (7 8))))
(T (reshape #(1 2 3 4 5 6 7 8) 2 2) #2((1 2) (3 4)))

; copy case.
(T (reshape #2((1 2) (3 4)) 4 3) #2((1 2 3) (4 1 2) (3 4 1) (2 3 4)))
; copy, but maybe not necessary.
(T (reshape #2((1 2) (3 4)) 4 4 2)
   #3(((1 2) (3 4) (1 2) (3 4))
      ((1 2) (3 4) (1 2) (3 4))
      ((1 2) (3 4) (1 2) (3 4))
      ((1 2) (3 4) (1 2) (3 4))))
; copy case, but maybe not necessary.
(T (reshape #2((1 2) (3 4)) 4 2) #2((1 2) (3 4) (1 2) (3 4)))

; same size case.
(T (reshape #(1 2 3 4) 1 2 2) #3(((1 2) (3 4))))
(T (reshape #(1 2 3 4) 2 2) #2((1 2) (3 4)))
(T (reshape #(1 2 3 4) 2 2 1) #3(((1) (2)) ((3) (4))))

; tile cases.
(T (reshape #2((1 2) (3 4)) 4 2 2) #3(((1 2) (3 4)) ((1 2) (3 4)) ((1 2) (3 4)) ((1 2) (3 4))))
(T (reshape #(1 2 3 4) 2 2 2) #3(((1 2) (3 4)) ((1 2) (3 4))))
(T (reshape #(1 2 3 4) 2 3 2) #3(((1 2) (3 4) (1 2)) ((3 4) (1 2) (3 4))))
(T (reshape #(9) 3) #(9 9 9))

; select cases,
(T (reshape #2((1 2) (3 4)) 2 2) #2((1 2) (3 4)))
(T (reshape #2((1 2) (3 4)) 2) #(1 2))

; take case.
(T (reshape #(1 2 3 4) 1 1 1) #3(((1))))
(T (reshape #(1 2 3 4 5 6 7 8) 2 3) #2((1 2 3) (4 5 6)))

; cases that need copy vs cases that do not. @TODO This for every case above.
(define A #2((1 2 3) (4 5 6) (7 8 9)))
(define B (from A (J 3) (J 2)))
(assert (eq? (shared-array-root A) (shared-array-root B)))
(T B #2((1 2) (4 5) (7 8)))
(define Arow (reshape A 9))
(T Arow #(1 2 3 4 5 6 7 8 9))
(assert (eq? (shared-array-root A) (shared-array-root Arow)))
(define Brow (reshape B 6))
(T Brow #(1 2 4 5 7 8))
(define A #(1 2 3 4))
(assert (eq? (shared-array-root A) (shared-array-root (reshape A 2 2 2))))

; was col->array, row->array
(let ((a #f64(1 2 3 4)))
  (T (reshape a #f 1) #2f64((1) (2) (3) (4)))
  (T (reshape a 1 #f) #2f64((1 2 3 4))))

; --------------------------------
; axis ops.
; --------------------------------

(T (rollaxis (i. 2 3 4 5)  0 -1)
   (transpose-array (i. 2 3 4 5) 3 0 1 2))
(T (rollaxis (i. 2 3 4 5) -1  0)
   (transpose-array (i. 2 3 4 5) 1 2 3 0))
(T (rollaxis (i. 2 3 4 5) 0 2)
   (transpose-array (i. 2 3 4 5) 2 0 1 3))
(T (rollaxis (i. 2 3 4 5) 3 1)
   (transpose-array (i. 2 3 4 5) 0 2 3 1))

(define (axes-to-front a . x)
  (let ((xy (sort (zip x (iota (length x)))
                  (lambda (a b) (< (car a) (car b))))))
    (apply transpose-array a
           (let loop ((y '()) (xy xy) (i (length x)) (j 0))
             (cond ((= j (rank a))
                    (reverse! y))
; one of x.
                   ((and (pair? xy) (= j (first (car xy))))
                    (loop (cons (second (car xy)) y)
                          (cdr xy)
                          i
                          (+ j 1)))
; one not of x.
                   (else
                    (loop (cons i y)
                          xy
                          (+ 1 i)
                          (+ 1 j))))))))

; @TODO Probably should test with perm of comb.
(define (combinations l k)
  (cond ((zero? k)
         '(()))
        ((null? l)
         '())
        (else
         (append
          (map (cute cons (car l) <>)
               (combinations (cdr l) (- k 1)))
          (combinations (cdr l) k)))))

(let ((I (i. 2 3 4 5)))
  (for-each
   (lambda (i)
     (T (map (cute list-ref ($ I) <>) (append i (lset-difference = (iota 4) i)))
        ($ (apply axes-to-front I i))))
   (append-map
    (cute combinations (iota (rank I)) <>)
    (iota (rank I)))))

; --------------------------------
; ply! with suffix matching. @TODO Do I want suffix matching?
; --------------------------------

(T (ply! (make-array 0 4) (verb (const 2) '()))
   #(2 2 2 2))
(T (ply! (make-array 0 4) (const 2))
   #(2 2 2 2))

(T (ply! (make-array 0 4) values 9)
   #(9 9 9 9))
(T (ply! (make-array 0 4 3) (cut iota. 3 <>) #(1 2 3 4))
   #2((1 2 3) (2 3 4) (3 4 5) (4 5 6)))

; suffix matching.
(T (ply! (make-array 0 4 3) values #(1 2 3))
   #2((1 2 3) (1 2 3) (1 2 3) (1 2 3)))

; prefix matching.
(T (ply!! (make-array 0 4 3) values #(1 2 3 4))
   #2((1 1 1) (2 2 2) (3 3 3) (4 4 4)))

(T (ply! (i. 4 3 3) (const #(1 2 3)))
   (reshape #(1 2 3) 4 3 3))
(T (ply! (i. 4 3 3) values (out * #(1 2 3) #(1 2 3)))
   (reshape (out * #(1 2 3) #(1 2 3)) 4 3 3))
(T (ply! (i. 4 3 3) (w/rank + 1 0)  #(1 2 3) #(10 20 30))
   (reshape (ply (w/rank + 1 0)  #(1 2 3) #(10 20 30)) 4 3 3))

; --------------------------------
; Regression test for match-frame (@TODO exact test).
; --------------------------------

(define q (reshape #(1 2 3) 6 2 8 3))
(define w (reshape #(10 20 30) 2 3))
(define y #2((10 20 30) (10 20 30)))
(T (ply (w/rank (w/rank + 1 1) -1 '_) q w)
   (ply (w/rank (w/rank + 1 1) -1 '_) q y))

; --------------------------------
; Regression test for array-map/frame (null? case)
; --------------------------------

(assert (equal? #f64(1 2) (ply/t 'f64 (verb identity #f 1) #(1 2)))
        "bad type conversion")

; --------------------------------
; There might be more bugs (@TODO research)
; --------------------------------

;; ; down to make-shared-array.

;; (import (srfi srfi-1))
;; (define a (make-shared-array #0(0) (lambda i '()) 6 2 3)) ; works, should it?
;; (define a (make-array 0 6 2 3)) ; doesn't work.
;; (apply make-shared-array a
;;        (lambda i (format #t "$a ~a\n" (array-dimensions a)) (pk 'attempt (append (take i 1) (take i 1) (take-right i 1))))
;;        '(6 2 8 3))

;; ; rank 0 arrays; (array-map/frame) has a case (null? f) that hands off #0() to op, which 0-rank verbs may not take. Probably should be handled in λ → verb conversion.

;; (ply + #0(99))
;; (from #(1 2 3) #0(1))

; --------------------------------
; experiments.
; --------------------------------

; (out op ...) supports variable arity but (ply (outv op) ...) doesn't.
; to fix that, have nested-op-frames actualize op with then-known arity.
(define (outv op)
  (let* ((op (if (verb? op) op (verb op)))
         (ri (verb-ri op))
         (n (begin (assert (list? ri) "outv requires known input ranks")
                   (length ri)))
         (infs (make-list n '_)))
    (let loop ((i 1))
      (if (>= i n)
        op
        (apply w/rank (loop (+ i 1)) (append (take ri i) (drop infs i)))))))

(T (ply/t 'f64 (outv (verb sqrm-reduce '() 1 1)) (i. 5 4 3) (i. 6 7 3))
   (out/t 'f64 (verb sqrm-reduce '() 1 1) (i. 5 4 3) (i. 6 7 3)))