tests.mojo

from memory.unsafe import Pointer
from python import Python

from mojograd import Value, Neuron, Layer, MLP, make_moons, plot

# Utils

fn plot_classifier_step(step: Int, inout model: MLP, inout X: Pointer[Pointer[Pointer[Value]]], n_samples: Int) raises:
    let outs = Pointer[Pointer[Value]].alloc(n_samples)

    for i in range(n_samples):
        var x = X.load(i)
        let packed = model.forward(x)
        let zero_or_one = 0 if packed.load(0).load().data.load() > 0 else 1
        let v = Value(zero_or_one)
        let ptr_v = Pointer[Value].alloc(1)
        ptr_v.store(v)

        # outs.store(i, packed.load(0))
        outs.store(i, ptr_v)

# Assets

fn asset_input(n_x: Int = 2) -> Pointer[Pointer[Value]]:
    let x = Pointer[Pointer[Value]].alloc(n_x)
    for i in range(n_x):
        let xi = Value(2.0)
        let ptr_xi = Pointer[Value].alloc(1)
        ptr_xi.store(xi)
        x.store(i, ptr_xi)

    return x

# Tests

fn test_simple_eq():
    var a = Value(2.0)
    var b = Value(3.0)
    var c: Float32 = 2.0
    var d = b**c
    var e = a + c
    e.backward()

    a.print()
    b.print()
    d.print()
    e.print()

fn test_autograd():
    var a = Value(2.0)
    let b = Value(1.0)
    var c = Value(4.0)
    let x = a + -b
    var x2 = c + c
    let x3 = x + x2
    var x4 = x3.relu()
    # var d = (a + -b + (c + c) / a * b**10 - b**2).relu()
    print(a.data.load(), b.data.load(), c.data.load(), x4.data.load())
    print(a.grad.load(), b.grad.load(), c.grad.load(), x4.grad.load())
    x4.backward()
    print(a.data.load(), b.data.load(), c.data.load(), x4.data.load())
    print(a.grad.load(), b.grad.load(), c.grad.load(), x4.grad.load())


# fn test_original_eq():
#     let a = Value(-4.0)
#     let b= Value(2.0)
#     var c = a + b
#     var d = a * b + b**3
#     c = c + c + 1.0
#     c = c + 1.0 + c + (-a)
#     d = d + d * 2.0 + (b + a).relu()
#     d = d + 3.0 * d + (b - a).relu()
#     var e = c - d
#     print(e.data.load())
#     e.backward()
#     print(a.grad.load())
#     print(b.grad.load())


fn test_neuron():
    var x = asset_input(2)

    var neuron = Neuron(2)
    let ptr_s = neuron.forward(x)
    print("s", ptr_s.load().data.load(), ptr_s.load().grad.load())
    var s = ptr_s.load()
    s.backward()
    print("s", ptr_s.load().data.load(), ptr_s.load().grad.load())
    
    for i in range(neuron.nin):
        let v = neuron.w.load(i).load()
        print("w", i, v.data.load(), v.grad.load())


fn test_layer():
    var x = asset_input(2)

    var l = Layer(2, 1)
    let res = l.forward(x)

    for i in range(l.nout):
        var v = res.load(i).load()
        print("v", i, v.data.load(), v.grad.load())
        v.backward()
        print("v", i, v.data.load(), v.grad.load())


fn test_mlp() raises:
    var x = asset_input()

    var nouts = DynamicVector[Int]()
    nouts.push_back(16)
    nouts.push_back(16)
    nouts.push_back(1)

    var m = MLP(2, nouts)

    let res = m.forward(x)

    var res_v = res.load(0).load()
    print("v", res_v.data.load(), res_v.grad.load())
    res_v.backward()
    print("v", res_v.data.load(), res_v.grad.load())

    # m.layers.load(0).load().neurons.load(0).load().parameters.load()

fn test_optmization() raises:
    ## Make datasets
    let time = Python.import_module("time")

    let n_samples = 30
    let n_dim = 2
    let out = make_moons(n_samples, 0.1)
    let X = out.get[0, Pointer[Pointer[Pointer[Value]]]]()
    let y = out.get[1, Pointer[Pointer[Value]]]()

    # print_datasets(X, y, n_samples)
    # plot(X, y, n_samples)

    ## Create MLP model

    var nouts = DynamicVector[Int]()
    nouts.push_back(16)
    nouts.push_back(16)
    nouts.push_back(1)

    var model = MLP(2, nouts)

    let num_epochs = 100
    let scores = Pointer[Pointer[Value]].alloc(n_samples)
    let losses = Pointer[Pointer[Value]].alloc(n_samples)
    for i in range(n_samples):
        let ptr_loss = Pointer[Value].alloc(1)
        losses.store(i, ptr_loss)

    for k in range(num_epochs):
        let scores = Pointer[Pointer[Value]].alloc(n_samples)
        # Forward
        for i in range(n_samples):
            var x = X.load(i)
            # For this example, it results only a "list" of one element,
            # so let's unpack it to the value
            let packed = model.forward(x)
            scores.store(i, packed.load(0))

        # SVM max-margin loss
        for i in range(n_samples):
            var scorei = scores.load(i)
            let yi = y.load(i).load()
            let prod = -yi * scorei # load?
            var one: Float32 = 1.0
            let loss = (one + prod).relu()
            let ptr_loss = losses.load(i)
            ptr_loss.store(loss)

        var sum_losses = Value(0.0)
        for i in range(n_samples):
            sum_losses = sum_losses + losses.load(i).load()
        var div: Float32 = 1.0 / n_samples
        var data_loss = sum_losses * div

        # Backward
        model.zero_grad()
        data_loss.backward()

        # SGD Update
        let params = model.parameters()
        let learning_rate = 1.0 - 0.9*k/100
        for i in range(len(params)):
            let param = params[i].load()
            let data = param.data
            data.store( data.load() - learning_rate * param.grad.load() )

        print("step", k, "loss", data_loss.data.load())

        # let outs = Pointer[Pointer[Value]].alloc(n_samples)

        # for i in range(n_samples):
        #     var x = X.load(i)
        #     let packed = model.forward(x)
        #     let zero_or_one = 0 if packed.load(0).load().data.load() > 0 else 1
        #     let v = Value(zero_or_one)
        #     let ptr_v = Pointer[Value].alloc(1)
        #     ptr_v.store(v)
        #     outs.store(i, ptr_v)

        # plot_classifier_step(k, model, X, n_samples)

fn main() raises:
    # TODO Assert expected outputs, for now only useful for dev
    # test_simple_eq()
    # test_autograd()
    # test_original_eq()
    # test_neuron()
    # test_layer()
    # test_mlp()
    test_optmization()