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

# # GRAPE calculation of control fields for cnot implementation

# Robert Johansson (robert@riken.jp)

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib.pyplot as plt
import time
import numpy as np


# In[2]:


from qutip import *
from qutip.control import *


# In[3]:


T = 2 * np.pi 
times = np.linspace(0, T, 500)


# In[4]:


U = cnot()
R = 500
H_ops = [tensor(sigmax(), identity(2)),
         tensor(sigmay(), identity(2)),
         tensor(sigmaz(), identity(2)),
         tensor(identity(2), sigmax()),
         tensor(identity(2), sigmay()),
         tensor(identity(2), sigmaz()),
         tensor(sigmax(), sigmax()) +
         tensor(sigmay(), sigmay()) +
         tensor(sigmaz(), sigmaz())]

H_labels = [r'$u_{1x}$', r'$u_{1y}$', r'$u_{1z}$',
            r'$u_{2x}$', r'$u_{1y}$', r'$u_{2z}$',
            r'$u_{xx}$',
            r'$u_{yy}$',
            r'$u_{zz}$',
        ]


# In[5]:


H0 = 0 * np.pi * (tensor(sigmax(), identity(2)) + tensor(identity(2), sigmax()))

c_ops = []

# This is the analytical result in the absense of single-qubit tunnelling
#g = pi/(4 * T)
#H = g * (tensor(sigmax(), sigmax()) + tensor(sigmay(), sigmay()))


# # GRAPE

# In[6]:


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


# In[7]:


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


# In[8]:


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

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

u_limits = None #[0, 1 * 2 * pi]
alpha = None


# In[9]:


result = cy_grape_unitary(U, H0, H_ops, R, times, u_start=u0, u_limits=u_limits,
                          eps=2*np.pi*1, alpha=alpha, phase_sensitive=False,
                          progress_bar=TextProgressBar())

U_f, H_list_func, u = grape_unitary_adaptive(U, H0, H_ops, R, times, u_start=u0, u_limits=u_limits,
                                    eps=2*pi*1, alpha=alpha, phase_sensitive=False,
                                    overlap_terminate=0.9999,
                                    progress_bar=TextProgressBar())
# ## Plot control fields for cnot gate in the presense of single-qubit tunnelling

# In[10]:


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


# ## Fidelity/overlap

# In[11]:


U


# In[12]:


result.U_f


# In[13]:


result.U_f/result.U_f[0,0]


# In[14]:


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


# ## Test numerical integration of GRAPE pulse

# In[15]:


U_f_numerical = propagator(result.H_t, times[-1], [], options=Odeoptions(nsteps=5000), args={})
U_f_numerical


# In[16]:


U_f_numerical / U_f_numerical[0,0]


# In[17]:


_overlap(result.U_f, U_f_numerical).real, abs(_overlap(result.U_f, U_f_numerical))**2


# # Process tomography

# ## Ideal cnot gate

# In[18]:


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


# In[19]:


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

U_i_s = to_super(U)

chi = qpt(U_i_s, op_basis)

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


# ## cnot gate calculated using GRAPE

# In[20]:


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

U_f_s = to_super(result.U_f)

chi = qpt(U_f_s, op_basis)

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


# ## Versions

# In[21]:


from qutip.ipynbtools import version_table

version_table()