%load_ext autoreload
%autoreload 2

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

def get_thetas_circuit(thetas, D2):
    qr = QuantumRegister(4, name="qubit")
    qc = QuantumCircuit(qr)
        
    for d in range(D2):
        qc.append(RXGate(thetas[0]), [qr[0]])
        qc.append(RXGate(thetas[1]), [qr[1]])
        qc.append(RXGate(thetas[2]), [qr[2]])
        qc.append(RXGate(thetas[3]), [qr[3]])
        
        qc.append(RZGate(thetas[4]), [qr[0]])
        qc.append(RZGate(thetas[5]), [qr[1]])
        qc.append(RZGate(thetas[6]), [qr[2]])
        qc.append(RZGate(thetas[7]), [qr[3]])
        
        qc.append(CZGate(), [qr[0], qr[1]])
        qc.append(CZGate(), [qr[1], qr[2]])
        qc.append(CZGate(), [qr[2], qr[3]])
        qc.barrier(qr)
    
    qc.append(RXGate(thetas[0]), [qr[0]])
    qc.append(RXGate(thetas[1]), [qr[1]])
    qc.append(RXGate(thetas[2]), [qr[2]])
    qc.append(RXGate(thetas[3]), [qr[3]])

    qc.append(RZGate(thetas[4]), [qr[0]])
    qc.append(RZGate(thetas[5]), [qr[1]])
    qc.append(RZGate(thetas[6]), [qr[2]])
    qc.append(RZGate(thetas[7]), [qr[3]])

    return qc

def get_phis_circuit(phis, D1, input_state):
    qr = QuantumRegister(4, name="qubit")
    qc = QuantumCircuit(qr)
    qc.initialize(input_state)
    
    for d in range(D1):
        qc.append(RXGate(phis[0]), [qr[2]])
        qc.append(RXGate(phis[1]), [qr[3]])
        
        qc.append(RZGate(phis[2]), [qr[2]])
        qc.append(RZGate(phis[3]), [qr[3]])
        
        qc.append(CZGate(), [qr[2], qr[3]])
        qc.barrier(qr)
    return qc

def get_full_variational_quantum_circuit(thetas, phis, D1, D2, input_state):
    thetas_quantum_circuit = get_thetas_circuit(thetas, D2)
    phis_quantum_circuit = get_phis_circuit(phis, D1, input_state)
    
    variational_quantum_circuit = phis_quantum_circuit.compose(thetas_quantum_circuit)
    
    return variational_quantum_circuit

