try:
    import recirq
except ImportError:
    !pip install --quiet git+https://github.com/quantumlib/ReCirq

import recirq.algorithmic_benchmark_library as algos
from IPython.display import display

for algo in algos.BENCHMARKS:
    display(algo)

"""
Explanation: Feature Testbed Benchmark Library
The Feature Testbed and Benchmarking Library (FTb/BL) consists of a series of workflow tools in cirq_google.workflow and a library of application-inspired algorithms conforming to those specifications living in ReCirq.
Computer code can serve as the intermediary language between algorithms researchers, software developers, and providers of quantum hardware -- each of whom have their own unique domain expertise. Higher-order abstraction in Quantum Executables can support higher-order functionality in our Quantum Runtime.
By architecting features, tools, and techniques to be part of a well specified runtime we can:
 - Easily run experiments ("A/B testing") on the impact of features by toggling them in the runtime and running against a full library of applications.
 - Create a clear specification for how to leverage a feature, tool, or technique in our library of applications.
 - Mock out runtime features for ease of testing without consuming the scarce resource of quantum computer time.

Algorithmic Benchmark Library
The algorithmic benchmark library is a collection of quantum executables that probe different aspects of a quantum computer's performance. Each benchmark is based off of an algorithm of interest and operates on > 2 qubits (in contrast to traditional 1- and 2-qubit fidelity metrics).
We use the library card catalog in recirq.algorithmic_benchmark_library to get a description of each algorithmic benchmark.
End of explanation
"""


algo = algos.get_algo_benchmark_by_executable_family(
    executable_family='recirq.otoc.loschmidt.tilted_square_lattice')
print(type(algo))
algo

from pathlib import Path

# In ReCirq, the `executable_family` is the Python module path where
# the relevant code can be found.
algo_src_dir = Path('..') / algo.executable_family.replace('.', '/')
algo_src_dir

"""
Explanation: Select an example benchmark using executable_family
Each benchmark has a name (e.g. "loschmidt.tilted_square_lattice") and a domain (e.g. this benchmark is inspired by the OTOC Experiment so it is given the domain of "recirq.otoc"). We combine these two properties to give the executable_family string, which serves as a globally-unique identifier for each benchmark. In ReCirq, the executable_family is the Python module path where the relevant code can be found.
We'll seelct the "recirq.otoc.loschmidt.tilted_square_lattice" benchmark by querying the card catalog using this unique executable_family identifier.
End of explanation
"""


for config in algo.configs:
    print(config.full_name)

"""
Explanation: Each AlgorithimBenchmark has a collection of BenchmarkConfigs
An algorithmic benchmark defines a class of quantum executables. Often times the specific size, shape, depth or other properties is left as a parameter. For each benchmark, we have a collection of BenchmarkConfigs that fully specify what to run and can be run repeatedly for controlled comparison over time or between processors.
End of explanation
"""


config = algo.get_config_by_full_name('loschmidt.tilted_square_lattice.small-cz-v1')
print(config.description)

"""
Explanation: We'll select the small-cz-v1 configuration, as described below.
End of explanation
"""


config

"""
Explanation: A benchmark has three steps:
 1. Executable generation
 2. Execution
 3. Analysis
Usually, these steps are done in order but independently and with differing frequencies. For a robust benchmark, executable generation should likely be done once and the serialized QuantumExecutableGroup cached and re-used for subsequent executions. Execution should happen on a regular cadence for historical data or as part of an A/B test for trialing different runtime configuraiton options. Analysis can happen at any moment and may incorperate the latest data or a collection of datasets across time, processors, or runtime configurations.
End of explanation
"""


# Helper function to show scripts from the ReCirq repository for comparison
from IPython.display import Code, HTML

def show_python_script(path: Path):
    with path.open() as f:
        contents = f.read()

    display(HTML(f"<u>The contents of {path}:</u>"))
    display(Code(contents[contents.find('import'):], language='python'))

show_python_script(algo_src_dir / config.gen_script)

"""
Explanation: Step 1: Executable Generation
Here, we generate a QuantumExecutableGroup for a given range of parameters. This step is usually done once for each BenchmarkConfig and the serialized result is saved and re-used for execution. We use a short python file more like a configuration file than a script to commit the parameters for a particular config. The filename can be found as config.gen_script.
End of explanation
"""


import numpy as np
from recirq.otoc.loschmidt.tilted_square_lattice import get_all_tilted_square_lattice_executables

exes = get_all_tilted_square_lattice_executables(
    min_side_length=2, max_side_length=3, side_length_step=1,
    n_instances=3,
    macrocycle_depths=np.arange(0, 4 + 1, 1),
    twoq_gate_name='cz',
    seed=52,
)
len(exes)

"""
Explanation: We've copied the important bit into the cell below so you can execute it within this notebook.
End of explanation
"""


from cirq_google.workflow import (
    QuantumRuntimeConfiguration, 
    SimulatedProcessorWithLocalDeviceRecord,
    EngineProcessorRecord,
    RandomDevicePlacer
)

rt_config = QuantumRuntimeConfiguration(
    processor_record=SimulatedProcessorWithLocalDeviceRecord('rainbow', noise_strength=0.005),
    qubit_placer=RandomDevicePlacer(),
    random_seed=52,
)
rt_config

"""
Explanation: Step 2: Execution
The QuantumExecutableGroup for our benchmark defines what to run. Now we configure how to run it. This is done with a QuantumRuntimeConfiguration.
We specify which processor to use (here: a simulated one. Try changing to EngineProcessorRecord to run on a real device), how to map the "logical" qubit identities in the problem to physical qubits (here: randomly), and how to set the random number generator's seed.
End of explanation
"""


show_python_script(algo_src_dir / 'run-simulator.py')

"""
Explanation: Again, we use a short Python file in the repository to commit the configuration options for a particular config.
End of explanation
"""


import tempfile
base_data_dir = tempfile.mkdtemp()
base_data_dir

"""
Explanation: Usually, we're very careful about saving everthing in a structured way relative to a base_data_dir. Since this notebook is run interactively, we'll make a temporary directory to serve as our base_data_dir.
End of explanation
"""


from cirq_google.workflow import execute
raw_results = execute(rt_config, exes, base_data_dir=base_data_dir)

"""
Explanation: Actual execution is as simple as calling execute once the runtime configuration and executables are created.
End of explanation
"""


run_id = raw_results.shared_runtime_info.run_id
run_id

"""
Explanation: Since we didn't input our own, a random run_id is generated for us. The run_ids must be unique within a data directory.
End of explanation
"""


import recirq.otoc.loschmidt.tilted_square_lattice.analysis as analysis
import cirq_google as cg
import pandas as pd

raw_results = cg.ExecutableGroupResultFilesystemRecord.from_json(run_id=run_id, base_data_dir=base_data_dir)\
                    .load(base_data_dir=base_data_dir)
df = analysis.loschmidt_results_to_dataframe(raw_results)
df.head()

vs_depth_df, vs_depth_gb_cols = analysis.agg_vs_macrocycle_depth(df)
fit_df, exp_ansatz = analysis.fit_vs_macrocycle_depth(df)
total_df = pd.merge(vs_depth_df, fit_df, on=vs_depth_gb_cols)

from matplotlib import pyplot as plt

colors = plt.get_cmap('tab10')

for i, row in total_df.iterrows():
    plt.errorbar(
        x=row['macrocycle_depth'],
        y=row['success_probability_mean'],
        yerr=row['success_probability_std'],
        marker='o', capsize=5, ls='',
        color=colors(i),
        label=f'{row["width"]}x{row["height"]} ({row["n_qubits"]}q) {row["processor_str"]}; f={row["f"]:.3f}'
    )
    
    xx = np.linspace(np.min(row['macrocycle_depth']), np.max(row['macrocycle_depth']))
    yy = exp_ansatz(xx, a=row['a'], f=row['f'])
    plt.plot(xx, yy, ls='--', color=colors(i))
    
plt.legend(loc='best')
plt.yscale('log')
plt.xlabel('Macrocycle Depth')
plt.ylabel('Success Probability')
plt.tight_layout()

"""
Explanation: Step 3: Analysis and Plotting
Finally, we can analyze one or more datasets and generate plots. Since we've decoupled problem generation and execution from this step you can slice and dice your data any way you want. Usually, analysis routines will use the accompanying analysis module for helper function and do much of the pd.DataFrame and matplotlib munging interactively in a Jupyter notebook. One of the plots from plots.ipynb is reproduced here.
End of explanation
"""


# Pick one `QuantumExecutable` from the `QuantumExecutableGroup` 
exe = exes.executables[0]

import dataclasses
print('exe fields:')
print([f.name for f in dataclasses.fields(exe)])

"""
Explanation: Appendix: QuantumExecutable and ExecutableSpec

A QuantumExecutable contains a complete description of what to run. Think of it as a souped-up version of a cirq.Circuit. 
An ExecutableSpec is a problem-specific dataclass minimally capturing the salient independent variables, usually with plain-old-datatypes suitable for databasing and plotting.

Each benchmark provides a problem-specific subclass of ExecutableSpec and a function that turns those specs into problem-agnostic, fully-specified QuantumExecutables.

QuantumExecutable
Each QuantumExecutable fully specifies the quantum program to be run, but as high-level as possible. The most familiar part is circuit: cirq.Circuit, but the executable also includes measurement (i.e. repetitions) information, sweep parameters, and other data. We'll look at the fields on one of our executables below:
End of explanation
"""


print(algo.spec_class)
print(exe.spec.__class__)

"""
Explanation: QuantumExectutable.spec is a reference to the ExecutableSpec used to create this executable. Here, it is a TiltedSquareLatticeLoschmidtSpec which derives from the ExecutableSpec base class. Each AlgorithmicBenchmark has its own class derived from ExecutableSpec. This correspondance is recorded as AlgorithmicBenchmark.spec_class.
End of explanation
"""


import cirq
from recirq.otoc.loschmidt.tilted_square_lattice import TiltedSquareLatticeLoschmidtSpec

spec = TiltedSquareLatticeLoschmidtSpec(
    topology=cirq.TiltedSquareLattice(width=2, height=2),
    macrocycle_depth=0,
    instance_i=0,
    n_repetitions=1_000,
    twoq_gate_name='cz'
)

"""
Explanation: ExecutableSpec
The ExecutableSpec is a problem-specific dataclass containing the relevant independent variables for a benchmark. Since each benchmark has its own subclass of ExecutableSpec, we'll continue using our example loschmidt benchmark and create an example TiltedSquareLatticeLoschmidtSpec:
End of explanation
"""


exe.spec == spec

"""
Explanation: I've chosen the parameters corresponding to our example executable, exe = exes.executables[0].
End of explanation
"""


from recirq.otoc.loschmidt.tilted_square_lattice import tilted_square_lattice_spec_to_exe

exe2 = tilted_square_lattice_spec_to_exe(exe.spec, rs=np.random.RandomState(52))
exe == exe2

"""
Explanation: Below, we re-create the executable using just the spec.
End of explanation
"""


