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

# # Examples of tables and plots available from a `Workspace`
# 
# PyGSTi's `Workspace` object is first a foremost a container and factory for plots and tables.  At the most basic level, it can be used to generate nice output based on quantities (e.g. `Model`, `DataSet`, etc. objects) that you've computed or loaded within a notebook.  For this, it's useful to call `init_notebook_mode` with `autodisplay=True` (see below) so that you don't have to `.display()` everything - `display()` gets called automatically when a plot or table is created.
# 
# ## Getting some results
# First, let's run Gate Set Tomography (GST) on the standard 1-qubit model to get some results to play with.  We generate a few `DataSet` objects and then call `run_long_sequence_gst` to run GST, generating a `ModelEstimateResults` object (essentially a container for `Model` objects).  For more details, see the tutorials [GST overview tutorial](../algorithms/GST-Overview.ipynb), the [tutorial on GST functions](../algorithms/GST-Driverfunctions.ipynb), and the [tutorial explaining the ModelEstimateResults object](../objects/advanced/Results.ipynb).   

# In[ ]:


import numpy as np
import pygsti
from pygsti.modelpacks import smq1Q_XYI


# In[ ]:


#The usual GST setup: we're going to run GST on the standard XYI 1-qubit model
target_model = smq1Q_XYI.target_model()
prep_fiducials = smq1Q_XYI.prep_fiducials()
meas_fiducials = smq1Q_XYI.meas_fiducials()
germs = smq1Q_XYI.germs()
maxLengths = [1,2]
listOfExperiments = pygsti.construction.create_lsgst_circuits(
    target_model.operations.keys(), prep_fiducials, meas_fiducials, germs, maxLengths)


# In[ ]:


#Create some datasets for analysis
mdl_datagen1 = target_model.depolarize(op_noise=0.1, spam_noise=0.02)
mdl_datagen2 = target_model.depolarize(op_noise=0.05, spam_noise=0.01).rotate(rotate=(0.01,0.01,0.01))

ds1 = pygsti.construction.simulate_data(mdl_datagen1, listOfExperiments, n_samples=1000,
                                            sample_error="binomial", seed=1234)
ds2 = pygsti.construction.simulate_data(mdl_datagen2, listOfExperiments, n_samples=1000,
                                            sample_error="binomial", seed=1234)
ds3 = ds1.copy_nonstatic(); ds3.add_counts_from_dataset(ds2); ds3.done_adding_data()


# In[ ]:


#Run GST on all three datasets
target_model.set_all_parameterizations("TP")
results1 = pygsti.run_long_sequence_gst(ds1, target_model, prep_fiducials, meas_fiducials, germs, maxLengths, verbosity=0)
results2 = pygsti.run_long_sequence_gst(ds2, target_model, prep_fiducials, meas_fiducials, germs, maxLengths, verbosity=0)
results3 = pygsti.run_long_sequence_gst(ds3, target_model, prep_fiducials, meas_fiducials, germs, maxLengths, verbosity=0)

#make some shorthand variable names for later
tgt = results1.estimates['GateSetTomography'].models['target']

ds1 = results1.dataset
ds2 = results2.dataset
ds3 = results3.dataset

mdl1 = results1.estimates['GateSetTomography'].models['go0']
mdl2 = results2.estimates['GateSetTomography'].models['go0']
mdl3 = results3.estimates['GateSetTomography'].models['go0']

gss = results1.circuit_lists['final']


# ## Gallery of `Workspace` plots and tables.
# Now that we have some results, let's create a `Workspace` and make some plots and tables.
# 
# To get tables and plots to display properly, one must run `init_notebook_mode`.  The `connected` argument indicates whether you want to rely on an active internet connection.  If `True`, then resources will be loaded from the web (e.g. a CDN), and if you save a notebook as HTML the file size may be smaller.  If `False`, then all the needed resources (except MathJax) are provided by pyGSTi, and an `offline` directory is automatically created in the same directory as your notebook.  This directory contains all the necessary resources, and must "tag along" with the notebook and any saved-as-HTML versions of it in order for everything to work.  The second argument, `autodisplay`, determines whether tables and plots are automatically displayed when they are created.  If `autodisplay=False`, one must call the `display()` member function of a table or plot to display it. 

