# unit test for the functions in quantum_states.py

import unittest
import numpy as np
from quantum_states import random_pure_state, von_neumann_entropy_bipartite_pure_state

class LearningCase(unittest.TestCase):
    def test_random_pure_state_shape_and_norm(self):
        dim_a = 2
        dim_b = 2
        state = random_pure_state(dim_a, dim_b)
        self.assertEqual(state.shape, (dim_a * dim_b,))
        self.assertAlmostEqual(np.linalg.norm(state), 1)

    def test_partial_trace_entropy(self):
        dim_a = 2
        dim_b = 2
        state = random_pure_state(dim_a, dim_b)
        self.assertAlmostEqual(von_neumann_entropy_bipartite_pure_state(state, dim_a, dim_b), von_neumann_entropy_bipartite_pure_state(state, dim_b, dim_a))

    def test_sample_uniformly(self):
        # calculate the distribution of the random pure state
        dim_a = 2
        dim_b = 2
        state = random_pure_state(dim_a, dim_b)
        

def main():
    unittest.main()

if __name__ == "__main__":
    main()