import unittest
import numpy as np
import sys
import os

sys.path.append(os.getcwd())
import maxwell_core

class TestMaxwellSOTA(unittest.TestCase):
    def setUp(self):
        self.M = 128
        self.K = 256
        self.N = 256
        self.A = np.random.randn(self.M, self.K).astype(np.float32)
        # Random Ternary Weights
        self.W_dense = np.random.randint(-1, 2, size=(self.K, self.N)).astype(np.int8)
        self.W_packed = maxwell_core.create_packed_weights(self.W_dense)
        
    def test_forward_cuda(self):
        print("\nTesting Forward CUDA...")
        C_cuda = np.zeros((self.M, self.N), dtype=np.float32)
        
        # Run CUDA
        maxwell_core.run_maxwell_cuda(self.A, self.W_packed, C_cuda)
        
        # Run Reference (NumPy)
        W_float = self.W_dense.astype(np.float32)
        C_ref = self.A @ W_float
        
        # Check tolerance (FP32 precision)
        # Note: CUDA atomicAdd order might cause slight deviation, but usually negligible in FP32
        diff = np.abs(C_cuda - C_ref)
        max_diff = np.max(diff)
        mean_diff = np.mean(diff)
        
        print(f"Max Diff: {max_diff}, Mean Diff: {mean_diff}")
        self.assertTrue(max_diff < 1e-3, f"Forward pass mismatch! Max diff: {max_diff}")

    def test_backward_weights_cuda(self):
        print("\nTesting Backward Weights CUDA...")
        grad_Y = np.random.randn(self.M, self.N).astype(np.float32)
        grad_W_cuda = np.zeros((self.K, self.N), dtype=np.float32)
        
        # Run CUDA
        maxwell_core.run_maxwell_backward_weights_cuda(self.A, grad_Y, grad_W_cuda)
        
        # Run Reference: grad_W = A.T @ grad_Y
        grad_W_ref = self.A.T @ grad_Y
        
        diff = np.abs(grad_W_cuda - grad_W_ref)
        max_diff = np.max(diff)
        
        print(f"Max Diff: {max_diff}")
        self.assertTrue(max_diff < 1e-3, f"Backward pass mismatch! Max diff: {max_diff}")

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