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

# # ModelMemberGraph and Serialization
# 
# Example notebook of ModelMemberGraph functionality

# In[ ]:


import numpy as np

import pygsti


# In[ ]:


from pygsti.modelpacks import smq2Q_XYICNOT


# # Similar/Equivalent

# In[ ]:


ex_mdl1 = smq2Q_XYICNOT.target_model()
ex_mdl2 = ex_mdl1.copy()


# In[ ]:


ex_mmg1 = ex_mdl1.create_modelmember_graph()


# In[ ]:


ex_mmg1.print_graph()


# In[ ]:


ex_mmg1.mm_nodes['operations']['Gxpi2', 0]


# In[ ]:


ex_mmg2 = ex_mdl2.create_modelmember_graph()
print(ex_mmg1.is_similar(ex_mmg2))
print(ex_mmg1.is_equivalent(ex_mmg2))


# In[ ]:


ex_mdl2.operations['Gxpi2', 0][0, 0] = 0.0
ex_mmg2 = ex_mdl2.create_modelmember_graph()
print(ex_mmg1.is_similar(ex_mmg2))
print(ex_mmg1.is_equivalent(ex_mmg2))


# In[ ]:


ex_mdl2.operations['Gxpi2', 0] = pygsti.modelmembers.operations.StaticArbitraryOp(ex_mdl2.operations['Gxpi2', 0])
ex_mmg2 = ex_mdl2.create_modelmember_graph()
print(ex_mmg1.is_similar(ex_mmg2))
print(ex_mmg1.is_equivalent(ex_mmg2))


# In[ ]:


pspec = pygsti.processors.QubitProcessorSpec(2, ['Gi', 'Gxpi2', 'Gypi2', 'mygate'], geometry='line', nonstd_gate_unitaries={'mygate': np.eye(2, dtype='complex')})

ln_mdl1 = pygsti.models.create_crosstalk_free_model(pspec,
                                                    depolarization_strengths={('Gxpi2', 0): 0.1, ('mygate', 0): 0.2},
                                                    lindblad_error_coeffs={('Gypi2', 1): {('H', 'X'): 0.2, ('S', 'Y'): 0.3}})
print(ln_mdl1)


# In[ ]:


ln_mmg1 = ln_mdl1.create_modelmember_graph()
ln_mmg1.print_graph()


# In[ ]:


# Should be exactly the same
ln_mdl2 = pygsti.models.create_crosstalk_free_model(pspec,
                                                    depolarization_strengths={('Gxpi2', 0): 0.1, ('mygate', 0): 0.2},
                                                    lindblad_error_coeffs={('Gypi2', 1): {('H', 'X'): 0.2, ('S', 'Y'): 0.3}})
ln_mmg2 = ln_mdl2.create_modelmember_graph()
print(ln_mmg1.is_similar(ln_mmg2))
print(ln_mmg1.is_equivalent(ln_mmg2))


# In[ ]:


# Should be similar if we change params
ln_mdl3 = pygsti.models.create_crosstalk_free_model(pspec,
                                                    depolarization_strengths={('Gxpi2', 0): 0.01, ('mygate', 0): 0.02},
                                                    lindblad_error_coeffs={('Gypi2', 1): {('H', 'X'): 0.5, ('S', 'Y'): 0.1}})
ln_mmg3 = ln_mdl3.create_modelmember_graph()
print(ln_mmg1.is_similar(ln_mmg3))
print(ln_mmg1.is_equivalent(ln_mmg3))


# In[ ]:


# Should fail both, depolarize is on different gate
ln_mdl4 = pygsti.models.create_crosstalk_free_model(pspec,
                                                    depolarization_strengths={('Gypi2', 0): 0.1},
                                                    lindblad_error_coeffs={('Gypi2', 1): {('H', "X"): 0.2, ('S', 'Y'): 0.3}})
ln_mmg4 = ln_mdl4.create_modelmember_graph()
print(ln_mmg1.is_similar(ln_mmg4))
print(ln_mmg1.is_equivalent(ln_mmg4))


# # Serialization

# In[ ]:


ex_mdl1.write('example_files/ex_mdl1.json')


# In[ ]:


ln_mdl1.write('example_files/ln_mdl1.json')


# In[ ]: