;;; Sly
;;; Copyright (C) 2014 David Thompson <davet@gnu.org>
;;;
;;; This program 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.
;;;
;;; This program is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;; General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with this program.  If not, see
;;; <http://www.gnu.org/licenses/>.

;;; Commentary:
;;
;; Vector math.
;;
;;; Code:

(define-module (sly math vector)
  #:use-module (ice-9 match)
  #:use-module (rnrs bytevectors)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-9)
  #:use-module (sly math)
  #:use-module (sly records)
  #:export (vector2 vector3 vector4
            vector2? vector3? vector4?
            vx vy vz vw
            vector2-x vector2-y
            vector3-x vector3-y vector3-z
            vector4-x vector4-y vector4-z vector4-w
            vmap v+ v- v* vdot vcross
            magnitude normalize vlerp))

(define-packed-f64-record-type <vector2>
  vector2
  bytevector->vector2 vector2->bytevector
  vector2?
  (x 0 vector2-x set-vector2-x!)
  (y 1 vector2-y set-vector2-y!))

(define-packed-f64-record-type <vector3>
  vector3
  bytevector->vector3 vector3->bytevector
  vector3?
  (x 0 vector3-x set-vector3-x!)
  (y 1 vector3-y set-vector3-y!)
  (z 2 vector3-z set-vector3-z!))

(define-packed-f64-record-type <vector4>
  vector4
  bytevector->vector4 vector4->bytevector
  vector4?
  (x 0 vector4-x set-vector4-x!)
  (y 1 vector4-y set-vector4-y!)
  (z 2 vector4-z set-vector4-z!)
  (w 3 vector4-w set-vector4-w!))

(define-inlinable (vx v)
  (cond
   ((vector2? v) (vector2-x v))
   ((vector3? v) (vector3-x v))
   ((vector4? v) (vector4-x v))))

(define-inlinable (vy v)
  (cond
   ((vector2? v) (vector2-y v))
   ((vector3? v) (vector3-y v))
   ((vector4? v) (vector4-y v))))

(define-inlinable (vz v)
  (cond
   ((vector3? v) (vector3-z v))
   ((vector4? v) (vector4-z v))))

(define-inlinable (vw v)
  (vector4-w v))

(define (vmap proc v)
  "Return a new vector that is the result of applying PROC to each
element of the 2D/3D/4D vector V."
  (match v
    (($ <vector2> x y)
     (vector2 (proc x) (proc y)))
    (($ <vector3> x y z)
     (vector3 (proc x) (proc y) (proc z)))
    (($ <vector4> x y z w)
     (vector4 (proc x) (proc y) (proc z) (proc w)))))

