;;;; math-utils.vector.scm  -*- Scheme -*-
;;;; Kon Lovett, Jul '23
;;;; Kon Lovett, Feb '24

;; Issues
;;
;; - Functor (VECTOR FIELD) - f32vector & fL | ...
;;
;; - SIMD

(module (math-utils vector)

(;export
  vector-sum vector-prod vector-min vector-max
  absolute-magnitude
  cosine-similarity
  vector-apply
  vector-reduce
  vector-compare vector-compare2
  dot-product
  cross-product
  #;vector-info
  vector-scale vector-shift
  normalize-vector-scale normalize-vector-shift
  softmax softmax*)

(import scheme)
(import (chicken base))
(import (chicken fixnum))
(import (chicken flonum))
(import (chicken type))
(import (only (srfi 1) reverse!))
(import (only vector-lib vector-map! vector-map))
(import math-utils)

(define-type real (or integer ratnum float))

(define-type boolean-vector (vector-of boolean))
(define-type number-vector (vector-of number))

(define-type nary-numeric (#!rest number -> number))

(: vector-sum (number-vector -> number))
(: vector-prod (number-vector -> number))
(: vector-min (number-vector -> number))
(: vector-max (number-vector -> number))

(: absolute-magnitude (number-vector -> number))

(: cosine-similarity (number-vector number-vector -> number))

(: vector-reduce (('a number -> 'a) 'a number-vector #!optional fixnum fixnum -> 'a))

(: vector-apply (nary-numeric #!rest number-vector -> number-vector))

(: vector-compare2 (number-vector number-vector -> number))
(: vector-compare (#!rest number-vector -> number))

(: dot-product (number-vector number-vector -> number))
(: cross-product (number-vector number-vector -> (or number number-vector)))

#;(: vector-info (number-vector boolean boolean boolean boolean boolean
                -> (vector number fixnum number fixnum fixnum fixnum fixnum)))

(: vector-scale (number-vector number -> number-vector))
(: vector-shift (number-vector number -> number-vector))

(: normalize-vector-scale (number-vector -> number-vector))
(: normalize-vector-shift (number-vector -> number-vector))

(: softmax (number-vector -> number-vector))
(: softmax* (number-vector #!optional real -> number-vector))

;;

#; ;UNUSED
(define-inline (exactess v ex)
  (cond ((exact? e) (inexact->exact v))
        (else       (exact->inexact v))) )

(define-inline (one? x) (= x 1))
(define-inline (small? x) (< x 1))

;;

;using `vector-fold' needs subvector for @st & @ed
(define (vector-reduce f id v #!optional (st 0) (ed (vector-length v)))
  (let loop ((i st) (acc id))
    (if (fx= i ed) acc
        (loop (fx+ i 1) (f acc (vector-ref v i)))) ) )

;always mixed-mode
(define (vector-sum v) (vector-reduce + 0 v))
(define (vector-prod v) (vector-reduce * 1 v))

#| ;always inexact result!
(define (vector-min v) (vector-reduce min +inf.0 v))
(define (vector-max v) (vector-reduce max -inf.0 v))
|#

(define (vector-min v)
  (let ((l (vector-length v)))
    (cond ((fx= 0 l)  +inf.0)
          ((fx= 1 l)  (vector-ref v 0))
          (else       (vector-reduce min (vector-ref v 0) v 1 l)) ) ) )

(define (vector-max v)
  (let ((l (vector-length v)))
    (cond ((fx= 0 l)  -inf.0)
          ((fx= 1 l)  (vector-ref v 0))
          (else       (vector-reduce max (vector-ref v 0) v 1 l)) ) ) )

(define (sumsqr acc elm) (+ (sqr elm) acc))

(define (absolute-magnitude v) (inexact->exact (sqrt (vector-reduce sumsqr 0 v))))

;https://en.wikipedia.org/wiki/Cosine_similarity
(define (cosine-similarity a b)
  #; ;unrolled is slightly faster
  (/ (dot-product a b)
     (* (absolute-magnitude a) (absolute-magnitude b)))
  (let ((len (vector-length a)))
    (let loop ((i 0) (dot 0) (maga 0) (magb 0))
      (if (fx= i len)
          (/ dot (inexact->exact (sqrt (* maga magb))))
          (let ((ai (vector-ref a i)) (bi (vector-ref b i)))
            (loop (fx+ i 1) (+ (* ai bi) dot) (sumsqr maga ai) (sumsqr magb bi)) ) ) ) ) )

; 1 vector only
(define-inline (*vector-apply/1 f v)
  (let* ((len (vector-length v))
         (r (make-vector len)) )
    (do ((i 0 (fx+ i 1)))
         ((fx= i len) r)
      (set! (vector-ref r i) (f (vector-ref v i))) ) ) )

; assumes all same length
(define (vector-apply f . vs)
  (let ((ln (length vs)))
    ;no data, then identity
    (if (fx= 0 ln) (f)
        (let* ((v (car vs))
               (len (vector-length v))
               (r (make-vector len)) )
          (if (= 1 ln)
              (do ((i 0 (fx+ i 1)))
                   ((fx= i len) r)
                (set! (vector-ref r i) (f (vector-ref v i))) )
              (do ((i 0 (fx+ i 1)))
                  ((fx= i len) r)
                (set! (vector-ref r i)
                      (apply f (map (cut vector-ref <> i) vs))) ) ) ) ) ) )

(define (vector-compare2 a b)
  (let* ((alen (vector-length a))
         (blen (vector-length b))
         (dlen (fx- alen blen)) )
    (if (not (fx= 0 dlen)) dlen
        (let loop ((i 0))
          (if (fx= i alen) 0
              (let ((d (- (vector-ref a i) (vector-ref b i))))
                (if (not (zero? d)) d
                    (loop (fx+ i 1)) ) ) ) ) ) ) )

(define (vector-compare . vs)
  (let loop ((vs vs))
    (if (or (null? vs) (null? (cdr vs))) 0
        (let ((d (vector-compare2 (car vs) (cadr vs))))
          (if (not (zero? d)) d
              (loop (cdr vs)) ) ) ) ) )

(define (dot-product a b)
  (let ((len (vector-length a)))
    (do ((i 0 (fx+ i 1))
         (acc 0 (+ acc (* (vector-ref a i) (vector-ref b i)))))
        ((fx= i len) acc)) ) )

;0 1 2 -> X, 3 -> #(X Y Z), 4 -> #(X Y Z W), 8 -> #(X Y Z W I J K L), * -> Error
(define (cross-product a b)
  (let ((alen (vector-length a)) (blen (vector-length b)))
    ;unequal dimension not supported
    (if (not (fx= alen blen))
        ;then problem
        (error 'cross-product "undefined for varying length" alen blen))
        ;else try
        (case alen
          ((0)
            0 )
          ((1)
            (* (vector-ref a 0) (vector-ref b 0)) )
          ((2)
            (- (* (vector-ref a 0) (vector-ref b 1))
               (* (vector-ref a 1) (vector-ref b 0))) )
          ((3)
            (vector
              ;0
              (- (* (vector-ref a 1) (vector-ref b 2))
                 (* (vector-ref a 2) (vector-ref b 1)))
              ;1
              (- (* (vector-ref a 2) (vector-ref b 0))
                 (* (vector-ref a 0) (vector-ref b 2)))
              ;2
              (- (* (vector-ref a 0) (vector-ref b 1))
                 (* (vector-ref a 1) (vector-ref b 0)))) )
          ((4)
            ;{ {a[0] * b[0]} - {a[1] * b[1]} - {a[2] * b[2]} - {a[3] * b[3]} }
            ;{ {a[0] * b[1]} + {a[1] * b[0]} + {a[2] * b[3]} - {a[3] * b[2]} }
            ;{ {a[0] * b[2]} - {a[1] * b[3]} + {a[2] * b[0]} + {a[3] * b[1]} }
            ;{ {a[0] * b[3]} + {a[1] * b[2]} - {a[2] * b[1]} + {a[3] * b[0]} }
            (vector
              ;0
              (- (* (vector-ref a 0) (vector-ref b 0))
                 (* (vector-ref a 1) (vector-ref b 1))
                 (* (vector-ref a 2) (vector-ref b 2))
                 (* (vector-ref a 3) (vector-ref b 3)))
              ;1
              (- (+ (+ (* (vector-ref a 0) (vector-ref b 1))
                       (* (vector-ref a 1) (vector-ref b 0)))
                    (* (vector-ref a 2) (vector-ref b 3)))
                 (* (vector-ref a 3) (vector-ref b 2)))
              ;2
              (+ (+ (- (* (vector-ref a 0) (vector-ref b 2))
                       (* (vector-ref a 1) (vector-ref b 3)))
                    (* (vector-ref a 2) (vector-ref b 0)))
                 (* (vector-ref a 3) (vector-ref b 1)))
              ;3
              (+ (- (+ (* (vector-ref a 0) (vector-ref b 3))
                       (* (vector-ref a 1) (vector-ref b 2)))
                    (* (vector-ref a 2) (vector-ref b 1)))
                 (* (vector-ref a 3) (vector-ref b 0)))) )
          ((8)
            ;{ {a[0] * b[0]} - {a[1] * b[1]} - {a[2] * b[2]} - {a[3] * b[3]} - {b[4] * a[4]} - {b[5] * a[5]} - {b[6] * a[6]} - {b[7] * a[7]} }
            ;{ {a[0] * b[1]} + {a[1] * b[0]} + {a[2] * b[3]} - {a[3] * b[2]} - {b[4] * a[5]} + {b[5] * a[4]} + {b[6] * a[7]} - {b[7] * a[6]} }
            ;{ {a[0] * b[2]} - {a[1] * b[3]} + {a[2] * b[0]} + {a[3] * b[1]} - {b[4] * a[6]} - {b[5] * a[7]} + {b[6] * a[4]} + {b[7] * a[5]} }
            ;{ {a[0] * b[3]} + {a[1] * b[2]} - {a[2] * b[1]} + {a[3] * b[0]} - {b[4] * a[7]} + {b[5] * a[6]} - {b[6] * a[5]} + {b[7] * a[4]} }
            ;{ {b[4] * a[0]} - {b[5] * a[1]} - {b[6] * a[2]} - {b[7] * a[3]} + {a[4] * b[0]} + {a[5] * b[1]} + {a[6] * b[2]} + {a[7] * b[3]} }
            ;{ {b[4] * a[1]} + {b[5] * a[0]} + {b[6] * a[3]} - {b[7] * a[2]} - {a[4] * b[1]} + {a[5] * b[0]} - {a[6] * b[3]} + {a[7] * b[2]} }
            ;{ {b[4] * a[2]} - {b[5] * a[3]} + {b[6] * a[0]} + {b[7] * a[1]} - {a[4] * b[2]} + {a[5] * b[3]} + {a[6] * b[0]} - {a[7] * b[1]} }
            ;{ {b[4] * a[3]} + {b[5] * a[2]} - {b[6] * a[1]} + {b[7] * a[0]} - {a[4] * b[3]} - {a[5] * b[2]} + {a[6] * b[1]} + {a[7] * b[0]} }

            (vector
              ;0
              (- (* (vector-ref a 0) (vector-ref b 0))
                 (* (vector-ref a 1) (vector-ref b 1))
                 (* (vector-ref a 2) (vector-ref b 2))
                 (* (vector-ref a 3) (vector-ref b 3))
                 (* (vector-ref b 4) (vector-ref a 4))
                 (* (vector-ref b 5) (vector-ref a 5))
                 (* (vector-ref b 6) (vector-ref a 6))
                 (* (vector-ref b 7) (vector-ref a 7)))
              ;1
              (- (+ (+ (- (- (+ (+ (* (vector-ref a 0) (vector-ref b 1))
                                   (* (vector-ref a 1) (vector-ref b 0)))
                                (* (vector-ref a 2) (vector-ref b 3)))
                             (* (vector-ref a 3) (vector-ref b 2)))
                          (* (vector-ref b 4) (vector-ref a 5)))
                       (* (vector-ref b 5) (vector-ref a 4)))
                    (* (vector-ref b 6) (vector-ref a 7)))
                 (* (vector-ref b 7) (vector-ref a 6)))
              ;2
              (+ (+ (- (- (+ (+ (- (* (vector-ref a 0) (vector-ref b 2))
                                   (* (vector-ref a 1) (vector-ref b 3)))
                                (* (vector-ref a 2) (vector-ref b 0)))
                             (* (vector-ref a 3) (vector-ref b 1)))
                          (* (vector-ref b 4) (vector-ref a 6)))
                       (* (vector-ref b 5) (vector-ref a 7)))
                    (* (vector-ref b 6) (vector-ref a 4)))
                 (* (vector-ref b 7) (vector-ref a 5)))
              ;3
              (+ (- (+ (- (+ (- (+ (* (vector-ref a 0) (vector-ref b 3))
                                   (* (vector-ref a 1) (vector-ref b 2)))
                                (* (vector-ref a 2) (vector-ref b 1)))
                             (* (vector-ref a 3) (vector-ref b 0)))
                          (* (vector-ref b 4) (vector-ref a 7)))
                       (* (vector-ref b 5) (vector-ref a 6)))
                    (* (vector-ref b 6) (vector-ref a 5)))
                 (* (vector-ref b 7) (vector-ref a 4)))
              ;4
              (+ (+ (+ (+ (- (- (- (* (vector-ref b 4) (vector-ref a 0))
                                   (* (vector-ref b 5) (vector-ref a 1)))
                                (* (vector-ref b 6) (vector-ref a 2)))
                             (* (vector-ref b 7) (vector-ref a 3)))
                          (* (vector-ref a 4) (vector-ref b 0)))
                       (* (vector-ref a 5) (vector-ref b 1)))
                    (* (vector-ref a 6) (vector-ref b 2)))
                 (* (vector-ref a 7) (vector-ref b 3)))
              ;5
              (+ (- (+ (- (- (+ (+ (* (vector-ref b 4) (vector-ref a 1))
                                   (* (vector-ref b 5) (vector-ref a 0)))
                                (* (vector-ref b 6) (vector-ref a 3)))
                             (* (vector-ref b 7) (vector-ref a 2)))
                          (* (vector-ref a 4) (vector-ref b 1)))
                       (* (vector-ref a 5) (vector-ref b 0)))
                    (* (vector-ref a 6) (vector-ref b 3)))
                 (* (vector-ref a 7) (vector-ref b 2)))
              ;6
              (- (+ (+ (- (+ (+ (- (* (vector-ref b 4) (vector-ref a 2))
                                   (* (vector-ref b 5) (vector-ref a 3)))
                                (* (vector-ref b 6) (vector-ref a 0)))
                             (* (vector-ref b 7) (vector-ref a 1)))
                          (* (vector-ref a 4) (vector-ref b 2)))
                       (* (vector-ref a 5) (vector-ref b 3)))
                    (* (vector-ref a 6) (vector-ref b 0)))
                 (* (vector-ref a 7) (vector-ref b 1)))
              ;7
              (+ (+ (- (- (+ (- (+ (* (vector-ref b 4) (vector-ref a 3))
                                   (* (vector-ref b 5) (vector-ref a 2)))
                                (* (vector-ref b 6) (vector-ref a 1)))
                             (* (vector-ref b 7) (vector-ref a 0)))
                          (* (vector-ref a 4) (vector-ref b 3)))
                       (* (vector-ref a 5) (vector-ref b 2)))
                    (* (vector-ref a 6) (vector-ref b 1)))
                 (* (vector-ref a 7) (vector-ref b 0)))) )
          (else
            (error 'cross-product "undefined for this length" alen blen) ) ) ) )

;; Normalize

(define-inline (zero~? n e) (or (zero? n) (< (abs n) e)))

;This is only useful as a wrapper around the data vector, holding the metadata
;vector
#; ; #(max xcnt min ncnt 0s +s -s)
(define (vector-info v mx? mn? 0s? +s? -s?)
  (define (vinfo mc a)
    ;maxmin
    (when mx?
      (let ((mx (vector-ref mc 0)))
        (cond ((= mx a) (set! (vector-ref mc 1) (+ (vector-ref mc 1) 1)))
              ((< mx a) (set! (vector-ref mc 0) a) (set! (vector-ref mc 1) 1))) ) )
    (when mn?
      (let ((mn (vector-ref mc 2)))
        (cond ((= mn a) (set! (vector-ref mc 3) (+ (vector-ref mc 3) 1)))
              ((> mn a) (set! (vector-ref mc 2) a) (set! (vector-ref mc 3) 1))) ) )
    ;signum
    (when (and 0s? (zero? a) (set! (vector-ref mc 4) (+ (vector-ref mc 4) 1))))
    (when (and +s? (positive? a) (set! (vector-ref mc 5) (+ (vector-ref mc 5) 1))))
    (when (and -s? (negative? a) (set! (vector-ref mc 6) (+ (vector-ref mc 6) 1))))
    mc )
  ;requires mixed-mode always
  (vector-reduce vinfo (vector -inf.0 0 +inf.0 0 0 0 0) v) )

;; simple coord algebra

;
(define (vector-scale v x) (*vector-apply/1 (cut * <> x) v))

;FIXME bad name
(define (vector-shift v x) (*vector-apply/1 (cut + <> x) v))

;scale by max - reduce range -> (... 1]
(define (normalize-vector-scale v) (vector-scale v (/ (vector-max v))))

;shift by min - bias towards 0
(define (normalize-vector-shift v) (vector-shift v (- (vector-min v))))

;; ranges -> scalers

(define-inline (clamp01 x)
  (cond ((exact? x) (if (eqv? x 1) 1 0))
        (else       (if (eqv? x 1.0) 1.0 0.0))) )

(define-inline (clamp0 x) (if (exact? x) 0 0.0))
(define-inline (clamp1 x) (if (exact? x) 1 1.0))

;; softargmax - normalized exponential function

;1 ~= summ(vmax(V)) but 1 ?= summ(V), normalize
(define (softmax v)
  (let* ((ve (*vector-apply/1 exp v))
         (sum-ve (vector-sum ve)) )
    (*vector-apply/1 (cut / <> sum-ve) ve) ) )

(define (softmax* v #!optional (t 1))
  (cond ((one? t)
          (softmax v))
        ((zero? t)      ;(zero~? t e)
          (let ((nv (normalize-vector-scale v)))
            (vector-map! (lambda (i a) (clamp01 a)) nv)
            nv ) )
        ((infinite? t)  ;(< (/ e) t)
          (vector-map (lambda (i a) (clamp0 a)) v) )
        (else
          (softmax (vector-scale v (/ t))) ) ) )

) ;(math-utils vector)