from qiskit import IBMQ
# IBMQ.save_account(MY_API_TOKEN)
import qiskit
qiskit.__version__


import scipy
import numpy as np
import random

from sklearn import preprocessing
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.quantum_info.operators import Operator
from qiskit import(QuantumCircuit, execute, Aer)
from qiskit.visualization import plot_histogram
from qiskit.extensions import Initialize # Import the Inititialize function
from qiskit.aqua.circuits.gates import multi_control_toffoli_gate
from qiskit.aqua.circuits.gates import multi_control_multi_target_gate

# JUŻ NIEPOTRZEBNE - TUTAJ BYŁY HIPERPARAMETRY OBWODU TESOWEGO
# #definicja rozmiaru sieci
# visible = 1
# hidden = 1
# ancilla = visible-1

# #definicja wejścia (x) oraz inicjalizacja macierzy wag
# x = np.array([random.uniform(0, 1) for n in range(visible)])
# weight_matrix = np.random.rand(visible, hidden) * np.pi

# JUŻ NIEPOTRZEBNE - TUTAJ BYŁ BUDOWANY OBWÓD TESTOWY
# # definicja obwodu (linii qubitów)
# # qc = QuantumCircuit(visible + hidden, visible)

# qr = QuantumRegister(visible + hidden + ancilla, 'q')
# cr = ClassicalRegister(visible + hidden, 'c')
# qc = QuantumCircuit(qr, cr)

# # inicjalizacja wartości qubitów wejściowych (x) na razie randomowa
# initial_state = [[np.sqrt(x[i]), np.sqrt(1-x[i])] for i in range(len(x))]

# # inicjalizacja wartości qubitów wejściowych i bramka Hadamarda
# for i in range(visible):
#     initialize_qubit = Initialize(initial_state[i])
#     qc.append(initialize_qubit, [i])
#     qc.h(i)
    
# # ciąg bramek CNOT i bramek rotacji R (zależnych od parametrów)
# for i in range(hidden):
#     for j in range(visible):
#         qc.rz(weight_matrix[j][i], j)
#     multi_control_toffoli_gate.mct(qc, [qr[k] for k in range(visible)], qr[visible + i], [qr[i] for i in range(visible + hidden, visible + hidden + ancilla)], mode='basic')
    
    
    
# # bramki Hadamarda na warstwie ukrytej
# for i in range(hidden):
#     qc.h(visible + i)

    
# # pomiar linii visible
# qc.measure(list(range(visible, visible+hidden)), list(range(visible, visible+hidden)))


# ### OBWÓD DRUGĄ STRONĘ

# # wyzerowanie wartości qubitów visible
# for i in range(visible):
#     initialize_qubit = Initialize([1, 0])
#     qc.append(initialize_qubit, [i])
    
# # bramka Hadamarda na qubitach hidden
# for i in range(visible, hidden + visible):
#     qc.h(i)
    
# #odwrócone rotacje i CNOTy
# for i in range(hidden):
#     for j in range(visible):
#         multi_control_toffoli_gate.mct(qc, [qr[k] for k in range(visible, visible + hidden)], qr[j], [qr[k] for k in range(visible + hidden, visible + hidden + ancilla)], mode='basic')
#         qc.rz(-weight_matrix[j][hidden - 1 - i], j)
        
# # bramka Hadamarda i pomiar na qubitach visible
# for i in range(visible):
#     qc.h(i)
#     qc.measure(list(range(visible)), list(range(visible)))
    

# EKSPERYMNET Z TESTOWYM UKŁADEM - JUŻ NIEPOTRZEBNY
# # Use Aer's qasm_simulator
# simulator = Aer.get_backend('qasm_simulator')

# # Execute the circuit on the qasm simulator
# job = execute(qc, simulator, shots=1000)

# # Grab results from the job
# result = job.result()

# # Returns counts
# counts = result.get_counts(qc)
# print("\nTotal count for 0 and 1 are:",counts)
# # qc.draw(output='mpl')
# qc.draw()

# # komórka do testowania wieloqubitowej bramki CNOT - JUŻ NIEPOTRZEBNA BO BRAMKA DZIAŁA
# # NIE ZALEŻNA OD WŁAŚCIWEJ SIECI

# # work = ancilla
# def n_control_gate(qc, target, control: list, work: list, gate_type):
#     if len(control) - len(work) != 1:
#         raise Exception("Wrong number of control or work qubits!")

#     qc.toffoli(control[0], control[1], work[0])
#     control_left = control[2:]
    
#     for i, c_bit in enumerate(control_left):
#         qc.toffoli(c_bit, work[i], work[i+1])
    
#     qc.gate_type(work[-1], target)
    
#     for i, c_bit in reversed(list(enumerate(control_left))):
#         qc.toffoli(c_bit, work[i], work[i+1])
    
