!pip install qiskit
!pip install pylatexenc

import matplotlib.pyplot as plt
import math
import matplotlib as mpl
import numpy as np
import pandas as pd
import qiskit as q
from qiskit import Aer, assemble, QuantumRegister, ClassicalRegister
from qiskit.circuit import QuantumCircuit
import pylatexenc

from qiskit.quantum_info import Statevector, PTM
from qiskit.visualization import array_to_latex

from qiskit import transpile
from qiskit.providers.aer import QasmSimulator
import random
from qiskit.quantum_info import partial_trace, Statevector, DensityMatrix
from numpy import kron, trace

backend = QasmSimulator()

def visualize_one_line(x1_vals, y1_vals, axis_labels, title, label):
  fig, ax = plt.subplots()
  plt.title(title)
  plt.xlabel(axis_labels[0])
  plt.ylabel(axis_labels[1])
  ax.plot(x1_vals, y1_vals, linestyle='-', color='red', label=label)
  plt.legend()

def stateTomography(p):
  I = (np.array([[1, 0],[0, 1]])* np.array(p)).trace()
  X = (np.array([[0, 1],[1, 0]])* np.array(p)).trace()
  Y = (np.array([[0, 0- 1j],[1j, 0]])* np.array(p)).trace()
  Z = (np.array([[1, 0],[0, -1]])* np.array(p)).trace()
  return np.array([I, X, Y, Z])

def normalNoise(circ, noise):
  circ.cry(noise, 0, 1)
  circ.cnot(1, 0)
  circ.reset(1)

def randomTwirl(circ, noise):
  x = random.randint(0,3)
  if x == 1:
    circ.x(0)
  if x == 2:
    circ.y(0)
  if x == 3:
    circ.z(0)
  circ.cry(noise, 0, 1)
  circ.cnot(1, 0)
  circ.reset(1)
  if x == 1:
    circ.x(0)
  if x == 2:
    circ.y(0)
  if x == 3:
    circ.z(0)

def PauliConjugation(circ, noise, gate):
  
  if gate == 'x':
    circ.x(0)
  if gate == 'y':
    circ.y(0)
  if gate == 'z':
    circ.z(0)
  circ.cry(noise, 0, 1)
  circ.cnot(1, 0)
  circ.reset(1)
  if gate == 'x':
    circ.x(0)
  if gate == 'y':
    circ.y(0)
  if gate == 'z':
    circ.z(0)

def pauli_conjug_x(circ, noise):
    PauliConjugation(circ, noise, 'x')

def pauli_conjug_y(circ, noise):
    PauliConjugation(circ, noise, 'y')

def pauli_conjug_z(circ, noise):
    PauliConjugation(circ, noise, 'z')

tvals = [0.7, 1, 7, 10, 70, 100, 700, 1000, 7000, 10000]

funcs = [normalNoise, randomTwirl, pauli_conjug_x, pauli_conjug_y, pauli_conjug_z]
strings_to_funcs = {str(noise_func).split(' ')[1]: noise_func for noise_func in funcs}
funcs_to_strings = {strings_to_funcs[a]: a for a in strings_to_funcs}

success_data = {noise_func: [] for noise_func in funcs}
state_tomography_data = {noise_func: [] for noise_func in funcs}

for noise_func in funcs:
    print(funcs_to_strings[noise_func])
    for val in tvals:
        coinQ = QuantumRegister(1)
        coinC = ClassicalRegister(1)
        t = val
        T_1 = 10000
        turn = math.pi/500
        noise = 2*math.asin(math.sqrt(1-math.e**(-t/T_1)))

        coinCircuit = QuantumCircuit(coinQ, coinC)
        coinCircuit.initialize([1, 0], 0)
        coinCircuit.ry(math.pi/3, 0)
        coinCircuit.measure(0, 0)

        coin_compiled = transpile(coinCircuit, backend)
        flip_sim = backend.run(coin_compiled, shots=1000)
        result_sim = flip_sim.result()
        counts = {}
        counts = result_sim.get_counts(coin_compiled)
        distribution = []
        for i in range(counts['1']):
            distribution.append(1)
        for i in range (counts['0']):
            distribution.append(0)
        random.shuffle(distribution)

        unfairQ = QuantumRegister(2)
        unfairC = ClassicalRegister(1)
        unfairCircuit = QuantumCircuit(unfairQ, unfairC)
        unfairCircuit.initialize([1, 0], 0)
        unfairCircuit.initialize([1, 0], 1)

        for i in distribution:
            if i == 1:
                unfairCircuit.ry(-turn, 0)
            if i == 0:
                unfairCircuit.ry(turn, 0)
            noise_func(unfairCircuit, noise)
            #For Pauli conjugation comment out randomTwirl and use the below with desired gate
            #PauliConjugation(unfairCircuit, noise, z)

            #For normal noise use below
            #normalNoise(unfairCircuit, noise)


        unfairCircuit.measure(0, 0)

        turn_compiled = transpile(unfairCircuit, backend)
        turn_sim = backend.run(turn_compiled, shots = 200)
        unfair_sim = turn_sim.result()

        count_totals = unfair_sim.get_counts(turn_compiled)

        unfairQ = QuantumRegister(2)
        unfairCircuit = QuantumCircuit(unfairQ)

        for i in distribution:
            if i == 1:
                unfairCircuit.ry(-turn, 0)
            if i == 0:
                unfairCircuit.ry(turn, 0)
            unfairCircuit.cry(noise, 0, 1)
            unfairCircuit.cnot(1, 0)
            unfairCircuit.reset(1)

        state = Statevector.from_int(0, 2**2)
        state = state.evolve(unfairCircuit)
        rho = DensityMatrix(state)
        qubitOut = partial_trace(rho, [1])

        qubitIn = DensityMatrix([1, 0])
        stateTomographyIn = stateTomography(qubitIn)
        stateTomographyOut = stateTomography(qubitOut)

        success_data[noise_func].append(count_totals['1'] / 200 if '1' in count_totals else 0)
        state_tomography_data[noise_func].append(stateTomographyOut[-1])


fig, axs = plt.subplots(1, 2, figsize=(13, 5))
fig.suptitle('error correction via twirling and Pauli conjugation')
axs[0].set_xscale('log')
axs[0].set_xlabel('t / T_1')
axs[0].set_ylabel('success probability')

axs[1].set_xscale('log')
axs[1].set_xlabel('t / T_1')
axs[1].set_ylabel('Pauli expectation out')

tvals = [t / T_1 for t in tvals]

for noise_func in funcs:
    axs[0].plot(tvals, success_data[noise_func], label=funcs_to_strings[noise_func])
    axs[1].plot(tvals, state_tomography_data[noise_func], label=funcs_to_strings[noise_func])

plt.legend()