#!/usr/bin/env python
# coding: utf-8

# # Easy experiment setup

# This notebook demonstrates some convenience functions for coming up with the
# experiments needed to do GST on a new gateset.

# In[1]:


import pygsti
import pygsti.algorithms.germselection as germsel
import pygsti.algorithms.fiducialselection as fidsel
import pygsti.construction as constr


# Let's construct a 1-qubit $X(\pi/2)$, $Y(\pi/2)$, $I$ gateset for which we will need to find germs and fiducials.

# In[2]:


gs_target = constr.build_gateset([2], [('Q0',)], ['Gi', 'Gx', 'Gy'],
                                 ["I(Q0)", "X(pi/2,Q0)", "Y(pi/2,Q0)"],
                                 prepLabels=['rho0'], prepExpressions=["0"],
                                 effectLabels=['E0'], effectExpressions=["1"],
                                 spamdefs={'plus': ('rho0', 'E0'),
                                           'minus': ('rho0', 'remainder')})


# ## Hands-off

# We begin by demonstrating the most hands-off approach.

# We can generate a germ set simply by providing the target gateset.

# In[3]:


germs = germsel.generate_germs(gs_target)


# In the same way we can generate preparation and measurement fiducials.

# In[4]:


prepFiducials, measFiducials = fidsel.generate_fiducials(gs_target)


# Now that we have germs and fiducials, we can construct the list of experiments we need to perform in
# order to do GST. The only new things to provide at this point are the sizes for the experiments we want
# to perform (in this case we want to perform between 0 and 256 gates between fiducial pairs, going up
# by a factor of 2 at each stage).

# In[5]:


maxLengths = [0] + [2**n for n in range(8 + 1)]
listOfExperiments = constr.make_lsgst_experiment_list(gs_target.gates.keys(), prepFiducials,
                                                      measFiducials, germs, maxLengths)


# The list of `GateString` that the previous function gave us isn't necessarily the most readable
# form to present the information in, so we can write the experiment list out to an empty data
# file to be filled in after the experiments are performed.

# In[6]:


pygsti.io.write_empty_dataset("tutorial_files/EasyDataTemplate.txt", listOfExperiments)


# That's it! You can now take the data this file is asking for and come back to analyze the data.

# ## More control

# There are many ways you can assume more control over the experiment design process. We'll only demonstrate
# a few here, but all options are discussed in the documentation for the various functions we've used.

# ### Different algorithms

# There are a number of different algorithms available for germ selection. You can choose a non-default
# algorithm by specifying the `algorithm` keyword argument.

# Each of the available algorithms has a set of keyword arguments of its own with which you
# can more precisely specify how you want it to behave. These keyword arguments can be passed
# as a dictionary to `generate_germs` through the keyword argument `algorithm_kwargs`.

# In[7]:


graspGerms = germsel.generate_germs(gs_target, algorithm='grasp', algorithm_kwargs={'iterations': 1})


# Fiducial selection can be controlled in much the same way.

# In[8]:


slackPrepFids, slackMeasFids = fidsel.generate_fiducials(gs_target, algorithm='slack',
                                                         algorithm_kwargs={'slackFrac': 0.25})


# ### Germ and fiducial lengths

# We can also adjust some algorithm-independent parameters for germ and fiducial selection. For instance, all
# of the algorithms currently rely on having a pool of gatestring from which they construct germs and fiducials.
# The size of this pool is set by specifying the longest germ or fiducial to include in this pool.

# For germ selection, the default maximum germ length is 6, and we can see that our original construction
# indeed makes use of germs of length 6.

# In[9]:


max([len(germ) for germ in germs])


# We can try and set the maximum germ length to 5 and see what we get.

# In[10]:


germsMaxLength5 = germsel.generate_germs(gs_target, maxGermLength=5)


# In[11]:


max([len(germ) for germ in germsMaxLength5])