;; Hoo boy, the things we do for efficiency. ;)
(define-syntax-rule (vector-arithmetic vectors op identity)
  (match vectors
    ;; Common cases: Adding just 2 vectors of the same type.
    ;; Matching against them here means avoiding the more expensive,
    ;; more general loops later on.
    (((and (? vector2?) (= vector2->bytevector bv1))
      (and (? vector2?) (= vector2->bytevector bv2)))
     (bytevector->vector2
      (f64vector (op (f64vector-ref bv1 0)
                     (f64vector-ref bv2 0))
                 (op (f64vector-ref bv1 1)
                     (f64vector-ref bv2 1)))))
    (((and (? vector3?) (= vector3->bytevector bv1))
      (and (? vector3?) (= vector3->bytevector bv2)))
     (bytevector->vector3
      (f64vector (op (f64vector-ref bv1 0)
                     (f64vector-ref bv2 0))
                 (op (f64vector-ref bv1 1)
                     (f64vector-ref bv2 1))
                 (op (f64vector-ref bv1 2)
                     (f64vector-ref bv2 2)))))
    (((and (? vector4?) (= vector4->bytevector bv1))
      (and (? vector4?) (= vector4->bytevector bv2)))
     (bytevector->vector4
      (f64vector (op (f64vector-ref bv1 0)
                     (f64vector-ref bv2 0))
                 (op (f64vector-ref bv1 1)
                     (f64vector-ref bv2 1))
                 (op (f64vector-ref bv1 2)
                     (f64vector-ref bv2 2))
                 (op (f64vector-ref bv1 3)
                     (f64vector-ref bv2 3)))))
    ;; Special cases for a list with a a single element, to handle use
    ;; with subtraction.
    (((and (? vector2?) (= vector2->bytevector head)))
     (vector2 (op (f64vector-ref head 0))
              (op (f64vector-ref head 1))))
    (((and (? vector3?) (= vector3->bytevector head)))
     (vector3 (op (f64vector-ref head 0))
              (op (f64vector-ref head 1))
              (op (f64vector-ref head 2))))
    (((and (? vector4?) (= vector4->bytevector head)))
     (vector4 (op (f64vector-ref head 0))
              (op (f64vector-ref head 1))
              (op (f64vector-ref head 2))
              (op (f64vector-ref head 3))))
    (((? number? x))
     (op x))
    ;; General case:
    (vectors*
     (let outer ((scalar-sum #f)
                 (vectors* vectors*))
       (match vectors*
         ;; First, add up all of the scalars that appear at the head of
         ;; the list, before we've been able to determine which vector
         ;; type to specialize on.
         (() (or scalar-sum identity))
         (((? number? x) . tail)
          (outer (if scalar-sum (op scalar-sum x) x) tail))
         ;; Specialize based on vector type once we actually encounter a
         ;; vector.
         ;;
         ;; 2D vectors, possibly mixed with scalars:
         (((and (? vector2?) (= vector2->bytevector head)) . tail)
          (let ((bv (if scalar-sum
                        (f64vector (op scalar-sum
                                       (f64vector-ref head 0))
                                   (op scalar-sum
                                       (f64vector-ref head 1)))
                        (bytevector-copy head))))
            (let inner ((vectors* tail))
              (match vectors*
                (() (bytevector->vector2 bv))
                (((? number? x) . tail)
                 (f64vector-set! bv 0 (op (f64vector-ref bv 0) x))
                 (f64vector-set! bv 1 (op (f64vector-ref bv 1) x))
                 (inner tail))
                (((and (? vector2?) (= vector2->bytevector head)) . tail)
                 (f64vector-set! bv 0 (op (f64vector-ref bv 0)
                                          (f64vector-ref head 0)))
                 (f64vector-set! bv 1 (op (f64vector-ref bv 1)
                                          (f64vector-ref head 1)))
                 (inner tail))))))
         ;; 3D vectors, possibly mixed with scalars:
         (((and (? vector3?) (= vector3->bytevector head)) . tail)
          (let ((bv (if scalar-sum
                        (f64vector (op scalar-sum
                                       (f64vector-ref head 0))
                                   (op scalar-sum
                                       (f64vector-ref head 1))
                                   (op scalar-sum
                                       (f64vector-ref head 2)))
                        (bytevector-copy head))))
            (let inner ((vectors* tail))
              (match vectors*
                (() (bytevector->vector3 bv))
                (((? number? x) . tail)
                 (f64vector-set! bv 0 (op (f64vector-ref bv 0) x))
                 (f64vector-set! bv 1 (op (f64vector-ref bv 1) x))
                 (f64vector-set! bv 2 (op (f64vector-ref bv 2) x))
                 (inner tail))
                (((and (? vector3?) (= vector3->bytevector head)) . tail)
                 (f64vector-set! bv 0 (op (f64vector-ref bv 0)
                                          (f64vector-ref head 0)))
                 (f64vector-set! bv 1 (op (f64vector-ref bv 1)
                                          (f64vector-ref head 1)))
                 (f64vector-set! bv 2 (op (f64vector-ref bv 2)
                                          (f64vector-ref head 2)))
                 (inner tail))))))
         ;; 4D vectors, possibly mixed with scalars:
         (((and (? vector4?) (= vector4->bytevector head)) . tail)
          (let ((bv (if scalar-sum
                        (f64vector (op scalar-sum
                                       (f64vector-ref head 0))
                                   (op scalar-sum
                                       (f64vector-ref head 1))
                                   (op scalar-sum
                                       (f64vector-ref head 2))
                                   (op scalar-sum
                                       (f64vector-ref head 3)))
                        (bytevector-copy head))))
            (let inner ((vectors* tail))
              (match vectors*
                (() (bytevector->vector4 bv))
                (((? number? x) . tail)
                 (f64vector-set! bv 0 (op (f64vector-ref bv 0) x))
                 (f64vector-set! bv 1 (op (f64vector-ref bv 1) x))
                 (f64vector-set! bv 2 (op (f64vector-ref bv 2) x))
                 (f64vector-set! bv 3 (op (f64vector-ref bv 3) x))
                 (inner tail))
                (((and (? vector4?) (= vector4->bytevector head)) . tail)
                 (f64vector-set! bv 0 (op (f64vector-ref bv 0)
                                          (f64vector-ref head 0)))
                 (f64vector-set! bv 1 (op (f64vector-ref bv 1)
                                          (f64vector-ref head 1)))
                 (f64vector-set! bv 2 (op (f64vector-ref bv 2)
                                          (f64vector-ref head 2)))
                 (f64vector-set! bv 3 (op (f64vector-ref bv 3)
                                          (f64vector-ref head 3)))
                 (inner tail)))))))))))

(define (v+ . vectors)
  "Compute the sum of VECTORS."
  (vector-arithmetic vectors + 0))

(define (v* . vectors)
  "Compute the product of VECTORS."
  (vector-arithmetic vectors * 1))

(define (v- vectors . rest)
  "Compute the difference of VECTORS."
  (vector-arithmetic (cons vectors rest) - 0))

(define (vdot v1 v2)
  "Compute the dot product of the vectors V1 and V2."
  (cond
   ((and (vector2? v1) (vector2? v2))
    (+ (* (vector2-x v1) (vector2-x v2))
       (* (vector2-y v1) (vector2-y v2))))
   ((and (vector3? v1) (vector3? v2))
    (+ (* (vector3-x v1) (vector3-x v2))
       (* (vector3-y v1) (vector3-y v2))
       (* (vector3-z v1) (vector3-z v2))))
   ((and (vector4? v1) (vector4? v2))
    (+ (* (vector4-x v1) (vector4-x v2))
       (* (vector4-y v1) (vector4-y v2))
       (* (vector4-z v1) (vector4-z v2))
       (* (vector4-w v1) (vector4-w v2))))))

(define (vcross v1 v2)
  "Compute the cross product of the 3D vectors V1 and V2."
  (vector3 (- (* (vector3-y v1) (vector3-z v2))
              (* (vector3-z v1) (vector3-y v2)))
           (- (* (vector3-z v1) (vector3-x v2))
              (* (vector3-x v1) (vector3-z v2)))
           (- (* (vector3-x v1) (vector3-y v2))
              (* (vector3-y v1) (vector3-x v2)))))

(define (magnitude v)
  "Return the magnitude of the vector V."
  (sqrt
   (cond
    ((vector2? v)
     (+ (square (vector2-x v))
        (square (vector2-y v))))
    ((vector3? v)
     (+ (square (vector3-x v))
        (square (vector3-y v))
        (square (vector3-z v))))
    ((vector4? v)
     (+ (square (vector4-x v))
        (square (vector4-y v))
        (square (vector4-z v))
        (square (vector4-w v)))))))

(define (normalize v)
  "Return the normalized form of the vector V."
  (let ((m (magnitude v)))
    (cond
     ((zero? m) v)
     ((vector2? v)
      (vector2 (/ (vector2-x v) m)
               (/ (vector2-y v) m)))
     ((vector3? v)
      (vector3 (/ (vector3-x v) m)
               (/ (vector3-y v) m)
               (/ (vector3-z v) m)))
     ((vector4? v)
      (vector4 (/ (vector4-x v) m)
               (/ (vector4-y v) m)
               (/ (vector4-z v) m)
               (/ (vector4-w v) m))))))

(define vlerp (make-lerp v+ v*))