%load_ext autoreload
%autoreload 2

from qiskit.circuit.library.standard_gates import RXGate, RZGate, RYGate, CXGate, CZGate, SGate, HGate
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister

import numpy as np

def get_linear_entangelment_ansatz(num_of_qubits, thetas, input_state, circuit_depth=3):
    quantum_register = QuantumRegister(num_of_qubits, name="qubit")
    quantum_circuit = QuantumCircuit(quantum_register)
    quantum_circuit.initialize(input_state)

    for iteration in range(circuit_depth):
        for qubit_index in range(num_of_qubits):
            RY_theta_index = iteration*2*num_of_qubits + qubit_index
            RZ_theta_index = RY_theta_index + num_of_qubits
            
            quantum_circuit.append(RYGate(thetas[RY_theta_index]), [quantum_register[qubit_index]])
            quantum_circuit.append(RZGate(thetas[RZ_theta_index]), [quantum_register[qubit_index]])
            
        for qubit_index in range(num_of_qubits - 1):
            quantum_circuit.append(CXGate(), [quantum_register[qubit_index], quantum_register[qubit_index + 1]])

    for qubit_index in range(num_of_qubits):
            RY_theta_index = 2*num_of_qubits*circuit_depth + qubit_index
            RZ_theta_index = RY_theta_index + num_of_qubits
            quantum_circuit.append(RYGate(thetas[RY_theta_index]), [quantum_register[qubit_index]])
            quantum_circuit.append(RZGate(thetas[RZ_theta_index]), [quantum_register[qubit_index]])

    return quantum_circuit

def get_full_entangelment_ansatz(num_of_qubits, thetas, input_state, circuit_depth=3):
    quantum_register = QuantumRegister(num_of_qubits, name="qubit")
    quantum_circuit = QuantumCircuit(quantum_register)
    quantum_circuit.initialize(input_state)

    for iteration in range(circuit_depth):
        for qubit_index in range(num_of_qubits):
            RY_theta_index = iteration*2*num_of_qubits + qubit_index
            RZ_theta_index = RY_theta_index + num_of_qubits
            
            quantum_circuit.append(RYGate(thetas[RY_theta_index]), [quantum_register[qubit_index]])
            quantum_circuit.append(RZGate(thetas[RZ_theta_index]), [quantum_register[qubit_index]])
        for qubit_index in range(num_of_qubits - 1):
            for target_qubit_index in range(qubit_index + 1, num_of_qubits):
                quantum_circuit.append(CXGate(), [quantum_register[qubit_index], quantum_register[target_qubit_index]])

    for qubit_index in range(num_of_qubits):
            RY_theta_index = 2*num_of_qubits*circuit_depth + qubit_index
            RZ_theta_index = RY_theta_index + num_of_qubits
            
            quantum_circuit.append(RYGate(thetas[RY_theta_index]), [quantum_register[qubit_index]])
            quantum_circuit.append(RZGate(thetas[RZ_theta_index]), [quantum_register[qubit_index]])

    return quantum_circuit

initial_thetas = np.random.uniform(low=0, high=360, size=32)
initial_eigenvector = np.identity(16)[0]
print(get_linear_entangelment_ansatz(4, initial_thetas, 3))

initial_thetas = np.random.uniform(low=0, high=2*np.pi, size=24)
initial_eigenvector = np.identity(8)[0]
print(get_linear_entangelment_ansatz(3, initial_thetas, 3))

initial_thetas = np.random.uniform(low=0, high=2*np.pi, size=16)
initial_eigenvector = np.identity(4)[0]
print(get_linear_entangelment_ansatz(2, initial_thetas, 3))

initial_thetas = np.random.uniform(low=0, high=2*np.pi, size=32)
initial_eigenvector = np.identity(16)[0]
print(get_full_entangelment_ansatz(4, initial_thetas, 3))

initial_thetas = np.random.uniform(low=0, high=2*np.pi, size=24)
initial_eigenvector = np.identity(8)[0]
print(get_full_entangelment_ansatz(3, initial_thetas, 3))

initial_thetas = np.random.uniform(low=0, high=2*np.pi, size=16)
initial_eigenvector = np.identity(4)[0]
print(get_full_entangelment_ansatz(2, initial_thetas, 3))

