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):
    qr = QuantumRegister(4, name="qubit")
    qc = QuantumCircuit(qr)
    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):
    thetas_quantum_circuit = get_thetas_circuit(thetas, D2)
    phis_quantum_circuit = get_phis_circuit(phis, D1)
    
    variational_quantum_circuit = phis_quantum_circuit.compose(thetas_quantum_circuit)
    return variational_quantum_circuit

import numpy as np

thetas = np.random.uniform(low=0, high=2*np.pi, size=8)
phis = np.random.uniform(low=0, high=2*np.pi, size=4)

D1 = 2
D2 = 6

qc1 = get_thetas_circuit(thetas,D2)
print(qc1.draw())

qc2 = get_phis_circuit(phis,D1)
print(qc2.draw())

print(get_full_variational_quantum_circuit(thetas, phis, D1, D2))