import logging
from itertools import product 

import networkx as nx
import matplotlib.pyplot as plt

from qiskit import IBMQ, QuantumCircuit, execute, transpile
from qiskit.test.mock import FakeMumbai as FakeDevice
from qiskit.providers.aer.noise import NoiseModel
from qiskit.visualization import plot_histogram


logging.getLogger().setLevel(logging.INFO)

IBMQ.load_account()
provider = IBMQ.get_provider(hub='ibm-q')
device = FakeDevice() 
noise_model = NoiseModel.from_backend(device)

conf = device.configuration()
print(conf.basis_gates)
G = nx.Graph()
G.add_nodes_from(range(0,conf.n_qubits))
G.add_edges_from(conf.coupling_map)
pos = nx.kamada_kawai_layout(G)
nx.draw_networkx(G, pos, with_labels = True)

simBackends = provider.backends(filters=lambda x: x.configuration().n_qubits >= conf.n_qubits
                                    and x.configuration().simulator
                                    and x.status().operational==True
                                    and all(basis_gate in x.configuration().basis_gates
                                     for basis_gate in conf.basis_gates if basis_gate!="reset"))
                                     # we exclude reset because this is not physically implemented yet
simBackends

n_entangled = 4
n_pairs = 5
n_qubits = n_entangled * n_pairs
qc = QuantumCircuit(n_qubits, n_qubits)
qubits_in_superpos = [i for i in range(0,n_qubits, n_entangled)]
qubits_to_entangle = {i:[i + j for j in range(1, n_entangled)] for i in qubits_in_superpos}
for control in qubits_in_superpos:
    qc.h(control)
    for target in qubits_to_entangle[control]:
        qc.cx(control, target)
qc.measure_all()
qc.draw(output='mpl')

circuit_transpiled = transpile(qc, backend=device, optimization_level=3)
circuit_transpiled.draw(output='mpl')

bell_states = [x * n_entangled for x in ["0", "1"]]
possible_states = ["".join(state) for state in product(bell_states, repeat=n_pairs)]
possible_states

backend = provider.get_backend('simulator_statevector')
shots = 4096
result = execute(qc, backend, shots=shots,
                 coupling_map=conf.coupling_map,
                 basis_gates=conf.basis_gates,
                 noise_model=noise_model).result() # took about 2 hours for problem with 26 qubits
counts = result.get_counts(0)

noise_ct = 0
for measurement, count in counts.items():
    measurement = measurement[:n_qubits] # drop ancillary qubits
    if measurement in possible_states:
        print(f"{measurement=} : {count=}")
    else: 
        noise_ct = noise_ct + count
print(f"{noise_ct=} valid states: {100 * (shots - noise_ct) / shots : 0.2f}%")