#     qc.toffoli(control[0], control[1], work[0])
    
    

        

# visible = 2
# hidden = 2
# ancilla = visible - 1

# # qr_visible = QuantumRegister(visible, 'visible')
# # qr_hidden = QuantumRegister(hidden, 'hidden')
# # cr = ClassicalRegister(3, 'c')
# # if visible > 2:
# #     anc = QuantumRegister(visible - 2, 'ancilla')
# #     qc = QuantumCircuit(qr_visible, qr_hidden, anc, cr)
# # else:
# #     qc = QuantumCircuit(qr_visible, qr_hidden, cr)


# # qr_visible = QuantumRegister(visible, 'visible')
# qr = QuantumRegister(visible + hidden + ancilla, 'q')
# cr = ClassicalRegister(3, 'c')
# qc = QuantumCircuit(qr, cr)


# for i in range(visible + hidden):
#     initialize_qubit = Initialize([0, 1])
#     qc.append(initialize_qubit, [i])
# # initialize_qubit = Initialize([1, 0])
# # qc.append(initialize_qubit, [2])
# # for i in range(hidden):
# #     for j in range(visible):
# #         qc.rz(weight_matrix[j][i], j)
# print(list(range(visible)) + [visible + 0])
# print(list(range(visible+hidden, visible+hidden+ancilla)))
# # multi_control_multi_target_gate.mcmt(qc, list(range(visible)), [], QuantumCircuit.cx, [visible])

# # print([qr_visible[i] for i in range(visible)])
# # print([anc[i] for i in range(ancilla)])
# # print([qr_hidden[i] for i in range(hidden)])
# for j in range(hidden):
#     multi_control_multi_target_gate.mcmt(qc, [qr[i] for i in range(visible)], [qr[i] for i in range(visible + hidden, visible + hidden + ancilla)], QuantumCircuit.cx, [qr[visible + j]])

# # multi_control_multi_target_gate.mcmt(qc, [qr_visible[i] for i in range(visible)], [anc[i] for i in range(ancilla)], QuantumCircuit.cx, [qr_hidden[0]])
# # qc.measure(list(range(visible+hidden)), list(range(visible + hidden)))
# # (qc, list(range(visible)) + [visible + 0], qancilla = 6)


# # Use Aer's qasm_simulator
# simulator = Aer.get_backend('qasm_simulator')

# # Execute the circuit on the qasm simulator
# job = execute(qc, simulator, shots=1000)

# # Grab results from the job
# result = job.result()

# # Returns counts
# counts = result.get_counts(qc)
# print("\nTotal count for 0 and 1 are:",counts)
# qc.draw()

# TO JUŻ NIEPOTRZEBNE - NIE DZIAŁA TEN SPOSÓB Z ARTYKUŁU

# print(weight_matrix, '\n')
# sinus_matrix = np.sin(weight_matrix + np.pi/4)
# print(sinus_matrix, '\n')
# print(sinus_matrix.prod(axis = 0), '\n')
# phi = np.arcsin(sinus_matrix.prod(axis = 0))
# print(phi, '\n')
# psi = phi + np.pi/4
# print(psi, '\n')
# ph = np.sin(psi)**2
# print(ph, '\n')
# # ph_normalized = preprocessing.normalize([ph], norm = 'l1')[0]
# # print(ph_normalized)

def exctract_single_qubit_measurment(dict_of_counts, qubit_range):
#     print(dict_of_counts)
#     print(len(list(dict_of_counts.keys())[0]))
    num_qubits = len(list(dict_of_counts.keys())[0])
#     result = np.zeros(len(qubit_range))
    result = np.zeros(num_qubits)
#     print(result)
    for el in dict_of_counts:
        for i in range(num_qubits):
#             print("i", i)
#             print("el[i]", el[i])
            if i in qubit_range and el[i] == '1':
                result[i] += dict_of_counts[el]
#     print(result)
#     print(result[qubit_range])
    return result[qubit_range]

def create_dataset(dataset_size):
    dataset = []
    for i in range(dataset_size):
        x = np.array([random.choice([0,1]), random.choice([0,1])])
        y = 1
        if np.array_equal(x, np.array([0,0])):
            y = 0
        dataset.append({"x": x, "y": y})
    return dataset

dataset = create_dataset(4)
print(dataset)

def circuit_function(qc, weight_matrix):
    qc.data = []

    # inicjalizacja wartości qubitów wejściowych (x)
    initial_state = [[np.sqrt(1-x[i]), np.sqrt(x[i])] for i in range(len(x))]

    # inicjalizacja wartości qubitów wejściowych i bramka Hadamarda
    for i in range(visible):
        initialize_qubit = Initialize(initial_state[i])
        qc.append(initialize_qubit, [i])
        qc.h(i)

    # ciąg bramek CNOT i bramek rotacji R (zależnych od parametrów)
    for i in range(hidden):
        for j in range(visible):
            qc.ry(weight_matrix[j][i], j)
        multi_control_toffoli_gate.mct(qc, [qr[k] for k in range(visible)], qr[visible + i], [qr[i] for i in range(visible + hidden, visible + hidden + ancilla)], mode='basic')

