; APL Runtime — array model + scalar primitives
;
; Array = SX dict {:shape (d1 d2 ...) :ravel (v1 v2 ...)}
; Scalar: rank 0, shape (), one element in ravel
; Vector: rank 1, shape (n), n elements in ravel
; Matrix: rank 2, shape (r c), r*c elements in ravel

; ============================================================
; Array constructors
; ============================================================

(define make-array (fn (shape ravel) {:ravel ravel :shape shape}))

(define apl-scalar (fn (v) {:ravel (list v) :shape (list)}))

(define apl-vector (fn (elems) {:ravel elems :shape (list (len elems))}))

; enclose — wrap any value in a rank-0 box
(define enclose (fn (v) (apl-scalar v)))

; disclose — unwrap rank-0 box, returning the first element
(define disclose (fn (arr) (first (get arr :ravel))))

; ============================================================
; Array accessors
; ============================================================

(define array-rank (fn (arr) (len (get arr :shape))))

(define scalar? (fn (arr) (= (len (get arr :shape)) 0)))

(define array-ref (fn (arr i) (nth (get arr :ravel) i)))

; ============================================================
; System variables
; ============================================================

(define apl-io 1)

; ============================================================
; Broadcast engine
; ============================================================

(define
  broadcast-monadic
  (fn (f arr) (make-array (get arr :shape) (map f (get arr :ravel)))))

(define
  broadcast-dyadic
  (fn
    (f a b)
    (cond
      ((and (scalar? a) (scalar? b))
        (apl-scalar (f (first (get a :ravel)) (first (get b :ravel)))))
      ((scalar? a)
        (let
          ((sv (first (get a :ravel))))
          (make-array
            (get b :shape)
            (map (fn (x) (f sv x)) (get b :ravel)))))
      ((scalar? b)
        (let
          ((sv (first (get b :ravel))))
          (make-array
            (get a :shape)
            (map (fn (x) (f x sv)) (get a :ravel)))))
      (else
        (if
          (equal? (get a :shape) (get b :shape))
          (make-array (get a :shape) (map f (get a :ravel) (get b :ravel)))
          (error "length error: shape mismatch"))))))

; ============================================================
; Arithmetic primitives
; ============================================================

; Monadic + : identity
(define apl-plus-m (fn (a) (broadcast-monadic (fn (x) x) a)))

; Dyadic +
(define apl-add (fn (a b) (broadcast-dyadic (fn (x y) (+ x y)) a b)))

; Monadic - : negate
(define apl-neg-m (fn (a) (broadcast-monadic (fn (x) (- 0 x)) a)))

; Dyadic -
(define apl-sub (fn (a b) (broadcast-dyadic (fn (x y) (- x y)) a b)))

; Monadic × : signum
(define
  apl-signum
  (fn
    (a)
    (broadcast-monadic
      (fn (x) (cond ((> x 0) 1) ((< x 0) -1) (else 0)))
      a)))

; Dyadic ×
(define apl-mul (fn (a b) (broadcast-dyadic (fn (x y) (* x y)) a b)))

; Monadic ÷ : reciprocal
(define apl-recip (fn (a) (broadcast-monadic (fn (x) (/ 1 x)) a)))

; Dyadic ÷
(define apl-div (fn (a b) (broadcast-dyadic (fn (x y) (/ x y)) a b)))

; Monadic ⌈ : ceiling
(define apl-ceil (fn (a) (broadcast-monadic (fn (x) (ceil x)) a)))

; Dyadic ⌈ : max
(define
  apl-max
  (fn (a b) (broadcast-dyadic (fn (x y) (if (>= x y) x y)) a b)))

; Monadic ⌊ : floor
(define apl-floor (fn (a) (broadcast-monadic (fn (x) (floor x)) a)))

; Dyadic ⌊ : min
(define
  apl-min
  (fn (a b) (broadcast-dyadic (fn (x y) (if (<= x y) x y)) a b)))

; Monadic * : e^x
(define apl-exp (fn (a) (broadcast-monadic (fn (x) (exp x)) a)))

; Dyadic * : power
(define apl-pow (fn (a b) (broadcast-dyadic (fn (x y) (pow x y)) a b)))

; Monadic ⍟ : natural log
(define apl-ln (fn (a) (broadcast-monadic (fn (x) (log x)) a)))

; Dyadic ⍟ : log base (a⍟b = log base a of b)
(define
  apl-log
  (fn (a b) (broadcast-dyadic (fn (x y) (/ (log y) (log x))) a b)))

; Monadic | : absolute value
(define
  apl-abs
  (fn (a) (broadcast-monadic (fn (x) (if (< x 0) (- 0 x) x)) a)))

; Dyadic | : modulo (a|b = b mod a)
(define
  apl-mod
  (fn
    (a b)
    (broadcast-dyadic
      (fn (x y) (if (= x 0) y (- y (* x (floor (/ y x))))))
      a
      b)))

; Monadic ! : factorial
(define
  apl-fact
  (fn
    (a)
    (broadcast-monadic
      (fn
        (n)
        (let
          ((loop nil))
          (begin
            (set!
              loop
              (fn (i acc) (if (> i n) acc (loop (+ i 1) (* acc i)))))
            (loop 1 1))))
      a)))

; Dyadic ! : binomial coefficient n!k  (a=n, b=k => a choose b)
(define
  apl-binomial
  (fn
    (a b)
    (broadcast-dyadic
      (fn
        (n k)
        (let
          ((loop nil))
          (begin
            (set!
              loop
              (fn
                (i num den)
                (if
                  (> i k)
                  (/ num den)
                  (loop (+ i 1) (* num (- (+ n 1) i)) (* den i)))))
            (loop 1 1 1))))
      a
      b)))

; Monadic ○ : pi times x
(define
  apl-pi-times
  (fn (a) (broadcast-monadic (fn (x) (* 3.14159 x)) a)))

; Dyadic ○ : trig functions (a○b, a=code, b=value)
(define
  apl-trig
  (fn
    (a b)
    (broadcast-dyadic
      (fn
        (n x)
        (cond
          ((= n 0) (pow (- 1 (* x x)) 0.5))
          ((= n 1) (sin x))
          ((= n 2) (cos x))
          ((= n 3) (tan x))
          ((= n -1) (asin x))
          ((= n -2) (acos x))
          ((= n -3) (atan x))
          (else (error "circle: unsupported trig code"))))
      a
      b)))

; ============================================================
; Comparison primitives (return 0 or 1)
; ============================================================

(define
  apl-lt
  (fn (a b) (broadcast-dyadic (fn (x y) (if (< x y) 1 0)) a b)))

(define
  apl-le
  (fn (a b) (broadcast-dyadic (fn (x y) (if (<= x y) 1 0)) a b)))

(define
  apl-eq
  (fn (a b) (broadcast-dyadic (fn (x y) (if (= x y) 1 0)) a b)))

(define
  apl-ge
  (fn (a b) (broadcast-dyadic (fn (x y) (if (>= x y) 1 0)) a b)))

(define
  apl-gt
  (fn (a b) (broadcast-dyadic (fn (x y) (if (> x y) 1 0)) a b)))

(define
  apl-ne
  (fn (a b) (broadcast-dyadic (fn (x y) (if (= x y) 0 1)) a b)))

; ============================================================
; Logical primitives
; ============================================================

; Monadic ~ : logical not
(define
  apl-not
  (fn (a) (broadcast-monadic (fn (x) (if (= x 0) 1 0)) a)))

; Dyadic ∧ : logical and
(define
  apl-and
  (fn
    (a b)
    (broadcast-dyadic
      (fn (x y) (if (and (not (= x 0)) (not (= y 0))) 1 0))
      a
      b)))

; Dyadic ∨ : logical or
(define
  apl-or
  (fn
    (a b)
    (broadcast-dyadic
      (fn (x y) (if (or (not (= x 0)) (not (= y 0))) 1 0))
      a
      b)))

; Dyadic ⍱ : logical nor
(define
  apl-nor
  (fn
    (a b)
    (broadcast-dyadic
      (fn (x y) (if (or (not (= x 0)) (not (= y 0))) 0 1))
      a
      b)))

; Dyadic ⍲ : logical nand
(define
  apl-nand
  (fn
    (a b)
    (broadcast-dyadic
      (fn (x y) (if (and (not (= x 0)) (not (= y 0))) 0 1))
      a
      b)))

; ============================================================
; Shape primitives
; ============================================================

; Monadic ⍴ : shape — returns shape as a vector array
(define apl-shape (fn (arr) (apl-vector (get arr :shape))))

; Monadic , : ravel — returns a rank-1 vector of all elements
(define apl-ravel (fn (arr) (apl-vector (get arr :ravel))))

; Monadic ≢ : tally — first dimension (1 for scalar)
(define
  apl-tally
  (fn
    (arr)
    (if
      (scalar? arr)
      (apl-scalar 1)
      (apl-scalar (first (get arr :shape))))))

; Monadic ≡ : depth
; simple number/string value → 0
; array containing only non-arrays → 0
; array containing arrays → 1 + max depth of elements
(define
  apl-depth
  (fn
    (arr)
    (define item-depth nil)
    (set!
      item-depth
      (fn
        (v)
        (if
          (and
            (dict? v)
            (not (= nil (get v :shape nil)))
            (not (= nil (get v :ravel nil))))
          (+ 1 (first (get (apl-depth v) :ravel)))
          0)))
    (let
      ((depths (map item-depth (get arr :ravel))))
      (apl-scalar (reduce (fn (a b) (if (> a b) a b)) 0 depths)))))

; Monadic ⍳ : iota — vector 1..n (with ⎕IO=1)
(define
  apl-iota
  (fn
    (n-arr)
    (let
      ((n (first (get n-arr :ravel))) (build nil))
      (begin
        (set!
          build
          (fn (i acc) (if (< i 1) acc (build (- i 1) (cons i acc)))))
        (apl-vector (build n (list)))))))