import os
from typing import Any, Dict, List, Optional, Union

import numpy as np
import matplotlib.pyplot as plt

from qiskit import IBMQ, QuantumCircuit, QuantumRegister, ClassicalRegister, quantum_info as qi
from qiskit.providers.ibmq import RunnerResult
from qiskit.result import marginal_counts
import qiskit.tools.jupyter

%matplotlib inline

import warnings
warnings.filterwarnings("ignore")


%load_ext autoreload
%autoreload 2

# Note: This can be any hub/group/project that has access to the required device and the Qiskit runtime.
# Verify that ``qasm3`` is present in ``backend.configuration().supported_features``.
hub = "ibm-q-community"
group = "ieee-demos"
project = "main"
backend_name = "ibm_peekskill"
hgp = f"{hub}/{group}/{project}"

from qiskit.circuit import Delay, Parameter
from qiskit_ibm_runtime import QiskitRuntimeService

service = QiskitRuntimeService(instance=hgp)
backend = service.backend(backend_name, instance=hgp)

# Addresses a bug in the current Qiskit runtime package
# that will be fixed on the next Qiskit release
target = backend.target
if "delay" not in target:
    target.add_instruction(
        Delay(Parameter("t")), {(bit,): None for bit in range(target.num_qubits)}
    )

from qiskit.providers.aer import AerSimulator, Aer
from qiskit.providers.aer.noise import NoiseModel

backend_noise_model = NoiseModel.from_backend(backend)
#backend_sim = AerSimulator(noise_model=backend_noise_model)
backend_sim = AerSimulator()
ideal_sim = Aer.get_backend('qasm_simulator')

import utils
backend

shots: int = 1024 # Number of shots to run each circuit for

init_num_resets: int = 3 # Set the number of resets to initialize qubits between circuits
init_delay: float = 0. # Set the initialization idle time for qubits between circuits

sim: bool = False # Set True to simulate our experiments and False to run in hardware
dynamical_decoupling: bool = False # Set True to enable dynamical decoupling
    
favourite_qubits: List[int] = [5, 3, 8, 2, 9]

# Some general helper routines that will be used throughout the notebook

import mapomatic as mm
from qiskit.circuit.library import XGate
from qiskit.transpiler import PassManager
from qiskit.transpiler.instruction_durations import InstructionDurations
from qiskit_ibm_runtime import RuntimeJob, IBMBackend

def convert_cycles(time_in_seconds: float, backend: IBMBackend) -> int:
    cycles = time_in_seconds / (backend.configuration().dt)
    return int(cycles + (16 - (cycles % 16)))

def execute(circuits: List[QuantumCircuit], verbose: bool = True, **kwargs) -> RuntimeJob:
    """A helper method to execute circuits with common settings in the notebook.
    
    Args:
        circuits: To execute.
        verbose: Emit execution information.
        sim: Run with simulator instead of hardware. By default reads global variable "sim"
    
    """
    if isinstance(circuits, QuantumCircuit):
        circuits = [circuits]

    if sim:
        return backend_sim.run(circuits, shots=shots, **kwargs)
    else:
    
        job = run_openqasm3(circuits, backend, verbose=False, shots=shots, init_num_resets=init_num_resets, init_delay=init_delay, **kwargs)
        if verbose:
            print(f"Running on qasm3, job id: {job.job_id}")

        return job
    
def calculate_initial_layout(
    circuit: QuantumCircuit, 
    mapomatic: Optional[bool] = None, 
    favourite_qubits: Optional[List[int]] = None, 
    blacklist_qubits: Optional[List[int]] = None
) -> List[int]:
    """Routine to help choose the ideal qubit layout given an input circuit.
    
    Args:
        circuit: To evaluate layout for.
        mapomatic: Use mapomatic to choose ideal qubits, otherwise naievly choose ``favourite_qubits``
        favourite_qubits: A list of favourite qubits to ensure are in the layout
        blacklist_qubits: Qubits that are not allowed in the layout.
    """
    deflated = mm.deflate_circuit(circuit)
    if mapomatic:
        deflated = mm.deflate_circuit(circuit)
        available_layouts = [layout for layout in mm.matching_layouts(deflated, backend.configuration().coupling_map, strict_direction=False) if layout[:len(favourite_qubits)] == favourite_qubits and not set(layout).intersection(set(blacklist_qubits))]
        scores = mm.evaluate_layouts(deflated, available_layouts, backend)
        return scores[0][0]
    
    if favourite_qubits is None:
        favourite_qubits = []
    
    if blacklist_qubits is None:
        blacklist_qubits = []
        
    return (favourite_qubits + list(set(range(backend.num_qubits)) - set(favourite_qubits)))[:deflated.num_qubits]

def apply_dynamical_decoupling(
    circuits: List[QuantumCircuit], 
    backend: IBMBackend, 
    initial_layout: List[int],
) -> List[QuantumCircuit]:
    """Apply dynamic circuit dynamical decoupling."""

    from qiskit_ibm_provider.transpiler.passes.scheduling import DynamicCircuitScheduleAnalysis, PadDynamicalDecoupling
    dd_sequence = [XGate(), XGate()]
    durations = InstructionDurations.from_backend(backend)
    pm = PassManager(
        [
            DynamicCircuitScheduleAnalysis(durations),
            PadDynamicalDecoupling(durations, dd_sequence, qubits=initial_layout),
        ]
    )
    return pm.run(circuits)

