# To use the package locally, add the C2QA repository's root folder to the path prior to importing c2qa.
import os
import sys
module_path = os.path.abspath(os.path.join("../.."))
if module_path not in sys.path:
    sys.path.append(module_path)

# Import packages that will be used
import c2qa
import numpy as np
import matplotlib.pyplot as plt
from qiskit import ClassicalRegister, visualization, AncillaRegister
from qiskit_aer.noise import NoiseModel, amplitude_damping_error
from qiskit.quantum_info import state_fidelity, Statevector
from qiskit.quantum_info.operators import Operator
from qiskit.quantum_info.operators.predicates import is_unitary_matrix
from math import pi, ceil
from scipy.optimize import curve_fit

# LaTeX typesetting for plt
plt.rc('font', **{'family': 'serif', 'serif': ['Computer Modern']})
plt.rcParams['text.usetex'] = True
plt.rc('text.latex', preamble=r'\usepackage{amssymb}')
plt.rc('text.latex', preamble=r'\usepackage{amsmath}')

## Demonstration of amplitude damping circuit
# Define how many qumodes we want, and how many qubits we should use to represent each. Our basis will consist of Fock states |0> to |7>
num_qumodes = 1
num_qubits_per_qumode = 3 # The photon number which the Hilbert space truncates at is 2 ** num_qubits_per_qumode = 8

# Create circuit
qmr = c2qa.QumodeRegister(num_qumodes, num_qubits_per_qumode)
cr = ClassicalRegister(num_qumodes * num_qubits_per_qumode)
circ = c2qa.CVCircuit(qmr, cr)

# To illustrate photon loss, first initialize the qumode in Fock state |7>
circ.cv_initialize(7, qmr[0])

# Now apply a delay gate and specify gate duration and units
gate_duration = 1
time_unit = "ms"
circ.cv_delay(duration=gate_duration, qumode=qmr[0], unit=time_unit)

# Now set a loss rate -- make it large enough such that we can actually see the loss
photon_loss_rate = 0.5 # This is a loss rate in units of 1/ms
noise_pass = c2qa.kraus.PhotonLossNoisePass(photon_loss_rates=photon_loss_rate, circuit=circ, time_unit=time_unit)

# To see the loss, we can make a measurement on the qumode and visualise counts with a histogram
circ.cv_measure(qmr[0], cr)

# Use noise pass with simulate
_, _, counts = c2qa.util.simulate(circ, noise_passes=noise_pass)

# Visualise photon loss using a histogram
visualization.plot_histogram(counts)

## Demonstration of photon parity measurement circuit
# Create new circuit with 1 qumode, 1 ancilla qubit, and 1 classical register for anc readout.
qmr = c2qa.QumodeRegister(num_qumodes, num_qubits_per_qumode)
anc = AncillaRegister(1)
cr = ClassicalRegister(1)
circ = c2qa.CVCircuit(qmr, anc, cr)

circ.initialize([1, 0], anc[0]) # Initialize ancilla qubit in |0> state

# Initialize qumode in either odd parity or even parity Fock state
#circ.cv_initialize(2, qmr) # |2> Fock state
circ.cv_initialize(3, qmr) # |3> Fock state

# Apply hadamard on ancilla qubit
circ.h(anc[0])

# Apply controlled phase space rotation gate on qumode
circ.cv_c_r(pi/2, qmr[0], anc[0])

# Apply unconditional phase space rotation gate to correct global phase
circ.cv_r(-pi/2, qmr[0])

# Apply hadamard on ancilla qubit again, and measure the ancilla qubit.
circ.h(anc[0])
circ.cv_measure(anc, cr)

_, _, counts = c2qa.util.simulate(circ)

# Measurement will return 0 for even parity and 1 for odd parity
print(counts)

# Create empty matrix
empty = np.zeros((2**num_qubits_per_qumode, 2**num_qubits_per_qumode))

## Defining our theoretical operation for U_odd
# Map logical 1's photon loss state |1> back to logical 1. (i.e. |2><1|, where fock_input = <1|, and fock_output = |2>)
u_odd = c2qa.util.fockmap(matrix=empty, fock_input=[1], fock_output=[2]) 

# Map logical 0's photon loss state |3> to logical 0
u_odd = c2qa.util.fockmap(u_odd, fock_input=[3, 3], fock_output=[0, 4], amplitude=[1/np.sqrt(2), 1/np.sqrt(2)]) 

# Mappings to satisfy unitary condition
u_odd = c2qa.util.fockmap(u_odd, [5, 6, 7], [5, 6, 7])
u_odd = c2qa.util.fockmap(u_odd, [2, 0], [3, 1]) 
u_odd = c2qa.util.fockmap(u_odd, [4, 4], [0, 4], 1/np.sqrt(2) * np.array([1, -1]))

# Check that our matrix is indeed unitary
assert(is_unitary_matrix(u_odd))

## Defining our theoretical operation for U_even. 
# We will delay the building of the gate until we have the value we need for kt
def u_even(kt):
    # Map logical 1 back to logical 1
    u_even = c2qa.util.fockmap(empty, 2, 2) 

    # Map logical 0's error state to logical 0
    u_even = c2qa.util.fockmap(u_even, [0, 0, 4, 4], [0, 4, 0, 4], 1/np.sqrt(1 + kt**2) * np.array([1, kt, -kt, 1])) 

    # Mappings to satisfy unitary condition
    u_even = c2qa.util.fockmap(u_even, [1, 3, 5, 6, 7], [1, 3, 5, 6, 7]) 

    # Check that our matrix is indeed unitary
    assert(is_unitary_matrix(u_even))
    return u_even

# Initialize new circuit 
qmr = c2qa.QumodeRegister(num_qumodes, num_qubits_per_qumode)
cr = ClassicalRegister(num_qumodes * num_qubits_per_qumode)
circ = c2qa.CVCircuit(qmr, cr)

# Initialize qumode in |1_L> error state
circ.cv_initialize([0, 1], qmr)

# Append the u_odd matrix to circuit and measure
circ.cv_gate_from_matrix(u_odd, qmr[0])
circ.cv_measure(qmr[0], cr)

_, _, counts = c2qa.util.simulate(circ)

# We measure |2> Fock state, demonstrating recovery of the logical 1 state
print(counts)

## Run this block to load in the functions we need for the simulation
# Runs circuit simulation
def error_correction_experiment(encoding: str, corrections: int, shots: int, additional_snapshots: int=10, total_circuit_time: int=1, apply_noise: bool=True, return_circuit: bool=False, had_damp: bool=False, identity: bool=False):
    multiplier = total_circuit_time
    total_circuit_time = int(total_circuit_time * 1000) # Code works with microsecond units, but input is in milliseconds
    
    if corrections > 333:
        raise ValueError("Please restrict number of correction procedures to 333 or less")
    photon_loss_rate = 0.0005

    num_qumodes = 1
    num_qubits_per_qumode = 3

    qmr = c2qa.QumodeRegister(num_qumodes, num_qubits_per_qumode)
    anc = AncillaRegister(1)
    cr = ClassicalRegister(1)
    circuit = c2qa.CVCircuit(qmr, anc, cr)

    ### 1. State Preparation
    perf_0 = np.array([1/np.sqrt(2), 0, 0, 0, 1/np.sqrt(2), 0, 0, 0])
    perf_1 = np.array([0, 0, 1, 0, 0, 0, 0, 0])
    
    if encoding == 'logi0':
        perf_state = perf_0
    elif encoding == 'logi1':
        perf_state = perf_1
    elif encoding == 'logi_xplus':
        perf_state = (perf_0 + perf_1)/np.sqrt(2)
    elif encoding == 'logi_xminus':
        perf_state = (perf_0 - perf_1)/np.sqrt(2)
    elif encoding == 'logi_yplus':
        perf_state = (perf_0 + 1j * perf_1)/np.sqrt(2)
    elif encoding == 'logi_yminus':
        perf_state = (perf_0 - 1j * perf_1)/np.sqrt(2)
    else:
        raise ValueError("Please input a valid state for the encoding")
    
    # Convert type to list
    perf_state = list(perf_state)
    #print(perf_state, "cv_initialzed this")
    
    # Initialize state
    circuit.cv_initialize(perf_state, qmr[0])

    ### 2. Using instructions stored in schedule, add snapshots, noise, and parity check + recovery operations
    schedule = _circuit_scheduler(corrections * multiplier, additional_snapshots * multiplier, total_circuit_time)
    schedule_copy = list(schedule.items())
    #print(schedule)

    adjustment = 0 # For calculating duration timing
    snapshot_counter = 0 # For tracking number of snapshots taken
    corr_position = [] # For tracking positions at which corrections were finished
    
    u_iden = c2qa.util.fockmap(empty, [0, 1, 2, 3, 4, 5, 6, 7], [0, 1, 2, 3, 4, 5, 6, 7])
    
    for index, (key, value) in enumerate(list(schedule.items())):
        # Instruction at 0µs should always be to take snapshot
        if (key == '0') & (value == 'snapshot'):
            circuit.save_statevector(label='snapshot_{}'.format(snapshot_counter), pershot=True)
            snapshot_counter += 1
            continue
        elif (key != '0'):
            pass
        else:
            raise Exception("Something seems to have gone wrong here")

        # Calculate delay gate durations using prevkey. Adjustment to duration timing is needed depending on previous instruction
        prevkey, prevvalue = schedule_copy[index - 1]
        delay_dur = int(key) - int(prevkey) + adjustment
        #print(delay_dur, key, value)
        adjustment = 0

        # Apply noise
        circuit.cv_delay(duration=delay_dur, qumode=qmr[0], unit="µs")

        # Apply snapshot only procedure
        if (value == 'snapshot'):
            # Don't take duplicate snapshots
            if (prevvalue == 'correction') & (int(key) == (int(prevkey) + 3)):
                pass
            else:
                circuit.save_statevector(label='snapshot_{}'.format(snapshot_counter), pershot=True)
                snapshot_counter += 1

        # Apply snapshot, parity check, recovery, snapshot 
        elif value == 'correction':
            # Find previous corr position, and determine kt_dur with that information
            while index > 0:
                k, v = list(schedule.items())[index-1]
                # Break loop if we iterated till start of dictionary
                if (v == 'snapshot') & (k == '0'):
                    kt_dur = int(key) - int(k)
                    break
                # Skip all snapshot instructions
                elif v == 'snapshot':
                    index -= 1
                # Find either snapshot+corr or corr instruction and extract kt_dur
                elif v != 'snapshot':
                    kt_dur = int(key) - int(k)
                    break
            #print(kt_dur)
            
            kt = photon_loss_rate * (kt_dur + 3)
            adjustment = -3

            # Parity check
            circuit.initialize([1, 0], anc[0])
            circuit.h(anc[0])
            circuit.cv_delay(duration=1, qumode=qmr[0], unit="µs")
            
            # Manually discretize gates by splitting 0.5µs pi/2 rotation into five 0.1µs pi/10 rotations
            for _ in range(5):
               circuit.cv_c_r(pi/10, qmr[0], anc[0], duration=0.1, unit="µs")
            for _ in range(5):
                circuit.cv_r(-pi/10, qmr[0], duration=0.1, unit="µs")
            # circuit.cv_c_r(pi/2, qmr[0], anc[0], duration=0.5, unit="µs")
            # circuit.cv_r(-pi/2, qmr[0], anc[0], duration=0.5, unit="µs")
            
            circuit.cv_delay(duration=1, qumode=qmr[0], unit="µs")
            circuit.h(anc[0])
            circuit.measure(anc[0], cr[0])

            # Apply recovery operation, without or with u_even
            if identity:
                
                with circuit.if_test((cr[0], 1), label='Recovery op') as else_:
                    circuit.cv_gate_from_matrix(u_odd, qmr[0], duration=0.1, unit="ns") #u_odd
                    circuit.reset(anc[0])
                with else_:
                    circuit.cv_gate_from_matrix(u_iden, qmr[0], duration=0.1, unit="ns") #u_iden
                
            else:
                with circuit.if_test((cr[0], 1), label='Recovery op') as else_:
                    circuit.cv_gate_from_matrix(u_odd, qmr[0], duration=0.1, unit="ns") #u_odd
                    circuit.reset(anc[0])
                with else_:
                    circuit.cv_gate_from_matrix(u_even(kt), qmr[0], duration=0.1, unit="ns") #u_even(kt)
            
            circuit.save_statevector(label='snapshot_{}'.format(snapshot_counter), pershot=True)
            snapshot_counter += 1
            corr_position.append(int(key) + 3)

        # Apply snapshot, parity check, recovery, followed by another snapshot
        else:
            circuit.save_statevector(label='snapshot_{}'.format(snapshot_counter), pershot=True)
            snapshot_counter += 1
            
            adjustment = -3
            # Find previous corr position, and determine kt_dur with that information
            while index > 0:
                k, v = list(schedule.items())[index-1]
                # Break loop if we iterated till start of dictionary
                if (v == 'snapshot') & (k == '0'):
                    kt_dur = int(key) - int(k)
                    break
                # Skip all snapshot instructions
                elif v == 'snapshot':
                    index -= 1
                # Find either snapshot+corr or corr instruction and extract kt_dur
                elif v != 'snapshot':
                    kt_dur = int(key) - int(k)
                    break
            #print(kt_dur)
            kt = photon_loss_rate * (kt_dur + 3)

            # Parity check
            circuit.initialize([1, 0], anc[0])
            circuit.h(anc[0])
            circuit.cv_delay(duration=1, qumode=qmr[0], unit="µs")
            
            # Manually discretize gates by splitting 0.5µs pi/2 rotation into five 0.1µs pi/10 rotations
            for _ in range(5): # Manually discretize cv_c_r gate by splitting one 1µs pi rotation into ten 0.1µs pi/10 rotations
               circuit.cv_c_r(pi/10, qmr[0], anc[0], duration=0.1, unit="µs")
            for _ in range(5):
                circuit.cv_r(-pi/10, qmr[0], duration=0.1, unit="µs")
            # circuit.cv_c_r(pi/2, qmr[0], anc[0], duration=0.5, unit="µs")
            # circuit.cv_r(-pi/2, qmr[0], anc[0], duration=0.5, unit="µs")
            
            circuit.cv_delay(duration=1, qumode=qmr[0], unit="µs")
            circuit.h(anc[0])
            circuit.measure(anc[0], cr[0])
            
            if identity:
                # Recovery operation
                with circuit.if_test((cr[0], 1), label='Recovery op') as else_:
                    circuit.cv_gate_from_matrix(u_odd, qmr[0], duration=0.1, unit="ns") #u_odd
                    circuit.reset(anc[0])
                with else_:
                    circuit.cv_gate_from_matrix(u_iden, qmr[0], duration=0.1, unit="ns") #u_even(kt)
                
            else:
                # Recovery operation
                with circuit.if_test((cr[0], 1), label='Recovery op') as else_:
                    circuit.cv_gate_from_matrix(u_odd, qmr[0], duration=0.1, unit="ns") #u_odd
                    circuit.reset(anc[0])
                with else_:
                    circuit.cv_gate_from_matrix(u_even(kt), qmr[0], duration=0.1, unit="ns") #u_even(kt)
                
            circuit.save_statevector(label='snapshot_{}'.format(snapshot_counter), pershot=True)
            snapshot_counter += 1
            corr_position.append(int(key) + 3)

    # Final delay gate
    finalkey, finalvalue = schedule_copy[-1]
    if finalvalue != 'snapshot':
        final_gate_delay = total_circuit_time - int(finalkey) - 3
    else:
        final_gate_delay = total_circuit_time - int(finalkey)
        
    circuit.cv_delay(duration=final_gate_delay, qumode=qmr[0], unit="µs")
    
    if return_circuit:
        return circuit
    
    ### 3. Define noise parameters
    # Noise pass for qumode photon loss
    noise_pass = c2qa.kraus.PhotonLossNoisePass(photon_loss_rates=photon_loss_rate, circuit=circuit, time_unit="µs")

    # Noise model for hadamard amplitude damping
    T1 = 10 # In µs
    hgate_duration = 1 # In µs
    error = amplitude_damping_error(1 - np.exp(-hgate_duration / T1))

    noise_model = NoiseModel()
    index, _ = circuit.find_bit(anc[0])
    noise_model.add_quantum_error(error, ['h'], [index])
    
    ### 4. Simulate circuit
    if apply_noise:
        if had_damp:
            state, result, _ = c2qa.util.simulate(circuit, shots=shots, per_shot_state_vector=True, return_fockcounts=False, noise_passes=noise_pass, noise_model=noise_model)
        else:
            state, result, _ = c2qa.util.simulate(circuit, shots=shots, per_shot_state_vector=True, return_fockcounts=False, noise_passes=noise_pass)
    elif not apply_noise:
        state, result, _ = c2qa.util.simulate(circuit, shots=shots, per_shot_state_vector=True, return_fockcounts=False)
    
    ### 5. Process data
    # Find fidelity and photon probability at each snapshot, across all shots
    traced_density_matrices = _state_traceout(circuit, state, result, snapshot_counter)
    expectation_list_tuple, fidelity_list, probability_list_tuple = _state_processor(encoding, perf_state, traced_density_matrices)
    
    # Find exact timing of each snapshot (we have to recalculate because no. of snapshots can change due to scheduling conflicts)
    snapshot_times = _snapshot_timings(schedule, total_circuit_time)
    
    return expectation_list_tuple, fidelity_list, probability_list_tuple, snapshot_times, corr_position#, circuit, schedule, snapshot_counter

# Trace out qubit for both sets of statevectors
def _state_traceout(circuit, state, result, actual_snapshots):
    # Process snapshot and final statevectors
    traced_snapshot_density_matrices = []
    for i in range(actual_snapshots):
        snapshot_list = result.data()['snapshot_{}'.format(i)]

        traced_snapshots = []
        for statevector in snapshot_list:
            traced_snapshots.append(c2qa.util.trace_out_qubits(circuit, statevector))

        traced_snapshot_density_matrices.append(traced_snapshots)

    traced_final_density_matrices = []
    for i in state:
        traced_final_density_matrices.append(c2qa.util.trace_out_qubits(circuit, i))
    
    traced_density_matrices = traced_snapshot_density_matrices + [traced_final_density_matrices]

    return traced_density_matrices

# Takes traced matrices and finds expectation, fidelity, and photon probability
def _state_processor(encoding, perf_state, traced_density_matrices):
    #print(perf_state, "compute fidelity with this")
    
    perf_0 = [1/np.sqrt(2), 0, 0, 0, 1/np.sqrt(2), 0, 0, 0]
    perf_1 = [0, 0, 1, 0, 0, 0, 0, 0]
    err = [1/np.sqrt(2), 0, 0, 0, -1/np.sqrt(2), 0, 0, 0]
    
    # Bring into qubit space
    projector = Operator(np.outer([1, 0], perf_0) + np.outer([0, 1], perf_1))
    
    # Z expectation
    if (encoding == 'logi0') or (encoding == 'logi1'):
        pauli = np.array([[1, 0], [0, -1]])
        
    # X expectation
    elif (encoding == 'logi_xplus') or (encoding == 'logi_xminus'):
        #projector = Operator(np.outer([1/np.sqrt(2), 1/np.sqrt(2)], (perf_0 + perf_1)/np.sqrt(2)) + np.outer([1/np.sqrt(2), -1/np.sqrt(2)], (perf_0 - perf_1)/np.sqrt(2)))
        pauli = np.array([[0, 1], [1, 0]]) 
    
    # Y expectation
    elif (encoding == 'logi_yplus') or (encoding == 'logi_yminus'):
        #projector = Operator(np.outer([1/np.sqrt(2), 1j/np.sqrt(2)], (perf_0 + 1j * perf_1)/np.sqrt(2)) + np.outer([1/np.sqrt(2), 1j/np.sqrt(2)], (perf_0 - 1j * perf_1)/np.sqrt(2)))
        pauli = np.array([[0, -1j], [1j, 0]])

    # Compute from each density matrix 1) Pauli expectation, 2) fidelity, 3) probability of measuring photon number
    expectation_list, exp_prob_list, fidelity_list, probability_listA, probability_listB = [], [], [], [], []
    for density_matrix_list in traced_density_matrices:
        
        cycle_list_expectation, cycle_list_exp_prob, cycle_list_fid, cycle_list_probA, cycle_list_probB = [], [], [], [], []
        for matrix in density_matrix_list:
            #print(matrix.trace())
            # Pauli Expectation
            qubit_space_projected_state = matrix.evolve(projector)
            projection_prob = qubit_space_projected_state.trace()
            
            expectation = (qubit_space_projected_state/projection_prob).expectation_value(pauli)
            
            cycle_list_exp_prob.append(projection_prob)
            cycle_list_expectation.append(expectation)
            
            # Fid
            cycle_list_fid.append(round(state_fidelity(matrix, Statevector(perf_state)), 6))
            
            # Photon meas prob
            cycle_list_probA.append(list(matrix.probabilities(decimals=5))) # Photon number
            
            fid0 = round(state_fidelity(matrix, perf_0), 6)
            fiderr = round(state_fidelity(matrix, err), 6)
            cycle_list_probB.append([fid0, fiderr]) # w/ respect to specific states
            
            # Avg photon num
            #cycle_list_avg_photon.append(c2qa.util.qumode_avg_photon_num(matrix))

        expectation_list.append(cycle_list_expectation)
        exp_prob_list.append(cycle_list_exp_prob)
        fidelity_list.append(cycle_list_fid)
        probability_listA.append(cycle_list_probA)
        probability_listB.append(cycle_list_probB)
        #avg_photon_list.append(cycle_list_avg_photon)

    return (expectation_list, exp_prob_list), fidelity_list, (probability_listA, probability_listB)

# Function used to determine order and timing in which snapshots and corrections are done.
def _circuit_scheduler(corrections, snapshots, total_time=1000):
    # Schedule snapshots
    if snapshots == 0:
        pass
    else:
        snapshot_schedule = {"{}".format(round(i * (total_time/snapshots))): 'snapshot' for i in range(snapshots)}

    # Schedule corrections
    corr_interval = total_time/(1 + corrections)
    correction_schedule = {"{}".format(int((i + 1) * corr_interval)): 'correction' for i in range(corrections)}
    
    
    # If no corrections, return snapshot_schedule
    if corrections == 0:
        return snapshot_schedule
    # If no snapshots, return correction_schedule with 1 snapshot at start
    elif snapshots == 0:
        correction_schedule = {**{'0' : 'snapshot'}, **correction_schedule}
        return correction_schedule
    else:
        pass
        
    # Merge the two dicts from above. If duplicate keys exist then concatenate values together.
    try:
        unsorted_merged_schedule = {**snapshot_schedule, **correction_schedule}

        for key, value in list(snapshot_schedule.items()):
            if key in list(correction_schedule.keys()):
                unsorted_merged_schedule[key] = '{}'.format(value + '+' + correction_schedule[key])
            else:
                pass
    except:
        print("_circuit_scheduler() has an error")
    
    # Sort merged dicts
    merged_schedule = {key:value for key, value in sorted(unsorted_merged_schedule.items(), key=lambda item: int(item[0]))}

    # Ensure that all corrections happen at least 3µs before each snapshot
    # but if two successive corrections happen within 3µs of each other, drop the second correction
    DictItems = list(merged_schedule.items())
    unsorted_cleaned_schedule = dict(merged_schedule)
    for index, (key, value) in enumerate(DictItems):
        # index == 0 case
        if index == 0:
            nextkey, _ = DictItems[index + 1]
            if (value == 'correction') & (int(key) + 3 > int(nextkey)):
                unsorted_cleaned_schedule['{}'.format(str(int(nextkey) - 3))] = unsorted_cleaned_schedule.pop('{}'.format(key))
        # index != 0 case
        else:
            # Automatically break loop when we get to last element of dictionary
            try:
                nextkey, _ = DictItems[index + 1]
                prevkey, _ = DictItems[index - 1]
            except:
                break
            
            if (value == 'correction') & (int(key) + 3 > int(nextkey)) & (int(prevkey) < int(key) - 3):
                unsorted_cleaned_schedule['{}'.format(str(int(nextkey) - 3))] = unsorted_cleaned_schedule.pop('{}'.format(key))
            elif (value == 'correction') & (int(key) + 3 > int(nextkey)):
                unsorted_cleaned_schedule.pop('{}'.format(key))
                print("Due to scheduling conflicts, a correction scheduled at {}µs will be dropped".format(key))

    # Sort overall_schedule
    cleaned_schedule = {key:value for key, value in sorted(unsorted_cleaned_schedule.items(), key=lambda item: int(item[0]))}

    return cleaned_schedule

# Function for extracting snapshot timings (in ms)
def _snapshot_timings(schedule, total_circuit_time):
    timings = []
    sch_copy = list(schedule.items())
    prevkey, prevvalue = 0, 0 # Just so that first iteration doesn't throw an error
    
    for index, (key, value) in enumerate(list(schedule.items())):
        if index != 0:
            prevkey, prevvalue = sch_copy[index - 1]
        
        if value == 'snapshot':
            if (prevvalue == 'correction') & (int(key) == (int(prevkey) + 3)):
                pass
            else:
                timings.append(int(key)/1000)
            
        elif value =='correction':
            timings.append((int(key) + 3)/1000)
        else:
            timings.append(int(key)/1000)
            timings.append((int(key) + 3)/1000)
            
    timings.append(total_circuit_time/1000)
    
    return timings

# Function for extracting positions where corrections were done
def correction_positions(corr_positions, correction_avg_fidelity, correction_snapshot_timings):
    corr_time = [key/1000 for key in corr_positions]
    corr_fid = [correction_avg_fidelity[correction_snapshot_timings.index(timing)] for timing in corr_time]
    
    return corr_time, corr_fid

# Function for fitting exponential curve
def curve_fitter(x, y, sign=1):
    popt, pcov = curve_fit(lambda x, b: sign * np.exp(-x/b), x, y)
    sigma = np.sqrt(np.diag(pcov))
    
    lifetime = round(popt[0], 2)
    uncertainty = round(sigma[0], 3)
    
    return lifetime, uncertainty

### Adjust these parameters ###
num_corr = 1 # per ms, number of correction circuits applied
additional_snaps = 2 # per ms, number of additional snapshots taken in between each correction circuit

noise_only_circ_time = 12 # ms, total duration of noise-only circuit
correction_circ_time = 12 # ms, total duration of correction circuit

shots = 1000 # Number of times circuit is repeated
encoding=['logi0', 'logi1', 'logi_xplus', 'logi_xminus'] # Input 'logi0', 'logi1', 'logi_xplus', 'logi_xminus', 'logi_yplus', or 'logi_yminus'
## 

### If you wish to generate the circuit diagram for a particular simulation, uncomment these lines. ###
# noise_circuit = error_correction_experiment(encoding[0], 0, shots, 10, noise_only_circ_time, return_circuit=True)
# correction_circuit = error_correction_experiment(encoding[0], num_corr, shots, additional_snaps, correction_circ_time, return_circuit=True)

# noise_circuit.draw()
# correction_circuit.draw()

## For pyplot formatting
# For formatting graphs
encoding_format = {'logi0': r'$0_L$', 'logi1': r'$1_L$', 'logi_xplus': r'$+_L$', 'logi_xminus':r'$-_L$', 'logi_yplus':r'$+i_L$', 'logi_yminus': r'$-i_L$'}

# Colors for plotting
colors= ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']

# Labels for logical states
logical_label = [r'$\lvert 0_L \rangle = (\lvert 0 \rangle + \lvert 4 \rangle)/\sqrt{2}$', 
                     r'$\lvert E_2 \rangle = (\lvert 0 \rangle - \lvert 4 \rangle)/\sqrt{2}$', 
                     r'$\lvert 1_L \rangle$']

# Run simulation with and without correction. Disable u_even
noise_only_simulation_list = []
correction_simulation_list = []

for i in encoding:
    noise_only_simulation_list.append(error_correction_experiment(i, 0, shots, additional_snaps, noise_only_circ_time))
    correction_simulation_list.append(error_correction_experiment(i, num_corr, shots, additional_snaps, correction_circ_time, identity=True))

## Fidelity visualisation
index = encoding.index('logi1')
value = 'logi1'

_, noise_only_fid, _, noise_only_snapshot_timings, _ = noise_only_simulation_list[index]
_, correction_fid, _, correction_snapshot_timings, corr_position_times = correction_simulation_list[index]

noise_only_avg_fid = [round(np.average(cycle), 4) for cycle in noise_only_fid]
noise_only_avg_fid_error = [round(np.std(cycle)/np.sqrt(shots), 4) for cycle in noise_only_fid]

correction_avg_fid = [round(np.average(cycle), 4) for cycle in correction_fid]
correction_avg_fid_error = [round(np.std(cycle)/np.sqrt(shots), 4) for cycle in correction_fid]

# Find correction positions
x, y = correction_positions(corr_position_times, correction_avg_fid, correction_snapshot_timings)

# Fit fidelity data using an exponential decay function
noise_only_lifetime, noise_lifetime_uncertainty = curve_fitter(noise_only_snapshot_timings, noise_only_avg_fid)
correction_lifetime_perf, correction_lifetime_uncertainty = curve_fitter(x, y)

# Round values to 2 decimal places
noise_lifetime_uncertainty = round(noise_lifetime_uncertainty, 2)
correction_lifetime_uncertainty = round (correction_lifetime_uncertainty, 2)

# Plot fidelity data w/ uncertainty
plt.figure(figsize=(15,6))

plt.errorbar(noise_only_snapshot_timings, noise_only_avg_fid, yerr=noise_only_avg_fid_error,
    label=r'No corrections, ' + '\n' + r'$\tau = $' + str(noise_only_lifetime) + r'ms $\pm$' + str(noise_lifetime_uncertainty) + r'ms', 
    marker='o', color='#5BCEFA', capsize=3, ecolor='#2a6278', linestyle='None')
plt.errorbar(correction_snapshot_timings, correction_avg_fid, yerr=correction_avg_fid_error,
    label=str(num_corr) + r' corrections/ms,' + '\n' + r'$\tau = $' + str(correction_lifetime_perf) + r'ms $\pm$' + str(correction_lifetime_uncertainty) + r'ms', 
    marker='o', color='#F5A9B8', capsize=3, ecolor='#9c6b75')

# Plot curve fits w/ uncertainty
noise_curve_timing = np.linspace(0, noise_only_circ_time, 100)
plt.plot(noise_curve_timing, np.exp(-noise_curve_timing/noise_only_lifetime))
plt.fill_between(noise_curve_timing, np.exp(-noise_curve_timing/(noise_only_lifetime + noise_lifetime_uncertainty)), 
            np.exp(-noise_curve_timing/(noise_only_lifetime - noise_lifetime_uncertainty)), color='#5BCEFA', alpha=0.3)

corr_curve_timing = np.linspace(0, correction_circ_time, 100)
plt.plot(corr_curve_timing, np.exp(-corr_curve_timing/correction_lifetime_perf))
plt.fill_between(corr_curve_timing, np.exp(-corr_curve_timing/(correction_lifetime_perf + correction_lifetime_uncertainty)), 
            np.exp(-corr_curve_timing/(correction_lifetime_perf - correction_lifetime_uncertainty)), color='#F5A9B8', alpha=0.3)

# Plot correction positions
label = 'Correction \n position'
for i in x:
    plt.axvline(x=i, color='grey', label=label, linestyle="--")
    label = ''

plt.xlabel("Time (ms)", size=20)
plt.ylabel("State fidelity, averaged over {} shots".format(shots), size=20)
plt.title(r'$\lvert$' + str(encoding_format.get(value)) + r'$\rangle$' + ' state fidelity without error correction and with ' + str(num_corr) + r' correction(s)/ms ($\hat{U}_{\text{odd}}$ only)', size=20)
plt.grid(color = 'lightgrey', linestyle = 'dashed', linewidth = 0.5)
plt.legend(loc='upper right', fontsize=20)
plt.xticks(fontsize=20)
plt.yticks(fontsize=20)
plt.show()

## Fidelity visualisation
for index, value in enumerate(encoding):
    if value == 'logi1':
        continue
    _, noise_only_fid_other, _, noise_only_snapshot_timings, _ = noise_only_simulation_list[index]
    _, correction_fid_other, _, correction_snapshot_timings, corr_position_times = correction_simulation_list[index]
    
    noise_only_avg_fid_other = [round(np.average(cycle), 4) for cycle in noise_only_fid_other]
    noise_only_avg_fid_error_other = [round(np.std(cycle)/np.sqrt(shots), 4) for cycle in noise_only_fid_other]
    
    correction_avg_fid_other = [round(np.average(cycle), 4) for cycle in correction_fid_other]
    correction_avg_fid_error_other = [round(np.std(cycle)/np.sqrt(shots), 4) for cycle in correction_fid_other]
    
    # Find correction positions
    x, y = correction_positions(corr_position_times, correction_avg_fid_other, correction_snapshot_timings)

    # Plot fidelity data w/ uncertainty
    plt.figure(figsize=(15,6))

    plt.errorbar(noise_only_snapshot_timings, noise_only_avg_fid_other, yerr=noise_only_avg_fid_error_other,
        label=r'No corrections', 
        marker='o', color='#5BCEFA', capsize=3, ecolor='#2a6278')
    plt.errorbar(correction_snapshot_timings, correction_avg_fid_other, yerr=correction_avg_fid_error_other,
        label=str(num_corr) + r' corrections/ms', 
        marker='o', color='#F5A9B8', capsize=3, ecolor='#9c6b75')

    # Plot an actual exponential decay as comparison
    noise_curve_timing = np.linspace(0, noise_only_circ_time, 100)
    plt.plot(noise_curve_timing, np.exp(-noise_curve_timing), label='Exponential curve')

    # Plot correction positions
    label = 'Correction \n position'
    for i in x:
        plt.axvline(x=i, color='grey', label=label, linestyle="--")
        label = ''

    plt.xlabel("Time (ms)", size=20)
    plt.ylabel("State fidelity, averaged over {} shots".format(shots), size=20)
    plt.title(r'$\lvert$' + str(encoding_format.get(value)) + r'$\rangle$' + ' state fidelity without error correction and with ' + str(num_corr) + r' correction(s)/ms ($\hat{U}_{\text{odd}}$ only)', size=20)
    plt.grid(color = 'lightgrey', linestyle = 'dashed', linewidth = 0.5)
    plt.legend(loc='upper right', fontsize=20)
    plt.xticks(fontsize=20)
    plt.yticks(fontsize=20)
    plt.show()

# Compute exponential fit across a variety of corrections/ms for |1_L> only
value = 'logi1'

_num_corr = [1, 2, 5, 10]
_correction_circ_time = [12 for _ in _num_corr] #[(i + 1) * 8 for i in range(5)]
_additional_snaps = 0
_shots = 1000
_encoding='logi1'

