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

#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

# 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)))
    

# 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
# 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()

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):
    num_qubits = int(np.log2(len(counts)))
#     result = np.zeros(len(qubit_range))
    result = np.zeros(num_qubits)
    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]

# template do uczenia

def update(exp_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 = 3
hidden = 2
ancilla = visible-1

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

#definicja parametrów uczenia
num_shots = 1000
num_epochs = 100

for epoch in range(num_epochs):
    ##########################
    # 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)))
    
    
    #################################
    #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)
    exp_ph = exctract_single_qubit_measurment(counts, list(range(visible, visible + hidden))) / num_shots
#     print("\nTotal count for hidden qubits '1' measurments are:",exctract_single_qubit_measurment(counts))
    print("\nProbabilities are:",exp_ph)
#     print("\nTotal count for 0 and 1 are:",counts)

    print(weight_matrix)
    qc.draw()    
    
    update(exp_ph)
    



