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

# Create a numpy array of 10 rows and 5 columns.
# Set the element at row i and column j to be i+j.
def grade_Q1():
    a = np.array(range(10)).reshape(-1,1)
    b = np.array(range(5)).reshape(1,-1)
    c = a+b
    if not hasattr(c,"shape"):
        return False
    if c.shape != (10,5):
        return False
    return all([c[i,j] == i+j for i in range(10) for j in range(5)])


# Add two numpy arrays together.
def grade_Q2():
    a = np.random.rand(100, 10)
    b = np.random.rand(100, 10)
    c = a + b
    if not hasattr(c, "shape"):
        return False
    if c.shape != (100, 10):
        return False
    return all([abs(c[i, j]-a[i, j]-b[i, j]) < 1e-9 for i in range(100) for j in range(10)])


# Multiply two 2D numpy arrays using matrix multiplication.
def grade_Q3():
    a = np.array([[i for i in range(10)]])
    b = a.transpose()
    c = a.dot(b)
    d = b.dot(a)
    if not hasattr(c, "shape") or not hasattr(d, "shape"):
        return False
    if c.shape != (1, 1) or d.shape != (10, 10):
        return False
    return c[0, 0] == 285 and all([d[i, j] == i*j for i in range(10) for j in range(10)])


# For each row of a 2D numpy array, find the column index
# with the maximum element. Return all these column indices.
def grade_Q4():
    a = np.random.rand(100, 10)
    b = np.argmax(a,axis=1)
    if len(b) != 100:
        return False
    return all([a[i, b[i]] >= a[i, j] for i in range(100) for j in range(10)])


# Solve Ax = b.
def grade_Q5():
    A = -np.random.rand(100, 100)+np.diag([100]*100)
    b = np.random.rand(100, 1)
    #以矩阵的形式解方程组
    x = np.linalg.solve(A, b)
    Ax = A.dot(x)
    return np.allclose(Ax, b)


# Return an nn expression for a+b.
def grade_Q6():
    a = np.random.rand(100, 10)
    b = np.random.rand(100, 10)
    expr_add_ab = nn.Input('a') + nn.Input('b')
    c = expr_add_ab.compile(golden.Builder())(a = a, b = b)
    return np.allclose(c, a+b)


# Return an nn expression for a+b*c.
def grade_Q7():
    a = np.random.rand(100, 50)
    b = np.random.rand(100, 10)
    c = np.random.rand(10, 50)        
    expr_add_mul_abc =  nn.Input('a') +  (nn.Input('b') * nn.Input('c'))
    d = expr_add_mul_abc.compile(golden.Builder())(a = a, b = b, c = c)
    return np.allclose(d, a+(b.dot(c)))


# Given A and b, return an nn expression for Ax+b.
def grade_Q8():
    A = np.random.rand(100, 100)
    x = np.random.rand(100, 1)
    b = np.random.rand(100, 1)
    expr_add_Axb = nn.Const(A)*nn.Input('x') + nn.Const(b)
    y = expr_add_Axb.compile(golden.Builder())(x = x)
    return np.allclose(y, A.dot(x)+b)


# Given n, return an nn expression for x**n.
def Q9(n):
    x = nn.Input('x')
    expression = nn.Const(1)
    while n:
        if n & 1:
            expression = expression * x
        x = x * x
        n >>= 1
        #n = n/2
    return expression

def grade_Q9():
    x = np.random.rand(100, 100)
    y = Q9(8).compile(golden.Builder())(x = x)
    return np.allclose(y, (((x.dot(x)).dot(x)).dot(x)).dot(((x.dot(x)).dot(x)).dot(x)))



# Return an nn expression to compute
# the element-wise absolute value |x|.
def Q10():
    x = nn.Input('x')
    relu = nn.ReLU()
    a = relu(x)
    b = relu(-x)
    return a+b

def grade_Q10():
    x = np.random.randn(100, 100)
    y = Q10().compile(golden.Builder())(x = x)
    return np.allclose(y, np.abs(x))



grade = 0

for q in range(1, 11):
    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 functions passed: %d" % grade)