# Colors for plotting
colors= ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']

# Plot noise_only trial as simulated previously
plt.figure(figsize=(15,6))
plt.scatter(noise_only_snapshot_timings, noise_only_avg_fid, marker='o', color='#1f77b4', zorder=3, alpha=0.6, linestyle='None', s= 50)
plt.plot(noise_curve_timing, np.exp(-noise_curve_timing/noise_only_lifetime), label=r'0 corrections, ' + '\n' + r'$\tau$ = ' + str(noise_only_lifetime) + r'ms $\pm$' + str(noise_lifetime_uncertainty) + r'ms', color='#1f77b4')

# Plot all other correction trials w/ lifetimes
for index, _corrections in enumerate(_num_corr):
    # print(corrections, "Run starting") # Debug
    _, _correction_fid, _, _correction_snapshot_timings, _corr_position_times = error_correction_experiment(_encoding, _corrections, _shots, _additional_snaps, _correction_circ_time[index], identity=True)
    _correction_avg_fid = [round(np.average(cycle), 4) for cycle in _correction_fid]
    
    _x, _y = correction_positions(_corr_position_times, _correction_avg_fid, _correction_snapshot_timings)
    _lifetime, _lifetime_uncertainty = curve_fitter(_x, _y)
    
    _curve_timing = np.linspace(0, _correction_circ_time[index], 100)
    if _additional_snaps != 0:
        plt.plot(correction_snapshot_timings, correction_avg_fid, marker='o', color=colors[index + 1], alpha=0.2)

    plt.plot(_curve_timing, np.exp(- _curve_timing/_lifetime), 
             label=str(_num_corr[index]) + r' corrections/ms, ' + '\n' + r'$\tau = $' + str(_lifetime) + r'ms $\pm$' + str(_lifetime_uncertainty) + r'ms', color=colors[index + 1])
    
    plt.scatter(_x, _y, marker='s', color=colors[index + 1], zorder=3, alpha=0.6, linestyle='None', s= 50)

plt.xlabel("Time (ms)", size=20)
plt.ylabel("State fidelity, averaged over {} shots".format(shots), size=20)
plt.title(r'$\lvert$' + str(encoding_format.get(value)) + r'$\rangle$' + r' state fidelity for varying circuit parameters ($\hat{U}_{\text{odd}}$ only)', size=20)
plt.grid(color = 'lightgrey', linestyle = 'dashed', linewidth = 0.5)
plt.legend(fontsize=20, bbox_to_anchor=(1.32, 0.5), loc='right')
plt.xticks(fontsize=20)
plt.yticks(fontsize=20)
plt.show()

### Adjust these parameters ###
run = [(14, 0), (1, 1), (3, 4), (0, 12)] # Choose single shot to graph. Tuples are formatted as (noise_only_shot, correction_shot).
## 

# Plot photon probability for all initial encodings
for index, value in enumerate(encoding):
    _, _, noise_only_prob_tuple, noise_only_snapshot_timings, _ = noise_only_simulation_list[index]
    _, _, correction_prob_tuple, correction_snapshot_timings, corr_position_times = correction_simulation_list[index]

    # Fock state probability visualization, single shot
    noise_only_prob, noise_only_prob_logical  = noise_only_prob_tuple
    correction_prob, correction_prob_logical = correction_prob_tuple

    # Runs to plot
    run_noise, run_corr = run[index]
    
    noise_only_fock_prob = [i[run_noise] for i in noise_only_prob]
    correction_fock_prob = [i[run_corr] for i in correction_prob]

    noise_only_fock_prob_logical = [i[run_noise] for i in noise_only_prob_logical]
    correction_fock_prob_logical = [i[run_corr] for i in correction_prob_logical]

    # Ignore trajectories of Fock states with prob below a certain value.
    prob_truncation = 0.005

    # Plot both graph in single image
    plt.figure(figsize=(15,10))
    plt.suptitle('Trajectories of Fock state probability simulations with ' + r'$\lvert$' + str(encoding_format.get(value)) + r'$\rangle$' + ' encoding', size=30)
    # plt.suptitle('No correction and ' + str(num_corr) + ' correction simulation, with ' +r'$\lvert$' + str(encoding_format.get(encoding)) + r'$\rangle$' + ' encoding', size=20)

    ## 1a. Noise only circuit plot (photon number)
    plt.subplot(2, 1, 1)
    plt.subplot(2, 1, 1).xaxis.set_tick_params(labelsize=20)
    plt.subplot(2, 1, 1).yaxis.set_tick_params(labelsize=20)
    n = len(noise_only_fock_prob[0])
    # n,  = noise_only_avg_prob[0].shape
    for index in range(n):
        trajectory = [value[index] for value in noise_only_fock_prob]
        # uncert = [value[index] for value in noise_only_fock_prob_error]
        if all(i <= prob_truncation for i in trajectory):
            continue
        
        state = index
        # Rename state |2> = |1_L>
        if index == 2:
            state=logical_label[2]
            plt.plot(noise_only_snapshot_timings, trajectory, marker='s', label=r'$\lvert 2 \rangle =$ ' + str(state), color=colors[index], markersize = 8, linewidth = 3)
        else:
            plt.plot(noise_only_snapshot_timings, trajectory, marker='o', label=r'$\lvert$' + str(state) + r'$\rangle$', color=colors[index])

    ## 1b. Noise only circuit plot (logical state)
    m = len(noise_only_fock_prob_logical[0])
    for index in range(m):
        trajectory = [value[index] for value in noise_only_fock_prob_logical]
        if all(i <= prob_truncation for i in trajectory):
            continue
        plt.plot(noise_only_snapshot_timings, trajectory, marker='s', label=str(logical_label[index]), color=colors[index+8], markersize = 8, linewidth = 3)

    # Subplot 1 formatting
    #plt.ylabel("Probability of measuring Fock state, averaged over over {} shots".format(shots))
    plt.ylabel(r'$P(\lvert \psi \rangle$)', size=20)
    plt.title(r"Without error correction", size=20)
    plt.grid(color = 'lightgrey', linestyle = 'dashed', linewidth = 0.5)

    ## 2a. Correction circuit plot (photon number)
    plt.subplot(2, 1, 2, sharex = plt.subplot(2, 1, 1))
    plt.subplot(2, 1, 2).xaxis.set_tick_params(labelsize=20)
    plt.subplot(2, 1, 2).yaxis.set_tick_params(labelsize=20)
    n = len(correction_fock_prob[0])
    # n,  = correction_avg_prob[0].shape
    for index in range(n):
        trajectory = [value[index] for value in correction_fock_prob]
        # print(round(np.max(trajectory),4), round(np.median(trajectory),4), round(np.ptp(trajectory),4), index) ### debug
        # uncert = [value[index] for value in correction_fock_prob_error]
        if all(i <= prob_truncation for i in trajectory):
            continue
        state = index
        # Rename state |2> = |1_L>
        if index == 2:
            state=logical_label[2]
            plt.plot(correction_snapshot_timings, trajectory, marker='s', label=r'$\lvert 2 \rangle =$ ' + str(state), color=colors[index], markersize = 8, linewidth = 3)
        else:
            plt.plot(correction_snapshot_timings, trajectory, marker='o', label=r'$\lvert$' + str(state) + r'$\rangle$', color=colors[index])

    ## 1b. Correction circuit plot (logical state)
    m = len(correction_fock_prob_logical[0])
    for index in range(m):
        trajectory = [value[index] for value in correction_fock_prob_logical]
        if all(i <= prob_truncation for i in trajectory):
            continue
        plt.plot(correction_snapshot_timings, trajectory, marker='s', label=str(logical_label[index]), color=colors[index+8], markersize = 8, linewidth = 3)

    # Plot correction positions
    x, y = correction_positions(corr_position_times, correction_fock_prob, correction_snapshot_timings)
    label = 'Correction \n position'
    for i in x:
        plt.axvline(x=i, color='grey', label=label, linestyle="--")
        label = ''

    # Subplot 2 formatting
    plt.xlabel("Time (ms)", size=20)
    #plt.ylabel("Probability of measuring Fock state, averaged over {} shots".format(shots))
    plt.ylabel(r'$P(\lvert \psi \rangle$)', size=20)
    plt.title("With " + str(num_corr) + r" correction(s)/ms ($\hat{U}_{\text{odd}}$ only)", size=20)
    plt.grid(color = 'lightgrey', linestyle = 'dashed', linewidth = 0.5)

    # Remove duplicate labels
    leg1 = {l:h for h,l in zip(*plt.subplot(2, 1, 1).get_legend_handles_labels())}
    leg2 = {l:h for h,l in zip(*plt.subplot(2, 1, 2).get_legend_handles_labels())}

    plt.figlegend(*[*zip(*{**leg1, **leg2}.items())][::-1], bbox_to_anchor=(1.15, 0.5), loc = 'right', fontsize=20)
    plt.show()