# In[ ]:


from pygsti.report import workspace
w = workspace.Workspace()
w.init_notebook_mode(connected=False, autodisplay=True) 


# Plots and tables are created via member functions of a `Workspace` (`w` in our case).  Note that you can start typing "`w.`" and TAB-complete to see the different things a `Workspace` can make for you.  Furthermore, pressing SHIFT-TAB after the opening parenthesis of a function,  e.g. after typing "`w.GatesVsTargetTable(`", will bring up Jupyter's help window showing you the function signature (the arguments you need to give the function).
# 
# #### The remainder of this tutorial demonstrates some of the tables and plots you can create. 
# Note that displayed objects have a resize handle in their lower right corner.

# In[ ]:


w.ColorBoxPlot(("logl",), gss, ds1, mdl1, typ='scatter')
w.ColorBoxPlot(("logl",), gss, ds1, mdl1, typ='boxes')
w.ColorBoxPlot(("logl",), gss, ds1, mdl1, typ='histogram')


# In[ ]:


w.FitComparisonBarPlot(gss, results1.circuit_lists['iteration'],
                     results1.estimates['GateSetTomography'].models['iteration estimates'], ds1)


# In[ ]:


w.GramMatrixBarPlot(ds1,tgt)


# In[ ]:


w.GatesVsTargetTable(mdl1, tgt)


# In[ ]:


w.SpamVsTargetTable(mdl2, tgt)


# In[ ]:


w.ColorBoxPlot(("chi2","logl"), gss, ds1, mdl1, box_labels=True)
  #Notice how long it takes to switch between "chi2" and "logl".  This 
  # is due to drawing all of the box labels (box_labels=True).


# In[ ]:


#This one requires knowng that each Results object holds a list of models
# from each GST intation along with the corresponding operation sequences that were used.
w.FitComparisonTable(gss, results1.circuit_lists['iteration'],
                     results1.estimates['GateSetTomography'].models['iteration estimates'], ds1)


# In[ ]:


# We can reuse 'gss' for all three since the operation sequences are the same.
w.FitComparisonTable(["GS1","GS2","GS3"], [gss, gss, gss], [mdl1,mdl2,mdl3], ds1, x_label="Model")


# In[ ]:


w.ChoiTable(mdl3, display=('matrix','barplot'))


# In[ ]:


w.GateMatrixPlot(mdl1[('Gxpi2',0)],scale=1.0, box_labels=True,ylabel="hello")
w.GateMatrixPlot(pygsti.tools.error_generator(mdl1[('Gxpi2',0)], tgt[('Gxpi2',0)], 'pp'), scale=1.5)


# In[ ]:


from pygsti.modelpacks import smq2Q_XYCNOT
w.GateMatrixPlot(smq2Q_XYCNOT.target_model()[('Gxpi2',0)],scale=1.0, box_labels=False,ylabel="hello",mx_basis="pp")


# In[ ]:


mx = np.array( 
[[ 7.3380823,   8.28446943,  7.4593754,   3.91256384,  0.68631199],
 [ 3.36139818,  7.42955114,  6.78516082,  0.35863173,  5.57713093],
 [ 2.61489939,  3.40182958,  6.77389064,  9.29736475,  0.33824271],
 [ 9.64258149,  9.45928809,  6.91516602,  5.61423854,  0.56480777],
 [ 2.15195669,  9.37588783,  5.1781991,   7.20087591,  1.46096288]], 'd')
cMap = pygsti.report.colormaps.LinlogColormap(vmin=0, vmax=10, n_boxes=25, pcntle=0.55, dof_per_box=1, color='blue')
w.MatrixPlot(mx, colormap=cMap, colorbar=False)


# In[ ]:


