# Importing Packages
from qiskit import QuantumCircuit
from qiskit import IBMQ, Aer, transpile, assemble
from qiskit.visualization import plot_histogram

style = {'backgroundcolor': 'lightyellow'} # Style of the circuits

def encode(qc, qubit, msg):
    if len(msg) != 2 or not set(msg).issubset({"0","1"}):
        raise ValueError(f"message '{msg}' is invalid")
    if msg[1] == "1":
        qc.x(qubit)
    if msg[0] == "1":
        qc.z(qubit)
    return qc

# Circuit for Superdense coding
qc = QuantumCircuit(2)
# Generating entangled state
qc.h(1)
qc.cx(1, 0)
qc.barrier()
#  Encoding message
message = '11'
qc = encode(qc, 1, message)
qc.barrier()
# Decoding Message
qc.cx(1, 0)
qc.h(1)
# Measuring qubits
qc.measure_all()
qc.draw(output="mpl", style=style)

aer_sim = Aer.get_backend('aer_simulator')
qobj = assemble(qc)
result = aer_sim.run(qobj).result()
counts = result.get_counts(qc)
print(counts)
plot_histogram(counts)