import unittest
import asyncio
import sys
sys.path.append('.')
sys.path.append('../')
import growingnn as gnn
import numpy as np
import time

# Import your script here, e.g., 
# from myscript import get_reshsper, Reshape, Reshape_forward_prop, Reshape_back_prop, eye_stretch

MAX_ENTRIES = 100  # Match with the script configuration


class TestReshapeFunctions(unittest.TestCase):

    def test_eye_stretch(self):
        a, b = 5, 10
        result = gnn.quaziIdentity.eye_stretch(a, b)
        self.assertEqual(result.shape, (a, b))
        
    def test_eye_stretch_different_sizes(self):
        """Test eye_stretch with various input and output sizes"""
        test_cases = [
            (1, 1),    # Minimal case
            (1, 10),   # Stretch from 1 to many
            (10, 1),   # Compress from many to 1
            (5, 5),    # Same size (identity)
            (10, 20),  # Moderate stretch
            (20, 10),  # Moderate compression
            (50, 100), # Large stretch
            (100, 50)  # Large compression
        ]
        
        for a, b in test_cases:
            result = gnn.quaziIdentity.eye_stretch(a, b)
            self.assertEqual(result.shape, (a, b))
            
            # Check that the matrix has the expected properties
            if a <= b:
                # For stretching, each row should have at least one non-zero element
                for i in range(a):
                    self.assertTrue(np.any(result[i, :] != 0))
            else:
                # For compression, each column should have at least one non-zero element
                for j in range(b):
                    self.assertTrue(np.any(result[:, j] != 0))

    def test_get_reshsper(self):
        size_from, size_to = 10, 20

        # First time, it should create a new resherper
        resherper = gnn.quaziIdentity.get_reshsper(size_from, size_to)
        self.assertIsNotNone(resherper)
        self.assertEqual(resherper.shape, (size_from, size_to))

        # Second time, it should fetch the same one
        resherper_again = gnn.quaziIdentity.get_reshsper(size_from, size_to)
        self.assertTrue(np.array_equal(resherper, resherper_again))
        
    def test_get_reshsper_cache(self):
        """Test that the cache works correctly for different sizes"""
        # Test with different sizes
        sizes = [(5, 10), (10, 5), (20, 40), (40, 20)]
        
        for size_from, size_to in sizes:
            # First call should create a new resherper
            resherper1 = gnn.quaziIdentity.get_reshsper(size_from, size_to)
            self.assertIsNotNone(resherper1)
            
            # Second call should return the same resherper
            resherper2 = gnn.quaziIdentity.get_reshsper(size_from, size_to)
            self.assertTrue(np.array_equal(resherper1, resherper2))
            
            # Different sizes should return different resherpers
            if size_from != size_to:
                resherper3 = gnn.quaziIdentity.get_reshsper(size_to, size_from)
                self.assertFalse(np.array_equal(resherper1, resherper3))

    def test_reshape(self):
        input_data = np.random.rand(10, 5)  # 10 features, 5 samples
        QIdentity = gnn.quaziIdentity.eye_stretch(10, 20)
        reshaped_data = gnn.quaziIdentity.Reshape(input_data, 20, QIdentity)

        self.assertEqual(reshaped_data.shape, (20, 5))
        self.assertTrue(np.allclose(np.dot(input_data[:, 0], QIdentity), reshaped_data[:, 0]))
        
    def test_reshape_different_sizes(self):
        """Test Reshape with different input and output sizes"""
        test_cases = [
            (5, 10, 3),    # Small stretch
            (10, 5, 3),    # Small compression
            (20, 40, 5),   # Moderate stretch
            (40, 20, 5),   # Moderate compression
            (50, 100, 2),  # Large stretch
            (100, 50, 2)   # Large compression
        ]
        
        for size_from, size_to, samples in test_cases:
            input_data = np.random.rand(size_from, samples)
            QIdentity = gnn.quaziIdentity.eye_stretch(size_from, size_to)
            reshaped_data = gnn.quaziIdentity.Reshape(input_data, size_to, QIdentity)
            
            self.assertEqual(reshaped_data.shape, (size_to, samples))
            
            # Check that the transformation is correct for the first sample
            self.assertTrue(np.allclose(np.dot(input_data[:, 0], QIdentity), reshaped_data[:, 0]))
            
    def test_reshape_with_zeros(self):
        """Test Reshape with input containing zeros"""
        size_from, size_to, samples = 10, 20, 5
        
        # Create input with some zeros
        input_data = np.zeros((size_from, samples))
        input_data[0, 0] = 1.0  # Set one element to 1
        
        QIdentity = gnn.quaziIdentity.eye_stretch(size_from, size_to)
        reshaped_data = gnn.quaziIdentity.Reshape(input_data, size_to, QIdentity)
        
        self.assertEqual(reshaped_data.shape, (size_to, samples))
        
        # Check that the transformation is correct
        self.assertTrue(np.allclose(np.dot(input_data[:, 0], QIdentity), reshaped_data[:, 0]))

    def test_reshape_forward_prop(self):
        input_data = np.random.rand(5, 3, 3, 2)  # Batch of 5, 3x3x2 feature maps
        flatten_size = 3 * 3 * 2
        QIdentity = gnn.quaziIdentity.eye_stretch(flatten_size, 10)
        reshaped_data = gnn.quaziIdentity.Reshape_forward_prop(input_data, 10, QIdentity)

        self.assertEqual(reshaped_data.shape, (10, 5))
        
    def test_reshape_forward_prop_different_shapes(self):
        """Test Reshape_forward_prop with different input shapes"""
        test_cases = [
            (2, 2, 2, 1),    # Small 3D
            (3, 3, 3, 2),    # Medium 3D
            (5, 5, 5, 3),    # Larger 3D
            (2, 2, 1, 1),    # 2D-like
            (1, 1, 1, 1)     # Minimal
        ]
        
        for h, w, c, batch in test_cases:
            input_data = np.random.rand(batch, h, w, c)
            flatten_size = h * w * c
            output_size = max(flatten_size // 2, 1)  # Ensure output size is at least 1
            
            QIdentity = gnn.quaziIdentity.eye_stretch(flatten_size, output_size)
            reshaped_data = gnn.quaziIdentity.Reshape_forward_prop(input_data, output_size, QIdentity)
            
            self.assertEqual(reshaped_data.shape, (output_size, batch))
            
    def test_reshape_forward_prop_with_zeros(self):
        """Test Reshape_forward_prop with input containing zeros"""
        batch, h, w, c = 3, 3, 3, 2
        input_data = np.zeros((batch, h, w, c))
        input_data[0, 0, 0, 0] = 1.0  # Set one element to 1
        
        flatten_size = h * w * c
        output_size = max(flatten_size // 2, 1)
        
        QIdentity = gnn.quaziIdentity.eye_stretch(flatten_size, output_size)
        reshaped_data = gnn.quaziIdentity.Reshape_forward_prop(input_data, output_size, QIdentity)
        
        self.assertEqual(reshaped_data.shape, (output_size, batch))

    def test_reshape_back_prop(self):
        output_size, input_shape = 10, (3, 3, 2)
        E = np.random.rand(10, 5)
        QIdentity = gnn.quaziIdentity.eye_stretch(output_size, np.prod(input_shape))
        back_prop_data = gnn.quaziIdentity.Reshape_back_prop(E, input_shape, QIdentity)

        self.assertEqual(back_prop_data.shape, (5, 3, 3, 2))
        
    def test_reshape_back_prop_different_shapes(self):
        """Test Reshape_back_prop with different input shapes"""
        test_cases = [
            (5, (2, 2, 2)),    # Small 3D
            (10, (3, 3, 3)),   # Medium 3D
            (20, (5, 5, 5)),   # Larger 3D
            (5, (2, 2, 1)),    # 2D-like
            (1, (1, 1, 1))     # Minimal
        ]
        
        for output_size, input_shape in test_cases:
            batch_size = 3
            E = np.random.rand(output_size, batch_size)
            QIdentity = gnn.quaziIdentity.eye_stretch(output_size, np.prod(input_shape))
            back_prop_data = gnn.quaziIdentity.Reshape_back_prop(E, input_shape, QIdentity)
            
            expected_shape = (batch_size,) + input_shape
            self.assertEqual(back_prop_data.shape, expected_shape)
            
    def test_reshape_back_prop_with_zeros(self):
        """Test Reshape_back_prop with error containing zeros"""
        output_size, input_shape = 10, (3, 3, 2)
        batch_size = 3
        
        E = np.zeros((output_size, batch_size))
        E[0, 0] = 1.0  # Set one element to 1
        
        QIdentity = gnn.quaziIdentity.eye_stretch(output_size, np.prod(input_shape))
        back_prop_data = gnn.quaziIdentity.Reshape_back_prop(E, input_shape, QIdentity)
        
        expected_shape = (batch_size,) + input_shape
        self.assertEqual(back_prop_data.shape, expected_shape)
        
    def test_reshape_back_prop_edge_cases(self):
        """Test Reshape_back_prop with edge cases"""
        # Test with minimal dimensions
        output_size, input_shape = 1, (1, 1, 1)
        batch_size = 1
        E = np.random.rand(output_size, batch_size)
        QIdentity = gnn.quaziIdentity.eye_stretch(output_size, np.prod(input_shape))
        back_prop_data = gnn.quaziIdentity.Reshape_back_prop(E, input_shape, QIdentity)
        self.assertEqual(back_prop_data.shape, (batch_size,) + input_shape)
        
        # Test with very large dimensions
        output_size, input_shape = 100, (10, 10, 10)
        batch_size = 2
        E = np.random.rand(output_size, batch_size)
        QIdentity = gnn.quaziIdentity.eye_stretch(output_size, np.prod(input_shape))
        back_prop_data = gnn.quaziIdentity.Reshape_back_prop(E, input_shape, QIdentity)
        self.assertEqual(back_prop_data.shape, (batch_size,) + input_shape)
        
    def test_reshape_roundtrip(self):
        """Test that reshaping forward and then back produces the expected result"""
        # Create input data
        batch, h, w, c = 3, 3, 3, 2
        input_data = np.random.rand(batch, h, w, c)
        
        # Flatten the input
        flatten_size = h * w * c
        output_size = max(flatten_size // 2, 1)
        
        # Forward reshape
        QIdentity = gnn.quaziIdentity.eye_stretch(flatten_size, output_size)
        reshaped_data = gnn.quaziIdentity.Reshape_forward_prop(input_data, output_size, QIdentity)
        
        # Backward reshape - we need to transpose QIdentity for the backward pass
        # to match the expected dimensions
        QIdentity_back = gnn.quaziIdentity.eye_stretch(output_size, flatten_size)
        back_prop_data = gnn.quaziIdentity.Reshape_back_prop(reshaped_data, (h, w, c), QIdentity_back)
        
        # Check shapes
        self.assertEqual(back_prop_data.shape, input_data.shape)
        
        # Note: We don't check for exact equality because the quasi-identity transformation
        # is not a perfect roundtrip (information is lost in the compression)
        
    def test_performance(self):
        """Test the performance of the reshape functions"""
        # Test with a moderately large input
        batch, h, w, c = 10, 32, 32, 3
        input_data = np.random.rand(batch, h, w, c)
        
        flatten_size = h * w * c
        output_size = flatten_size // 4  # Compress to 1/4 size
        
        # Measure time for eye_stretch
        start_time = time.time()
        QIdentity = gnn.quaziIdentity.eye_stretch(flatten_size, output_size)
        eye_stretch_time = time.time() - start_time
        
        # Measure time for forward reshape
        start_time = time.time()
        reshaped_data = gnn.quaziIdentity.Reshape_forward_prop(input_data, output_size, QIdentity)
        forward_time = time.time() - start_time
        
        # Measure time for backward reshape
        # We need to create a new QIdentity for the backward pass with the correct dimensions
        QIdentity_back = gnn.quaziIdentity.eye_stretch(output_size, flatten_size)
        start_time = time.time()
        back_prop_data = gnn.quaziIdentity.Reshape_back_prop(reshaped_data, (h, w, c), QIdentity_back)
        backward_time = time.time() - start_time
        
        # Print performance metrics
        print(f"Performance test results:")
        print(f"  eye_stretch: {eye_stretch_time:.6f} seconds")
        print(f"  forward reshape: {forward_time:.6f} seconds")
        print(f"  backward reshape: {backward_time:.6f} seconds")
        
        # We don't assert specific time limits as they depend on the hardware
        # But we can check that the operations complete successfully
        self.assertIsNotNone(QIdentity)
        self.assertEqual(reshaped_data.shape, (output_size, batch))
        self.assertEqual(back_prop_data.shape, (batch, h, w, c))


if __name__ == '__main__':
    unittest.main()