def apply_transpile(
    circuits: List[QuantumCircuit], 
    backend: IBMBackend, 
    initial_layout: List[int], 
    dynamical_decoupling: bool = dynamical_decoupling, 
    **transpile_kwargs,
) -> List[QuantumCircuit]:
    """Fixed transpile routine to be used for consistency throughout the notebook."""

    transpiled_circuits = transpile(circuits, backend, initial_layout=initial_layout, optimization_level=3, **transpile_kwargs)
    if dynamical_decoupling:
        return apply_dynamical_decoupling(transpiled_circuits, backend, initial_layout)
    return transpiled_circuits

qubit = favourite_qubits[0]
print(f"We will use qubit {qubit}")

# Exercise

from utils import run_openqasm3

from qiskit import QuantumRegister, ClassicalRegister, transpile

qr = QuantumRegister(1)
crx = ClassicalRegister(1, name="xresult")
crm = ClassicalRegister(1, name="measureresult")
qc_reset = QuantumCircuit(qr, crx, crm, name="Reset")

# Complete the circuit to measure and conditionally reset the target
# qubit.

qc_reset = transpile(qc_reset, backend)
qc_reset.draw(output="mpl", idle_wires=False)

# Solution

from utils import run_openqasm3

from qiskit import QuantumRegister, ClassicalRegister, transpile

qr = QuantumRegister(1)
crx = ClassicalRegister(1, name="xresult")
crm = ClassicalRegister(1, name="measureresult")
qc_reset = QuantumCircuit(qr, crx, crm, name="Reset")
qc_reset.x(0)
qc_reset.measure(0, crx)
qc_reset.x(0).c_if(crx, 1)
qc_reset.measure(0, crm)

qc_reset = transpile(qc_reset, backend)
qc_reset.draw(output="mpl", idle_wires=False)

reset_job = run_openqasm3(qc_reset, backend, verbose=True, shots=shots, init_num_resets=0, init_delay=0) # Turn off automatic init


from qiskit.result import marginal_counts

reset_result = reset_job.result()
reset_counts = reset_result.get_counts(0)
# Look at only the final measurement result
initial_measurement_result = marginal_counts(reset_counts, indices=[0])
mitigated_reset_results = marginal_counts(reset_counts, indices=[1])
print(f"Full counts including reset: {reset_counts}")
print(f"Conditional X gate applied on: {initial_measurement_result['1']} shots")
print(f"Results from our reset - |0\rangles prepared: {mitigated_reset_results.get('0')}, |1\rangles prepared: {mitigated_reset_results.get('1', 0)}"
)

from qiskit import qasm3

basis_gates = backend.configuration().basis_gates

def dump_qasm3(circuit: QuantumCircuit, backend: IBMBackend = backend) -> str:
    return qasm3.Exporter(includes=[],
                          basis_gates=backend.configuration().basis_gates,
                          disable_constants=True).dumps(circuit)

qasm3_reset = dump_qasm3(qc_reset)
print(qasm3_reset)

runtime_job = run_openqasm3(qasm3_reset, backend, verbose=True, shots=shots, init_num_resets=0, init_delay=0)

classical_compute_qasm3 = """
OPENQASM 3.0;

bit a;
bit b;
bit c;

x $0;
x $2;

a = measure $0; // expected "1"
b = measure $1; // expected "0"
c = measure $2; // expected "1"

bit[3] d = "100";

if (((a | b) & c) == 1) {
    // Path will nearly always execute 
    d[0] = 1;
} else {
    // Path will rarely execute outside of SPAM errors
    d[0] = 0;
}

bit[3] e = "101";
// Conditionally execute based on classical bit array comparison
// Should always execute
if (d == e) {
    x $1;
}

bit final;
final = measure $1; // expected "1"
"""

classical_compute_job = run_openqasm3(classical_compute_qasm3, backend, verbose=False, shots=shots)

qr = QuantumRegister(3)
cr = ClassicalRegister(3, name="result")
crz = cr[0]
crx = cr[1]
result = cr[2]

qc_teleport = QuantumCircuit(qr, cr, name="Teleport")

# Apply teleportation circuit
qc_teleport.h(qr[1])
qc_teleport.cx(qr[1], qr[2])
qc_teleport.cx(qr[0], qr[1])
qc_teleport.h(qr[0])
qc_teleport.measure(qr[0], crz)
qc_teleport.measure(qr[1], crx)
qc_teleport.z(qr[2]).c_if(crz, 1)
qc_teleport.x(qr[2]).c_if(crx, 1)
qc_teleport.draw(output="mpl")

# Prepare |1> for teleportation
qc_state_prep = QuantumCircuit(1)
qc_state_prep.x(0)
print(qc_state_prep)

target_state_prep = qi.Statevector.from_instruction(qc_state_prep)
target_state_prep.draw(output="bloch")

teleport_qubits = favourite_qubits[:3]

qc_teleport_state = QuantumCircuit(qr, cr, name="Teleport Hadamard")

# Prepare state to teleport
qc_teleport_state.compose(qc_state_prep, [qr[0]], inplace=True)
qc_teleport_state.barrier(qr)

# Compose with teleportation circuit
qc_teleport_state.compose(qc_teleport, inplace=True)

qc_teleport_state.draw(output="mpl")