# Sure enough, we now have a germ set with a shorter longest germ. If we get too ambitious in shortening the maximum
# germ length, germ selection won't be able to find an amplificationally complete germ set. It will send a warning
# message to `stderr` if this happens and return `None`.

# In[12]:


germsMaxLength3 = germsel.generate_germs(gs_target, maxGermLength=3)
print(germsMaxLength3)


# Fiducial selection defaults to a maximum fiducial length of 2. This allows us to construct an informationally
# complete set of states and measurements, but for this gateset we know that there is a uniformly
# informationally complete set of states and measurements that require fiducials of up to length 3. We can find
# that set of fiducials by telling `generate_fiducials` to consider longer fiducials.

# In[13]:


uniformPrepFids, uniformMeasFids = fidsel.generate_fiducials(gs_target, maxFidLength=3,
                                                             algorithm='grasp',
                                                             algorithm_kwargs={'iterations': 100})


# As was the case with germ selection, if you are too aggressive in limiting fiducial length you may
# constrain the algorithm to the extent that it cannot even find a set of fiducials to generate an
# informationally complete set of states and measurements. In that case, it will also send a warning
# message to `stderr` and return `None` for the preparation and measurement fiducial sets.

# In[14]:


incompletePrepFids, incompleteMeasFids = fidsel.generate_fiducials(gs_target, maxFidLength=1)


# In[15]:


print(incompleteMeasFids, incompletePrepFids)


# ### Set requirements

# There are several natural things to require of the returned germ and fiducial sets. For germ sets, you will usually
# want the individual gates to be included as germs. If for some reason you don't want this, you can set the
# *forceSingletons* keyword argument to `False`.

# In[16]:


nonSingletonGerms = germsel.generate_germs(gs_target, force=None, maxGermLength=4,
                                           algorithm='grasp', algorithm_kwargs={'iterations': 5})


# In fiducial selection, it is likewise natural to require the empty gate string to be in the
# fiducial set. This requirement may be disabled by setting *forceEmpty* to `False`. It is also
# often desireable for identity gates to be left out of fiducials, since they add no diversity
# to the set of states and measurements generated. You can allow identity gates in fiducials by
# setting *omitIdentity* to `False`.

# A more common modification to the fiducial set requirements is to leave out additional gates from fiducials.
# This might be desireable if you have a multi-qubit system and you expect your 2-qubit gates to be of lower
# fidelity than your single-qubit gates. In this case you might want to construct fiducials from only
# single-qubit gates. A list of gates that you would like to omit from your fiducials can be provided as a
# list of gate labels to the *gatesToOmit* keyword argument.

# Our gateset doesn't have multi-qubit gates, but we can demonstrate several pieces of this
# functionality by setting *omitIdentity* to `False` and omitting the identity manually using
# *gatesToOmit*.

# In[17]:


omitIdentityPrepFids, omitIdentityMeasFids = fidsel.generate_fiducials(gs_target, omitIdentity=False,
                                                                       gatesToOmit=['Gi'])


# In[18]:


omitIdentityPrepFids


# In[19]:


omitIdentityMeasFids


# ### Verbosity

# The various algorithms can tell you something of what's going on with them while they're running. By default,
# this output is silenced, but it can be turned on using the *verbosity* keyword argument.

# A verbosity level of 1 is the default. This prints out what algorithm is being used, the returned set, and
# the score of that set.

# A verbosity level of 0 silences all output (other than warnings that things have gone wrong).

# A verbosity level of $n+1$ where $n\geq0$ prints the output of verbosity level 1 in addition to
# the output that the current algorithm displays when its own verbosity is set to $n$.

# In[20]:


verbosePrepFids, verboseMeasFids = fidsel.generate_fiducials(gs_target, verbosity=2)


# In[21]:


silentGerms = germsel.generate_germs(gs_target, algorithm='slack', algorithm_kwargs={'maxIter': 5},
                                      verbosity=0)


# In[ ]: