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

# # An example of how to run GST on a 2-qubit system
# This tutorial gives an overview of the typical steps used to perform an end-to-end (i.e. experimental-data-to-report) Gate Set Tomography analysis on a 2-qubit system.  The steps are very similar to the single-qubit case; the main differences are:
# - the use of more general syntax when constructing 2-qubit gate sets
# - the increased number of fiducial and germ gate sequences
# - the increased run time required to compute GST estimates
# - a different report style 

# In[1]:


from __future__ import print_function
import matplotlib
matplotlib.use('Agg')


# In[2]:


import pygsti


# ### Step 1: Construct the desired 2-qubit gateset
# There are several ways to do this, as outlined by the comments in the cell below.

# In[3]:


# via build_gateset:
#   [4] = a 4-dimensional Hilbert (state) space
#   [('Q0','Q1')] = interpret this 4-d space as that of two qubits 'Q0', and 'Q1' (note these labels *must* begin with 'Q'!)
#   "Gix" = gate label; can be anything that begins with 'G' and is followed by lowercase letters
#   "X(pi/2,Q1)" = pi/2 single-qubit x-rotation gate on the qubit labeled Q1
#   "CX(pi,Q0,Q1)" = controlled pi x-rotation using qubits Q0 (control) and Q1 (target) 
#   "rho0" = prep label; can be anything that begins with "rho"
#   "E1" = effect label; can be anything that begins with "E"
#   "2" = a prep or effect expression indicating a projection/preparation of the 3rd (b/c 0-based) computational basis element
#   'dnup': ('rho0','E2') = designate the SPAM label "dnup" to mean preparation using "rho0" (a prep label) and measuring the outcome "E2" (an effect label)
#   "pp" = create all of these gate & SPAM operators in the Pauli-product basis.
gs_target = pygsti.construction.build_gateset( 
            [4], [('Q0','Q1')],['Gix','Giy','Gxi','Gyi','Gcnot'], 
            [ "X(pi/2,Q1)", "Y(pi/2,Q1)", "X(pi/2,Q0)", "Y(pi/2,Q0)", "CX(pi,Q0,Q1)" ],
            prepLabels=['rho0'], prepExpressions=["0"],
            effectLabels=['E0','E1','E2'], effectExpressions=["0","1","2"], 
            spamdefs={'upup': ('rho0','E0'), 'updn': ('rho0','E1'),
                      'dnup': ('rho0','E2'), 'dndn': ('rho0','remainder') }, basis="pp")

# Note that you can also explicity add identity operations, e.g. "I(Q0)" to get the same gate set,
#  and that this same syntax could be used for non-entangling 2-qubit gates, e.g. "X(pi/2,Q0):X(pi/2,Q1)".
gs_targetB = pygsti.construction.build_gateset( 
            [4], [('Q0','Q1')],['Gix','Giy','Gxi','Gyi','Gcnot'], 
            [ "I(Q0):X(pi/2,Q1)", "I(Q0):Y(pi/2,Q1)", "X(pi/2,Q0):I(Q1)", "Y(pi/2,Q0):I(Q1)", "CX(pi,Q0,Q1)" ],
            prepLabels=['rho0'], prepExpressions=["0"], 
            effectLabels=['E0','E1','E2'], effectExpressions=["0","1","2"], 
            spamdefs={'upup': ('rho0','E0'), 'updn': ('rho0','E1'),
                      'dnup': ('rho0','E2'), 'dndn': ('rho0','remainder') }, basis="pp")

# If you're lucky and your gateset is one of pyGSTi's "standard" gate sets, you can just import it.
from pygsti.construction import std2Q_XYCNOT
gs_targetC = std2Q_XYCNOT.gs_target

#check that these are all the same
assert(abs(gs_target.frobeniusdist(gs_targetB)) < 1e-6)
assert(abs(gs_target.frobeniusdist(gs_targetC)) < 1e-6)