# Run simulation with and without correction. Enable u_even
noise_only_simulation_list = []
correction_simulation_list = []

for i in encoding:
    noise_only_simulation_list.append(error_correction_experiment(i, 0, shots, additional_snaps, noise_only_circ_time))
    correction_simulation_list.append(error_correction_experiment(i, num_corr, shots, additional_snaps, correction_circ_time))

## Fidelity visualisation
index = encoding.index('logi1')
value = 'logi1'

_, noise_only_fid, _, noise_only_snapshot_timings, _ = noise_only_simulation_list[index]
_, correction_fid, _, correction_snapshot_timings, corr_position_times = correction_simulation_list[index]

noise_only_avg_fid = [round(np.average(cycle), 4) for cycle in noise_only_fid]
noise_only_avg_fid_error = [round(np.std(cycle)/np.sqrt(shots), 4) for cycle in noise_only_fid]

correction_avg_fid = [round(np.average(cycle), 4) for cycle in correction_fid]
correction_avg_fid_error = [round(np.std(cycle)/np.sqrt(shots), 4) for cycle in correction_fid]

# Find correction positions
x, y = correction_positions(corr_position_times, correction_avg_fid, correction_snapshot_timings)

# Plot fidelity data w/ uncertainty
plt.figure(figsize=(15,6))

plt.errorbar(noise_only_snapshot_timings, noise_only_avg_fid, yerr=noise_only_avg_fid_error,
    label=r'No corrections', marker='o', color='#5BCEFA', capsize=3, ecolor='#2a6278')
plt.errorbar(correction_snapshot_timings, correction_avg_fid, yerr=correction_avg_fid_error,
    label=str(num_corr) + r' corrections/ms', marker='o', color='#F5A9B8', capsize=3, ecolor='#9c6b75')

# Plot correction positions
label = 'Correction \n position'
for i in x:
    plt.axvline(x=i, color='grey', label=label, linestyle="--")
    label = ''

# Plot previous logi1 lifetime
plt.plot(noise_curve_timing, np.exp(-noise_curve_timing/correction_lifetime_perf), label=r'Simulation 1 result' + '\n' + r'$\tau = $' + str(correction_lifetime_perf) + r'ms')
    
plt.xlabel("Time (ms)", size=20)
plt.ylabel("State fidelity, averaged over {} shots".format(shots), size=20)
plt.title(r'$\lvert$' + str(encoding_format.get(value)) + r'$\rangle$' + ' state fidelity without error correction and with ' + str(num_corr) + r' correction(s)/ms ($\hat{U}_{\text{odd}}$ and $\hat{U}_{\text{even}}$)', size=20)
plt.grid(color = 'lightgrey', linestyle = 'dashed', linewidth = 0.5)
plt.legend(loc='upper right', fontsize=20)
plt.xticks(fontsize=20)
plt.yticks(fontsize=20)
plt.show()

## Fidelity visualisation
for index, value in enumerate(encoding):
    if value == 'logi1':
        continue
    _, noise_only_fid_other, _, noise_only_snapshot_timings, _ = noise_only_simulation_list[index]
    _, correction_fid_other, _, correction_snapshot_timings, corr_position_times = correction_simulation_list[index]
    
    noise_only_avg_fid_other = [round(np.average(cycle), 4) for cycle in noise_only_fid_other]
    noise_only_avg_fid_error_other = [round(np.std(cycle)/np.sqrt(shots), 4) for cycle in noise_only_fid_other]
    
    correction_avg_fid_other = [round(np.average(cycle), 4) for cycle in correction_fid_other]
    correction_avg_fid_error_other = [round(np.std(cycle)/np.sqrt(shots), 4) for cycle in correction_fid_other]
    
    # Find correction positions
    x, y = correction_positions(corr_position_times, correction_avg_fid_other, correction_snapshot_timings)

    # Plot fidelity data w/ uncertainty
    plt.figure(figsize=(15,6))

    plt.errorbar(noise_only_snapshot_timings, noise_only_avg_fid_other, yerr=noise_only_avg_fid_error_other,
        label=r'No corrections', 
        marker='o', color='#5BCEFA', capsize=3, ecolor='#2a6278')
    plt.errorbar(correction_snapshot_timings, correction_avg_fid_other, yerr=correction_avg_fid_error_other,
        label=str(num_corr) + r' corrections/ms', 
        marker='o', color='#F5A9B8', capsize=3, ecolor='#9c6b75')

    # Plot an actual exponential decay as comparison
    noise_curve_timing = np.linspace(0, noise_only_circ_time, 100)
    plt.plot(noise_curve_timing, np.exp(-noise_curve_timing), label='Exponential curve')

    # Plot correction positions
    label = 'Correction \n position'
    for i in x:
        plt.axvline(x=i, color='grey', label=label, linestyle="--")
        label = ''

    plt.xlabel("Time (ms)", size=20)
    plt.ylabel("State fidelity, averaged over {} shots".format(shots), size=20)
    plt.title(r'$\lvert$' + str(encoding_format.get(value)) + r'$\rangle$' + ' state fidelity without error correction and with ' + str(num_corr) + r' correction(s)/ms ($\hat{U}_{\text{odd}}$ and $\hat{U}_{\text{even}}$)', size=20)
    plt.grid(color = 'lightgrey', linestyle = 'dashed', linewidth = 0.5)
    plt.legend(loc='upper right', fontsize=20)
    plt.xticks(fontsize=20)
    plt.yticks(fontsize=20)
    plt.show()

### Adjust these parameters ###
run = [(41, 24), (3, 3), _, _] # Choose single shot to graph. Tuples are formatted as (noise_only_shot, correction_shot).
##

# Plot photon probability for all initial encodings
for index, value in enumerate(encoding):
    if index not in [0, 1]:
        continue
    _, _, noise_only_prob_tuple, noise_only_snapshot_timings, _ = noise_only_simulation_list[index]
    _, _, correction_prob_tuple, correction_snapshot_timings, corr_position_times = correction_simulation_list[index]

    # Fock state probability visualization, single shot
    noise_only_prob, noise_only_prob_logical  = noise_only_prob_tuple
    correction_prob, correction_prob_logical = correction_prob_tuple

    # Runs to plot
    run_noise, run_corr = run[index]
    
    noise_only_fock_prob = [i[run_noise] for i in noise_only_prob]
    correction_fock_prob = [i[run_corr] for i in correction_prob]

    noise_only_fock_prob_logical = [i[run_noise] for i in noise_only_prob_logical]
    correction_fock_prob_logical = [i[run_corr] for i in correction_prob_logical]

    # Ignore trajectories of Fock states with prob below a certain value.
    prob_truncation = 0.005

    # Plot both graph in single image
    plt.figure(figsize=(15,10))
    plt.suptitle('Trajectories of Fock state probability simulations with ' + r'$\lvert$' + str(encoding_format.get(value)) + r'$\rangle$' + ' encoding', size=30)
    # plt.suptitle('No correction and ' + str(num_corr) + ' correction simulation, with ' +r'$\lvert$' + str(encoding_format.get(encoding)) + r'$\rangle$' + ' encoding', size=20)

    ## 1a. Noise only circuit plot (photon number)
    plt.subplot(2, 1, 1)
    plt.subplot(2, 1, 1).xaxis.set_tick_params(labelsize=20)
    plt.subplot(2, 1, 1).yaxis.set_tick_params(labelsize=20)
    n = len(noise_only_fock_prob[0])
    # n,  = noise_only_avg_prob[0].shape
    for index in range(n):
        trajectory = [value[index] for value in noise_only_fock_prob]
        # uncert = [value[index] for value in noise_only_fock_prob_error]
        if all(i <= prob_truncation for i in trajectory):
            continue
        
        state = index
        # Rename state |2> = |1_L>
        if index == 2:
            state=logical_label[2]
            plt.plot(noise_only_snapshot_timings, trajectory, marker='s', label=r'$\lvert 2 \rangle =$ ' + str(state), color=colors[index], markersize = 8, linewidth = 3)
        else:
            plt.plot(noise_only_snapshot_timings, trajectory, marker='o', label=r'$\lvert$' + str(state) + r'$\rangle$', color=colors[index])

    ## 1b. Noise only circuit plot (logical state)
    m = len(noise_only_fock_prob_logical[0])
    for index in range(m):
        trajectory = [value[index] for value in noise_only_fock_prob_logical]
        if all(i <= prob_truncation for i in trajectory):
            continue
        plt.plot(noise_only_snapshot_timings, trajectory, marker='s', label=str(logical_label[index]), color=colors[index+8], markersize = 8, linewidth = 3)

    # Subplot 1 formatting
    #plt.ylabel("Probability of measuring Fock state, averaged over over {} shots".format(shots))
    plt.ylabel(r'$P(\lvert \psi \rangle$)', size=20)
    plt.title(r"Without error correction", size=20)
    plt.grid(color = 'lightgrey', linestyle = 'dashed', linewidth = 0.5)

    ## 2a. Correction circuit plot (photon number)
    plt.subplot(2, 1, 2, sharex = plt.subplot(2, 1, 1))
    plt.subplot(2, 1, 2).xaxis.set_tick_params(labelsize=20)
    plt.subplot(2, 1, 2).yaxis.set_tick_params(labelsize=20)
    n = len(correction_fock_prob[0])
    # n,  = correction_avg_prob[0].shape
    for index in range(n):
        trajectory = [value[index] for value in correction_fock_prob]
        # print(round(np.max(trajectory),4), round(np.median(trajectory),4), round(np.ptp(trajectory),4), index) ### debug
        # uncert = [value[index] for value in correction_fock_prob_error]
        if all(i <= prob_truncation for i in trajectory):
            continue
        state = index
        # Rename state |2> = |1_L>
        if index == 2:
            state=logical_label[2]
            plt.plot(correction_snapshot_timings, trajectory, marker='s', label=r'$\lvert 2 \rangle =$ ' + str(state), color=colors[index], markersize = 8, linewidth = 3)
        else:
            plt.plot(correction_snapshot_timings, trajectory, marker='o', label=r'$\lvert$' + str(state) + r'$\rangle$', color=colors[index])

    ## 1b. Correction circuit plot (logical state)
    m = len(correction_fock_prob_logical[0])
    for index in range(m):
        trajectory = [value[index] for value in correction_fock_prob_logical]
        if all(i <= prob_truncation for i in trajectory):
            continue
        plt.plot(correction_snapshot_timings, trajectory, marker='s', label=str(logical_label[index]), color=colors[index+8], markersize = 8, linewidth = 3)

    # Plot correction positions
    x, y = correction_positions(corr_position_times, correction_fock_prob, correction_snapshot_timings)
    label = 'Correction \n position'
    for i in x:
        plt.axvline(x=i, color='grey', label=label, linestyle="--")
        label = ''

    # Subplot 2 formatting
    plt.xlabel("Time (ms)", size=20)
    #plt.ylabel("Probability of measuring Fock state, averaged over {} shots".format(shots))
    plt.ylabel(r'$P(\lvert \psi \rangle$)', size=20)
    plt.title("With " + str(num_corr) + r" correction(s)/ms ($\hat{U}_{\text{odd}}$ and $\hat{U}_{\text{even}}$)", size=20)
    plt.grid(color = 'lightgrey', linestyle = 'dashed', linewidth = 0.5)

    # Remove duplicate labels
    leg1 = {l:h for h,l in zip(*plt.subplot(2, 1, 1).get_legend_handles_labels())}
    leg2 = {l:h for h,l in zip(*plt.subplot(2, 1, 2).get_legend_handles_labels())}

    plt.figlegend(*[*zip(*{**leg1, **leg2}.items())][::-1], bbox_to_anchor=(1.15, 0.5), loc = 'right', fontsize=20)
    plt.show()

