import pennylane as qml
from pennylane import numpy as np
from sklearn import datasets

dev = qml.device("lightning.qubit", wires=2)
#dev = qml.device("default.qubit", wires=2)


@qml.qnode(dev)
def circuit(weights, features):
    qml.AngleEmbedding(features, wires=[0, 1])
    qml.StronglyEntanglingLayers(weights, wires=[0, 1])
    return qml.expval(qml.PauliZ(0))


def variational_classifier(var, features):
    weights = var[0]
    bias = var[1]
    return circuit(weights, features) + bias


def square_loss(labels, predictions):
    loss = 0
    for l, p in zip(labels, predictions):
        loss = loss + (l - p) ** 2

    loss = loss / len(labels)
    return loss


def cost(weights, features, labels):
    predictions = [variational_classifier(weights, f) for f in features]
    return square_loss(labels, predictions)


# prepare data
X, y = datasets.make_moons(noise=0.05)
y = y * 2 - np.ones(len(y))

num_qubits = 2
num_layers = 6
var_init = (np.zeros((num_layers, num_qubits, 3)), 0.0)

opt = qml.NesterovMomentumOptimizer(0.01)
batch_size = 5

# train the variational classifier
var = var_init
for it in range(20000):
    # Update the weights by one optimizer step
    batch_index = np.random.randint(0, len(y), (batch_size,))
    X_batch = X[batch_index]
    y_batch = y[batch_index]
    opt.step(lambda v: cost(v, X_batch, y_batch), var)