# ### Step 2: Obtain lists of fiducial and germ gate sequences
# These are the building blocks of the gate sequences performed in the experiment. Typically, these lists are either given to you by the folks at Sandia National Labs (email pygsti@sandia.gov), provided by pyGSTi because you're using a "standard" gate set, or computed using "fiducial selection" and "germ selection" algorithms (which are a part of pyGSTi, but not covered in this tutorial).

# In[4]:


#If you know the fiducial strings you can create a list manually.  Note
# that in general there can be different "preparation" and "measurement"
# (or "effect") fiducials.
prep_fiducials = pygsti.construction.gatestring_list( \
[ (), ('Gix',), ('Giy',), ('Gix','Gix'), 
('Gxi',), ('Gxi','Gix'), ('Gxi','Giy'), ('Gxi','Gix','Gix'), 
('Gyi',), ('Gyi','Gix'), ('Gyi','Giy'), ('Gyi','Gix','Gix'), 
('Gxi','Gxi'), ('Gxi','Gxi','Gix'), ('Gxi','Gxi','Giy'), ('Gxi','Gxi','Gix','Gix') ] )

effect_fiducials = pygsti.construction.gatestring_list( \
[(), ('Gix',), ('Giy',), 
 ('Gix','Gix'), ('Gxi',),
 ('Gyi',), ('Gxi','Gxi'),
 ('Gxi','Gix'), ('Gxi','Giy'),
 ('Gyi','Gix'), ('Gyi','Giy')] )

#Or, if you're lucky, you can just import them
prep_fiducialsB = std2Q_XYCNOT.prepStrs
effect_fiducialsB = std2Q_XYCNOT.effectStrs

#check that these are the same
assert(prep_fiducials == prep_fiducialsB)
assert(effect_fiducials == effect_fiducialsB)

#Use fiducial sequences to create a "spam specifiers" object, telling
# GST which preparation and measurement fiducials to follow and precede which
# state preparation and effect operators, respectively.
specs = pygsti.construction.build_spam_specs(
    prepStrs=prep_fiducials,
    effectStrs=effect_fiducials,
    prep_labels=gs_target.get_prep_labels(),
    effect_labels=gs_target.get_effect_labels() )

#Alternatively, if you're lucky, you can grab the specs directly:
specsB = std2Q_XYCNOT.specs
assert(specs[0] == specsB[0])


# In[5]:


#germ lists can be specified in the same way.  In this case, there are
# 71 germs required to do honest GST.  Since this would crowd this tutorial
# notebook, we create some smaller lists of germs manually and import the
# full 71-germ list from std2Q_XYCNOT
germs4 = pygsti.construction.gatestring_list(
    [ ('Gix',), ('Giy',), ('Gxi',), ('Gyi',) ] )

germs11 = pygsti.construction.gatestring_list(
    [ ('Gix',), ('Giy',), ('Gxi',), ('Gyi',), ('Gcnot',), ('Gxi','Gyi'), ('Gix','Giy'),
      ('Gix','Gcnot'), ('Gxi','Gcnot'), ('Giy','Gcnot'), ('Gyi','Gcnot') ] )

germs71 = std2Q_XYCNOT.germs


# ### Step 3: Data generation
# Now that fiducial and germ strings have been found, we can generate the list of experiments needed to run GST, just like in the 1-qubit case.  As an additional input we'll need a list of lengths indicating the maximum length strings to use on each successive GST iteration.

# In[6]:


#A list of maximum lengths for each GST iteration
maxLengths = [1,2,4]

#Create a list of GST experiments for this gateset, with
#the specified fiducials, germs, and maximum lengths.  We use
#"germs4" here so that the tutorial runs quickly; really, you'd
#want to use germs71!
listOfExperiments = pygsti.construction.make_lsgst_experiment_list(gs_target.gates.keys(), prep_fiducials,
                                                                   effect_fiducials, germs4, maxLengths)

#Create an empty dataset file, which stores the list of experiments
# and zerod-out columns where data should be inserted.  Note the use of the SPAM
# labels in the "Columns" header line.
pygsti.io.write_empty_dataset("tutorial_files/My2QDataTemplate.txt", listOfExperiments,
                              "## Columns = upup count, updn count, dnup count, dndn count")