qc_teleport_experiment = qc_teleport_state.copy()
qc_teleport_experiment.measure(qr[2], result)
qc_teleport_experiment = apply_transpile(qc_teleport_experiment, backend, initial_layout=teleport_qubits)
qc_teleport_experiment.draw(output="mpl", idle_wires=True)

teleport_job = execute(qc_teleport_experiment, verbose=False)

teleport_result = teleport_job.result()
print(f"All teleportation counts: {teleport_result.get_counts(0)}")

marginal_teleport_counts = marginal_counts(teleport_result.get_counts(0), indices=[2])
print(f"Marginalized teleportation counts: {marginal_teleport_counts}")

p1 = 0.01
p3 = 3 * p1**2 * (1-p1) + p1**3 # probability of 2 or 3 errors
print('Probability of a single reply being garbled: {}'.format(p1))
print('Probability of a majority of the three replies being garbled: {:.4f}'.format(p3))

# Approximate duration of the measurement processing / conditional latency
block_branch_duration = 0.5e-6 
block_branch_cycles = convert_cycles(block_branch_duration, backend)

# Setup a base quantum circuit for our experiments
qreg_data = QuantumRegister(3)
qreg_measure = QuantumRegister(2)
creg_data = ClassicalRegister(3)
creg_syndrome = ClassicalRegister(2)
state_data = qreg_data[0]
ancillas_data = qreg_data[1:]

def build_qc() -> QuantumCircuit:
    return QuantumCircuit(qreg_data, qreg_measure, creg_data, creg_syndrome)

# Exercise

qc_init = build_qc()

# Complete to initialize the |1> state
qc_init.barrier(qreg_data)

qc_init.draw(output="mpl")

# Solution

qc_init = build_qc()

qc_init.x(qreg_data[0])
qc_init.barrier(qreg_data)

qc_init.draw(output="mpl")

# Exercise

# Complete the method below to encode the bit-flip code
def encode_bit_flip(qc, state, ancillas):
    # Complete
    qc.barrier(state, *ancillas)
    return qc

qc_encode_bit = build_qc()

encode_bit_flip(qc_encode_bit, state_data, ancillas_data)

qc_encode_bit.draw(output="mpl")

# Solution

def encode_bit_flip(qc, state, ancillas):
    control = state
    for ancilla in ancillas:
        qc.cx(control, ancilla)
    qc.barrier(state, *ancillas)
    return qc

qc_encode_bit = build_qc()

encode_bit_flip(qc_encode_bit, state_data, ancillas_data)

qc_encode_bit.draw(output="mpl")

# Exercise

# Complete the routine below to decode from the bit-flip codespace
def decode_bit_flip(qc, state, ancillas):
    #complete
    
    return qc

qc_decode_bit = build_qc()

qc_decode_bit = decode_bit_flip(qc_decode_bit, state_data, ancillas_data)

qc_decode_bit.draw(output="mpl")

qc_encoded_state_bit = qc_init.compose(qc_encode_bit)
qc_encoded_state_bit.draw(output="mpl")

# Solution

def decode_bit_flip(qc, state, ancillas):
    inv = qc_encode_bit.inverse()
    return qc.compose(inv)

qc_decode_bit = build_qc()

qc_decode_bit = decode_bit_flip(qc_decode_bit, state_data, ancillas_data)

qc_decode_bit.draw(output="mpl")

# Exercise

# Complete the syndome measurement routine below
def measure_syndrome_bit(qc, qreg_data, qreg_measure, creg_measure):
    # Complete
    qc.barrier(*qreg_data, *qreg_measure)
    return qc

qc_syndrome_bit = measure_syndrome_bit(build_qc(), qreg_data, qreg_measure, creg_syndrome)
qc_syndrome_bit.draw(output="mpl")    

qc_measure_syndrome_bit = qc_encoded_state_bit.compose(qc_syndrome_bit)
qc_measure_syndrome_bit.draw(output="mpl")

# Solution

def measure_syndrome_bit(qc, qreg_data, qreg_measure, creg_measure):
    def branch_delay():
        if sim:
            qc.barrier(qreg_data)
            qc.delay(block_branch_cycles, qreg_data)
            qc.barrier(qreg_data)
  
    qc.cx(qreg_data[0], qreg_measure[0])
    qc.cx(qreg_data[1], qreg_measure[0])
    qc.cx(qreg_data[0], qreg_measure[1])
    qc.cx(qreg_data[2], qreg_measure[1])
    qc.measure(qreg_measure, creg_measure)
    branch_delay() # grouped in hardware
    qc.x(qreg_measure[0]).c_if(creg_measure[0], 1)
    qc.x(qreg_measure[1]).c_if(creg_measure[1], 1)
    qc.barrier(*qreg_data, *qreg_measure)
    return qc

qc_syndrome_bit = measure_syndrome_bit(build_qc(), qreg_data, qreg_measure, creg_syndrome)
qc_syndrome_bit.draw(output="mpl")    

# Exercise

# Complete the routine below to apply our decoding and correction sequence
def apply_correction_bit(qc, qreg_data, creg_syndrome):
    # Complete
    qc.barrier(qreg_data)
    return qc
    
qc_correction_bit = apply_correction_bit(build_qc(), qreg_data, creg_syndrome)
qc_correction_bit.draw(output="mpl")

