;    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 tuples "generics.ss"
  (require (all-except (lib "43.ss" "srfi") vector-fill! vector->list))
    
  (provide <tuple> <up> <down> 
           tuple? up? down? tuple-same-type?
           (rename my-up up) (rename my-down down) 
           up->down down->up vector->up vector->down flip
           tuple->vector tuple->list
           tuple-map tuple-fold contraction-compatible?
           same-structure/fill tuple-apply
           up-unfold down-unfold)
  
  (defclass <tuple> ()
    elts
    :auto #t
    :print #t)
  
  (define (tuple-equals? t1 t2)
    (and (or (and (up? t1) (up? t2))
             (and (down? t1) (down? t2)))
         (vector= equals? (slot-ref t1 'elts) (slot-ref t2 'elts))))
  
  (add-equals?-method <tuple> tuple-equals?)
  
  (defmethod (zero? (t <tuple>))
    (let/ec k
      (vector-for-each (lambda (i x)
                         (if (not (zero? x))
                             (k #f)))
                       (slot-ref t 'elts))
      #t))
                             
  
  (defclass <up> (<tuple>)
    :auto #t
    :print #t)
  (defclass <down> (<tuple>)
    :auto #t
    :print #t)
  
  (define (unfold-helper class fn n . starts)
    (make class 
          :elts (apply vector-unfold fn n starts)))
  
  (define (up-unfold fn n . starts)
    (apply unfold-helper <up> fn n starts))
  
  (define (down-unfold fn n . starts)
    (apply unfold-helper <down> fn n starts))
  
  (define (my-up . vals)
    (make-up (list->vector vals)))
  (define (my-down . vals)
    (make-down (list->vector vals)))
  
  (define (vector->up v)
    (make-up (vector-copy v)))
  (define (vector->down v)
    (make-down (vector-copy v)))
  
  (define (up->down t)
    (make-down (vector-copy (tuple-elts t))))
  (define (down->up t)
    (make-up (vector-copy (tuple-elts t))))
  
  (define (flip t)
    (if (up? t)
        (up->down t)
        (down->up t)))
  
  (define (tuple-same-type? t1 t2)
    (or (and (up? t1)
             (up? t2))
        (and (down? t1)
             (down? t2))))
  
  (define (tuple->vector t)
    (slot-ref t 'elts))
  
  (define (tuple->list t)
    (vector->list
     (tuple->vector t)))
  
  (defmethod (norm-squared (t <tuple>))
    (* t (flip t)))
    
  (define (tuple-map fn . tups)
    (if (not (or (every up? tups)
                 (every down? tups)))
        (raise-mismatch-error 'tuple-map "must have tuples of the same type (up or down)" tups))
    (let ((elts (map tuple-elts tups)))
      (if (and (not (null? (cdr elts)))
               (not (apply = (map vector-length elts))))
          (raise-mismatch-error 'tuple-map "must have tuples of equal lengths" tups))
      (let ((new-elts (apply vector-map fn elts)))
        (if (up? (car tups))
            (make-up new-elts)
            (make-down new-elts)))))
  
  (define (same-structure/fill tup fill)
    (tuple-map (lambda (i x)
                 (if (tuple? x)
                     (same-structure/fill x fill)
                     fill))
               tup))
  
  (define (tuple-apply fn tup)
    (apply fn (vector->list (slot-ref tup 'elts))))
  
  (define (tuple-fold fn start . tups)
    (let ((vecs (map tuple-elts tups)))
      (apply vector-fold fn start vecs)))
  
  (define tuple-ref 
    (case-lambda 
      ((t i) (vector-ref (tuple-elts t) i))
      ((t i . js) (apply tuple-ref (tuple-ref t i) js))))
  (define tuple-put!
    (case-lambda
      ((t v i) (vector-set! (tuple-elts t) i v))
      ((t v i . js) (apply tuple-put! (tuple-ref t i) v js))))
  
  (add-len-method <tuple> (lambda (t) (vector-length (tuple-elts t))))
  (add-ref-method <tuple> tuple-ref)
  (add-put!-method <tuple> tuple-put!)
  
  (define (contraction-compatible? t1 t2)
    (and (or (and (up? t1) (down? t2))
             (and (down? t1) (up? t2)))
         (= (vector-length (tuple-elts t1))
            (vector-length (tuple-elts t2)))))
  
  ;; Two conformable tuples act elementwise.  (Multiplication will be addressed later.)
  (defmethod (plus (t <tuple>))
    (tuple-map (lambda (i x) (plus x)) t))
  (defmethod (add (t1 <tuple>) (t2 <tuple>))
    (tuple-map (lambda (i x y) (add x y)) t1 t2))
  (defmethod (minus (t <tuple>))
    (tuple-map (lambda (i x) (minus x)) t))
  (defmethod (sub (t1 <tuple>) (t2 <tuple>))
    (tuple-map (lambda (i x y) (sub x y)) t1 t2))
  (defmethod (times (t <tuple>))
    (tuple-map (lambda (i x) (times x)) t))
  (defmethod (inverse (t <tuple>))
    (tuple-map (lambda (i x) (inverse x)) t))
  (defmethod (div (t1 <tuple>) (t2 <tuple>))
    (tuple-map (lambda (i x y) (div x y)) t1 t2))

  ;; Any single number is conformable to a tuple (can't remember if this is true in scmutils, too)
  (defmethod (add (t <tuple>) (x <number>))
    (tuple-map (lambda (i y) (add y x)) t))
  (defmethod (add (x <number>) (t <tuple>))
    (tuple-map (lambda (i y) (add x y)) t))
  (defmethod (sub (t <tuple>) (x <number>))
    (tuple-map (lambda (i y) (sub y x)) t))
  (defmethod (sub (x <number>) (t <tuple>))
    (tuple-map (lambda (i y) (sub x y)) t))
  (defmethod (mul (x <number>) (t <tuple>))
    (tuple-map (lambda (i y) (mul x y)) t))
  (defmethod (mul (t <tuple>) (x <number>))
    (tuple-map (lambda (i y) (mul y x)) t))
  (defmethod (div (t <tuple>) (x <number>))
    (tuple-map (lambda (i y) (div y x)) t))
  (defmethod (div (x <number>) (t <tuple>))
    (tuple-map (lambda (i y) (div x y)) t))
  
  ;; Multiplying two tuples does different things depending on whether the tuples are compatible for contration or not.
  (defmethod (mul (t1 <tuple>) (t2 <tuple>))
    (if (contraction-compatible? t1 t2)
        (let ((elts1 (tuple-elts t1))
              (elts2 (tuple-elts t2)))
          (vector-fold (lambda (i sum e1 e2) (+ sum (* e1 e2))) 0 elts1 elts2))
        (tuple-map (lambda (i x y) (mul x y)) t1 t2)))
  
  ;; The transcendental functions
  (define-syntax lift
    (syntax-rules ()
      ((_ fn)
       (defmethod (fn (t <tuple>))
         (tuple-map (lambda (i x) (fn x)) t)))))

  (lift exp)
  (lift log)
  (lift sin)
  (lift cos)
  (lift tan)
  (lift asin)
  (lift acos)
  (lift atan1)
  (lift sqrt)

  (defmethod (atan2 (t1 <tuple>) (t2 <tuple>))
    (tuple-map (lambda (i x y) (atan2 x y)) t1 t2))
  
  (defmethod (expt (t1 <tuple>) (x <number>))
    (tuple-map (lambda (i y) (expt y x)) t1))
  (defmethod (expt (x <number>) (t <tuple>))
    (tuple-map (lambda (i y) (expt x y)) t))
  (defmethod (expt (t1 <tuple>) (t2 <tuple>))
    (tuple-map (lambda (i x y) (expt x y)) t1 t2)))