#     bramki Hadamarda na warstwie ukrytej
#     for i in range(hidden):
#         qc.h(visible + i)

    # pomiar linii visible
    qc.measure(list(range(visible, visible+hidden)), list(range(hidden)))
    
    #eksperyment:
    #symylator
    simulator = Aer.get_backend('qasm_simulator')

    # Execute the circuit on the qasm simulator
    job = execute(qc, simulator, shots=num_shots)

    # Grab results from the job
    result = job.result()

    # Returns counts
    counts = result.get_counts(qc)
#     print(counts)
    ph = exctract_single_qubit_measurment(counts, list(range(hidden))) / num_shots
#     print("\nProbabilities are:",ph)
    return ph[0]


epsilon = 0.01
def update2(ph, expected_ph, weight_matrix, lr):
    gradient = []
    for i, row in enumerate(weight_matrix):
        gradient_row = []
        for j, el in enumerate(row):
            weight_matrix[i][j] += epsilon
            result_plus = circuit_function(qc, weight_matrix)

            weight_matrix[i][j] -= 2*epsilon
            result_minus = circuit_function(qc, weight_matrix)

            weight_matrix[i][j] += epsilon
#             result_zero = circuit_function(qc, weight_matrix)
#             print("ph", result_zero)
#             print("exp_ph", expected_ph)
#             print("el final:", el)
#             print(result_plus - result_minus)
            result = (result_plus - result_minus)/(2*epsilon) * lr * (ph - expected_ph)
            gradient_row.append(result)
        gradient.append(gradient_row)
#     print("gradient", gradient)

    weight_matrix = weight_matrix - gradient
    return weight_matrix
    
    
# update2()

# template do uczenia
import matplotlib.pyplot as plt
%matplotlib inline

def update(exp_ph, ph):
    global weight_matrix
#     #obliczanie prawdopodobieństwa p(|h>) ze wzoru
#     sinus_matrix = np.sin(weight_matrix + np.pi/4)
#     phi = np.arcsin(sinus_matrix.prod(axis = 0))
# #     print('phi', phi)
#     psi = phi + np.pi/4
# #     print('psi', psi)
#     ph = np.sin(psi)**2
# #     print('ph', ph)

#obliczanie gradientu i,k
    S = 1/(np.sqrt(1-(np.sin(phi))**2))
#     print("S", S)
    T = np.sin(phi) * np.sum(1/np.tan(weight_matrix + np.pi/4), axis = 0)
#     print("T", T)
#     print("sum products", (exp_ph - ph) * ph * np.cos(2*psi) * S * T)
    gradient = np.zeros(hidden)
    for k in range(hidden):
        gradient[k] = np.sum( (exp_ph - ph) * ph * np.cos(2*psi) * S * T[k])
#     print("gradient = ", gradient)
#     print([gradient] * len(weight_matrix))
    weight_matrix = weight_matrix + ([gradient] * len(weight_matrix))



#definicja rozmiaru sieci
visible = 2
hidden = 1
ancilla = visible-1

#definicja wejścia (x)oraz inicjalizacja macierzy wag
# x = np.array([random.uniform(0, 1) for n in range(visible)])

dataset = create_dataset(10)
print(dataset[0]["x"][0])
print(dataset[0]["x"][1])
print([n for n in range(visible)])


weight_matrix = np.random.rand(visible, hidden) * np.pi

#definicja parametrów uczenia
num_shots = 1000
num_epochs = 100
qr = QuantumRegister(visible + hidden + ancilla, 'q')
cr = ClassicalRegister(hidden, 'c')
qc = QuantumCircuit(qr, cr)

cost_function_data = []
lr = 0.05

for epoch in range(num_epochs):
    print("epoch: ", epoch)
    for i, element in enumerate(dataset):
#         print(element)
        x = np.array([dataset[i]["x"][n] for n in range(visible)])
        exp_ph = dataset[i]["y"]
        ph = circuit_function(qc, weight_matrix)
        weight_matrix = update2(ph, exp_ph, weight_matrix, lr)
#         print("exp_ph", exp_ph, "ph", ph, "weight_matrix", weight_matrix, "cost_function", 0.5 * (ph - exp_ph)**2)   
    cost_function_data.append(0.5 * (ph - exp_ph)**2)
qc.draw()
plt.xlabel('number of epochs')
plt.ylabel('cost')
plt.plot(cost_function_data)


qc.draw()



