# AUTOGENERATED! DO NOT EDIT! File to edit: 04_tests.ipynb (unless otherwise specified).

__all__ = ['zero_ham', 'u0', 'evolution_time', 'no_time_step_for_eval', 'zero_ham', 'u0', 'evolution_time',
           'no_time_step_for_eval', 'identity_ham', 'u0', 'evolution_time', 'no_time_step_for_eval', 'H', 'psi0',
           'times', 'mere_result', 'q_ham', 'u0', 'evolution_time', 'no_time_step_for_eval', 'psi0_ndarray', 'upsi']

# Cell
from numpy.random import random_sample
from numpy import array, identity, zeros, linspace, pi, empty
from numpy.testing import assert_allclose
from qutip import mesolve
from qutip import basis
from qutip import sigmax, sigmay, sigmaz, qeye
from qutip.ui.progressbar import BaseProgressBar, TextProgressBar
from math import sin, cos
from .core import ihu, y_der, check_for_stupidity
from .herm import rand_herm_ndarray, test_herm_ndarray
from .unitary import rand_unitary_ndarray, test_unitary_ndarray

# Cell
def zero_ham(t):
    zerodha = zeros((2,2), dtype=complex)
    return zerodha

# Cell
# u0 = rand_unitary_ndarray((2,2))
u0 = array([[1.+0.j, 0.+0.j],
           [0.+0.j, 1.+0.j]])
evolution_time = 5
no_time_step_for_eval = 100

# Cell
u_series, result = ihu(zero_ham, u0, evolution_time,
                       no_time_step_for_eval,)

# Cell
assert_allclose(u_series[:, :, -1], u0)

# Cell
for i in range(u_series.shape[-1]):
    assert_allclose(u_series[:, :, i], u0)

# Cell
def zero_ham(t):
    zerodha = zeros((2,2), dtype=complex)
    return zerodha

# Cell
u0 = rand_unitary_ndarray((2,2))
#u0 = array([[1.+0.j, 0.+0.j],
#           [0.+0.j, 1.+0.j]])
evolution_time = 5
no_time_step_for_eval = 100

# Cell
u_series, result = ihu(zero_ham, u0, evolution_time,
                       no_time_step_for_eval,)

# Cell
assert_allclose(u_series[:, :, -1], u0)

# Cell
for i in range(u_series.shape[-1]):
    assert_allclose(u_series[:, :, i], u0)

# Cell
def identity_ham(t):
    iden = identity(2, dtype=complex)
    return iden

# Cell
u0 = rand_unitary_ndarray((2,2))
evolution_time = 5
no_time_step_for_eval = 100

# Cell
u_series, result = ihu(identity_ham, u0, evolution_time,
                       no_time_step_for_eval,)

# Cell
H = 2*pi*0.1*sigmax()

# Cell
psi0 = basis(2, 0)

# Cell
times = linspace(0.0, 10, 20)

# Cell
mere_result = mesolve(H, psi0, times)

# Cell
def q_ham(t):
    qutham = H.full()
    return qutham

# Cell
u0 = identity(2, dtype=complex)
evolution_time = 10
no_time_step_for_eval = 20

# Cell
u_series, result = ihu(q_ham, u0, evolution_time,
                       no_time_step_for_eval,)

# Cell
psi0_ndarray = psi0.full().ravel()

# Cell
upsi = empty((2, u_series.shape[-1]), dtype=complex)
for i in range(u_series.shape[-1]):
    upsi[:, i] = u_series[:, :, i]@psi0_ndarray


# Cell
assert_allclose(upsi[:, 0], mere_result.states[0].full().ravel())

# Cell
for i in range(u_series.shape[-1]):
    assert_allclose(upsi[:, i], mere_result.states[i].full().ravel(),
                   rtol=1e-01, atol=1e-02)

# Cell
for i in range(u_series.shape[-1]):
    assert_allclose(upsi[:, i], mere_result.states[i].full().ravel(),
                   rtol=1e-07, atol=1e-02)

# Cell
for i in range(u_series.shape[-1]):
    assert_allclose(upsi[:, i], mere_result.states[i].full().ravel(),
                   rtol=1e-08, atol=1e-02)