#|
  This file is a part of stoe project.
  Copyright (c) 2015 Renaud Casenave-Péré (renaud@casenave-pere.fr)
|#

(uiop:define-package :stoe/maths/matrix
    (:use :cl :alexandria
          :stoe/core/utils
          :stoe/maths/types
          :stoe/maths/vector)
  (:import-from :stoe/maths/vector
                #:fill-vector #:make-displaced-vector)
  (:export #:mref
           #:mat-null #:mat-id
           #:make-matrix
           #:mat #:mat2 #:mat3 #:mat4
           #:transpose
           #:m+ #:m- #:m*))
(in-package :stoe/maths/matrix)

(defmethod element-type ((m matrix)) (array-element-type (raw-data m)))

(defun mref (m &rest subscripts)
  (let ((len (length subscripts))
        (dim-x (first (dimensions m)))
        (dim-y (second (dimensions m))))
    (assert (< len 3))
    (case len
      (2 (progn
           (assert (< (first subscripts) dim-x))
           (assert (< (second subscripts) dim-y))
           (aref (slot-value m 'array) (+ (* (first subscripts) dim-y)
                                          (second subscripts)))))
      (1 (progn
           (assert (< (first subscripts) dim-x))
           (make-displaced-vector (slot-value m 'array)
                                  (* (first subscripts) dim-y)
                                  dim-y))))))

(defun set-mref (m &rest subscripts-or-value)
  (let ((len (length subscripts-or-value)))
    (assert (< len 4))
    (case len
      (3 (setf (aref (slot-value m 'array) (+ (* (first subscripts-or-value)
                                                 (second (dimensions m)))
                                              (second subscripts-or-value)))
               (third subscripts-or-value)))
      (2 (let* ((dim (second (dimensions m)))
                (offset (* (first subscripts-or-value) dim))
                (v (second subscripts-or-value)))
           (assert (= dim (dimensions v)))
           (loop for i below dim
              do (setf (aref (slot-value m 'array) (+ i offset)) (vref v i))
              finally (return (second subscripts-or-value))))))))

(defsetf mref set-mref)

(defun matrix-type (dim-x dim-y type)
  (if (/= dim-x dim-y)
      'matrix
      (case dim-x
        (2 (case type (single-float 'float22) (fixnum 'int22) (otherwise 'matrix)))
        (3 (case type (single-float 'float33) (fixnum 'int33) (otherwise 'matrix)))
        (4 (case type (single-float 'float44) (fixnum 'int44) (otherwise 'matrix)))
        (otherwise 'matrix))))

(defun mat-null (dims type)
  (let ((dim-x (if (listp dims) (first dims) dims))
        (dim-y (if (listp dims) (second dims) dims)))
    (make-instance (matrix-type dim-x dim-y type)
                   :dims (list dim-x dim-y)
                   :array (make-array (* dim-x dim-y) :element-type type))))

(defun mat-id (dim type)
  (let ((m (mat-null dim type)))
    (loop for i below dim
       do (setf (mref m i i) (coerce 1 type)))
    m))

(defun clone-matrix (dims type mat)
  (let ((m (mat-null dims type)))
    (loop for i below (first (dimensions mat))
       do (loop for j below (second (dimensions mat))
             do (setf (mref m i j) (mref mat i j))))
    m))

(defun make-matrix (dims type &rest attribs)
  (let* ((m (mat-null dims type))
         (v (make-displaced-vector (slot-value m 'array))))
    (loop with i = 0
       for attr in attribs
       do (setf i (fill-vector v attr i)))
    m))

(defmacro mat (&rest attribs)
  (once-only ((attrib (first attribs)))
    (if (= (length attribs) 1)
         `(clone-matrix (dimensions ,attrib) (element-type ,attrib) ,attrib)
         (let ((dim (list '+ 0)) type)
           (loop for attr in attribs
              do (progn
                   (unless type
                     (setf type (cond
                                  ((floatp attr) 'single-float)
                                  ((integerp attr) 'fixnum))))
                   (if (numberp attr)
                       (setf (cadr dim) (1+ (cadr dim)))
                       (setf dim (append dim (list `(dimensions ,attr)))))))
           `(let* ((len ,dim)
                   (dim-x (floor (sqrt len)))
                   (dim-y (if (= (* dim-x dim-x) len) dim-x (/ len dim-x))))
              (make-matrix (list dim-x dim-y)
                           ,(if type
                                `',type
                                `(element-type ,attrib))
                           ,@attribs))))))

(defmacro mat2 (&rest attribs)
  (once-only ((attrib (first attribs)))
    (if (= (length attribs) 1)
      `(clone-matrix '(2 2) (element-type ,attrib) ,attrib)
      (let (type)
        (loop for attr in attribs
           do (unless type
                (setf type (cond
                             ((floatp attr) 'single-float)
                             ((integerp attr) 'fixnum)))))
        `(make-matrix '(2 2)
                      ,(if type
                           `',type
                           `(element-type ,attrib))
                      ,@attribs)))))

(defmacro mat3 (&rest attribs)
  (once-only ((attrib (first attribs)))
    (if (= (length attribs) 1)
      `(clone-matrix '(3 3) (element-type ,attrib) ,attrib)
      (let (type)
        (loop for attr in attribs
           do (unless type
                (setf type (cond
                             ((floatp attr) 'single-float)
                             ((integerp attr) 'fixnum)))))
        `(make-matrix '(3 3)
                      ,(if type
                           `',type
                           `(element-type ,attrib))
                      ,@attribs)))))

(defmacro mat4 (&rest attribs)
  (once-only ((attrib (first attribs)))
    (if (= (length attribs) 1)
      `(clone-matrix '(4 4) (element-type ,attrib) ,attrib)
      (let (type)
        (loop for attr in attribs
           do (unless type
                (setf type (cond
                             ((floatp attr) 'single-float)
                             ((integerp attr) 'fixnum)))))
        `(make-matrix '(4 4)
                      ,(if type
                           `',type
                           `(element-type ,attrib))
                      ,@attribs)))))

(defgeneric transpose (m))
(defmethod transpose ((m matrix))
  (let* ((dim-x (first (dimensions m)))
         (dim-y (second (dimensions m)))
         (transposed (mat-null (list dim-y dim-x) (element-type m))))
    (loop for i below dim-x
       do (loop for j below dim-y
             do (setf (mref transposed j i) (mref m i j))))
    transposed))

(defgeneric madd (m1 m2))
(defmethod madd ((m1 matrix) (m2 matrix))
  (let ((dim (dimensions m1))
        (type (element-type m1)))
    (assert (equal dim (dimensions m2)))
    (assert (eq type (element-type m2)))
    (let ((mat (mat-null dim (element-type m1))))
      (loop for i below (apply #'* dim)
         do (setf (aref (slot-value mat 'array) i)
                  (+ (aref (slot-value m1 'array) i)
                     (aref (slot-value m2 'array) i))))
      mat)))

(defun m+ (&rest m-list)
  (reduce #'madd m-list))

(defgeneric msub (m1 m2))
(defmethod msub ((m1 matrix) (m2 matrix))
  (let ((dim (dimensions m1))
        (type (element-type m1)))
    (assert (equal dim (dimensions m2)))
    (assert (eq type (element-type m2)))
    (let ((mat (mat-null dim (element-type m1))))
      (loop for i below (apply #'* dim)
         do (setf (aref (slot-value mat 'array) i)
                  (- (aref (slot-value m1 'array) i)
                     (aref (slot-value m2 'array) i))))
      mat)))

(defun m- (&rest m-list)
  (reduce #'msub m-list))

(defgeneric mmul (m1 m2))
(defmethod mmul ((m1 matrix) (m2 number))
  (let* ((dim (dimensions m1))
         (mat (mat-null dim (element-type m1))))
    (loop for i below (apply #'* dim)
       do (setf (aref (slot-value mat 'array) i)
                (* (aref (slot-value m1 'array) i) m2)))
    mat))

(defmethod mmul ((m1 number) (m2 matrix))
  (mmul m2 m1))

(defmethod mmul ((m1 matrix) (m2 matrix))
  (let ((dim-1 (dimensions m1))
        (dim-2 (dimensions m2))
        (type (element-type m1)))
    (assert (= (first dim-1) (second dim-2)))
    (assert (eq type (element-type m2)))
    (let ((mat (mat-null (list (first dim-2) (second dim-1)) type)))
      (loop for i below (first dim-2)
         do (loop for j below (second dim-1)
               do (setf (mref mat i j)
                        (loop for k below (first dim-1)
                           for l below (second dim-2)
                           sum (* (mref m1 k j) (mref m2 i l))))))
      mat)))

(defmethod mmul ((m1 matrix) (m2 vect))
  (let* ((dim-1 (dimensions m1))
         (dim-2 (dimensions m2))
         (type (element-type m1)))
    (assert (= (first dim-1) dim-2))
    (assert (eq type (element-type m2)))
    (let ((vec (vec-null (second dim-1) type)))
      (loop for i below (second dim-1)
         do (setf (vref vec i)
                  (loop for j below (second dim-1)
                     sum (* (mref m1 j i) (vref m2 j)))))
      vec)))

(defmethod mmul ((m1 vect) (m2 matrix))
  (let* ((dim-1 (dimensions m1))
         (dim-2 (dimensions m2))
         (type (element-type m1)))
    (assert (= dim-1 (second dim-2)))
    (assert (eq type (element-type m2)))
    (let ((vec (vec-null (first dim-2) type)))
      (loop for i below (first dim-2)
         do (setf (vref vec i)
                  (loop for j below (second dim-2)
                     sum (* (vref m1 j) (mref m2 i j)))))
      vec)))

(defun m* (&rest mat-list)
  (reduce #'mmul mat-list))