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

//export
import Path
import TensorFlow

import FastaiNotebook_00_load_data

let zeros = Tensor<Float>(zeros: [1,4,5])
let ones  = Tensor<Float>(ones: [12,4,5])
let twos  = Tensor<Float>(repeating: 2.0, shape: [2,3,4,5])
let range = Tensor<Int32>(rangeFrom: 0, to: 32, stride: 1)

let xTrain = Tensor<Float>(randomNormal: [5, 784])
var weights = Tensor<Float>(randomNormal: [784, 10]) / sqrt(784)
print(weights[0])

// a and b are the flattened array elements, aDims/bDims are the #rows/columns of the arrays.
func swiftMatmul(a: [Float], b: [Float], aDims: (Int,Int), bDims: (Int,Int)) -> [Float] {
    assert(aDims.1 == bDims.0, "matmul shape mismatch")
    
    var res = Array(repeating: Float(0.0), count: aDims.0 * bDims.1)
    for i in 0 ..< aDims.0 {
        for j in 0 ..< bDims.1 {
            for k in 0 ..< aDims.1 {
                res[i*bDims.1+j] += a[i*aDims.1+k] * b[k*bDims.1+j]
            }
        }
    }
    return res
}

let flatA = xTrain[0..<5].scalars
let flatB = weights.scalars
let (aDims,bDims) = ((5, 784), (784, 10))

var resultArray = swiftMatmul(a: flatA, b: flatB, aDims: aDims, bDims: bDims)

time(repeating: 100) {
    _ = swiftMatmul(a: flatA, b: flatB, aDims: aDims, bDims: bDims)
}

// a and b are the flattened array elements, aDims/bDims are the #rows/columns of the arrays.
func swiftMatmulUnsafe(a: UnsafePointer<Float>, b: UnsafePointer<Float>, aDims: (Int,Int), bDims: (Int,Int)) -> [Float] {
    assert(aDims.1 == bDims.0, "matmul shape mismatch")
    
    var res = Array(repeating: Float(0.0), count: aDims.0 * bDims.1)
    res.withUnsafeMutableBufferPointer { res in 
        for i in 0 ..< aDims.0 {
            for j in 0 ..< bDims.1 {
                for k in 0 ..< aDims.1 {
                    res[i*bDims.1+j] += a[i*aDims.1+k] * b[k*bDims.1+j]
                }
            }
        }
    }
    return res
}

time(repeating: 100) {
    _ = swiftMatmulUnsafe(a: flatA, b: flatB, aDims: aDims, bDims: bDims)
}

import Glibc

let ptr : UnsafeMutableRawPointer = malloc(42)

print("☠️☠️ Uninitialized garbage =", ptr.load(as: UInt8.self))

free(ptr)

import Python
let np = Python.import("numpy")
let pickle = Python.import("pickle")
let sys = Python.import("sys")

print("🐍list =    ", pickle.dumps([1, 2, 3]))
print("🐍ndarray = ", pickle.dumps(np.array([[1, 2], [3, 4]])))


var bias = Tensor<Float>(zeros: [10])

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

print("m1: ", m1.shape)
print("m2: ", m2.shape)

let small = Tensor<Float>([[1, 2],
                           [3, 4]])

print("🔢2x2:\n", small)

print("⊞ matmul:\n",  matmul(small, small))
print("\n⊞ again:\n", small • small)


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

sqrt((m * m).sum())

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

print("add:  ", a + b)
print("mul:  ", a * b)
print("sqrt: ", sqrt(a))
print("pow:  ", pow(a, b))

a < b

a .< b

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

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

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

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)

print(c.expandingShape(at: 0).shape)

print(c.expandingShape(at: 1).shape)

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

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

func tensorMatmul(_ 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] {
        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
}

_ = tensorMatmul(m1, m2)

time { 
    let tmp = tensorMatmul(m1, m2)
    
    // Copy a scalar back to the host to force a GPU sync.
    _ = tmp[0, 0].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 {
        let row = a[i]
        for j in 0 ..< bc {
            res[i, j] = (row * b.slice(lowerBounds: [0,j], upperBounds: [ac,j+1]).squeezingShape(at: 1)).sum()
        }
    }
    return res
}

_ = elementWiseMatmul(m1, m2)

time { 
    let tmp = elementWiseMatmul(m1, m2)

    // Copy a scalar back to the host to force a GPU sync.
    _ = tmp[0, 0].scalar
}

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
}

_ = broadcastMatmult(m1, m2)

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

    // Copy a scalar back to the host to force a GPU sync.
    _ = tmp[0, 0].scalar
}

time(repeating: 100) { _ = m1 • m2 }

func timeMatmulTensor(size: Int) {
    var matrix = Tensor<Float>(randomNormal: [size, size])
    print("\n\(size)x\(size):\n  ⏰", terminator: "")
    time(repeating: 10) { 
        let matrix = matrix • matrix 
        _ = matrix[0, 0].scalar
    }
}

timeMatmulTensor(size: 1)     // Tiny
timeMatmulTensor(size: 10)    // Bigger
timeMatmulTensor(size: 100)   // Even Bigger
timeMatmulTensor(size: 1000)  // Biggerest
timeMatmulTensor(size: 5000)  // Even Biggerest

func timeMatmulSwift(size: Int, repetitions: Int = 10) {
    var matrix = Tensor<Float>(randomNormal: [size, size])
    let matrixFlatArray = matrix.scalars

    print("\n\(size)x\(size):\n  ⏰", terminator: "")
    time(repeating: repetitions) { 
       _ = swiftMatmulUnsafe(a: matrixFlatArray, b: matrixFlatArray, aDims: (size,size), bDims: (size,size))
    }
}

timeMatmulSwift(size: 1)     // Tiny
timeMatmulSwift(size: 10)    // Bigger
timeMatmulSwift(size: 100)   // Even Bigger
timeMatmulSwift(size: 1000, repetitions: 1)  // Biggerest

print("\n5000x5000: skipped, it takes tooo long!")

withDevice(.cpu) {
    timeMatmulTensor(size: 1)     // Tiny
    timeMatmulTensor(size: 10)    // Bigger
    timeMatmulTensor(size: 100)   // Even Bigger
    timeMatmulTensor(size: 1000)  // Biggerest
    timeMatmulTensor(size: 5000)  // Even Biggerest
}

// Explore the contents of the Raw namespace by typing Raw.<tab>
print(Raw.zerosLike(c))

// Raw.

//export
public extension StringTensor {
    // Read a file into a Tensor.
    init(readFile filename: String) {
        self.init(readFile: StringTensor(filename))
    }
    init(readFile filename: StringTensor) {
        self = Raw.readFile(filename: filename)
    }

    // Decode a StringTensor holding a JPEG file into a Tensor<UInt8>.
    func decodeJpeg(channels: Int = 0) -> Tensor<UInt8> {
        return Raw.decodeJpeg(contents: self, channels: Int64(channels), dctMethod: "") 
    }
}


import NotebookExport
let exporter = NotebookExport(Path.cwd/"01_matmul.ipynb")
print(exporter.export(usingPrefix: "FastaiNotebook_"))