%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import datetime

from qutip import Qobj, identity, sigmax, sigmay, sigmaz, tensor
from qutip.qip import hadamard_transform
import qutip.logging as logging
logger = logging.get_logger()
#Set this to None or logging.WARN for 'quiet' execution
log_level = logging.INFO
#QuTiP control modules
import qutip.control.pulseoptim as cpo

example_name = 'Lindblad'

Sx = sigmax()
Sy = sigmay()
Sz = sigmaz()
Si = identity(2)

Sd = Qobj(np.array([[0, 1],
             [0, 0]]))
Sm = Qobj(np.array([[0, 0],
             [1, 0]]))
Sd_m = Qobj(np.array([[1, 0],
              [0, 0]]))
Sm_d = Qobj(np.array([[0, 0],
              [0, 1]]))

#Amplitude damping#
#Damping rate:
gamma = 0.1
L0_Ad = gamma*(2*tensor(Sm, Sd.trans()) - 
            (tensor(Sd_m, Si) + tensor(Si, Sd_m.trans())))
#sigma X control
LC_x = -1j*(tensor(Sx, Si) - tensor(Si, Sx))
#sigma Y control
LC_y = -1j*(tensor(Sy, Si) - tensor(Si, Sy.trans()))
#sigma Z control
LC_z = -1j*(tensor(Sz, Si) - tensor(Si, Sz))

#Drift
drift = L0_Ad
#Controls
ctrls = [LC_z, LC_x]
# Number of ctrls
n_ctrls = len(ctrls)

initial = identity(4)
#Target
#Hadamard gate
had_gate = hadamard_transform(1)
target_DP = tensor(had_gate, had_gate)

# Number of time slots
n_ts = 200
# Time allowed for the evolution
evo_time = 2

# Fidelity error target
fid_err_targ = 1e-10
# Maximum iterations for the optisation algorithm
max_iter = 200
# Maximum (elapsed) time allowed in seconds
max_wall_time = 30
# Minimum gradient (sum of gradients squared)
# as this tends to 0 -> local minima has been found
min_grad = 1e-20

# pulse type alternatives: RND|ZERO|LIN|SINE|SQUARE|SAW|TRIANGLE|
p_type = 'RND'

#Set to None to suppress output files
f_ext = "{}_n_ts{}_ptype{}.txt".format(example_name, n_ts, p_type)

# Note that this call will take the defaults
#    dyn_type='GEN_MAT'
# This means that matrices that describe the dynamics are assumed to be
# general, i.e. the propagator can be calculated using:
# expm(combined_dynamics*dt)
#    prop_type='FRECHET'
# and the propagators and their gradients will be calculated using the
# Frechet method, i.e. an exact gradent
#    fid_type='TRACEDIFF'
# and that the fidelity error, i.e. distance from the target, is give
# by the trace of the difference between the target and evolved operators 
result = cpo.optimize_pulse(drift, ctrls, initial, target_DP, n_ts, evo_time, 
                fid_err_targ=fid_err_targ, min_grad=min_grad, 
                max_iter=max_iter, max_wall_time=max_wall_time, 
                out_file_ext=f_ext, init_pulse_type=p_type, 
                log_level=log_level, gen_stats=True)

result.stats.report()
print("Final evolution\n{}\n".format(result.evo_full_final))
print("********* Summary *****************")
print("Final fidelity error {}".format(result.fid_err))
print("Final gradient normal {}".format(result.grad_norm_final))
print("Terminated due to {}".format(result.termination_reason))
print("Number of iterations {}".format(result.num_iter))
print("Completed in {} HH:MM:SS.US".format(datetime.timedelta(seconds=result.wall_time)))

t = result.time[:n_ts]

fig1 = plt.figure()
ax1 = fig1.add_subplot(2, 1, 1)
ax1.set_title("Initial Control amps")
ax1.set_xlabel("Time")
ax1.set_ylabel("Control amplitude")
for j in range(n_ctrls):
    amps = result.initial_amps[:, j]
    ax1.plot(t, amps)

ax2 = fig1.add_subplot(2, 1, 2)
ax2.set_title("Optimised Control Amplitudes")
ax2.set_xlabel("Time")
ax2.set_ylabel("Control amplitude")
for j in range(n_ctrls):
    amps = result.final_amps[:, j]
    ax2.plot(t, amps)

plt.show()

from qutip.ipynbtools import version_table

version_table()