import traceback
import numpy as np
import easynn as nn
import easynn_golden as golden
import easynn_cpp as cpp

def random_kwargs(kwargs):
    return {k: np.random.random(shape) if shape != None else np.random.random() for k, shape in kwargs.items()}

def is_same(p, n, **kwargs):
    e0 = p.compile(golden.Builder())
    e1 = p.compile(cpp.Builder())
    nkwargs = [random_kwargs(kwargs) for i in range(n)]
    return all([np.allclose(e0(**nkwargs[i]), e1(**nkwargs[i])) for i in range(n)])


def grade_Q1():
    relu = nn.ReLU()
    x = relu(nn.Input("x"))
    return is_same(x, 1, x = (10, 11, 12, 13))

def grade_Q2():
    flatten = nn.Flatten()
    x = flatten(nn.Input("x"))
    return is_same(x, 1, x = (10, 11, 12, 13))

def grade_Q3():
    x = nn.Input2d("images", 10, 11, 3)
    return is_same(x, 1, images = (50, 10, 11, 3))

def grade_Q4():
    f = nn.Linear("f", 100, 10)
    x = f(nn.Input("x"))
    x.resolve({
        "f.weight": np.random.random((10, 100)),
        "f.bias": np.random.random((10,))})
    return is_same(x, 1, x = (50, 100))

def grade_Q5():
    relu = nn.ReLU()
    flatten = nn.Flatten()
    f = nn.Linear("f", 16*32*3, 8)
    x = nn.Input2d("images", 16, 32, 3)
    x = relu(f(flatten(x)))
    x.resolve({
        "f.weight": np.random.random((8, 16*32*3)),
        "f.bias": np.random.random((8,))})
    return is_same(x, 1, images = (20, 16, 32, 3))

def grade_Q6():
    relu = nn.ReLU()
    flatten = nn.Flatten()
    f1 = nn.Linear("f1", 28*28, 100)
    f2 = nn.Linear("f2", 100, 10)
    x = nn.Input2d("images", 28, 28, 1)
    x = flatten(x)
    x = f2(relu(f1(x)))
    x.resolve({
        "f1.weight": np.random.random((100, 28*28)),
        "f1.bias": np.random.random((100,)),
        "f2.weight": np.random.random((10, 100)),
        "f2.bias": np.random.random((10))})
    return is_same(x, 1, images = (20, 28, 28, 1))

def grade_Q7():
    relu = nn.ReLU()
    flatten = nn.Flatten()
    f1 = nn.Linear("f1", 28*28, 100)
    f2 = nn.Linear("f2", 100, 10)
    x = nn.Input2d("images", 28, 28, 1)
    x = flatten(x)
    x = f2(relu(f1(x)))
    x.resolve(np.load("msimple_params.npz"))
    mnist_test = np.load("mnist_test.npz")
    images = mnist_test["images"][:20]

    infer0 = x.compile(golden.Builder())
    infer1 = x.compile(cpp.Builder())
    logit0 = infer0(images = images)
    logit1 = infer1(images = images)
    return np.allclose(logit0, logit1)

def grade_Q8():
    relu = nn.ReLU()
    flatten = nn.Flatten()
    f1 = nn.Linear("f1", 28*28, 100)
    f2 = nn.Linear("f2", 100, 10)
    x = nn.Input2d("images", 28, 28, 1)
    x = flatten(x)
    x = f2(relu(f1(x)))
    x.resolve(np.load("msimple_params.npz"))
    mnist_test = np.load("mnist_test.npz")
    images = mnist_test["images"]

    infer0 = x.compile(golden.Builder())
    infer1 = x.compile(cpp.Builder())
    label0 = infer0(images = images).argmax(axis = 1)
    label1 = infer1(images = images).argmax(axis = 1)
    return np.allclose(label0, label1)

grade = 0
for q in range(1, 9):
    func = globals()["grade_Q%d" % q]
    try:
        if func():
            grade += 1
        else:
            print("============Q%d failed!============\n" % q)
    except Exception as e:
        print("============Q%d failed!============" % q)
        print(traceback.format_exc())

print("Total questions passed: %d" % grade)