from singa import tensor, device, layer

#help(layer.Layer)
layer.engine='singacpp'

"""
Explanation: SINGA Model Classes
<img src="http://singa.apache.org/en/_static/images/singav1-sw.png" width="500px"/>
Layer
Typically, the life cycle of a layer instance includes:
 1. construct layer without input_sample_shapes, goto 2; or,
construct layer with input_sample_shapes, goto 3;


call setup to create the parameters and setup other meta fields;
initialize the parameters of the layer
call forward or access layer members
call backward and get parameters for update
End of explanation
"""


from singa.layer import Dense, Conv2D, MaxPooling2D, Activation, BatchNormalization, Softmax

"""
Explanation: Common layers
End of explanation
"""


dense = Dense('dense', 3, input_sample_shape=(2,))
#dense.param_names()
w, b = dense.param_values()
print(w.shape, b.shape)

w.gaussian(0, 0.1)
b.set_value(0)

x  = tensor.Tensor((2,2))
x.uniform(-1, 1)
y = dense.forward(True, x)
tensor.to_numpy(y)


gx, [gw, gb] = dense.backward(True, y)
print(gx.shape, gw.shape, gb.shape)

"""
Explanation: Dense Layer
End of explanation
"""


conv = Conv2D('conv', 4, 3, 1, input_sample_shape=(3, 6, 6))
print(conv.get_output_sample_shape())

"""
Explanation: Convolution Layer
End of explanation
"""


pool = MaxPooling2D('pool', 3, 2, input_sample_shape=(4, 6, 6))
print(pool.get_output_sample_shape())

"""
Explanation: Pooling Layer
End of explanation
"""


from singa.layer import Split, Merge, Slice, Concat

split = Split('split', 2, input_sample_shape=(4, 6, 6))
print(split.get_output_sample_shape())

merge = Merge('merge', input_sample_shape=(4, 6, 6))
print(merge.get_output_sample_shape())

sli = Slice('slice', 1, [2], input_sample_shape=(4, 6, 6))
print(sli.get_output_sample_shape())

concat = Concat('concat', 1, input_sample_shapes=[(3, 6, 6), (1, 6, 6)])
print(concat.get_output_sample_shape())

"""
Explanation: Branch layers
End of explanation
"""


from singa import metric
import numpy as np

x = tensor.Tensor((3, 5))
x.uniform(0, 1)  # randomly genearte the prediction activation
x = tensor.softmax(x)  # normalize the prediction into probabilities
print(tensor.to_numpy(x))
y = tensor.from_numpy(np.array([0, 1, 3], dtype=np.int))  # set the truth

f = metric.Accuracy()
acc = f.evaluate(x, y)  # averaged accuracy over all 3 samples in x
print(acc)

from singa import loss

x = tensor.Tensor((3, 5))
x.uniform(0, 1)  # randomly genearte the prediction activation
y = tensor.from_numpy(np.array([0, 1, 3], dtype=np.int))  # set the truth

f = loss.SoftmaxCrossEntropy()
l = f.forward(True, x, y)  # l is tensor with 3 loss values
g = f.backward()  # g is a tensor containing all gradients of x w.r.t l
print(l.l1())
print(tensor.to_numpy(g))

"""
Explanation: Metric and Loss
End of explanation
"""


from singa import optimizer

sgd = optimizer.SGD(lr=0.01, momentum=0.9, weight_decay=1e-4)
p = tensor.Tensor((3,5))
p.uniform(-1, 1)
g = tensor.Tensor((3,5))
g.gaussian(0, 0.01)

sgd.apply(1, g, p, 'param')  # use the global lr=0.1 for epoch 1
sgd.apply_with_lr(2, 0.03, g, p, 'param')  # use lr=0.03 for epoch 2

"""
Explanation: Optimizer
End of explanation
"""


from singa import net as ffnet
layer.engine = 'singacpp'
net = ffnet.FeedForwardNet(loss.SoftmaxCrossEntropy(), metric.Accuracy())
net.add(layer.Conv2D('conv1', 32, 5, 1, input_sample_shape=(3,32,32,)))
net.add(layer.Activation('relu1'))
net.add(layer.MaxPooling2D('pool1', 3, 2))
net.add(layer.Flatten('flat'))
net.add(layer.Dense('dense', 10))

# init parameters
for p in net.param_values():
    if len(p.shape) == 0:
        p.set_value(0)
    else:
        p.gaussian(0, 0.01)
print(net.param_names())

layer.engine = 'cudnn'
net = ffnet.FeedForwardNet(loss.SoftmaxCrossEntropy(), metric.Accuracy())
net.add(layer.Conv2D('conv1', 32, 5, 1, input_sample_shape=(3,32,32,)))
net.add(layer.Activation('relu1'))
net.add(layer.MaxPooling2D('pool1', 3, 2))
net.add(layer.Flatten('flat'))
net.add(layer.Dense('dense', 10))

# init parameters
for p in net.param_values():
    if len(p.shape) == 0:
        p.set_value(0)
    else:
        p.gaussian(0, 0.01)
        
# move net onto gpu
dev = device.create_cuda_gpu()
net.to_device(dev)

"""
Explanation: FeedForwardNet
End of explanation
"""


