;; lib/apl/runtime.sx — APL primitives on SX
;;
;; APL vectors are represented as SX lists (functional, immutable results).
;; Operations are rank-polymorphic: scalar/vector arguments both accepted.
;; Index origin: 1 (traditional APL).
;;
;; Primitives used:
;;   map (multi-arg, Phase 1)
;;   bitwise-and/or/xor/not/arithmetic-shift (Phase 7)
;;   make-set/set-member?/set-add!/set->list (Phase 18)

;; ---------------------------------------------------------------------------
;; 1. Core vector constructors
;; ---------------------------------------------------------------------------

;; ⍳N — iota: generate integer vector 1, 2, ..., N
(define
  (apl-iota n)
  (letrec
    ((go (fn (i acc) (if (< i 1) acc (go (- i 1) (cons i acc))))))
    (go n (list))))

;; ⍴A — shape (length of a vector)
(define (apl-rho v) (if (list? v) (len v) 1))

;; A[I] — 1-indexed access
(define (apl-at v i) (nth v (- i 1)))

;; Scalar predicate
(define (apl-scalar? v) (not (list? v)))

;; ---------------------------------------------------------------------------
;; 2. Rank-polymorphic helpers
;;    dyadic: scalar/vector × scalar/vector → scalar/vector
;;    monadic: scalar/vector → scalar/vector
;; ---------------------------------------------------------------------------

(define
  (apl-dyadic op a b)
  (cond
    ((and (list? a) (list? b)) (map op a b))
    ((list? a) (map (fn (x) (op x b)) a))
    ((list? b) (map (fn (y) (op a y)) b))
    (else (op a b))))

(define (apl-monadic op a) (if (list? a) (map op a) (op a)))

;; ---------------------------------------------------------------------------
;; 3. Arithmetic (element-wise, rank-polymorphic)
;; ---------------------------------------------------------------------------

(define (apl-add a b) (apl-dyadic + a b))
(define (apl-sub a b) (apl-dyadic - a b))
(define (apl-mul a b) (apl-dyadic * a b))
(define (apl-div a b) (apl-dyadic / a b))
(define (apl-mod a b) (apl-dyadic modulo a b))
(define (apl-pow a b) (apl-dyadic pow a b))
(define (apl-max a b) (apl-dyadic (fn (x y) (if (> x y) x y)) a b))
(define (apl-min a b) (apl-dyadic (fn (x y) (if (< x y) x y)) a b))

(define (apl-neg a) (apl-monadic (fn (x) (- 0 x)) a))
(define (apl-abs a) (apl-monadic abs a))
(define (apl-floor a) (apl-monadic floor a))
(define (apl-ceil a) (apl-monadic ceil a))
(define (apl-sqrt a) (apl-monadic sqrt a))
(define (apl-exp a) (apl-monadic exp a))
(define (apl-log a) (apl-monadic log a))

;; ---------------------------------------------------------------------------
;; 4. Comparison (element-wise, returns 0/1 booleans)
;; ---------------------------------------------------------------------------

(define (apl-bool v) (if v 1 0))

(define (apl-eq a b) (apl-dyadic (fn (x y) (apl-bool (= x y))) a b))
(define
  (apl-neq a b)
  (apl-dyadic (fn (x y) (apl-bool (not (= x y)))) a b))
(define (apl-lt a b) (apl-dyadic (fn (x y) (apl-bool (< x y))) a b))
(define (apl-le a b) (apl-dyadic (fn (x y) (apl-bool (<= x y))) a b))
(define (apl-gt a b) (apl-dyadic (fn (x y) (apl-bool (> x y))) a b))
(define (apl-ge a b) (apl-dyadic (fn (x y) (apl-bool (>= x y))) a b))

;; Boolean logic (0/1 vectors)
(define
  (apl-and a b)
  (apl-dyadic
    (fn
      (x y)
      (if
        (and (not (= x 0)) (not (= y 0)))
        1
        0))
    a
    b))
(define
  (apl-or a b)
  (apl-dyadic
    (fn
      (x y)
      (if
        (or (not (= x 0)) (not (= y 0)))
        1
        0))
    a
    b))
(define
  (apl-not a)
  (apl-monadic (fn (x) (if (= x 0) 1 0)) a))

;; ---------------------------------------------------------------------------
;; 5. Bitwise operations (element-wise)
;; ---------------------------------------------------------------------------

(define (apl-bitand a b) (apl-dyadic bitwise-and a b))
(define (apl-bitor a b) (apl-dyadic bitwise-or a b))
(define (apl-bitxor a b) (apl-dyadic bitwise-xor a b))
(define (apl-bitnot a) (apl-monadic bitwise-not a))
(define
  (apl-lshift a b)
  (apl-dyadic (fn (x n) (arithmetic-shift x n)) a b))
(define
  (apl-rshift a b)
  (apl-dyadic (fn (x n) (arithmetic-shift x (- 0 n))) a b))

;; ---------------------------------------------------------------------------
;; 6. Reduction (fold) and scan
;; ---------------------------------------------------------------------------

