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

# In[1]:


import numpy as np
import pygsti
from pygsti.extras import rb
from pygsti.baseobjs import Label

n_1 = 4
glist = ['Gx','Gy','Gcnot']
pspec_1 = pygsti.obj.ProcessorSpec(n_1,glist,verbosity=0)

n_2 = 3
glist = ['Gxpi','Gypi','Gzpi','Gh','Gp','Gcphase']
availability = {'Gcphase':[(0,1),(1,2)]}
pspec_2 = pygsti.obj.ProcessorSpec(n_2,glist,availability=availability,verbosity=0)

#Fix processor specs to work with n=1 qubits.
#n = 1
#glist = ['Gh','Gp']
#pspec_3 = pygsti.obj.ProcessorSpec(n=1,glist,verbosity=0)


# In[ ]:


#pspec_1.models['clifford'].gates.keys()


# In[ ]:


# --------------------------------------- #
# --- Test the circuit layer samplers --- #
# --------------------------------------- #

# Tests for the sampling by pairs function
layer = rb.sample.circuit_layer_by_pairings(pspec_1, twoQprob=0.0, oneQgatenames='all', twoQgatenames='all', 
                                  gatesetname = 'clifford')
assert(len(layer) == n_1)
layer = rb.sample.circuit_layer_by_pairings(pspec_1, twoQprob=1.0, oneQgatenames='all', twoQgatenames='all', 
                                  gatesetname = 'clifford')
assert(len(layer) == n_1//2)
layer = rb.sample.circuit_layer_by_pairings(pspec_1, twoQprob=0.0, oneQgatenames=['Gx',], twoQgatenames='all', 
                                  gatesetname = 'target')
assert(layer[0].name == 'Gx')

layer = rb.sample.circuit_layer_by_Qelimination(pspec_2, twoQprob=0.0, oneQgates='all', twoQgates='all',
                                                gatesetname='clifford')
assert(len(layer) == n_2)
layer = rb.sample.circuit_layer_by_Qelimination(pspec_2, twoQprob=1.0, oneQgates='all', twoQgates='all',
                                                gatesetname='clifford')
assert(len(layer) == (n_2 % 2) + n_2//2)
layer = rb.sample.circuit_layer_by_pairings(pspec_1, twoQprob=0.0, oneQgatenames=['Gxpi'], twoQgatenames='all', 
                                  gatesetname = 'target')
assert(layer[0].name == 'Gxpi')

# Tests for the sampling by sectors function
C01 = Label('Gcnot',(0,1))
C23 = Label('Gcnot',(2,3))
sectors = [[],[C01,C23]]
layer = rb.sample.circuit_layer_by_sectors(pspec_1, sectors, sectorsprob='uniform', twoQprob=1.0, 
                                           oneQgatenames='all', gatesetname='clifford')
assert(len(layer) == n_1 or len(layer) == n_1//2)
layer = rb.sample.circuit_layer_by_sectors(pspec_1, sectors, sectorsprob=[0.,1.], twoQprob=1.0, 
                                           oneQgatenames='all', gatesetname='clifford')
assert(len(layer) == n_1//2)
layer = rb.sample.circuit_layer_by_sectors(pspec_1, sectors, sectorsprob=[1.,0.], twoQprob=1.0, 
                                           oneQgatenames=['Gx',], gatesetname='clifford')
assert(len(layer) == n_1)
assert(layer[0].name == 'Gx')
layer = rb.sample.circuit_layer_by_sectors(pspec_1, sectors, sectorsprob=[0.25,0.75], twoQprob=0.5, 
                                           oneQgatenames='all', gatesetname='clifford')

# Tests for the sampling a layer of 1Q gates.
layer = rb.sample.circuit_layer_of_1Q_gates(pspec_1, oneQgatenames='all', pdist='uniform',
                                            gatesetname='clifford')
assert(len(layer) == n_1)
layer = rb.sample.circuit_layer_of_1Q_gates(pspec_1, oneQgatenames=['Gx','Gy'], pdist=[1.,0.],
                                            gatesetname='clifford')
assert(layer[0].name == 'Gx')
layer = rb.sample.circuit_layer_of_1Q_gates(pspec_1, oneQgatenames=['Gx'], pdist=[3.,],
                                            gatesetname='clifford')
assert(layer[0].name == 'Gx')
layer = rb.sample.circuit_layer_of_1Q_gates(pspec_1, oneQgatenames=['Gx'], pdist='uniform',
                                            gatesetname='clifford')


# In[ ]:


circuit = rb.sample.circuit(pspec_1, length=100, sampler='Qelimination')
assert(circuit.depth() == 100)
circuit = rb.sample.circuit(pspec_2, length=100, sampler='Qelimination', samplerargs=[0.1,], addlocal = True)
assert(circuit.depth() == 201)
assert(len(circuit.get_circuit_layer(0)) == n_2)
circuit = rb.sample.circuit(pspec_1, length=100, sampler='pairings')
circuit = rb.sample.circuit(pspec_1, length=10, sampler='pairings', samplerargs=[0.1,['Gx',]])

circuit = rb.sample.circuit(pspec_1, length=100, sampler='sectors', samplerargs=[sectors])
circuit = rb.sample.circuit(pspec_1, length=100, sampler='sectors', samplerargs=[sectors,[0.1,0.2],0.1], 
                            addlocal = True, lsargs=[['Gx',]])
assert(circuit.depth() == 201)
circuit = rb.sample.circuit(pspec_1, length=5, sampler='local')
assert(circuit.size() == n_1*5)
circuit = rb.sample.circuit(pspec_1, length=5, sampler='local',samplerargs=[['Gx']])
assert(circuit.line_items[0][0].name == 'Gx')


# In[ ]:


circuit = rb.sample.circuit(pspec_1, length=5, sampler='local')


# In[ ]:


print(circuit)


# In[ ]:


order = [1,0,3,2]


# In[ ]:


circuit.relabel_qubits(order)


# In[ ]:


print(circuit)


# In[ ]: