import unittest
import numpy as np

from gradescope_utils.autograder_utils.decorators import weight, number

try:
    from programming.vlm import (
        patch_count,
        extract_patches,
        add_cls_token,
        add_positional_embedding,
        linear_projection,
        vit_preprocess,
    )
except ImportError:
    from vlm import (
        patch_count,
        extract_patches,
        add_cls_token,
        add_positional_embedding,
        linear_projection,
        vit_preprocess,
    )


class TestVLM(unittest.TestCase):

    # patch_count tests

    @weight(5)
    @number("2.1")
    def test_patch_count(self):
        """Test patch_count returns the correct number of patches."""
        self.assertEqual(patch_count(224, 224, 16), 196)
        self.assertEqual(patch_count(32, 32, 8), 16)
        self.assertEqual(patch_count(16, 32, 8), 8)

    # extract_patches tests

    @weight(5)
    @number("2.2")
    def test_extract_patches_shape(self):
        """Test extract_patches produces the correct output shape."""
        image = np.zeros((3, 32, 32))
        patches = extract_patches(image, patch_size=8)
        # N = (32//8)*(32//8) = 16 patches, each of dimension 3*8*8 = 192
        self.assertEqual(patches.shape, (16, 192))

    @weight(5)
    @number("2.3")
    def test_extract_patches_values(self):
        """Test extract_patches preserves pixel values."""
        image = np.ones((1, 4, 4))
        patches = extract_patches(image, patch_size=2)
        # 4 patches, each of dimension 1*2*2 = 4, all ones
        self.assertEqual(patches.shape, (4, 4))
        np.testing.assert_array_equal(patches, np.ones((4, 4)))

    @weight(5)
    @number("2.4")
    def test_extract_patches_order(self):
        """Test extract_patches returns patches in raster-scan order."""
        image = np.zeros((1, 4, 4))
        image[0, 0:2, 0:2] = 1.0  # top-left patch
        image[0, 0:2, 2:4] = 2.0  # top-right patch
        image[0, 2:4, 0:2] = 3.0  # bottom-left patch
        image[0, 2:4, 2:4] = 4.0  # bottom-right patch

        patches = extract_patches(image, patch_size=2)
        self.assertEqual(patches.shape, (4, 4))
        np.testing.assert_array_equal(patches[0], np.ones(4) * 1.0)
        np.testing.assert_array_equal(patches[1], np.ones(4) * 2.0)
        np.testing.assert_array_equal(patches[2], np.ones(4) * 3.0)
        np.testing.assert_array_equal(patches[3], np.ones(4) * 4.0)

    # add_cls_token tests

    @weight(5)
    @number("2.5")
    def test_add_cls_token_shape(self):
        """Test add_cls_token produces the correct output shape."""
        patches = np.zeros((16, 64))
        cls_token = np.ones((1, 64))
        result = add_cls_token(patches, cls_token)
        self.assertEqual(result.shape, (17, 64))

    @weight(5)
    @number("2.6")
    def test_add_cls_token_prepended(self):
        """Test CLS token is at position 0 with correct values."""
        patches = np.zeros((16, 64))
        cls_token = np.ones((1, 64)) * 9.0
        result = add_cls_token(patches, cls_token)
        np.testing.assert_array_equal(result[0], np.ones(64) * 9.0)
        np.testing.assert_array_equal(result[1:], np.zeros((16, 64)))

    # add_positional_embedding tests

    @weight(5)
    @number("2.7")
    def test_add_positional_embedding(self):
        """Test positional embeddings are added element-wise."""
        tokens = np.ones((17, 64)) * 2.0
        pos_embedding = np.ones((17, 64)) * 3.0
        result = add_positional_embedding(tokens, pos_embedding)
        self.assertEqual(result.shape, (17, 64))
        np.testing.assert_array_equal(result, np.ones((17, 64)) * 5.0)

    # linear_projection tests

    @weight(5)
    @number("2.8")
    def test_linear_projection_shape(self):
        """Test linear_projection produces the correct output shape."""
        patches = np.ones((16, 192), dtype=np.float32) * 2.0
        weight = np.ones((192, 64), dtype=np.float32) * 0.1
        result = linear_projection(patches, weight)
        self.assertEqual(result.shape, (16, 64))

    @weight(5)
    @number("2.9")
    def test_linear_projection_values(self):
        """Test linear_projection computes the correct matrix product."""
        patches = np.ones((4, 12))
        weight = np.ones((12, 8)) * 0.5
        result = linear_projection(patches, weight)
        self.assertEqual(result.shape, (4, 8))
        # Each output element = sum of 12 ones × 0.5 = 6.0
        np.testing.assert_allclose(result, np.ones((4, 8)) * 6.0)


    # vit_preprocess tests

    @weight(5)
    @number("2.10")
    def test_vit_preprocess_shape(self):
        """Test vit_preprocess returns the correct output shape."""
        C, H, W, P, D = 3, 8, 8, 4, 16
        N = (H // P) * (W // P)  # 4
        image = np.ones((C, H, W))
        projection_weight = np.ones((C * P * P, D))
        cls_token = np.zeros((1, D))
        pos_embedding = np.zeros((N + 1, D))
        result = vit_preprocess(image, P, projection_weight, cls_token, pos_embedding)
        self.assertEqual(result.shape, (N + 1, D))  # (5, 16)

    @weight(5)
    @number("2.11")
    def test_vit_preprocess_cls_token(self):
        """Test CLS token value is preserved at position 0 through the pipeline."""
        C, H, W, P, D = 1, 4, 4, 2, 8
        N = (H // P) * (W // P)  # 4
        image = np.ones((C, H, W))
        projection_weight = np.zeros((C * P * P, D))  # projects all patches to zero
        cls_token = np.ones((1, D)) * 7.0
        pos_embedding = np.zeros((N + 1, D))           # no positional shift
        result = vit_preprocess(image, P, projection_weight, cls_token, pos_embedding)
        np.testing.assert_array_equal(result[0], np.ones(D) * 7.0)

    @weight(5)
    @number("2.12")
    def test_vit_preprocess_values(self):
        """Test vit_preprocess computes correct end-to-end values."""
        C, H, W, P, D = 1, 4, 4, 2, 4
        N = (H // P) * (W // P)  # 4
        image = np.ones((C, H, W))
        # Each flattened patch is [1,1,1,1]; weight 0.5 → projection = [2,2,2,2]
        projection_weight = np.ones((C * P * P, D)) * 0.5
        cls_token = np.ones((1, D))                    # [1,1,1,1]
        pos_embedding = np.ones((N + 1, D))            # adds 1 everywhere
        result = vit_preprocess(image, P, projection_weight, cls_token, pos_embedding)
        # CLS (pos 0): [1,1,1,1] + pos[0,1,1,1,1] = [2,2,2,2]
        # Patches (pos 1-4): [2,2,2,2] + pos[1,1,1,1] = [3,3,3,3]
        expected = np.ones((N + 1, D))
        expected[0] = 2.0
        expected[1:] = 3.0
        np.testing.assert_allclose(result, expected)


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