import numpy as np

#Importing standard Qiskit libraries
from qiskit import *
from qiskit.tools.jupyter import *
from qiskit.visualization import *
from qiskit.providers import *
from qiskit.circuit import Parameter
import matplotlib.pyplot as plt

#Loading IBM Quantum account(s)
provider = IBMQ.load_account()

q = QuantumRegister(1,'q') #System register
e = QuantumRegister(1,'e') #Environment register
c = ClassicalRegister(1,'c')

def dephasing_M(steps):
    
    qc = QuantumCircuit(q,e,c)
    theta = np.pi/5

    qc.h(0)
    qc.barrier()

    for i in range(steps):
        if (i>0):
            qc.reset(e)
        qc.ry(theta,1)
        qc.cz(0,1)
        qc.measure(e,c)
        qc.barrier()
    
    qc.h(q)
    qc.measure(q,c)
    
    sim = Aer.get_backend('qasm_simulator')
    job = execute(qc, sim, shots=1024)
    res = job.result()
    counts = res.get_counts()['0']/1024
    
    return counts


qc = dephasing_M(1)
qc.draw('mpl')

steps = np.linspace(0,100,21)
prob0 = []

for st in steps:
    prob0.append(dephasing_M(int(st)))

fig = plt.figure(figsize=(8,5))

plt.scatter(steps,prob0)
plt.plot(steps,prob0)

plt.xticks(size=12)
plt.yticks(size=12)

plt.title('\nDephasing of a qubit', size=14)
plt.xlabel(r'Steps $\longrightarrow$', size=13)
plt.ylabel(r'Population in Hadamard basis $\longrightarrow$', size=13)

plt.show()

q = QuantumRegister(1,'q') #System register
e = QuantumRegister(3,'e') #Environment register
c = ClassicalRegister(1,'c')

def dephasingNM(steps):
    
    qc = QuantumCircuit(q,e,c)
    t1 = np.pi/5
    t2 = np.pi/4
    t3 = np.pi/2

    qc.h(0)
    qc.barrier()

    for i in range(steps):
        qc.ry(t1,e[0])
        qc.ry(t2,e[1])
        qc.ry(t3,e[2])
        qc.cz(q[0],e[0])
        qc.measure(e[0],c)
        
        qc.reset(e[0])
        qc.swap(e[0],e[1])
        qc.swap(e[1],e[2])
        qc.barrier()
    
    qc.h(q)
    qc.measure(q,c)
    

    sim = Aer.get_backend('qasm_simulator')
    job = execute(qc, sim, shots=1024)
    res = job.result()
    counts = res.get_counts()['0']/1024
    
    return counts


qc = dephasingNM(1)
qc.draw('mpl')

steps = np.linspace(0,100,21)
prob0_NM = []

for st in steps:
    prob0_NM.append(dephasingNM(int(st)))

fig = plt.figure(figsize=(8,5))

plt.scatter(steps,prob0_NM)
plt.plot(steps,prob0_NM)

plt.xticks(size=12)
plt.yticks(size=12)

plt.title('\nDephasing of a qubit', size=14)
plt.xlabel(r'Steps $\longrightarrow$', size=13)
plt.ylabel(r'Population in Hadamard basis $\longrightarrow$', size=13)

plt.show()

fig = plt.figure(figsize=(8,5))

plt.plot(steps,prob0)
plt.plot(steps,prob0_NM)

plt.xticks(size=12)
plt.yticks(size=12)

plt.legend(['Markovian','Non-Markovian'])
plt.title('\nDephasing of a qubit', size=14)
plt.xlabel(r'Steps $\longrightarrow$', size=13)
plt.ylabel(r'Population in Hadamard basis $\longrightarrow$', size=13)

plt.show()