def apply_final_readout(qc, qreg_data, creg_data):
    """Apply inverse mapping so that we always try and measure |1> in the computational basis.
    
    TODO: The above is just a stand in for proper measurement basis measurement
    """
    qc.barrier(qreg_data)
    qc.measure(qreg_data, creg_data)
    return qc

qc_final_measure = apply_final_readout(build_qc(), qreg_data, creg_data)
qc_final_measure.draw(output="mpl")

bit_code_circuit = qc_measure_syndrome_bit.compose(qc_correction_bit).compose(qc_final_measure)
bit_code_circuit.draw(output="mpl")

# Solution

def apply_correction_bit(qc, qreg_data, creg_syndrome):
    # If simulating we need to insert a delay to mirror the hardware
    def branch_delay():
        if sim:
            qc.barrier(qreg_data)
            qc.delay(block_branch_cycles, qreg_data)
            qc.barrier(qreg_data)
    
    branch_delay()
    qc.x(qreg_data[0]).c_if(creg_syndrome, 3)
    branch_delay()
    qc.x(qreg_data[1]).c_if(creg_syndrome, 1)
    branch_delay()
    qc.x(qreg_data[2]).c_if(creg_syndrome, 2)
    qc.barrier(qreg_data)
    return qc
    
qc_correction_bit = apply_correction_bit(build_qc(), qreg_data, creg_syndrome)
qc_correction_bit.draw(output="mpl")

def build_error_correction_sequence(
    qc_base: QuantumCircuit,
    qc_init: Optional[QuantumCircuit], 
    qc_encode: QuantumCircuit, 
    qc_channels: List[QuantumCircuit],
    qc_syndrome: QuantumCircuit,
    qc_correct: QuantumCircuit,
    qc_decode: Optional[QuantumCircuit] = None,
    qc_final: Optional[QuantumCircuit] = None,
    name=None,
) -> QuantumCircuit:
    """Build a typical error correction circuit"""
    qc = qc_base
    
    if qc_init:
        qc = qc.compose(
            qc_init
        )
    
    qc = qc.compose(
            qc_encode
        )
    if name is not None:
        qc.name = name
    
    if not qc_channels:
        qc_channels = [QuantumCircuit(*qc.qregs)]
        
    for qc_channel in qc_channels:
        qc = qc.compose(
                qc_channel
            ).compose(
                qc_syndrome
            ).compose(
                qc_correct
            )
    if qc_decode:
        qc = qc.compose(qc_decode)
    
    if qc_final:
        qc = qc.compose(qc_final)
    
    return qc

bit_code_circuit = build_error_correction_sequence(
    build_qc(),
    qc_init,
    qc_encode_bit,
    [],
    qc_syndrome_bit,
    qc_correction_bit,
    None,
    qc_final_measure,
)
bit_code_circuit.draw(output="mpl")

layout_circuit = transpile(bit_code_circuit, backend, optimization_level=3)
print(initial_layout := calculate_initial_layout(layout_circuit, False, favourite_qubits))

transpiled_bit_code_circuit = apply_transpile(bit_code_circuit, backend, initial_layout=initial_layout)
transpiled_bit_code_circuit.draw(output="mpl")

result = execute(transpiled_bit_code_circuit).result()

def decode_result(data_counts, syndrome_counts, verbose=True, indent=0):
    shots = sum(data_counts.values())
    success_trials = data_counts.get('000', 0) + data_counts.get('111', 0)
    failed_trials = shots-success_trials
    error_correction_events = shots-syndrome_counts.get('00', 0)
    
    if verbose:
        print(f"{' ' * indent}Bit flip errors were corrected on {error_correction_events}/{shots} trials")
        print(f"{' ' * indent}A final parity error was detected on {failed_trials}/{shots} trials")
        print(f"{' ' * indent}No error was detected on {success_trials}/{shots} trials")
    return error_correction_events, failed_trials

data_indices = list(range(len(qreg_data)))
syndrome_indices = list(range(data_indices[-1]+1, len(qreg_data) + len(qreg_measure) ))
                    
marginalized_data_result = marginal_counts(result, data_indices)
marginalized_syndrome_result = marginal_counts(result, syndrome_indices)

print(f'Completed bit code experiment data measurement counts {marginalized_data_result.get_counts(0)}')
print(f'Completed bit code experiment syndrome measurement counts {marginalized_syndrome_result.get_counts(0)}')
decode_result(marginalized_data_result.get_counts(0), marginalized_syndrome_result.get_counts(0));

from qiskit.circuit.library import IGate, XGate, ZGate

qreg_error_ancilla = QuantumRegister(1)
creg_error_ancilla = ClassicalRegister(1)

def build_random_error_channel(gate, ancilla, creg_ancilla, error_qubit):
    """Build an error channel that randomly applies a single-qubit gate based on an ancilla qubit measurement result"""
    qc = build_qc()
    qc.add_register(qreg_error_ancilla)
    qc.add_register(creg_error_ancilla)
    qc.barrier(ancilla, error_qubit.register)
    # 50-50 chance of applying a bit-flip
    qc.h(ancilla)
    qc.measure(ancilla, creg_ancilla)
    qc.append(gate, [error_qubit]).c_if(creg_ancilla, 1)
    qc.barrier(ancilla, error_qubit.register)
    return qc

qc_id_error_channel = build_random_error_channel(IGate(), qreg_error_ancilla, creg_error_ancilla, qreg_data[0])
print("Identity error channel")
print(qc_id_error_channel.draw(idle_wires=False,output='text',fold=-1))

