# Superdense Coding Implementation    Ben Foxman
import matplotlib.pyplot as plt
import numpy as np
from qiskit import *


bits = input("Enter 2 bits.") # the 2 bits to encode/decode
assert bits in ['00', '01', '10', '11']

# p1 encodes by creating map from each of the 4 bell states to each of the 4 2-bit codes
def encode(target):
  bell = QuantumCircuit(2, 2, name="bell")
  bell.h([0, 1]) # this entangles p1 and p2's qubit (on measurement, they both evaluate to 0 or 1)
  if target == '01': 
    bell.x([0, 1]) 
  elif target == '10':
    bell.z(0)
  elif target == '11':
    bell.z(0)
    bell.x([0, 1])
  # '00'= identity, keep the initial bell states bell state
  return bell



# now, p2 decodes the circuit 
def decode(circuit):
  # step 1: cnot it 
  circuit.cx(0, 1)
  #step 2: h-gate and measure
  circuit.h(0)
  circuit.measure([0, 1], [0, 1])
  return circuit

backend = Aer.get_backend('statevector_simulator')
# encode bits
circ = encode(bits)
# decode bits
message = decode(circ)
job = execute(message, backend)
result = job.result()
# print out the statevector. There should be a 1 at index n, where n (base 10) is the code 
print(np.round(result.get_statevector(), 2))
