#lang scheme
;;; matrix.ss

;;; (struct matrix (rows cols data))
;;;   rows : exact-positive-integer?
;;;   cols : exact-positive-integer?
;;;   data : (vectorof any/c)
;;; Defines the matrix structure. We use the lower-level make-struct-type so we
;;; can write our own make-matrix.
(define-values (struct:matrix
                matrix-constructor
                matrix?
                matrix-field-ref
                matrix-field-set!)
  (make-struct-type 'matrix #f 3 0))

(define matrix-rows (make-struct-field-accessor matrix-field-ref 0))

(define matrix-cols (make-struct-field-accessor matrix-field-ref 1))

(define matrix-data (make-struct-field-accessor matrix-field-ref 2))

;;; (make-matrix rows cols v) -> matrix?
;;;   rows : exact-positive-integer?
;;;   cols : exact-positive-integer?
;;;   v : any/c = 0
;;; Returns a new rows x cols matrix with each element initialized to v, which
;;; defaults to 0.
(define (make-matrix rows cols (v 0))
  (matrix-constructor rows cols (make-vector (* rows cols) v)))

;;; (matrix-area matrix) -> exct-positive-integer?
;;;   matrix : matrix?
;;; Returns the area, which is rows x cols, for the matrix.
(define (matrix-area matrix)
  (* (matrix-rows matrix) (matrix-cols matrix)))

;;; (matrix-ref matrix i j) -> any/c
;;;   matrix : matrix?
;;;   i : exact-nonnegative-integer?
;;;   j : exact-nonnegative-integer?
;;; Returns the (i,j)th element of the matrix. An error is raised if the
;;; coordinates are outside of the matrix bounds.
(define (matrix-ref matrix i j)
  (unless (and (< i (matrix-rows matrix)) (< j (matrix-cols matrix)))
    (error 'matrix-ref
           "expected exact integers less than ~a and ~a, given ~a ~a"
           (matrix-rows matrix) (matrix-cols matrix) i j))
  (vector-ref (matrix-data matrix) (+ (* i (matrix-cols matrix)) j)))

;;; (matrix-set! matrix i j v) -> void
;;;   matrix : matrix?
;;;   i : exact-nonnegative-integer?
;;;   j : exact-nonnegative-integer?
;;;   v : any/c
;;; Sets the (i,j)th element of the matrix to v. An error is raised if the
;;; coordinates are outside of the matrix bounds.
(define (matrix-set! matrix i j v)
  (unless (and (< i (matrix-rows matrix)) (< j (matrix-cols matrix)))
    (error 'matrix-set!
           "expected exact integers less than ~a and ~a, given ~a ~a"
           (matrix-rows matrix) (matrix-cols matrix) i j))
  (vector-set! (matrix-data matrix) (+ (* i (matrix-cols matrix)) j) v))

;;; (matrix-fill! matrix v) -> void?
;;;   matrix : matrix?
;;;   v : any/c
;;; Fills the matrix with the specifief value, v.
;;; Note that the PLT Scheme documentation says that vector-fill returns void,
;;; but it doesn't. Thus the explicit call to (void) to match the contract.
(define (matrix-fill! matrix v)
  (vector-fill! (matrix-data matrix) v)
  (void))

;;; (build-matrix rows cols proc) -> matrix?
;;;   rows : exact-positive-integer?
;;;   cols : exact-positive-integer?
;;;   proc : (-> exact-nonnegative-integer? exact-nonnegative-integer? any/c)
;;; Returns a rows x cols matrix, where each element is initialized by calling
;;; (proc i j).
(define (build-matrix rows cols proc)
  (make-matrix rows cols
               (build-vector
                (* rows cols)
                (lambda (ii)
                  (let-values (((i j) (quotient/remainder ii cols)))
                    (proc i j))))))

;;; (matrix->list matrix) -> (listof (listof any/c))
;;;   matrix : matrix?
;;; Returns a list corresponding to the matrix. The list has rows sublists that
;;; are each cols long.
(define (matrix->list matrix)
  (for/list ((i (in-range (matrix-rows matrix))))
    (for/list ((j (in-range (matrix-cols matrix))))
      (matrix-ref matrix i j))))

;;; (list->matrix list) -> matrix?
;;;    list : (listof (listof any/c))
;;; Return a matrix corresponding to the list with rows equal to the length of
;;; the list and cols equal to the maximum length of the sublists.
(define (list->matrix list)
  (let* ((rows (length list))
         (cols (for/fold ((max-row-cols 0))
                         ((row (in-list list)))
                 (let ((row-cols (length row)))
                   (max row-cols max-row-cols))))
         (matrix (make-matrix rows cols)))
    (for ((row (in-list list))
          (i (in-naturals)))
      (for ((col (in-list row))
            (j (in-naturals)))
        (matrix-set! matrix i j col)))
    matrix))

;;; Module contracts

(provide/contract
 (matrix?
  (-> any/c boolean?))
 (make-matrix
  (->* (exact-positive-integer? exact-positive-integer?)
       (any/c)
       matrix?))
 (matrix-rows
  (-> matrix? exact-positive-integer?))
 (matrix-cols
  (-> matrix? exact-positive-integer?))
 (matrix-data
  (-> matrix? (vectorof any/c)))
 (matrix-area
  (-> matrix? exact-positive-integer?))
; (matrix-ref
;  (->d ((matrix matrix?)
;        (i exact-nonnegative-integer?)
;        (j exact-nonnegative-integer?))
;       ()
;       #:pre-cond (and (< i (matrix-rows matrix))
;                       (< j (matrix-cols matrix)))
;       (result any/c)))
; (matrix-ref
;  (->d ((matrix matrix?)
;        (i (and/c exact-nonnegative-integer? (</c (matrix-rows matrix))))
;        (j (and/c exact-nonnegative-integer? (</c (matrix-cols matrix)))))
;       ()
;       (result any/c)))
 (matrix-ref
  (-> matrix? exact-nonnegative-integer? exact-nonnegative-integer? any/c))
; (matrix-set!
;  (->d ((matrix matrix?)
;        (i exact-nonnegative-integer?)
;        (j exact-nonnegative-integer?)
;        (value any/c))
;       ()
;       #:pre-cond (and (< i (matrix-rows matrix))
;                       (< j (matrix-cols matrix)))
;       (result void?)))
; (matrix-set!
;  (->d ((matrix matrix?)
;        (i (and/c exact-nonnegative-integer? (</c (matrix-rows matrix))))
;        (j (and/c exact-nonnegative-integer? (</c (matrix-cols matrix))))
;        (value any/c))
;       ()
;       (result void?)))
 (matrix-set!
  (-> matrix? exact-nonnegative-integer? exact-nonnegative-integer? any/c void?))
 (matrix-fill!
  (-> matrix? any/c void?))
 (build-matrix
  (-> exact-positive-integer?
      exact-positive-integer?
      (-> exact-nonnegative-integer?
          exact-nonnegative-integer?
          any/c)
      matrix?))
 (list->matrix
  (-> (listof (listof any/c)) matrix?))
 (matrix->list
  (-> matrix? (listof (listof any/c)))))