from qiskit import *
from qiskit.tools.parallel import parallel_map
from qiskit.tools.events import TextProgressBar
from qiskit.tools.jupyter import *  # Needed to load the Jupyter HTMLProgressBar

num_circuits = 1000
width = 4
depth = 4

import copy
import math
import numpy as np
from qiskit.tools.qi.qi import random_unitary_matrix
from qiskit.mapper import two_qubit_kak

def build_qv_circuit(idx, seeds, width, depth):
    """Builds a single Quantum Volume circuit.  Two circuits,
    one with measurements, and one widthout, are returned.

    The model circuits consist of layers of Haar random
    elements of SU(4) applied between corresponding pairs
    of qubits in a random bipartition.
    
    See: https://arxiv.org/abs/1811.12926
    """
    np.random.seed(seeds[idx])
    q = QuantumRegister(width, "q")
    c = ClassicalRegister(width, "c")
    # Create measurement subcircuit
    qc = QuantumCircuit(q,c)
    # For each layer
    for j in range(depth):
        # Generate uniformly random permutation Pj of [0...n-1]
        perm = np.random.permutation(width)
        # For each pair p in Pj, generate Haar random SU(4)
        # Decompose each SU(4) into CNOT + SU(2) and add to Ci
        for k in range(math.floor(width/2)):
            qubits = [int(perm[2*k]), int(perm[2*k+1])]
            U = random_unitary_matrix(4)
            for gate in two_qubit_kak(U):
                i0 = qubits[gate["args"][0]]
                if gate["name"] == "cx":
                    i1 = qubits[gate["args"][1]]
                    qc.cx(q[i0], q[i1])
                elif gate["name"] == "u1":
                    qc.u1(gate["params"][2], q[i0])
                elif gate["name"] == "u2":
                    qc.u2(gate["params"][1], gate["params"][2],
                                 q[i0])
                elif gate["name"] == "u3":
                    qc.u3(gate["params"][0], gate["params"][1],
                                 gate["params"][2], q[i0])
                elif gate["name"] == "id":
                    pass  # do nothing
    qc_no_meas = copy.deepcopy(qc)
    # Create circuit with final measurement
    qc.measure(q,c)
    return qc, qc_no_meas

num_circuits = 1000
seeds = np.random.randint(np.iinfo(np.int32).max, size=num_circuits)
TextProgressBar()
parallel_map(build_qv_circuit, np.arange(num_circuits), task_args=(seeds, width, depth));

seeds = np.random.randint(np.iinfo(np.int32).max, size=num_circuits)
HTMLProgressBar()
parallel_map(build_qv_circuit, np.arange(num_circuits), task_args=(seeds, width, depth));

