;;; Port of quaternion.scm to Chicken Scheme, by Peter Lane, 2010.

;;; Original version Copyright (c) 2001-2002, Dorai Sitaram.
;;; All rights reserved.

;;; Permission to distribute and use this work for any
;;; purpose is hereby granted provided this copyright
;;; notice is included in the copy.  This work is provided
;;; as is, with no warranty of any kind.

(module
  quaternions
  (export 
    make-rectangular
    make-polar
    real-part
    imag-part
    magnitude
    angle
    number?
    quaternion?
    =
    +
    -
    *
    /
    exp
    log
    expt
    sqrt
    sin
    cos
    tan
    asin
    acos
    atan
    jmag-part
    kmag-part
    vector-part
    colatitude
    longitude
    conjugate
    unit-vector
    dot-product
    cross-product
    )

  (import chicken 
          (except scheme ; except for those that are redefined
                  make-rectangular
                  make-polar
                  real-part
                  imag-part
                  magnitude
                  angle
                  number?
                  =
                  +
                  -
                  *
                  /
                  exp
                  log
                  expt
                  sqrt
                  sin
                  cos
                  tan
                  asin
                  acos
                  atan
                  )
          (only extras fprintf)
          ) 

  (require-extension numbers)
  (define c-make-rectangular make-rectangular)
  (define c-real-part real-part)
  (define c-imag-part imag-part)
  (define c= =)
  (define c+ +)
  (define c- -)
  (define c* *)
  (define c/ /)
  (define cexp exp)
  (define clog log)
  (define cexpt expt)
  (define csqrt sqrt)
  (define csin sin)
  (define ccos cos)
  (define catan atan)

  (define-record-type :quaternion 
                      (make-quaternion w x y z)
                      q-quaternion?
                      (w quaternion-w)
                      (x quaternion-x)
                      (y quaternion-y)
                      (z quaternion-z))

  (define-record-printer (:quaternion q out)
                         (fprintf out "~A~A~Ai~A~Aj~A~Ak"
                                  (quaternion-w q)
                                  (if (positive? (quaternion-x q)) "+" "")
                                  (quaternion-x q)
                                  (if (positive? (quaternion-y q)) "+" "")
                                  (quaternion-y q)
                                  (if (positive? (quaternion-z q)) "+" "")
                                  (quaternion-z q)))

  ;make-rectangular returns a plain real or
  ;a plain complex if some arguments are zero
  (define make-rectangular
    (lambda args
      (let ((w 0) (x 0) (y 0) (z 0)
                  (n (length args)))
        (when (> n 0)
          (set! w (list-ref args 0)))
        (when (> n 1)
          (set! x (list-ref args 1)))
        (when (> n 2)
          (set! y (list-ref args 2)))
        (when (> n 3)
          (set! z (list-ref args 3)))
        (when (> n 4)
          (error 'make-rectangular "more than 4 args supplied"))
        (if (c= y z 0) 
          (c-make-rectangular w x)
          (make-quaternion w x y z)))))

  ;real-part, etc, operate on reals and complexes too
  (define real-part 
    (lambda (q)
      (cond ((complex? q) (c-real-part q))
            ((quaternion? q) (quaternion-w q))
            (else (error 'real-part "~a not a quaternion" q)))))

  (define imag-part
    (lambda (q)
      (cond ((complex? q) (c-imag-part q))
            ((quaternion? q) (quaternion-x q))
            (else (error 'imag-part "~a not a quaternion" q)))))

  (define jmag-part
    (lambda (q)
      (cond ((complex? q) 0)
            ((quaternion? q) (quaternion-y q))
            (else (error 'jmag-part "~a not a quaternion" q)))))

  (define kmag-part
    (lambda (q)
      (cond ((complex? q) 0)
            ((quaternion? q) (quaternion-z q))
            (else (error 'kmag-part "~a not a quaternion" q)))))

  ;vector part is quaternion minus its real part
  (define vector-part 
    (lambda (q)
      (make-rectangular 0 (imag-part q) (jmag-part q) (kmag-part q))))

  ;make-polar doesn't need colatitude (phi) and
  ;longitude (psi), in which case it returns plain
  ;complex
  (define make-polar
    (lambda (mu theta . phipsi)
      (let ((phi 0) (psi 0) (n (length phipsi)))
        (when (> n 0) (set! phi (list-ref phipsi 0)))
        (when (> n 1) (set! psi (list-ref phipsi 1)))
        (when (> n 2) (error 'make-polar
                             "more than 4 args supplied"))
        (make-rectangular
          (c* mu (ccos theta))
          (c* mu (csin theta) (ccos phi))
          (c* mu (csin theta) (csin phi) (ccos psi))
          (c* mu (csin theta) (csin phi) (csin psi))))))

  (define magnitude
    (lambda (q)
      (csqrt (c+ (cexpt (real-part q) 2)
                 (cexpt (imag-part q) 2)
                 (cexpt (jmag-part q) 2)
                 (cexpt (kmag-part q) 2)))))

  (define angle
    (lambda (q)
      (catan (csqrt (c+ (cexpt (imag-part q) 2)
                        (cexpt (jmag-part q) 2)
                        (cexpt (kmag-part q) 2)))
             (real-part q))))

  (define colatitude
    (lambda (q)
      (catan (csqrt (c+ (cexpt (jmag-part q) 2)
                        (cexpt (kmag-part q) 2)))
             (imag-part q))))

  (define longitude
    (lambda (q)
      (catan (kmag-part q)
             (jmag-part q))))

  ;the new number? must succeed on quaternions too
  (define number?
    (lambda (q)
      (or (complex? q)
          (q-quaternion? q))))

  ;in the quaternion language, quaternion?
  ;should return #t for all numbers, not just
  ;the quaternion structs
  (define quaternion? number?)

  (define =
    (lambda (q . qq)
      (let ((w1 (real-part q))
            (x1 (imag-part q))
            (y1 (jmag-part q))
            (z1 (kmag-part q)))
        (let loop ((qq qq))
          (or (null? qq) 
              (let ((q (car qq)))
                (and (c= (real-part q) w1)
                     (c= (imag-part q) x1)
                     (c= (jmag-part q) y1)
                     (c= (kmag-part q) z1)
                     (loop (cdr qq)))))))))

  (define +
    (lambda qq 
      (let loop ((qq qq) (w1 0) (x1 0) (y1 0) (z1 0)) 
        (if (null? qq) (make-rectangular w1 x1 y1 z1)
          (let ((q (car qq)))
            (loop (cdr qq)
                  (c+ w1 (real-part q))
                  (c+ x1 (imag-part q))
                  (c+ y1 (jmag-part q))
                  (c+ z1 (kmag-part q))))))))

  (define -
    (lambda (q . qq)
      (let ((w1 (real-part q))
            (x1 (imag-part q))
            (y1 (jmag-part q))
            (z1 (kmag-part q)))
        (if (null? qq)
          (make-rectangular (c- w1) (c- x1) (c- y1) (c- z1))
          (let loop ((qq qq) (w1 w1) (x1 x1) (y1 y1) (z1 z1)) 
            (if (null? qq)
              (make-rectangular w1 x1 y1 z1)
              (let ((q (car qq)))
                (loop (cdr qq)
                      (c- w1 (real-part q))
                      (c- x1 (imag-part q))
                      (c- y1 (jmag-part q))
                      (c- z1 (kmag-part q))))))))))

  (define *
    (lambda qq
      (let loop ((qq qq) (w1 1) (x1 0) (y1 0) (z1 0))
        (if (null? qq)
          (make-rectangular w1 x1 y1 z1)
          (let ((q (car qq)))
            (let ((w2 (real-part q))
                  (x2 (imag-part q))
                  (y2 (jmag-part q))
                  (z2 (kmag-part q)))
              (loop (cdr qq)
                    (c- (c* w1 w2) (c* x1 x2) (c* y1 y2) (c* z1 z2))
                    (c+ (c* w1 x2) (c* x1 w2) (c* y1 z2) (c- (c* z1 y2)))
                    (c+ (c* w1 y2) (c* y1 w2) (c* z1 x2) (c- (c* x1 z2)))
                    (c+ (c* w1 z2) (c* z1 w2) (c* x1 y2) (c- (c* y1 x2))))))))))

  ;this will work on complexes too
  (define conjugate
    (lambda (q)
      (make-rectangular (real-part q)
                        (c- (imag-part q))
                        (c- (jmag-part q))
                        (c- (kmag-part q)))))

  ;note we use "right" division 
  (define /
    (lambda (q . qq)
      (if (null? qq)
        (cond ((complex? q) (c/ q))
              ((quaternion? q) 
               (let ((n (let ((m (magnitude q))) (c* m m))))
                 (make-rectangular
                   (c/ (real-part q) n)
                   (c- (c/ (imag-part q) n))
                   (c- (c/ (jmag-part q) n))
                   (c- (c/ (kmag-part q) n)))))
              (else (error '/ "~a not a quaternion" q)))
        (let loop ((qq qq) (r q)) 
          (if (null? qq) 
            r
            (let ((q (car qq)))
              (loop (cdr qq)
                    (* r (/ q)))))))))

  (define unit-vector
    (lambda (q)
      (if (real? q) (c-make-rectangular 0 1) ; 0+1i
        (let ((v (vector-part q)))
          (* v 
             (/ (magnitude v)))))))

  ;dot- and cross-product only accept 
  ;vector quaternion arguments
  (define dot-product
    (lambda (q1 q2)
      (unless (and (c= (real-part q1) 0)
                   (c= (real-part q2) 0))
        (error 'dot-product "arguments ~a, ~a, are not vector quaternions"
               q1 q2))
      (c- (real-part (* q1 q2)))))

  (define cross-product
    (lambda (q1 q2)
      (unless (and (c= (real-part q1) 0)
                   (c= (real-part q2) 0))
        (error 'dot-product "arguments ~a, ~a, are not vector quaternions"
               q1 q2))
      (vector-part (* q1 q2))))

  ;the following are all based on Maclaurin's series
  (define exp
    (lambda (q)
      (let ((w (real-part q))
            (u (unit-vector q))
            (v (magnitude (vector-part q))))
        (* (cexp w)
           (+ (ccos v) (* u (csin v)))))))

  (define log
    (lambda (q)
      (let ((w (real-part q))
            (u (unit-vector q))
            (v (magnitude (vector-part q))))
        (+ (c* 1/2 (clog (c+ (cexpt w 2) (cexpt v 2)))) 
           (* u (catan v w))))))

  (define expt
    (lambda (q1 q2)
      (exp (* (log q1) q2))))

  (define sqrt
    (lambda (q)
      (expt q 1/2)))

  ;real-valued sinh and cosh, which aren't in standard Scheme,
  ;are needed to define quat trig
  (define sinh
    (lambda (x)
      (c* 1/2 (c- (cexp x) (cexp (c- x))))))

  (define cosh
    (lambda (x)
      (c* 1/2 (c+ (cexp x) (cexp (c- x))))))

  (define sin
    (lambda (q)
      (let ((w (real-part q))
            (u (unit-vector q))
            (v (magnitude (vector-part q))))
        (+ (c* (csin w) (cosh v))
           (* u (c* (ccos w) (sinh v)))))))

  (define cos
    (lambda (q)
      (let ((w (real-part q))
            (u (unit-vector q))
            (v (magnitude (vector-part q))))
        (- (c* (ccos w) (cosh v))
           (* u (c* (csin w) (sinh v)))))))

  (define tan
    (lambda (q)
      (* (sin q) (/ (cos q)))))

  (define asin
    (lambda (q)
      (let ((u (unit-vector q)))
        (- (* u (log (+ (* u q)
                        (sqrt (- 1 (* q q))))))))))

  (define acos
    (lambda (q)
      (let ((u (unit-vector q)))
        (- (* u (log (+ q
                        (sqrt (- (* q q) 1)))))))))

  ;2-argument atan remains same as before 
  (define atan
    (lambda (q . qq)
      (let ((u (unit-vector q)))
        (if (null? qq)
          (* 1/2 u (log (* (+ u q)
                           (/ (- u q)))))
          (catan 
            ;note: args must be real here!
            q (car qq))))))

  ) ; end of library