author: Eduardo Villasenor
created-on: 20/11/17
"""

import qutip as qt
import numpy as np
import stabilizer

# Determine parameters
ps = 0.006
pm = 0.006
pg = 0.006

# Initialize  objects
stab = stabilizer.Stabilizer(ps, pm, pg)
stab_ideal = stabilizer.Stabilizer(0, 0, 0)

print("------------------TEST SWAP PAIR--------------------------")
pair = qt.bell_state('00') * qt.bell_state('00').dag()
pair = qt.tensor(pair, qt.basis(2, 0) * qt.basis(2, 0).dag())

pair_swapped = stab._swap_pair(pair, [0, 1])
print(pair)
print(pair_swapped)
print(qt.fidelity(pair, pair_swapped)**2)

print("------------------TEST SWAP GHZ--------------------------")
ghz = qt.ghz_state(4) * qt.ghz_state(4).dag()

ghz_swapped = stab._swap_ghz(ghz)
        ghz = np.sum(ghzs)

    ghz = ghz / (N - ignore_number)
    TIME += [t_max]
    # print("F: ", qt.fidelity(rho, rho_ref))
    print("_____________________________________________________")
    print("N: ", N)
    print("T_avg: ", tavg, tstd)
    print("TIME_MAX:", t_max)
    ################################################
    # CHOI decomposition
    for parity in ["X", "Z"]:
        # Initialize objects
        model = noise_modeling.NoiseModel(stab_size, parity)
        model.separate_basis_parity()
        stab = stabilizer.Stabilizer(ps=ps, pm=pm, pg=pg)

        ghz = stab.twirl_ghz(ghz)

        # Choi state for noise noise modeling
        choi = model._choi_state_ket(stab_size)
        choi = choi * choi.dag()
        targets = list(range(stab_size))

        extra_name = ""
        if len(ghz.dims[0]) == 2:
            extra_name = "_extra"
        protocol_chi2 = protocol_chi + extra_name

        p_res, rhos = stab.measure_ghz_stabilizer(choi, ghz, targets, parity)
cb = circb.Blocks(ps, pm, pg, eta, a0, a1, theta)

# First assemeble the small single selection circuit
single_sel = circ.Circuit(a0=a0, a1=a1, circuit_block=cb.start_epl)
single_sel.add_circuit(circuit_block=cb.swap_pair, pair=[0, 1])
single_sel.add_circuit(circuit_block=cb.single_selection,
                       operation_qubits=[0, 1],
                       sigma="Z")

# Initialize objects and define references
rho_ref = qt.bell_state('00') * qt.bell_state('00').dag()
rho_ref2 = qt.bell_state('01') * qt.bell_state('01').dag()
ghz_ref = qt.ghz_state(4) * qt.ghz_state(4).dag()
ghz3_ref = qt.ghz_state(3) * qt.ghz_state(3).dag()
stab = stabilizer.Stabilizer(0, 0, 0)

# Lists to save results
FIDELITY = []
TIMES = []

for s in [0]:
    print("------> Var=", a0)

    circuit = protocols.pair_single_sel(ps, pm, pg, eta, a0, a1, theta)

    # Get average number of steps
    fidelity = []
    times = []
    fidelity2 = []
    times2 = []