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

# # Workspace Pickling
# This example notebook shows how to generate reports with a pickled workspace object.  
# First, we run GST:

# In[1]:


import pygsti
import pickle
from pygsti.construction import std1Q_XYI

gs_target = std1Q_XYI.gs_target
fiducials = std1Q_XYI.fiducials
germs = std1Q_XYI.germs
maxLengths = [1,2,4]
#maxLengths = [1, 2, 4, 8, 16, 32, 64]

#Generate some data
gs_datagen = gs_target.depolarize(gate_noise=0.1, spam_noise=0.001)
gs_datagen = gs_datagen.rotate(rotate=0.04)
listOfExperiments = pygsti.construction.make_lsgst_experiment_list(gs_target, fiducials, fiducials, germs, maxLengths)
ds = pygsti.construction.generate_fake_data(gs_datagen, listOfExperiments, nSamples=1000,
                                            sampleError="binomial", seed=1234)
#Run GST
gs_target.set_all_parameterizations("TP") #TP-constrained
results = pygsti.do_long_sequence_gst(ds, gs_target, fiducials, fiducials, germs,
                                      maxLengths, verbosity=0)
with open('example_files/example_report_results.pkl', 'wb') as outfile:
    pickle.dump(results, outfile, protocol=2)

# Case1: TP-constrained GST
tpTarget = gs_target.copy()
tpTarget.set_all_parameterizations("TP")
results_tp = pygsti.do_long_sequence_gst(ds, tpTarget, fiducials, fiducials, germs,
                                      maxLengths, gaugeOptParams=False, verbosity=0)
# Gauge optimize
est = results_tp.estimates['default']
gsFinal = est.gatesets['final iteration estimate']
gsTarget = est.gatesets['target']
for spamWt in [1e-4,1e-3,1e-2,1e-1,1.0]:
    gs = pygsti.gaugeopt_to_target(gsFinal,gsTarget,{'gates':1, 'spam':spamWt})
    est.add_gaugeoptimized({'itemWeights': {'gates':1, 'spam':spamWt}}, gs, "Spam %g" % spamWt) 

#Case2: "Full" GST
fullTarget = gs_target.copy()
fullTarget.set_all_parameterizations("full")
results_full = pygsti.do_long_sequence_gst(ds, fullTarget, fiducials, fiducials, germs,
                                      maxLengths, gaugeOptParams=False, verbosity=0)
#Gauge optimize
est = results_full.estimates['default']
gsFinal = est.gatesets['final iteration estimate']
gsTarget = est.gatesets['target']
for spamWt in [1e-4,1e-3,1e-2,1e-1,1.0]:
    gs = pygsti.gaugeopt_to_target(gsFinal,gsTarget,{'gates':1, 'spam':spamWt})
    est.add_gaugeoptimized({'itemWeights': {'gates':1, 'spam':spamWt}}, gs, "Spam %g" % spamWt)

with open('example_files/full_report_results.pkl', 'wb') as outfile:
    pickle.dump((results_tp, results_full), outfile, protocol=2)


# ## Workspace Pickling
# Now that we've saved some GST data, we can begin generating reports.  
# The smart cache will track which plots we are unable to pickle.  
# At the time of writing, these are only plotly plots. `ReportTable` objects should delete their `ws` and `figs` members, allowing themselves to be pickled. 

# In[2]:


from pygsti.tools import timed_block

with timed_block('pickle loading'):
    with open('example_files/full_report_results.pkl', 'rb') as infile:
        results_tp, results_full = pickle.load(infile)
with timed_block('TP/Full multi report'):
    ws = pygsti.report.create_general_report({'TP': results_tp, "Full": results_full},
                                             "tutorial_files/exampleMultiGenReport.html",
                                             "Workspace Pickling Example Report1", 
                                             verbosity=3, auto_open=False)
ws.save_cache('example_files/wscache.pkl', True)


# In the Cell below, we create the same report, but use the existing workspace object.

# In[3]:


with timed_block('TP/Full multi report'):
    ws = pygsti.report.create_general_report({'TP': results_tp, "Full": results_full},
                                             "tutorial_files/exampleMultiGenReport.html",
                                             "Workspace Pickling Example Report2", 
                                             verbosity=3, auto_open=False, ws=ws)


# However, most of the time, we care about the time of the full report generation progress:

# In[4]:


del ws
with timed_block('Full report creation from pickle'):
    ws = pygsti.report.Workspace('example_files/wscache.pkl')
    pygsti.report.create_general_report({'TP': results_tp, "Full": results_full},
                                        "tutorial_files/exampleMultiGenReport.html",
                                        "Workspace Pickling Example Report3", 
                                        verbosity=3, auto_open=False, ws=ws)


# In[ ]: