;    Copyright (C) 2006  Will M. Farr <farr@mit.edu>
;
;    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 2 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, write to the Free Software Foundation, Inc.,
;    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.


(module quaternions "generics.ss"
  (require "tuples.ss")
  
  (provide <quaternion> quaternion? conjugate normalize
           vector->quaternion quaternion->vector 
           vector->quaternion-rotation quaternion-rotation->vector
           quaternion-rotation->euler-angles euler-angles->quaternion-rotation
           quaternion-rotation->rectangular-angles rectangular-angles->quaternion-rotation
           quaternion-rotate quaternion-identity-rotation)
  
  (defclass <quaternion> ()
    real-part
    imag-part
    :auto #t
    :print #t)
  
  (add-equals?-method <quaternion> class+slots-equals?)
  
  (define (vector->quaternion-rotation v)
    (let* ((theta (norm v))
           (theta/2 (/ theta 2)))
      (make <quaternion>
            :real-part (cos theta/2)
            :imag-part (* (sin theta/2)
                          (/ v theta)))))
  
  (define (quaternion-rotation->vector q)
    (* (slot-ref q 'imag-part) 2))
  
  (define (euler-angles->quaternion-rotation ea)
    (let ((phi (ref ea 0))
          (theta (ref ea 1))
          (psi (ref ea 2)))
      (let ((phi2 (/ phi 2))
            (theta2 (/ theta 2))
            (psi2 (/ psi 2)))
        (let ((cphi2 (cos phi2))
              (sphi2 (sin phi2))
              (ctheta2 (cos theta2))
              (stheta2 (sin theta2))
              (cpsi2 (cos psi2))
              (spsi2 (sin psi2)))
      (make <quaternion>
            :real-part (+ (* cphi2 ctheta2 cpsi2)
                          (* sphi2 stheta2 spsi2))
            :imag-part (up (- (* sphi2 ctheta2 cpsi2)
                              (* cphi2 stheta2 spsi2))
                           (+ (* cphi2 stheta2 cpsi2)
                              (* sphi2 ctheta2 spsi2))
                           (- (* cphi2 ctheta2 spsi2)
                              (* sphi2 stheta2 cpsi2))))))))
  
  (define (quaternion-rotation->euler-angles q)
    (let ((q0 (slot-ref q 'real-part))
          (i (slot-ref q 'imag-part)))
      (let ((q1 (ref i 0))
            (q2 (ref i 1))
            (q3 (ref i 2)))
        (up (atan (/ (* 2 (+ (* q0 q1) 
                             (* q2 q3)))
                     (- 1 (* 2 (+ (square q1)
                                  (square q2))))))
            (asin (* 2 (- (* q0 q2)
                          (* q3 q1))))
            (atan (/ (* 2 (+ (* q0 q3)
                             (* q1 q2)))
                     (- 1 (* 2 (+ (square q2)
                                  (square q3))))))))))
  
  ;; rectangular angles are given by (x,y,z) |-> Rz(z)*Ry(y)*Rx(x)
  (define (quaternion-rotation->rectangular-angles q)
    (let ((first-column (quaternion-rotate (up 1 0 0) q)))
      (let ((a (ref first-column 1))
            (b (ref first-column 2))
            (c (ref (quaternion-rotate (up 0 1 0) q) 2)))
        (let* ((y (- (asin b)))
               (c-y (cos y)))
          (up (asin (/ c c-y))
              y
              (asin (/ a c-y)))))))
          
  
  (define (rectangular-angles->quaternion-rotation r)
    (let ((x/2 (/ (ref r 0) 2))
          (y/2 (/ (ref r 1) 2))
          (z/2 (/ (ref r 2) 2)))
      (* (make <quaternion>
               :real-part (cos z/2)
               :imag-part (up 0 0 (sin z/2)))
         (make <quaternion>
               :real-part (cos y/2)
               :imag-part (up 0 (sin y/2) 0))
         (make <quaternion>
               :real-part (cos x/2)
               :imag-part (up (sin x/2) 0 0)))))
  
  (define quaternion-identity-rotation
    (make <quaternion>
          :real-part 1
          :imag-part (up 0 0 0)))
  
  (define (vector->quaternion v)
    (make <quaternion>
          :real-part 0
          :imag-part v))
  
  (define (quaternion->vector q)
    (slot-ref q 'imag-part))
  
  (define (quaternion-rotate v q)
    (quaternion->vector 
     (* q (vector->quaternion v) (conjugate q))))
  
  (defmethod (norm-squared (q <quaternion>))
    (with-slots q ((r 'real-part) (i 'imag-part))
      (+ (square r)
         (square (norm i)))))
  
  (define (normalize q)
    (let ((nq (norm q)))
      (if (zero? nq)
          (error 'normalize "trying to normalize zero vector!")
          (/ q nq))))
  
  (define (conjugate q)
    (with-slots q ((r 'real-part) (i 'imag-part))
      (make <quaternion> 
            :real-part r
            :imag-part (- i))))
  
  (defmethod (plus (q <quaternion>)) q)
  (defmethod (add (q1 <quaternion>) (q2 <quaternion>))
    (with-slots q1 ((r1 'real-part) (i1 'imag-part))
      (with-slots q2 ((r2 'real-part) (i2 'imag-part))
        (make <quaternion>
              :real-part (+ r1 r2)
              :imag-part (+ i1 i2)))))
  
  (defmethod (minus (q <quaternion>))
    (with-slots q ((r 'real-part) (i 'imag-part))
      (make <quaternion>
            :real-part (- r)
            :imag-part (- i))))
  (defmethod (sub (q1 <quaternion>) (q2 <quaternion>))
    (with-slots q1 ((r1 'real-part) (i1 'imag-part))
      (with-slots q2 ((r2 'real-part) (i2 'imag-part))
        (make <quaternion>
              :real-part (- r1 r2)
              :imag-part (- i1 i2)))))
  
  (defmethod (times (q <quaternion>)) q)
  (defmethod (mul (q <quaternion>) (x <number>))
    (with-slots q ((r 'real-part) (i 'imag-part))
      (make <quaternion>
            :real-part (* x r)
            :imag-part (* x i))))
  (defmethod (mul (x <number>) (q <quaternion>))
    (mul q x))
  (defmethod (mul (q1 <quaternion>) (q2 <quaternion>))
    (with-slots q1 ((r1 'real-part) (i1 'imag-part))
      (with-slots q2 ((r2 'real-part) (i2 'imag-part))
        (let ((x1 (ref i1 0))
              (y1 (ref i1 1))
              (z1 (ref i1 2))
              (x2 (ref i2 0))
              (y2 (ref i2 1))
              (z2 (ref i2 2)))
          (make <quaternion>
                :real-part (- (* r1 r2)
                              (* (flip i1) i2))
                :imag-part (+ (* r1 i2)
                              (* r2 i1)
                              (up (- (* y1 z2)
                                     (* z1 y2))
                                  (- (* z1 x2)
                                     (* z2 x1))
                                  (- (* x1 y2)
                                     (* x2 y1)))))))))
  
  (defmethod (div (q <quaternion>) (x <number>))
    (with-slots q ((r 'real-part) (i 'imag-part))
      (make <quaternion>
            :real-part (/ r x)
            :imag-part (/ i x)))))