;;;;;
;;;;;
;;;;; Tensor
;;;;;
;;;;;

(define $T.map
  (lambda [$fn $t]
    (tensor-map fn t)))

(define $T.map2
  (lambda [$fn $t1 $t2]
    (tensor-map2 fn t1 t2)))

(define $clear-index
  (lambda [$t]
    (| (tensor-size t)
       (tensor-to-list t) |)))

(define $unit-tensor
  (lambda [$ns]
    (generate-tensor kronecker-delta ns)))

(define $scalar-to-tensor
  (lambda [$x $ns]
    (T.map (* x $) (unit-tensor ns))))

(define $zero-tensor
  (lambda [$ns]
    (generate-tensor (cambda $ns 0) ns)))

(define $T.unit (unit-tensor $))
(define $T.zero (zero-tensor $))

;;
;; Arithmetic
;;
(define $T.arith
  (lambda [$op]
    (lambda [$t1 $t2]
      (match [(tensor? t1) (tensor? t2)] [bool bool]
        {[[,#t ,#t] (T.map2 op t1 t2)]
         [[,#t ,#f] (T.map2 op t1 (scalar-to-tensor t2 (tensor-size t1)))]
         [[,#f ,#t] (T.map2 op (scalar-to-tensor t1 (tensor-size t2)) t2)]
         }))))

(define $T.+ (T.arith +))
(define $T.- (T.arith -))

;;
;; Vectors
;;
(define $V.*
  (lambda [$v1 $v2]
    (. v1_i v2_i)))

;;
;; Matrices
;;
(define $M.*
  (cambda $ms
    (foldl M.*' (car ms) (cdr ms))))

(define $M.*'
  (lambda [$m1 $m2]
    (clear-index (. m1_i_j m2_j_k))))

(define $M.inverse
  (lambda [$m]
    (match (tensor-size m) (list integer)
      {[<cons ,2 <cons ,2 <nil>>>
        (T.map (/ $ (M.det m)) (| {2 2} {m_2_2 (* -1 m_1_2) (* -1 m_2_1) m_1_1} |))]
       [_ undefined]})))

;;
;; Linear algebra
;;
(define $M.LU
  (lambda [$x]
    (match (tensor-size x) (list integer)
      {[<cons ,2 <cons ,2 <nil>>>
        (let* {[$L (generate-tensor 2#(match (compare %1 %2) ordering {[<less> 0] [<equal> 1] [<greater> b_%1_%2]}) {2 2})]
               [$U (generate-tensor 2#(match (compare %1 %2) ordering {[<greater> 0] [_ c_%1_%2]}) {2 2})]
               [$m (M.* L U)]
               [$ret (solve {[m_1_1 x_1_1 c_1_1] [m_1_2 x_1_2 c_1_2]
                             [m_2_1 x_2_1 b_2_1] [m_2_2 x_2_2 c_2_2]})]}
          [(substitute ret L) (substitute ret U)])]
       [<cons ,3 <cons ,3 <nil>>>
        (let* {[$L (generate-tensor 2#(match (compare %1 %2) ordering {[<less> 0] [<equal> 1] [<greater> b_%1_%2]}) {3 3})]
               [$U (generate-tensor 2#(match (compare %1 %2) ordering {[<greater> 0] [_ c_%1_%2]}) {3 3})]
               [$m (M.* L U)]
               [$ret (solve {[m_1_1 x_1_1 c_1_1] [m_1_2 x_1_2 c_1_2] [m_1_3 x_1_3 c_1_3]
                             [m_2_1 x_2_1 b_2_1] [m_2_2 x_2_2 c_2_2] [m_2_3 x_2_3 c_2_3]
                             [m_3_1 x_3_1 b_3_1] [m_3_2 x_3_2 b_3_2] [m_3_3 x_3_3 c_3_3]})]}
          [(substitute ret L) (substitute ret U)])]
       [_ undefined]})))

;;
;; Determinant
;;
(define $even-and-odd-permutations
  (lambda [$n]
    (match n integer
      {[,2 [{{1 2}} {{2 1}}]]
       [_ (let* {[[$es $os] (even-and-odd-permutations (- n 1))]
                 [$es' (map 1#{@%1 n} es)]
                 [$os' (map 1#{@%1 n} os)]}
            [{@es'
              @(concat (map (lambda [$i] (map (permutate i n $) os')) (between 1 (- n 1))))
              }
             {@os'
              @(concat (map (lambda [$i] (map (permutate i n $) es')) (between 1 (- n 1))))
              }
             ]
            )]})))

(define $permutate
  (lambda [$x $y $xs]
    (match xs (list eq)
      {[<join $hs <cons ,x <join $ms <cons ,y $ts>>>>
        {@hs y @ms x @ts}]
       [<join $hs <cons ,y <join $ms <cons ,x $ts>>>>
        {@hs x @ms y @ts}]})))

(define $M.determinant
  (lambda [$m]
    (match (tensor-size m) (list integer)
      {[<cons $n <cons ,n <nil>>>
        (let {[[$es $os] (even-and-odd-permutations n)]}
          (- (sum (map (lambda [$e]
                         (product (map2 (lambda [$i $j] m_i_j)
                                        (between 1 n)
                                        e)))
                       es))
             (sum (map (lambda [$o]
                         (product (map2 (lambda [$i $j] m_i_j)
                                        (between 1 n)
                                        o)))
                       os))))]
       [_ undefined]})))

(define $M.det M.determinant)

;;;
;;; Eigenvalue
;;;
(define $M.eigenvalues
  (lambda [$m]
    (match (tensor-size m) (list integer)
      {[<cons ,2 <cons ,2 <nil>>>
        (let {[[$e1 $e2] (q-f (M.det (T.- m (scalar-to-tensor x {2 2}))) x)]}
          {e1 e2})]
       [_ undefined]})))

(define $M.eigenvectors
  (lambda [$m]
    (match (tensor-size m) (list integer)
      {[<cons ,2 <cons ,2 <nil>>>
        (let {[[$e1 $e2] (q-f (M.det (T.- m (scalar-to-tensor x {2 2}))) x)]}
          {[e1 (clear-index (T.- m (scalar-to-tensor e1 {2 2}))_i_1)]
           [e2 (clear-index (T.- m (scalar-to-tensor e2 {2 2}))_i_1)]})
        ]
       [_ undefined]})))