from __future__ import absolute_import, print_function

import tvm
import sys
import numpy as np
import verify

# Config
display_on = False
iter_time = 100
in_n = 20000
in_d = 784

# NOTE: This learning rate (2.42e-7) was tuned for the data generated by
# numpy with random seed 0. Change it to a suitable number if you changed
# the random seed or using your own data set.
learning_rate = 2.5e-3

if len(sys.argv) > 1:
    data_file = sys.argv[1]
else:
    data_file = None

# === Start computation
N = tvm.var('N') # Data set size
D = tvm.var('D') # Feature number

label = tvm.placeholder((N, ), name='label')
data = tvm.placeholder((N, D), name='data')
weight = tvm.placeholder((D + 1, ), name='weight')

data_expand = tvm.compute((N, D + 1), lambda n, d:
        tvm.select((d < D), data[n, d], tvm.const(1, dtype=data.dtype)),
        name='data_expand')

rd = tvm.reduce_axis((0, D + 1), name='rd')
dot = tvm.compute((N, ), lambda n:
        tvm.sum(weight[rd] * data_expand[n, rd], axis=rd),
        name='dot')

rn = tvm.reduce_axis((0, N), name='rn')
gradient = tvm.compute((D + 1, ), lambda d:
        tvm.sum(
            tvm.select((label[rn] * dot[rn] < 1.0),
                -label[rn] * data_expand[rn, d] * (1.0 / N),
                tvm.const(0, dtype='float32')),
            axis=rn),
        name='gradient')

new_weight = tvm.compute((D + 1, ), lambda d:
        weight[d] - learning_rate * gradient[d],
        name='new_weight')

# === End computation

# Scheduling
s = tvm.create_schedule(new_weight.op)

# Compilation
#print(tvm.lower(s, [data, label, weight, data_expand, dot, scale,
#    gradient, new_weight], simple_mode=True))
func = tvm.build(s, [data, label, weight, new_weight])
assert func

#print("------func code------")
#print(func.imported_modules[0].get_source())

# Generate data
np.random.seed(0)
np_data = np.random.uniform(size=(in_n, in_d), low=-1, high=3)
golden_weight = np.random.uniform(size=in_d, low=-1, high=1)
noise = np.random.uniform(size=(in_n), low=-0.2, high=0.2)
np_label = []
for n in range(in_n):
    factor = golden_weight.dot(np_data[n]) + noise[n]
    np_label.append(1 if factor > 0 else -1)
np_label = np.array(np_label)

in_label = tvm.nd.array(np_label.astype(label.dtype), tvm.cpu(0))
in_data = tvm.nd.array(np_data.astype(data.dtype), tvm.cpu(0))
in_weight = tvm.nd.array(np.zeros(in_d + 1, dtype=weight.dtype),
        tvm.cpu(0))

# Evaluation
err_rate_history = []
for i in range(iter_time):
    out_weight = tvm.nd.array(np.zeros(in_d + 1, dtype=weight.dtype),
            tvm.cpu(0))
    func(in_data, in_label, in_weight, out_weight)
    in_weight = out_weight
    out_err = tvm.nd.array(np.zeros(1, dtype='int32'), tvm.cpu(0))
    verify.calc(in_data, in_label, in_weight, out_err)
    err_rate = float(out_err.asnumpy()[0]) / in_n
    print("Round {0}, error rate {1}%".format(i, 100.0 * err_rate))
    err_rate_history.append(err_rate)

if display_on:
    print("Generating learning trend graph...")
    import matplotlib.pyplot as plt
    from matplotlib import animation
    fig = plt.figure()
    ax = plt.axes()
    ax.plot(np.arange(0, iter_time), err_rate_history)
    plt.show()

print("Finished")
