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

# # Example for qutrit GST
# This notebook demonstrates how to construct the operation sequences and perform the analysis for qutrit GST when the model consists of symmetric $\pi/2$-rotations on each single qubit separately, `X`, `Y` and a 2-qubit Molmer-Sorenson gate which rotates around the `XX` axis by $\pi/2$.

# In[1]:


import pygsti
import pygsti.construction as pc
from pygsti.construction import qutrit

from numpy import pi, array
import pickle

import numpy as np


# First, we construct the target model.  This functionality is built into pyGSTi, so we just need to specify the single-qubit and M-S angles.

# In[2]:


target_model = qutrit.make_qutrit_model(errorScale=0, Xangle=pi/2, Yangle=pi/2, MSglobal=pi/2, MSlocal=0, basis="qt")
#print(target_model)


# Now construct the operation sequences needed by GST.  These fiducials and germs have been computed ahead of time and the results are used to construct the operation sequence lists below.  Then we construct an empty dataset containing all of the necessary experimental sequences which can serve as a template for the actual experimental results.

# In[3]:


fiducialPrep = pc.circuit_list(
    [(),('Gy',),('Gx',),('Gm',),
    ('Gx','Gx'), ('Gm','Gy'),('Gm','Gx'),
    ('Gy','Gy','Gy'),('Gx','Gx','Gx')])

fiducialMeasure = pc.circuit_list(
    [(),('Gy',),('Gx',),('Gm',),
    ('Gy','Gm'),('Gx','Gm')])

maxLengths = [1,2,4]

germs = pygsti.construction.circuit_list(
[('Gi',),
 ('Gy',),
 ('Gx',),
 ('Gm',),
 ('Gi', 'Gy'),
 ('Gi', 'Gx'),
 ('Gi', 'Gm'),
 ('Gy', 'Gx'),
 ('Gy', 'Gm'),
 ('Gx', 'Gm'),
 ('Gi', 'Gi', 'Gy'),
 ('Gi', 'Gi', 'Gx'),
 ('Gi', 'Gi', 'Gm'),
 ('Gi', 'Gy', 'Gy'),
 ('Gi', 'Gy', 'Gx'),
 ('Gi', 'Gy', 'Gm'),
 ('Gi', 'Gx', 'Gy'),
 ('Gi', 'Gx', 'Gx'),
 ('Gi', 'Gx', 'Gm'),
 ('Gi', 'Gm', 'Gy'),
 ('Gi', 'Gm', 'Gx'),
 ('Gi', 'Gm', 'Gm'),
 ('Gy', 'Gy', 'Gx'),
 ('Gy', 'Gy', 'Gm'),
 ('Gy', 'Gx', 'Gx'),
 ('Gy', 'Gx', 'Gm'),
 ('Gy', 'Gm', 'Gx'),
 ('Gy', 'Gm', 'Gm'),
 ('Gx', 'Gx', 'Gm'),
 ('Gx', 'Gm', 'Gm')])


# In[4]:


#Note above construction is now a "standard" qutrit model
from pygsti.construction import stdQT_XYIMS
target_model = stdQT_XYIMS.target_model()
fiducialPrep = stdQT_XYIMS.prepStrs
fiducialMeasure = stdQT_XYIMS.effectStrs
germs = stdQT_XYIMS.germs_lite
maxLengths = [1,2,4]


# In[5]:


print("%d prep fiducials" % len(fiducialPrep))
print("%d meas fiducials" % len(fiducialMeasure))
print("%d germs" % len(germs))


# In[6]:


#generate data template
expList = pygsti.construction.make_lsgst_experiment_list(target_model.operations.keys(), fiducialPrep, fiducialMeasure, germs,  maxLengths)
pygsti.io.write_empty_dataset("example_files/dataTemplate_qutrit_maxL=4.txt", expList, "## Columns = 0bright count, 1bright count, 2bright count")


# At this point **STOP** and create/fill a dataset file using the template written in the above cell.  Then proceed with the lines below to run GST on the data and create (hopefully useful) reports telling you about your gates.

# In[7]:


mdl_datagen = target_model.depolarize(op_noise=0.05)
DS = pygsti.construction.generate_fake_data(mdl_datagen, expList, 500, sampleError='multinomial', seed=2018)


# In[8]:


#DS = pygsti.io.load_dataset('PATH_TO_YOUR_DATASET',cache=True) # (cache=True speeds up future loads)


# In[9]:


#Run qutrit GST... which could take a while on a single CPU.  Please adjust memLimit to machine specs 
# (now 3GB; usually set to slightly less than the total machine memory)
result = pygsti.do_stdpractice_gst(DS,target_model,fiducialPrep,fiducialMeasure,germs,maxLengths,
                                     verbosity=2, comm=None, memLimit=3*(1024)**3, modes="TP,CPTP")


# In[10]:


#Create a report
ws = pygsti.report.create_standard_report(result, "example_files/sampleQutritReport",
                                         "Example Qutrit Report", verbosity=3, auto_open=True)


# In[11]:


print(target_model.basis)


# In[ ]: