;; matrix.hs — transpose and 2D list operations.

(define
  hk-prog-val
  (fn
    (src name)
    (hk-deep-force (get (hk-eval-program (hk-core src)) name))))

(define
  hk-as-list
  (fn
    (xs)
    (cond
      ((and (list? xs) (= (first xs) "[]")) (list))
      ((and (list? xs) (= (first xs) ":"))
        (cons (nth xs 1) (hk-as-list (nth xs 2))))
      (:else xs))))

(define
  hk-mat-src
  "transpose [] = []\ntranspose ([] : _) = []\ntranspose xss = map head xss : transpose (map tail xss)\n\nmatAdd xss yss = zipWith (zipWith (+)) xss yss\n\ndiagonal [] = []\ndiagonal xss = head (head xss) : diagonal (map tail (tail xss))\n\nrowSum = map sum\ncolSum xss = map sum (transpose xss)\n")

(hk-test
  "transpose 2x2"
  (hk-deep-force
    (hk-prog-val (str hk-mat-src "r = transpose [[1,2],[3,4]]\n") "r"))
  (list
    ":"
    (list ":" 1 (list ":" 3 (list "[]")))
    (list ":" (list ":" 2 (list ":" 4 (list "[]"))) (list "[]"))))

(hk-test
  "transpose 1x3"
  (hk-deep-force
    (hk-prog-val (str hk-mat-src "r = transpose [[1,2,3]]\n") "r"))
  (list
    ":"
    (list ":" 1 (list "[]"))
    (list
      ":"
      (list ":" 2 (list "[]"))
      (list ":" (list ":" 3 (list "[]")) (list "[]")))))

(hk-test
  "transpose empty = []"
  (hk-as-list (hk-prog-val (str hk-mat-src "r = transpose []\n") "r"))
  (list))

(hk-test
  "rowSum [[1,2],[3,4]] = [3,7]"
  (hk-as-list (hk-prog-val (str hk-mat-src "r = rowSum [[1,2],[3,4]]\n") "r"))
  (list 3 7))

(hk-test
  "colSum [[1,2],[3,4]] = [4,6]"
  (hk-as-list (hk-prog-val (str hk-mat-src "r = colSum [[1,2],[3,4]]\n") "r"))
  (list 4 6))

(hk-test
  "matAdd [[1,2],[3,4]] [[5,6],[7,8]] = [[6,8],[10,12]]"
  (hk-deep-force
    (hk-prog-val
      (str hk-mat-src "r = matAdd [[1,2],[3,4]] [[5,6],[7,8]]\n")
      "r"))
  (list
    ":"
    (list ":" 6 (list ":" 8 (list "[]")))
    (list ":" (list ":" 10 (list ":" 12 (list "[]"))) (list "[]"))))

(hk-test
  "diagonal [[1,2],[3,4]] = [1,4]"
  (hk-as-list
    (hk-prog-val (str hk-mat-src "r = diagonal [[1,2],[3,4]]\n") "r"))
  (list 1 4))

(hk-test
  "diagonal 3x3"
  (hk-as-list
    (hk-prog-val
      (str hk-mat-src "r = diagonal [[1,2,3],[4,5,6],[7,8,9]]\n")
      "r"))
  (list 1 5 9))

{:fails hk-test-fails :pass hk-test-pass :fail hk-test-fail}