(define (apl-reduce-add v) (reduce + 0 v))
(define (apl-reduce-mul v) (reduce * 1 v))
(define
  (apl-reduce-max v)
  (reduce (fn (acc x) (if (> acc x) acc x)) (first v) (rest v)))
(define
  (apl-reduce-min v)
  (reduce (fn (acc x) (if (< acc x) acc x)) (first v) (rest v)))
(define
  (apl-reduce-and v)
  (reduce
    (fn
      (acc x)
      (if
        (and (not (= acc 0)) (not (= x 0)))
        1
        0))
    1
    v))
(define
  (apl-reduce-or v)
  (reduce
    (fn
      (acc x)
      (if
        (or (not (= acc 0)) (not (= x 0)))
        1
        0))
    0
    v))

;; Scan: prefix reduction (yields a vector of running totals)
(define
  (apl-scan op v)
  (if
    (= (len v) 0)
    (list)
    (letrec
      ((go (fn (xs acc result) (if (= (len xs) 0) (reverse result) (let ((next (op acc (first xs)))) (go (rest xs) next (cons next result)))))))
      (go (rest v) (first v) (list (first v))))))

(define (apl-scan-add v) (apl-scan + v))
(define (apl-scan-mul v) (apl-scan * v))

;; ---------------------------------------------------------------------------
;; 7. Vector manipulation
;; ---------------------------------------------------------------------------

;; ⌽A — reverse
(define (apl-reverse v) (reverse v))

;; A,B — catenate
(define
  (apl-cat a b)
  (cond
    ((and (list? a) (list? b)) (append a b))
    ((list? a) (append a (list b)))
    ((list? b) (cons a b))
    (else (list a b))))

;; ↑N A — take first N elements (negative: take last N)
(define
  (apl-take n v)
  (if
    (>= n 0)
    (letrec
      ((go (fn (xs i) (if (or (= i 0) (= (len xs) 0)) (list) (cons (first xs) (go (rest xs) (- i 1)))))))
      (go v n))
    (apl-reverse (apl-take (- 0 n) (apl-reverse v)))))

;; ↓N A — drop first N elements
(define
  (apl-drop n v)
  (if
    (>= n 0)
    (letrec
      ((go (fn (xs i) (if (or (= i 0) (= (len xs) 0)) xs (go (rest xs) (- i 1))))))
      (go v n))
    (apl-reverse (apl-drop (- 0 n) (apl-reverse v)))))

;; Rotate left by n positions
(define
  (apl-rotate n v)
  (let ((m (modulo n (len v)))) (append (apl-drop m v) (apl-take m v))))

;; Compression: A/B — select elements of B where A is 1
(define
  (apl-compress mask v)
  (if
    (= (len mask) 0)
    (list)
    (let
      ((rest-result (apl-compress (rest mask) (rest v))))
      (if
        (not (= (first mask) 0))
        (cons (first v) rest-result)
        rest-result))))

;; Indexing: A[B] — select elements at indices B (1-indexed)
(define (apl-index v indices) (map (fn (i) (apl-at v i)) indices))

;; Grade up: indices that would sort the vector ascending
(define
  (apl-grade-up v)
  (let
    ((indexed (map (fn (x i) (list x i)) v (apl-iota (len v)))))
    (map (fn (p) (nth p 1)) (sort indexed))))

;; ---------------------------------------------------------------------------
;; 8. Set operations (∊ ∪ ∩ ~)
;; ---------------------------------------------------------------------------

;; Membership ∊: for each element in A, is it in B? → 0/1 vector
(define
  (apl-member a b)
  (let
    ((bset (let ((s (make-set))) (for-each (fn (x) (set-add! s x)) b) s)))
    (if
      (list? a)
      (map (fn (x) (apl-bool (set-member? bset x))) a)
      (apl-bool (set-member? bset a)))))

;; Nub ∪A — unique elements, preserving order
(define
  (apl-nub v)
  (let
    ((seen (make-set)))
    (letrec
      ((go (fn (xs acc) (if (= (len xs) 0) (reverse acc) (if (set-member? seen (first xs)) (go (rest xs) acc) (begin (set-add! seen (first xs)) (go (rest xs) (cons (first xs) acc))))))))
      (go v (list)))))

;; Union A∪B — nub of concatenation
(define (apl-union a b) (apl-nub (apl-cat a b)))

;; Intersection A∩B
(define
  (apl-intersect a b)
  (let
    ((bset (let ((s (make-set))) (for-each (fn (x) (set-add! s x)) b) s)))
    (filter (fn (x) (set-member? bset x)) a)))

;; Without A~B
(define
  (apl-without a b)
  (let
    ((bset (let ((s (make-set))) (for-each (fn (x) (set-add! s x)) b) s)))
    (filter (fn (x) (not (set-member? bset x))) a)))

;; ---------------------------------------------------------------------------
;; 9. Format (⍕) — APL-style display
;; ---------------------------------------------------------------------------

(define
  (apl-format v)
  (if
    (list? v)
    (letrec
      ((go (fn (xs acc) (if (= (len xs) 0) acc (go (rest xs) (str acc (if (= acc "") "" " ") (str (first xs))))))))
      (go v ""))
    (str v)))