### Adjust these parameters ###
run = [_, _, (0, 3), (2, 8)] # Choose single shot to graph. Tuples are formatted as (noise_only_shot, correction_shot).
##

# Plot photon probability for all initial encodings
for index, value in enumerate(encoding):
    if index not in [2, 3]:
        continue
    _, _, noise_only_prob_tuple, noise_only_snapshot_timings, _ = noise_only_simulation_list[index]
    _, _, correction_prob_tuple, correction_snapshot_timings, corr_position_times = correction_simulation_list[index]

    # Fock state probability visualization, single shot
    noise_only_prob, noise_only_prob_logical  = noise_only_prob_tuple
    correction_prob, correction_prob_logical = correction_prob_tuple

    # Runs to plot
    run_noise, run_corr = run[index]
    
    noise_only_fock_prob = [i[run_noise] for i in noise_only_prob]
    correction_fock_prob = [i[run_corr] for i in correction_prob]

    noise_only_fock_prob_logical = [i[run_noise] for i in noise_only_prob_logical]
    correction_fock_prob_logical = [i[run_corr] for i in correction_prob_logical]

    # Ignore trajectories of Fock states with prob below a certain value.
    prob_truncation = 0.005

    # Plot both graph in single image
    plt.figure(figsize=(15,10))
    plt.suptitle('Trajectories of Fock state probability simulations with ' + r'$\lvert$' + str(encoding_format.get(value)) + r'$\rangle$' + ' encoding', size=30)
    # plt.suptitle('No correction and ' + str(num_corr) + ' correction simulation, with ' +r'$\lvert$' + str(encoding_format.get(encoding)) + r'$\rangle$' + ' encoding', size=20)

    ## 1a. Noise only circuit plot (photon number)
    plt.subplot(2, 1, 1)
    plt.subplot(2, 1, 1).xaxis.set_tick_params(labelsize=20)
    plt.subplot(2, 1, 1).yaxis.set_tick_params(labelsize=20)
    n = len(noise_only_fock_prob[0])
    # n,  = noise_only_avg_prob[0].shape
    for index in range(n):
        trajectory = [value[index] for value in noise_only_fock_prob]
        # uncert = [value[index] for value in noise_only_fock_prob_error]
        if all(i <= prob_truncation for i in trajectory):
            continue
        
        state = index
        # Rename state |2> = |1_L>
        if index == 2:
            state=logical_label[2]
            plt.plot(noise_only_snapshot_timings, trajectory, marker='s', label=r'$\lvert 2 \rangle =$ ' + str(state), color=colors[index], markersize = 8, linewidth = 3)
        else:
            plt.plot(noise_only_snapshot_timings, trajectory, marker='o', label=r'$\lvert$' + str(state) + r'$\rangle$', color=colors[index])

    ## 1b. Noise only circuit plot (logical state)
    m = len(noise_only_fock_prob_logical[0])
    for index in range(m):
        trajectory = [value[index] for value in noise_only_fock_prob_logical]
        if all(i <= prob_truncation for i in trajectory):
            continue
        plt.plot(noise_only_snapshot_timings, trajectory, marker='s', label=str(logical_label[index]), color=colors[index+8], markersize = 8, linewidth = 3)

    # Subplot 1 formatting
    #plt.ylabel("Probability of measuring Fock state, averaged over over {} shots".format(shots))
    plt.ylabel(r'$P(\lvert \psi \rangle$)', size=20)
    plt.title(r"Without error correction", size=20)
    plt.grid(color = 'lightgrey', linestyle = 'dashed', linewidth = 0.5)

    ## 2a. Correction circuit plot (photon number)
    plt.subplot(2, 1, 2, sharex = plt.subplot(2, 1, 1))
    plt.subplot(2, 1, 2).xaxis.set_tick_params(labelsize=20)
    plt.subplot(2, 1, 2).yaxis.set_tick_params(labelsize=20)
    n = len(correction_fock_prob[0])
    # n,  = correction_avg_prob[0].shape
    for index in range(n):
        trajectory = [value[index] for value in correction_fock_prob]
        # print(round(np.max(trajectory),4), round(np.median(trajectory),4), round(np.ptp(trajectory),4), index) ### debug
        # uncert = [value[index] for value in correction_fock_prob_error]
        if all(i <= prob_truncation for i in trajectory):
            continue
        state = index
        # Rename state |2> = |1_L>
        if index == 2:
            state=logical_label[2]
            plt.plot(correction_snapshot_timings, trajectory, marker='s', label=r'$\lvert 2 \rangle =$ ' + str(state), color=colors[index], markersize = 8, linewidth = 3)
        else:
            plt.plot(correction_snapshot_timings, trajectory, marker='o', label=r'$\lvert$' + str(state) + r'$\rangle$', color=colors[index])

    ## 1b. Correction circuit plot (logical state)
    m = len(correction_fock_prob_logical[0])
    for index in range(m):
        trajectory = [value[index] for value in correction_fock_prob_logical]
        if all(i <= prob_truncation for i in trajectory):
            continue
        plt.plot(correction_snapshot_timings, trajectory, marker='s', label=str(logical_label[index]), color=colors[index+8], markersize = 8, linewidth = 3)

    # Plot correction positions
    x, y = correction_positions(corr_position_times, correction_fock_prob, correction_snapshot_timings)
    label = 'Correction \n position'
    for i in x:
        plt.axvline(x=i, color='grey', label=label, linestyle="--")
        label = ''

    # Subplot 2 formatting
    plt.xlabel("Time (ms)", size=20)
    #plt.ylabel("Probability of measuring Fock state, averaged over {} shots".format(shots))
    plt.ylabel(r'$P(\lvert \psi \rangle$)', size=20)
    plt.title("With " + str(num_corr) + r" correction(s)/ms ($\hat{U}_{\text{odd}}$ and $\hat{U}_{\text{even}}$)", size=20)
    plt.grid(color = 'lightgrey', linestyle = 'dashed', linewidth = 0.5)

    # Remove duplicate labels
    leg1 = {l:h for h,l in zip(*plt.subplot(2, 1, 1).get_legend_handles_labels())}
    leg2 = {l:h for h,l in zip(*plt.subplot(2, 1, 2).get_legend_handles_labels())}

    plt.figlegend(*[*zip(*{**leg1, **leg2}.items())][::-1], bbox_to_anchor=(1.15, 0.5), loc = 'right', fontsize=20)
    plt.show()

