%install '.package(path: "$cwd/FastaiNotebook_00_load_data")' FastaiNotebook_00_load_data

import FastaiNotebook_00_load_data

import Path
import TensorFlow

let (xTrain, yTrain, xValid, yValid) = loadMNIST(path: mnistPath, flat: true)

//export
import Python

//export
public let np = Python.import("numpy")
public let plt = Python.import("matplotlib.pyplot")

%include "EnableIPythonDisplay.swift"
IPythonDisplay.shell.enable_matplotlib("inline")

let img = np.array(xTrain[0].array.scalars).reshape(28,28)

plt.figure(figsize: [5,5])
let fig = plt.imshow(X:img, cmap:"gray")
plt.show(fig)

var weights = Tensor<Float>(randomNormal: [784, 10]) / sqrt(784)
var bias = Tensor(repeating: 0.0, shape: [10])

let (arc,brc) = ((5, 784), (784, 10))

func swiftMatmul(a:[Float], b:[Float], arc:(Int,Int), brc:(Int,Int)) -> [Float] {
    var res:[Float] = Array(repeating: 0.0, count: arc.0*brc.1)
    for i in 0..<arc.0 {
        for j in 0..<brc.1 {
            for k in 0..<arc.1 {
                res[i+arc.0*j] += a[i+arc.0*k] * b[k+brc.0*j]
            }
        }
    }
    return res
}

let flatA = xTrain[0..<5].array.scalars
let flatB = weights.array.scalars

let res = swiftMatmul(a: flatA, b:flatB, arc:arc, brc:brc)

time(repeating: 100) {
    let _ = swiftMatmul(a: flatA, b:flatB, arc:arc, brc:brc) }

func TensorMatmul(_ a:Tensor<Float>, _ b:Tensor<Float>) -> Tensor<Float>{
    var res:Tensor<Float> = Tensor(repeating: 0.0, shape: [a.shape[0], b.shape[1]])
    for i in 0..<a.shape[0]{
        for j in 0..<b.shape[1]{
            for k in 0..<a.shape[1]{
                res[i][j] += a[i][k] * b[k][j]
            }
        }
    }
    return res
}

let m1 = Tensor<Float>(randomNormal: [5, 784])
let m2 = Tensor<Float>(randomNormal: [784, 10])

time() { let res = TensorMatmul(m1, m2)}

var a = Tensor([10.0, 6, -4])
var b = Tensor([2.0, 8, 7])
(a,b)

a + b

a < b

a .< b

var m = Tensor([1.0, 2, 3, 4, 5, 6, 7, 8, 9]).reshaped(to: [3,3])
m

2 * m

sqrt((m * m).sum().scalar!)

func elementWiseMatmul(_ a:Tensor<Float>, _ b:Tensor<Float>) -> Tensor<Float>{
    let (ar,ac) = (a.shape[0],a.shape[1])
    let (br,bc) = (b.shape[0],b.shape[1])
    var res = Tensor<Float>(zeros: [ac, br])
    for i in 0..<ar {
        for j in 0..<bc {
            res[i][j] = (a[i] * b.slice(lowerBounds: [0,j], upperBounds: [ac,j+1]).squeezingShape(at: 1)).sum()
        }
    }
    return res
}

let res = elementWiseMatmul(m1, m2)

time() { let _ = elementWiseMatmul(m1, m2)}

var a = Tensor([10.0, 6.0, -4.0])

print(a .> 0)

print((a .> 0).all())

print((a .> 0).any())

print(a+1)

2 * m

let c = Tensor([10.0,20.0,30.0])

m + c

c + m

m + c.expandingShape(at: 1)

c.expandingShape(at: 1)

func broadcastMatmult(_ a:Tensor<Float>, _ b:Tensor<Float>) -> Tensor<Float>{
    var res = Tensor<Float>(zeros: [a.shape[0], b.shape[1]])
    for i in 0..<a.shape[0]{
        res[i] = (a[i].expandingShape(at: 1) * b).sum(squeezingAxes: 0)
    }
    return res
}

let res = broadcastMatmult(m1, m2)

time(repeating: 100) { let _ = broadcastMatmult(m1, m2)}

c.expandingShape(at: 0).shape

c.expandingShape(at: 1).shape

c.expandingShape(at: 0) * c.expandingShape(at: 1)

c.expandingShape(at: 0) .> c.expandingShape(at: 1)

time(repeating: 100) { let _ = matmul(m1, m2)}

notebookToScript(fname: (Path.cwd / "01_matmul.ipynb").string)