mx = np.identity(3,'d')
mx[0,1] = 2.1
mx[2,2] = 4.0
mx[2,0] = 3.0
mx[0,2] = 7.0
mx[2,1] = 10.0
mx[0,0] = np.nan
cMap = pygsti.report.colormaps.PiecewiseLinearColormap(
            [[0,(0,0.5,0)],[1,(0,1.0,0)],[2,(1.0,1.0,0)],
             [4,(1.0,0.5,0)],[10,(1.0,0,0)]])
#print(cMap.colorscale())
w.MatrixPlot(mx, colormap=cMap, colorbar=False, grid="white:1", box_labels=True, prec=2,
             xlabels=('TP',"CPTP","full"),ylabels=("DS0","DS1","DS2"))


# In[ ]:


w.ErrgenTable(mdl3,tgt)


# In[ ]:


w.PolarEigenvaluePlot([np.linalg.eigvals(mdl2[('Gxpi2',0)])],["purple"],scale=1.5)


# In[ ]:


w.GateEigenvalueTable(mdl2, display=('evals','polar'))


# In[ ]:


w.GateDecompTable(mdl1,target_model)
#w.old_GateDecompTable(gs1) #historical; 1Q only


# In[ ]:


#Note 2Q angle decompositions
from pygsti.modelpacks import smq2Q_XXYYII
from pygsti.modelpacks import smq2Q_XYCNOT

w.GateDecompTable(smq2Q_XXYYII.target_model(), smq2Q_XXYYII.target_model())

import scipy
I = np.array([[1,0],[0,1]],'complex')
X = np.array([[0,1],[1,0]],'complex')
Y = np.array([[0,1j],[-1j,0]],'complex')
XX = np.kron(X,X)
YY = np.kron(Y,Y)
IX = np.kron(I,X)
XI = np.kron(X,I)
testU = scipy.linalg.expm(-1j*np.pi/2*XX)
testS = pygsti.unitary_to_process_mx(testU)
testS = pygsti.change_basis(testS,"std","pp")

#mdl_decomp = std2Q_XYCNOT.target_model()
#mdl_decomp.operations['Gtest'] = testS
#w.GateDecompTable(mdl_decomp, mdl_decomp)


# In[ ]:


dsLabels = ["A","B","C"]
datasets = [ds1, ds2, ds3]
dscmps = {}
for i,ds_a in enumerate(datasets):
    for j,ds_b in enumerate(datasets[i+1:],start=i+1):
        dscmps[(i,j)] = pygsti.objects.DataComparator([ds_a, ds_b])

w.DatasetComparisonSummaryPlot(dsLabels, dscmps)


# In[ ]:


w.DatasetComparisonHistogramPlot(dscmps[(1,2)])


# ### Saving figures to file
# You can also save plot and figures to separate files using their `saveas` method.  The output format is determined by the file extension, and allowed extensions are:
# 
# - 'pdf': Adobe portable document format
# - 'tex': LaTeX source (uncompiled, *tables only*)
# - 'pkl': Python pickle (of a pandas `DataFrame` for tables, a dict for plots)
# - 'html': A stand-alone HTML document

# In[ ]:


import os
if not os.path.exists("../tutorial_files/tempTest"):
    os.mkdir("../tutorial_files/tempTest")

obj = w.GatesVsTargetTable(mdl1, tgt)
#obj = w.ErrgenTable(mdl3,tgt)
#obj = w.ColorBoxPlot(("logl",), gss, ds1, mdl1, typ='boxes')

obj.saveas("../tutorial_files/tempTest/testSave.pdf")
obj.saveas("../tutorial_files/tempTest/testSave.tex")
obj.saveas("../tutorial_files/tempTest/testSave.pkl")
obj.saveas("../tutorial_files/tempTest/testSave.html")


# ## Exporting notebooks to HTML
# If you want, you can save figure-containing notebooks (like this one) as an HTML file by going to **File => Download As => HTML** in the Jupyter menu.  The resulting file will retain all of the plot interactivity, so long as its in a directory with an `offline` folder (because we set `connected=False` above).