### Adjust these parameters ###
num_corr = 2 # per ms, number of correction circuits applied
additional_snaps = 0 # per ms, number of additional snapshots taken in between each correction circuit

noise_only_circ_time = 2 # ms, total duration of noise-only circuit
correction_circ_time = 5 # ms, total duration of correction circuit

shots = 1000 # Number of times circuit is repeated
encoding=['logi_xplus', 'logi_yplus', 'logi1'] # Input 'logi0', 'logi1', 'logi_xplus', 'logi_xminus', 'logi_yplus', or 'logi_yminus'
##

# Run simulation with and without correction. Enable u_even
noise_only_simulation_list = []
correction_simulation_list = []

for i in encoding:
    noise_only_simulation_list.append(error_correction_experiment(i, 0, shots, 2, noise_only_circ_time))
    correction_simulation_list.append(error_correction_experiment(i, num_corr, shots, additional_snaps, correction_circ_time))

# Backup chunk
## Pauli Expectation Value visualisation
colors= ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']
label = 'Correction \n position'

plt.figure(figsize=(15,6))

# Plot with correction
for index, value in enumerate(encoding):
    if value in ['logi0', 'logi_xplus', 'logi_yplus']:
        sign = 1
    else:
        sign = -1
        
    # Plot without correction
    (noise_only_pauli_expectation, _), _, _, noise_only_snapshot_timings, _ = noise_only_simulation_list[index]
    
    noise_only_expectation_avg = [np.round(np.average(cycle), 4) for cycle in noise_only_pauli_expectation]
    noise_only_expectation_avg_error = [round(np.std(cycle)/np.sqrt(shots), 4) for cycle in noise_only_pauli_expectation]

    noise_only_lifetime, noise_lifetime_uncertainty = curve_fitter(noise_only_snapshot_timings, noise_only_expectation_avg, sign)
    plt.errorbar(noise_only_snapshot_timings, noise_only_expectation_avg, yerr=noise_only_expectation_avg_error,
        label=r'$\lvert$' + str(encoding_format.get(value)) + r'$\rangle$' + r' no corrections,' + '\n' + r'$\tau$ = ' + str(noise_only_lifetime) + r'ms $\pm$' + str(round(noise_lifetime_uncertainty, 2)) + r'ms',
        marker='o', color=colors[-index - 1], capsize=3, ecolor='#2a6278')

    noise_curve_timing = np.linspace(0, noise_only_circ_time, 100)
    plt.plot(noise_curve_timing, sign * np.exp(-noise_curve_timing/noise_only_lifetime), alpha=0.5, color=colors[-index - 1])
    
    # Plot with correction
    (correction_pauli_expectation, _), _, _, correction_snapshot_timings, corr_position_times = correction_simulation_list[index]

    correction_expectation_avg = [np.round(np.average(cycle), 4) for cycle in correction_pauli_expectation]
    correction_expectation_avg_error = [round(np.std(cycle)/np.sqrt(shots), 4) for cycle in correction_pauli_expectation]
    
    correction_lifetime, correction_lifetime_uncertainty = curve_fitter(correction_snapshot_timings, correction_expectation_avg, sign)
    plt.errorbar(correction_snapshot_timings, correction_expectation_avg, yerr=correction_expectation_avg_error,
        label=r'$\lvert$' + str(encoding_format.get(value)) + r'$\rangle$ ' + str(num_corr) + r' corrections/ms,'  + '\n' + r'$\tau$ = ' + str(correction_lifetime) + r'ms $\pm$' + str(round(correction_lifetime_uncertainty, 2)) + r'ms', 
        marker='o', color=colors[index], capsize=3, ecolor='#9c6b75')

    corr_curve_timing = np.linspace(0, correction_circ_time, 100)
    plt.plot(corr_curve_timing, sign * np.exp(-corr_curve_timing/correction_lifetime), alpha=0.5, color=colors[index])

    # Plot error correction positions
    x, y = correction_positions(corr_position_times, correction_pauli_expectation, correction_snapshot_timings)
    for i in x:
        plt.axvline(x=i, color='grey', label=label, linestyle="--")
        label = ''
    
plt.xlabel("Time (ms)", size=20)
plt.ylabel("Pauli expectation, averaged over {} shots".format(shots), size=20)
plt.title(r'Pauli expectation for varying logical states ($\hat{U}_{\text{odd}}$ and $\hat{U}_{\text{even}}$)', size=20)
plt.grid(color = 'lightgrey', linestyle = 'dashed', linewidth = 0.5)
plt.legend(fontsize=20, bbox_to_anchor=(1.35, 0.5), loc='right')
plt.xticks(fontsize=20)
plt.yticks(fontsize=20)
plt.show()

## All six states on separate charts
encoding = ['logi_xplus', 'logi_xminus', 'logi_yplus', 'logi_yminus', 'logi0', 'logi1']
sign = [1, -1, 1, -1, 1, -1]
for index, i in enumerate(encoding):
    noise_only_pauli_tuple, noise_only_fid, noise_only_prob, noise_only_snapshot_timings, _ = error_correction_experiment(i, 0, shots, 3, noise_only_circ_time)
    correction_pauli_tuple, correction_fid, correction_prob, correction_snapshot_timings, corr_position_times = error_correction_experiment(i, num_corr, shots, additional_snaps, correction_circ_time)

    ## Pauli operator decay visualization
    noise_only_pauli_expectation, _ = noise_only_pauli_tuple #[[matrix.real for matrix in cycle] for cycle in noise_only_pauli]
    correction_pauli_expectation, _ = correction_pauli_tuple

    # Compute average expectation and prob across all shots
    noise_only_expectation_avg = [np.round(np.average(cycle), 4) for cycle in noise_only_pauli_expectation]
    correction_expectation_avg = [np.round(np.average(cycle), 4) for cycle in correction_pauli_expectation]

    ## Expectation plots
    # Fit expectation data using an exponential decay function
    noise_only_lifetime, noise_lifetime_uncertainty = curve_fitter(noise_only_snapshot_timings, noise_only_expectation_avg, sign[index])
    correction_lifetime, correction_lifetime_uncertainty = curve_fitter(correction_snapshot_timings, correction_expectation_avg, sign[index])

    # Plot expectation data
    plt.figure(figsize=(15,6))
    plt.plot(noise_only_snapshot_timings, noise_only_expectation_avg,
            label=r'$\lvert$' + str(encoding_format.get(value)) + r'$\rangle$' + r' no corrections,' + '\n' + r'$\tau$ = ' + str(noise_only_lifetime) + r'ms $\pm$' + str(round(noise_lifetime_uncertainty, 2)) + r'ms', 
            marker='o', color='#5BCEFA')
    plt.plot(correction_snapshot_timings, correction_expectation_avg,
            label=r'$\lvert$' + str(encoding_format.get(value)) + r'$\rangle$ ' + str(num_corr) + r' corrections/ms,'  + '\n' + r'$\tau$ = ' + str(correction_lifetime) + r'ms $\pm$' + str(round(correction_lifetime_uncertainty, 2)) + r'ms', 
            marker='o', color='#F5A9B8')

    # Plot expectation data curve fits
    noise_curve_timing = np.linspace(0, noise_only_circ_time, 100)
    plt.plot(noise_curve_timing, sign[index] * np.exp(-noise_curve_timing/noise_only_lifetime))
    corr_curve_timing = np.linspace(0, correction_circ_time, 100)
    plt.plot(corr_curve_timing, sign[index] * np.exp(-corr_curve_timing/correction_lifetime))

    # Plot correction positions
    x, y = correction_positions(corr_position_times, correction_expectation_avg, correction_snapshot_timings)
    label = 'Correction \n position'
    for i in x:
        plt.axvline(x=i, color='grey', label=label, linestyle="--")
        label = ''

    plt.xlabel("Time (ms)", size=20)
    plt.ylabel("Pauli expectation, averaged over {} shots".format(shots), size=20)
    plt.title(r'Pauli expectation for varying logical states ($\hat{U}_{\text{odd}}$ and $\hat{U}_{\text{even}}$)', size=20)
    plt.grid(color = 'lightgrey', linestyle = 'dashed', linewidth = 0.5)
    plt.legend(fontsize=20, bbox_to_anchor=(1.35, 0.5), loc='right')
    plt.xticks(fontsize=20)
    plt.yticks(fontsize=20)
    plt.show()