# In[7]:


#Generate some "fake" (simulated) data based on a depolarized version of the target gateset
gs_datagen = gs_target.depolarize(gate_noise=0.1, spam_noise=0.001)
ds = pygsti.construction.generate_fake_data(gs_datagen, listOfExperiments, nSamples=1000,
                                            sampleError="multinomial", seed=2016)

#if you have a dataset file with real data in it, load it using something like:
#ds = pygsti.io.load_dataset("mydir/My2QDataset.txt")


# ### Step 4: Run GST using `do_long_sequence_gst`
# Just like for 1-qubit GST, we call the driver routine `do_long_sequence_gst` to compute the GST estimates.  Usually for two qubits this could take a long time (hours) based on the number of gate sequences used.  However, since we chose an incomplete set of only 4 germs and set our maximum max-length to 4, this will run fairly quickly (~20min).
# 
# Some notes about the options/arguments to `do_long_sequence_gst` that are particularly relevant to 2-qubit GST:
# - `mxBasis` indicates which basis the target gateset matrices are in -- this should be the same as that used to create the gateset.  Here `'pp'` denotes the Pauli-product basis.
# - `advancedOptions` expects a dictionary of with option names as the keys.  Relevant options are:
#   - `memoryLimitInBytes` gives an estimate of how much memory is available to use on your system (in bytes).  This is currently *not* a hard limit, and pyGSTi may require 50-100% more memory than this "limit".  So you'll need to be conservative in the value you place here: if your machine has 10GB of RAM, set this to 3 or 5 GB initially and increase it as you see how much memory is actually used using a separate OS performance monitor tool.
#   - `depolarizeLGST` gives an amount (in [0,1]) to depolarize the initial LGST estimate that is used as the initial guess for long-sequence GST.  In practice, we find that in the larger 2-qubit Hilbert space, the LGST estimate may be so poor as to adversely affect the subsequent long-sequence GST (e.g. very slow convergence).  Depolarizing the LGST estimate remedies this.  If you're unsure what to put here, either don't specify `depolarizeLGST` at all (the same as using 0.0), or just use 0.1.
#   - `verbosity` tells the routine how much detail to print to stdout.  If you don't mind waiting a while without getting any output, you can leave this at its default value (2).  If you can't standing wondering whether GST is still running or has locked up, set this to 3.

# In[8]:


import time
start = time.time()
results = pygsti.do_long_sequence_gst(ds, gs_target, prep_fiducials, effect_fiducials, germs4,
                                    maxLengths, gaugeOptParams={'itemWeights': {'spam':0.1,'gates': 1.0}},
                                    advancedOptions={ 'depolarizeStart' : 0.1 }, memLimit=3*(1024)**3,
                                    verbosity=3 )
end = time.time()
print("Total time=%f hours" % ((end - start) / 3600.0))

#If you wanted to, you could pickle the results for later analysis:
#pickle.dump(results, open("MySavedResults.pkl", "w"))


# ### Step 5: Create report(s) using the returned `Results` object
# The `Results` object returned from `do_long_sequence_gst` is able to generate several different types of reports.  Most of these are designed to display single-qubit results (for historical reasons).  The "general"-type report was designed with 2-qubit data presentation in mind, so creating reports with `create_general_report` will be the most useful.

# In[9]:


results.gatesets['final estimate'].set_basis("pp", [4])
print(results.gatesets['final estimate'].get_basis_name())
print(results.gatesets['final estimate'].get_basis_dimension())


# In[10]:


results.create_general_report_pdf(filename="tutorial_files/easy_2q_general.pdf",verbosity=2)


# Now open [tutorial_files/easy_2q_general.pdf](tutorial_files/easy_2q_general.pdf) to see the results.  You've run 2-qubit GST!

# In[11]:


import pickle
with open("tutorial_files/easy_2q_results.pkl","wb") as pklfile:
        pickle.dump(results, pklfile)


# In[ ]: