from qiskit import QuantumCircuit
from qiskit.quantum_info import Statevector, DensityMatrix
from qiskit.tools.visualization import plot_state_qsphere, plot_histogram

def init_message(circuit,n = 1,m = 0,hads = 0):
    if isinstance(m,int):
        m = bin(m)[2:]
    m = m.zfill(n)
    if isinstance(hads,int):
        hads = bin(hads)[2:]
    hads = hads.zfill(n)
    for i in range(0,n):
        if m[i] == '1':
            circuit.x(i)
        if hads[i] == '1':
            circuit.h(i)

def encrypt(n,key1 = 0,key2 = 0):
    circuit = QuantumCircuit(n)
    if isinstance(key1,int):
        k1 = bin(key1)[2:]
    k1 = k1.zfill(n)
    if isinstance(key2,int):
        k2 = bin(key2)[2:]
    k2 = k2.zfill(n)
    for i in range(0,n):
        if k1[i] == '1':
            circuit.x(i)
        if k2[i] == '1':
            circuit.z(i)
    return circuit.to_gate()

def decrypt(n,key1 = 0,key2 = 0):
    return encrypt(n,key1,key2).inverse()
    

n = 2
k1,k2 = 3,3
m = 0
hads = 3
circuit = QuantumCircuit(n)
init_message(circuit,n,m,hads)
#encryption
enc = encrypt(n,k1,k2)
enc.label = 'Enc({},{},{})'.format(n,k1,k2)
circuit.append(enc,list(range(0,n)))
# #decryption: uncomment to apply
# dec = decrypt(n,k1,k2)
# dec.label = 'Dec({},{},{})'.format(n,k1,k2)
# circuit.append(dec,list(range(0,n)))

%matplotlib inline
circuit.draw(output = 'mpl')
#Uncomment to view contents of the encryption and/or decryption circuits
#circuit.decompose().draw(output = 'mpl')

sv = Statevector.from_label("00")
dm = DensityMatrix(sv)
nsv = sv.evolve(circuit)
ndm = dm.evolve(circuit)
print(nsv.data)
print(ndm.data)