qc_bit_flip_error_channel = build_random_error_channel(XGate(), qreg_error_ancilla, creg_error_ancilla, qreg_data[0])
print("Bit flip error channel")
print(qc_bit_flip_error_channel.draw(idle_wires=False,output='text',fold=-1))

qc_phase_flip_error_channel = build_random_error_channel(ZGate(), qreg_error_ancilla, creg_error_ancilla, qreg_data[0])
print("Phase flip error channel")
print(qc_phase_flip_error_channel.draw(idle_wires=False,output='text',fold=-1))

def build_error_channel_base():
    qc = build_qc()
    qc.add_register(qreg_error_ancilla)
    qc.add_register(creg_error_ancilla)
    return qc

qc_id_error_bit_flip_code = build_error_correction_sequence(
    build_error_channel_base(),
    qc_init,
    qc_encode_bit,
    [qc_id_error_channel],
    qc_syndrome_bit,
    qc_correction_bit,
    None,
    qc_final_measure,
    "Identity error channel"
)

qc_bit_flip_error_bit_flip_code = build_error_correction_sequence(
    build_error_channel_base(),
    qc_init,
    qc_encode_bit,
    [qc_bit_flip_error_channel],
    qc_syndrome_bit,
    qc_correction_bit,
    None,
    qc_final_measure,
    "Bit flip error channel"
)

qc_phase_flip_error_bit_flip_code = build_error_correction_sequence(
    build_error_channel_base(),
    qc_init,
    qc_encode_bit,
    [qc_phase_flip_error_channel],
    qc_syndrome_bit,
    qc_correction_bit,
    None,
    qc_final_measure,
    "Phase flip error channel"
)

circuits_error_channels_bit_flip_code = [qc_id_error_bit_flip_code, qc_bit_flip_error_bit_flip_code, qc_phase_flip_error_bit_flip_code]

qc_bit_flip_error_bit_flip_code.draw(output="mpl")

# We need to add an extra ancilla qubit to our layout
# It doesn't matter which qubit for the must part as we are using as a
# source of random information
error_channel_layout = error_channel_layout = initial_layout + list(set(range(circuits_error_channels_bit_flip_code[0].num_qubits)) - set(initial_layout))[:1]

transpiled_circuits_error_channels_bit_flip_code = apply_transpile(circuits_error_channels_bit_flip_code, backend, initial_layout=error_channel_layout)
result_error_channels_bit_flip_code = execute(transpiled_circuits_error_channels_bit_flip_code).result()

from qiskit.quantum_info.analysis import hellinger_fidelity

def decode_error_channel_result(qc_init, data_counts, syndrome_counts, verbose=True, indent=0):
    shots = sum(data_counts.values())
    logical_zero = data_counts.get('000', 0)
    logical_one = data_counts.get('111', 0)
    success_trials = logical_zero + logical_one
    failed_trials = shots-success_trials
    logical_counts = {"0": logical_zero, "1": logical_one}
    
    
    ideal_transpiled = apply_transpile(qc_init_outcome, backend, initial_layout)
    counts_ideal = marginal_counts(ideal_sim.run(ideal_transpiled, shots=success_trials).result().get_counts(0), indices=[0])
    
    fidelity = hellinger_fidelity(counts_ideal, logical_counts)
    
    error_correction_events = shots-syndrome_counts.get('00', 0)
    
    if verbose:
        print(f"{' ' * indent}Bit flip errors were corrected on {error_correction_events}/{shots} trials")
        print(f"{' ' * indent}A final parity error was detected on {failed_trials}/{shots} trials")
        print(f"{' ' * indent}For the successful trials the Hellinger fidelity is {fidelity}")
    return error_correction_events, failed_trials

qc_init_outcome = qc_init.copy()
qc_init_outcome.measure(qreg_data[0], 0)

qreg_indices = list(range(len(qreg_data)))
data_indices = qreg_indices[:1]
syndrome_indices = list(range(qreg_indices[-1]+1, len(qreg_data) + len(qreg_measure) ))

result_decoded_data_qubit_marginal_err_ch = marginal_counts(result_error_channels_bit_flip_code, data_indices)
result_data_marginal_err_ch = marginal_counts(result_error_channels_bit_flip_code, qreg_indices)
result_syndrome_marginal_err_ch = marginal_counts(result_error_channels_bit_flip_code, syndrome_indices)

for i, qc in enumerate(transpiled_circuits_error_channels_bit_flip_code):
    print(f"For {qc.name} with bit flip code")
    print(f'    Completed bit code experiment decoded data qubit measurement counts {result_decoded_data_qubit_marginal_err_ch.get_counts(i)}')
    print(f'    Completed bit code experiment data qubits measurement counts {result_data_marginal_err_ch.get_counts(i)}')
    print(f'    Completed bit code experiment syndrome measurement counts {result_data_marginal_err_ch.get_counts(i)}')
    decode_error_channel_result(qc_init_outcome, result_data_marginal_err_ch.get_counts(i), result_syndrome_marginal_err_ch.get_counts(i), indent=4);
    print("")

# Exercise

# Complete the routine below to apply an "identity" correction rather than
# the standard bit-flip correction sequence
def apply_no_correction_bit(qc, qreg_data, creg_syndrome):
    """Complete the method below to apply an "identity correction"
    
    Make sure to still apply the conditional gates so that the comparison to the bit-flip
    code is faithful.
    """
    # Complete
    qc.barrier(qreg_data)
    return qc

    
