matrixmultiply.swift

//
//  matrixmultiply.swift
//  AISafety
//
//  Created by Andrew Zheng (github.com/aheze) on 4/12/24.
//  Copyright © 2024 Andrew Zheng. All rights reserved.
//

import Foundation

// MARK: Structures

struct Matrix {
    var dimensions: Int

    // stored row by row
    var contents = [Int]()
}

// MARK: - Main process

print("Enter matrix dimensions: ")
guard let line = readLine(), let dimensions = Int(line) else { fatalError("Please enter an int") }

print("Enter two matrices: ")
var matrixA = Matrix(dimensions: dimensions)
var matrixB = Matrix(dimensions: dimensions)

for index in 0 ..< dimensions * 2 {
    let row = readLine()!.split(separator: " ").compactMap { Int($0) }
    if index < dimensions {
        matrixA.contents += row
    } else {
        matrixB.contents += row
    }
}

func multiply(a: Matrix, b: Matrix) -> Matrix {
    let dimensions = a.dimensions
    var result = Matrix(dimensions: a.dimensions)

    // pre-populate with dummy values
    result.contents = Array(repeating: -1, count: dimensions * dimensions)

    for i in 0 ..< dimensions {
        for j in 0 ..< dimensions {
            var sum = 0
            for k in 0 ..< dimensions {
                sum += a[i, k] * b[k, j]
            }
            result[i, j] = sum
        }
    }

    return result
}

/* Sample input:
   Enter matrix dimensions: 3
   Enter two matrixes:
   0 1 2
   3 4 5
   6 7 8
   1 0 0
   0 1 0
   0 0 1

   Sample output:
   0 1 2
   3 4 5
   6 7 8
 */
let result = multiply(a: matrixA, b: matrixB)
result.prettyPrint()

// MARK: - Extensions

extension Matrix {
    // get index inside `contents` for a r and c
    func contentsIndex(r: Int, c: Int) -> Int {
        r * dimensions + c
    }

    subscript(_ r: Int, _ c: Int) -> Int {
        get {
            let index = contentsIndex(r: r, c: c)
            return contents[index]
        } set {
            let index = contentsIndex(r: r, c: c)
            contents[index] = newValue
        }
    }

    func prettyPrint() {
        for r in 0 ..< dimensions {
            let row = (0 ..< dimensions).map { "\(self[r, $0])" }.joined(separator: " ")
            print(row)
        }
    }
}