import unittest

class TestClass(unittest.Testcase):

    def setUp(self):
        pass

    def test_nonzero_infidelity_gradient_with_degenerate_hamiltonian():
        """
        Tests that the infidelity gradient is non-zero for a degenerate non-trivial noisy system.
        """
        with tf.Graph().as_default() as g:
            pulse = tf.compat.v1.placeholder(shape=(2,), dtype=tf.float64)
            hamiltonian_pwc = TensorPwc([1, 1],
                                        (tf.cast(pulse[:, None, None], dtype=tf.complex128)
                                         * tf.constant(np.array([np.diag([1, 1, 2, 2]),
                                                                 np.diag([1, 1, 2, 2])]),
                                                       dtype=tf.complex128)))
            target = Target(np.eye(4))

            infidelity = create_infidelity(hamiltonian_pwc=hamiltonian_pwc,
                                           noise_operator_pwcs=[],
                                           target=target)

            noise_operator_pwc = TensorPwc([2], tf.constant(np.diag([0, 0, 1, 1])[None],
                                                            dtype=tf.complex128))
            infidelity_noise = create_infidelity(hamiltonian_pwc=hamiltonian_pwc,
                                                 noise_operator_pwcs=[noise_operator_pwc],
                                                 target=target)

            grad = tf.gradients(infidelity, pulse)[0]
            grad_noise = tf.gradients(infidelity_noise, pulse)[0]

            with tf.compat.v1.Session() as session:
                grad_value = session.run(grad, feed_dict={pulse: np.array([2, 1])})
                grad_noise_value = session.run(grad_noise, feed_dict={pulse: np.array([2, 1])})

        # Changing the pulse should affect the infidelities, so the gradients should be non-zero.
        self.assertNotIn(np.Nan, grad_value)
        for val in grad_value:
            self.assertAlmostEqual(val, 0., delta=0.001)
        self.assertNotIn(np.NaN, grad_noise_value)