import numpy as np
import pennylane as qml
from pennylane.optimize import NesterovMomentumOptimizer

n = 6

backend = qml.device("default.qubit", wires = n)

def layerBlock(wgt):
    for q_i in range(n):
        qml.Rot(wgt[q_i, 0], wgt[q_i, 1], wgt[q_i, 2], wires = q_i)
    
    for q_i in range(n):
        qml.CNOT(wires = [q_i, (q_i+1)%n])

def preparation(x_n):
    qb = [i for i in range(n)]
    qml.BasisState(x_n, wires=qb)

@qml.qnode(backend)
def circuit(weights, x_n = None):
    
    preparation(x_n)
    
    for w in weights:
        layerBlock(w)
    
    return qml.expval(qml.PauliZ(0))

def variationalCircuit(param, x_n = None):
    weights = param[0]
    bias = param[1]
    
    return circuit(weights, x_n = x_n) + bias

def squareLoss(target, predictions):
    loss = 0
    for l, p in zip(target, predictions):
        loss += (l - p) ** 2
    
    loss /= len(target)
    
    return loss

def accuracy(target, predictions):
    loss = 0
    for l, p in zip(target, predictions):
        if abs(l - p) < 1e-5:
            loss += 1
    
    loss /= len(target)
    
    return loss

def cost(param, x_n, y_n):
    predictions = [variationalCircuit(param, x_n = x) for x in x_n]
    return squareLoss(y_n, predictions)

data = np.loadtxt("data/dataset1.txt")

dataset = data[:, :n]
target = data[:, n]

#maps {0, 1} -> {-1, 1}
target = target * 2 - np.ones(len(target))

np.random.seed(0)
blocks = 3
m = 4

initialParams = (0.01 * np.random.randn(blocks, n, m), 0.0)

optimizer = NesterovMomentumOptimizer(0.5)
batchSize = 5

params = initialParams
steps = 15

for i in range(steps):
    
    batchIdx = np.random.randint(0, len(dataset), (batchSize,))
    batchX = dataset[batchIdx]
    batchY = target[batchIdx]
    
    params = optimizer.step(lambda f  :cost(f, batchX, batchY), params)
    
    predictions = [np.sign(variationalCircuit(params, x)) for x in dataset]
    acc = accuracy(target, predictions)
    
    print(
        "Step: {:5d} | Cost: {:0.7f} | Accuracy: {:0.7f} ".format(
            i + 1, cost(params, dataset, target), acc
        )
    )