%matplotlib inline
import matplotlib.pyplot as plt
import time

from qutip import *
qutip.settings.auto_tidyup = False

T = 1
times = np.linspace(0, T, 500)

U = toffoli()

R = 5000
H_ops = [# qubit 1: single-qubit control
         tensor(sigmax(), identity(2), identity(2)),
         tensor(sigmay(), identity(2), identity(2)),
         tensor(sigmaz(), identity(2), identity(2)),
         # qubit 1: single-qubit control
         tensor(identity(2), sigmax(), identity(2)),
         tensor(identity(2), sigmay(), identity(2)),
         tensor(identity(2), sigmaz(), identity(2)),
         # qubit 3: single-qubit control
         tensor(identity(2), identity(2), sigmax()),
         tensor(identity(2), identity(2), sigmay()),
         tensor(identity(2), identity(2), sigmaz()),
         # pairwise X-X interactions
         tensor(sigmax(), sigmax(), identity(2)),
         tensor(identity(2), sigmax(), sigmax()),
         tensor(sigmax(), identity(2), sigmax()),
         # pairwise Y-Y interactions
         tensor(sigmay(), sigmay(), identity(2)),
         tensor(identity(2), sigmay(), sigmay()),
         tensor(sigmay(), identity(2), sigmay()),
         # pairwise Z-Z interactions
         tensor(sigmaz(), sigmaz(), identity(2)),
         tensor(identity(2), sigmaz(), sigmaz()),
         tensor(sigmaz(), identity(2), sigmaz()),
         ]

H_labels = [r'$u_{1x}$',
            r'$u_{1y}$',
            r'$u_{1z}$',
            
            r'$u_{2x}$',
            r'$u_{2y}$',
            r'$u_{2z}$',
            
            r'$u_{3x}$',
            r'$u_{3y}$',
            r'$u_{3z}$',
            
            r'$u_{xxi}$',
            r'$u_{ixx}$',
            r'$u_{xix}$',
            
            r'$u_{yyi}$',
            r'$u_{iyy}$',
            r'$u_{yiy}$',
            
            r'$u_{zzi}$',
            r'$u_{izz}$',
            r'$u_{ziz}$',
        ]

H0 = 2 * pi * (tensor(sigmaz(), identity(2), identity(2)) + 
               tensor(identity(2), sigmaz(), identity(2)) + 
               tensor(identity(2), identity(2), sigmaz()))

c_ops = []

from qutip.control.grape import cy_grape_unitary, grape_unitary_adaptive, plot_grape_control_fields, _overlap

from scipy.interpolate import interp1d
from qutip.ui.progressbar import TextProgressBar

u0 = np.array([np.random.rand(len(times)) * 2 * pi * 0.01 for _ in range(len(H_ops))])

u0 = [np.convolve(np.ones(10)/10, u0[idx,:], mode='same') for idx in range(len(H_ops))]

result = cy_grape_unitary(U, H0, H_ops, R, times, phase_sensitive=False,
                          u_start=u0, progress_bar=TextProgressBar(),
                          eps=2*pi*5)

plot_grape_control_fields(times, result.u / (2 * pi), H_labels, uniform_axes=True);

U

result.U_f.tidyup(1e-1)

result.U_f / result.U_f[0,0] #.tidyup(1e-1)

abs(_overlap(U, result.U_f))**2

op_basis = [[qeye(2), sigmax(), sigmay(), sigmaz()]] * 3
op_label = [["i", "x", "y", "z"]] * 3

fig = plt.figure(figsize=(16,12))

SU = spre(U) * spost(U.dag())

chi = qpt(SU, op_basis)

fig = qpt_plot_combined(chi, op_label, fig=fig, threshold=0.001)

fig = plt.figure(figsize=(16,12))

SU = spre(result.U_f) * spost(result.U_f.dag())

chi = qpt(SU, op_basis)

fig = qpt_plot_combined(chi, op_label, fig=fig, threshold=0.001)

from qutip.ipynbtools import version_table

version_table()