qc_no_correction_bit = apply_no_correction_bit(build_qc(), qreg_data, creg_syndrome)
qc_no_correction_bit.draw(output="mpl")

# Exercise

# Complete the sequence using `build_error_correction_sequence` to evaluate the
# performance of the identity correction sequence

# Complete the following routines
qc_id_error_no_correct = QuantumCircuit() # build_error_correction_sequence(...)
qc_bit_flip_error_no_correct = QuantumCircuit() # build_error_correction_sequence(...)
qc_phase_flip_error_no_correct = QuantumCircuit() # build_error_correction_sequence(...)


circuits_error_channels_no_correct = [qc_id_error_no_correct, qc_bit_flip_error_no_correct, qc_phase_flip_error_no_correct]

qc_id_error_no_correct.draw(output="mpl")

# We need to add an extra ancilla qubit to our layout
# It doesn't matter which qubit for the must part as we are using as a
# source of random information
error_channel_layout = error_channel_layout = initial_layout + list(set(range(circuits_error_channels_bit_flip_code[0].num_qubits)) - set(initial_layout))[:1]

transpiled_circuits_error_channels_no_correct = apply_transpile(circuits_error_channels_no_correct, backend, initial_layout=error_channel_layout)
result_error_channels_no_correct = execute(transpiled_circuits_error_channels_no_correct).result()

qc_init_outcome = qc_init.copy()
qc_init_outcome.measure(qreg_data[0], 0)

qreg_indices = list(range(len(qreg_data)))
data_indices = qreg_indices[:1]
syndrome_indices = list(range(qreg_indices[-1]+1, len(qreg_data) + len(qreg_measure) ))

result_decoded_data_qubit_marginal_no_correct = marginal_counts(result_error_channels_no_correct, data_indices)
result_data_marginal_no_correct = marginal_counts(result_error_channels_no_correct, qreg_indices)
result_syndrome_marginal_no_correct = marginal_counts(result_error_channels_no_correct, syndrome_indices)

for i, qc in enumerate(transpiled_circuits_error_channels_no_correct):
    print(f"For {qc.name} with bit flip code")
    print(f'    Completed bit code experiment decoded data qubit measurement counts {result_decoded_data_qubit_marginal_no_correct.get_counts(i)}')
    print(f'    Completed bit code experiment data qubits measurement counts {result_data_marginal_no_correct.get_counts(i)}')
    print(f'    Completed bit code experiment syndrome measurement counts {result_data_marginal_no_correct.get_counts(i)}')
    decode_error_channel_result(qc_init_outcome, result_data_marginal_no_correct.get_counts(i), result_syndrome_marginal_no_correct.get_counts(i), indent=4);
    print("")

def apply_final_readout_invert(qc, qc_init, qreg_data, creg_data):
    """Apply inverse mapping so that we always try and measure |1> in the computational basis."""
    qc = qc.compose(qc_init.inverse())
    qc.x(qreg_data[0])
    qc.barrier(qreg_data)
    qc.measure(qreg_data, creg_data)
    return qc

qc_final_measure_invert = apply_final_readout_invert(build_qc(), qc_init, qreg_data, creg_data)
qc_final_measure_invert.draw(output="mpl")

# Solution

qc_id_error_no_correct = build_error_correction_sequence(
    build_error_channel_base(),
    qc_init,
    qc_encode_bit,
    [qc_id_error_channel],
    qc_syndrome_bit,
    qc_no_correction_bit,
    None,
    qc_final_measure,
    "Identity error channel"
)

qc_bit_flip_error_no_correct = build_error_correction_sequence(
    build_error_channel_base(),
    qc_init,
    qc_encode_bit,
    [qc_bit_flip_error_channel],
    qc_syndrome_bit,
    qc_no_correction_bit,
    None,
    qc_final_measure,
    "Bit flip error channel"
)

qc_phase_flip_error_no_correct = build_error_correction_sequence(
    build_error_channel_base(),
    qc_init,
    qc_encode_bit,
    [qc_phase_flip_error_channel],
    qc_syndrome_bit,
    qc_no_correction_bit,
    None,
    qc_final_measure,
    "Phase flip error channel"
)

circuits_error_channels_no_correct = [qc_id_error_no_correct, qc_bit_flip_error_no_correct, qc_phase_flip_error_no_correct]

qc_id_error_no_correct.draw(output="mpl")

from collections import defaultdict

from qiskit.transpiler import PassManager, InstructionDurations
from qiskit.transpiler.passes import ASAPSchedule, PadDelay
from qiskit_ibm_provider.transpiler.passes.scheduling import DynamicCircuitScheduleAnalysis, PadDelay

def get_circuit_duration_(qc: QuantumCircuit) -> int:
    """Get duration of circuit in hardware cycles."""
    durations = InstructionDurations.from_backend(backend)
    pm = PassManager([DynamicCircuitScheduleAnalysis(durations)])
    pm.run(qc)
    node_start_times = pm.property_set["node_start_time"]
    
    block_durations = defaultdict(int)
    for inst, (block, t0) in node_start_times.items():
        block_durations[block] = max(block_durations[block], t0+inst.op.duration)
        
    duration = sum(block_durations.values())
    
    # If we are running on real hardware the delays have not been appended to the
    # circuit directly and are instead built into the conditional operations.
    # We account for them manually here
    if not sim:
        duration += block_branch_cycles * (len(block_durations) - 1)
    return duration

def build_idle_error_correction_sequence(
    qc_base: QuantumCircuit,
    qc_init: Optional[QuantumCircuit], 
    qc_encode: QuantumCircuit, 
    qc_channels: List[QuantumCircuit],
    qc_syndrome: QuantumCircuit,
    qc_correct: QuantumCircuit,
    qc_decode: Optional[QuantumCircuit] = None,
    qc_final: Optional[QuantumCircuit] = None,
    initial_layout=initial_layout,
    name: str = None,
) -> QuantumCircuit:
    """Build a quantum circuit that idles for the period of the input error correction sequence."""
    qc = qc_base
    
    if qc_init:
        qc = qc.compose(
            qc_init
        )
    if name is not None:
        qc.name = name

    qc_idle_region = qc_base.copy()
    
    qc_idle_region.compose(
            qc_encode
        )
  
    if not qc_channels:
        qc_channels = [QuantumCircuit(*qc.qregs)]
        
    for qc_channel in qc_channels:
        qc_idle_region = qc_idle_region.compose(
                qc_channel
            ).compose(
                qc_syndrome
            ).compose(
                qc_correct
            )
    
    if qc_decode:
        qc_idle_region = qc_idle_region.compose(qc_decode)
    
    qc_idle_transpiled = apply_transpile(qc_idle_region, backend, initial_layout=initial_layout, scheduling_method=None)
    idle_duration = get_circuit_duration_(qc_idle_transpiled)
    
    qc_idle = qc_base.copy()
    qc_idle.barrier()
    for qubit in qc_idle.qubits:
        qc_idle.delay(idle_duration, qubit, unit="dt")
    qc_idle.barrier()
    qc = qc.compose(qc_idle)
    
    if qc_final:
        qc = qc.compose(qc_final_measure)
    
    return qc

# Exercise

def build_idle_error_channel(time_in_seconds, qreg):
    idle_cycles = convert_cycles(time_in_seconds, backend)
    qc_idle = build_qc()
    # Complete
    qc_idle.barrier()
    return qc_idle

# Solution

def build_idle_error_channel(time_in_seconds, qreg):
    idle_cycles = convert_cycles(time_in_seconds, backend)
    qc_idle = build_qc()
    qc_idle.delay(idle_cycles, qreg, "dt")
    qc_idle.barrier()
    return qc_idle

# Exercise

num_rounds = 5

# Idle for a specified period in seconds
# This is how we build an idle error channel
idle_period = 5e-6

# Use the circuit we created above
qc_idle = build_idle_error_channel(idle_period, qreg_data)


qcs_corr_bit = []
qcs_no_corr_bit = []
qcs_idle_equiv_bit = []
for round_ in range(num_rounds):
    qc_error_channels = [qc_idle] * (round_ + 1)
    
    # Complete bit flip code iteration
    qcs_corr_bit.append(QuantumCircuit())
    # Complete no correction bit flip code iteration
    qcs_no_corr_bit.append(QuantumCircuit())
    # Complete idle equivalent bit flip code iteration
    qcs_idle_equiv_bit.append(QuantumCircuit())



qcs_corr_bit[0].draw(output="mpl", idle_wires=False)

transpiled_qcs_corr_bit = apply_transpile(qcs_corr_bit, backend, initial_layout=initial_layout)
job_qcs_corr_bit = execute(transpiled_qcs_corr_bit)

transpiled_qcs_no_corr_bit = apply_transpile(qcs_no_corr_bit, backend, initial_layout=initial_layout)
job_qcs_no_corr_bit = execute(transpiled_qcs_no_corr_bit)

transpiled_qcs_idle_equiv_bit = apply_transpile(qcs_idle_equiv_bit, backend, initial_layout=initial_layout, dynamical_decoupling=True)
job_qcs_idle_equiv_bit = execute(transpiled_qcs_idle_equiv_bit)

result_qcs_corr_bit = job_qcs_corr_bit.result()
result_qcs_no_corr_bit = job_qcs_no_corr_bit.result()
result_qcs_idle_equiv_bit = job_qcs_idle_equiv_bit.result()

transpiled_qcs_idle_equiv_bit[-1].draw(output="mpl")

from qiskit.quantum_info.analysis import hellinger_fidelity

qc_init_outcome = qc_init.copy()
qc_init_outcome = qc_init_outcome.compose(qc_final_measure)
result_ideal = ideal_sim.run(apply_transpile(qc_init_outcome, backend, initial_layout)).result()

def calculate_hellinger_fidelity(result_ideal, result_experiment, data_qubit=0):
    result_ideal = marginal_counts(result_ideal, indices=[data_qubit])
    result_experiment = marginal_counts(result_experiment, indices=[data_qubit])
    
    counts_ideal = result_ideal.get_counts(0)

    hellinger_fidelities = []
    for circuit_idx in range(len(result_experiment.results)):
        hellinger_fidelities.append(hellinger_fidelity(counts_ideal, result_experiment.get_counts(circuit_idx)))

    return hellinger_fidelities

# Evaluate the Hellinger fidelity using the routine above
fidelities_corr_bit = 0.
fidelities_no_corr_bit = 0.
fidelities_idle_equiv_bit = 0.

# Solution

num_rounds = 5

# Idle for a specified period in seconds
# This is how we build an idle error channel
idle_period = 5e-6

# Use the circuit we created above
qc_idle = build_idle_error_channel(idle_period, qreg_data)


qcs_corr_bit = []
qcs_no_corr_bit = []
qcs_idle_equiv_bit = []
for round_ in range(num_rounds):
    qc_error_channels = [qc_idle] * (round_ + 1)
    
    #qcs_corr_bit.append(QuantumCircuit())

    qcs_corr_bit.append(
        build_error_correction_sequence(
            build_qc(),
            qc_init,
            qc_encode_bit,
            qc_error_channels,
            qc_syndrome_bit,
            qc_correction_bit,
            qc_decode_bit,
            qc_final_measure_invert,
            f"With Correction {round_}"
        )
    )
    #qcs_no_corr_bit.append(QuantumCircuit())
    qcs_no_corr_bit.append(
        build_error_correction_sequence(
            build_qc(),
            qc_init,
            qc_encode_bit,
            qc_error_channels,
            qc_syndrome_bit,
            qc_no_correction_bit,
            qc_decode_bit,
            qc_final_measure_invert,
            name=f"Without Correction {round_}"
        )
    )
    
    #qcs_idle_equiv_bit.append(QuantumCircuit())
    qcs_idle_equiv_bit.append(
        build_idle_error_correction_sequence(
            build_qc(),
            qc_init,
            qc_encode_bit,
            qc_error_channels,
            qc_syndrome_bit,
            qc_correction_bit,
            qc_decode_bit,
            qc_final_measure_invert,
            initial_layout=initial_layout,
            name=f"Idle {round_}"
        )
    )


fidelities_corr_bit = calculate_hellinger_fidelity(result_ideal, result_qcs_corr_bit)
fidelities_no_corr_bit = calculate_hellinger_fidelity(result_ideal, result_qcs_no_corr_bit)
fidelities_idle_equiv_bit = calculate_hellinger_fidelity(result_ideal, result_qcs_idle_equiv_bit)

import matplotlib.pyplot as plt
plt.rcParams["figure.figsize"] = (12, 6)

iters = range(1, num_rounds+1)
plt.plot(iters, fidelities_corr_bit, label="Bit flip code - correction")
plt.plot(iters, fidelities_no_corr_bit, label="Bit flip code - no correction")
plt.plot(iters, fidelities_idle_equiv_bit, label="Idle equivalent circuit")
plt.ylabel("Hellinger Fidelity")
plt.xlabel("Correction cycles")
plt.xticks(iters)
plt.suptitle("Comparing the performance of error-correction strategies for multiple correction cycles")
plt.title(f"Idle period: {idle_period*1e6}us - Qubit layout: {initial_layout} - Dynamical decoupling: {dynamical_decoupling}")
plt.legend(loc="upper right")

result_qcs_corr_bit = job_qcs_corr_bit.result()
result_qcs_no_corr_bit = job_qcs_no_corr_bit.result()
result_qcs_idle_equiv_bit = job_qcs_idle_equiv_bit.result()

def build_412(circuit,qubits,parity,layout='zxz') -> QuantumCircuit:
    ### qubits
    # qubits: list of 7 qubits 
    # ex. [q[0],q[1],q[2],q[3],q[4],q[5],q[6],q[7]] 
    #where first 4 are data qubits, 5 and 6 are flags, and 7 is ancilla (see image from heavy-hex paper)
    ### parity
    # parity is 'x' or 'z'
    ### initialized_ancilla = True or False. 
    # If True, will prepare f1,f2,a1 in correct basis and apply post rotation before measurement 
    # Will include barrier after prep and before post rot if True
    [d1,d2,d3,d4,f1,f2,a1]=qubits
    
    if layout == 'zxz':
        if parity == 'x':
            circuit.h(a1)
            circuit.cx(a1,f2)
            circuit.cx(a1,f1)
            circuit.cx(f2,d1)
            circuit.cx(f1,d4)
            circuit.cx(f2,d2)
            circuit.cx(f1,d3)
            circuit.cx(a1,f2)
            circuit.cx(a1,f1)
            circuit.h(a1)           
        if parity == 'z':
            circuit.cx(d2,f2)
            circuit.cx(d1,f2)
            circuit.cx(d3,f1)
            circuit.cx(d4,f1)
        return circuit
    
    elif layout == 'xzx':
        if parity == 'z':
            circuit.h(f1)
            circuit.h(f2)

            circuit.cx(f2,a1)
            circuit.cx(f1,a1)
            circuit.cx(d1,f2)        
            circuit.cx(d3,f1)
            circuit.cx(d2,f2)
            circuit.cx(d4,f1)
            circuit.cx(f2,a1)
            circuit.cx(f1,a1)

            circuit.h(f1)
            circuit.h(f2)

        if parity == 'x':
            circuit.h(f1)
            circuit.h(f2)
            circuit.cx(f2,d2)
            circuit.cx(f2,d1)
            circuit.cx(f1,d3)
            circuit.cx(f1,d4)

            circuit.h(f1)
            circuit.h(f2)
        return circuit        

# Complete - how do you do the Z and X checks for this code? How do you decode the errors?
# See: https://arxiv.org/abs/2110.04285


import qiskit.tools.jupyter
%qiskit_version_table
%qiskit_copyright

