import numpy as np
import unittest
from unittest.mock import MagicMock, patch
from notscared.cpa import CPA


class TestCPA(unittest.TestCase):

    def setUp(self):
        # Create mock Tracehandler object
        self.trace_handler = MagicMock()

        # Create mock Bytes
        self.bytes = [0, 1, 2]

        # Create CPA instance
        self.cpa = CPA(self.trace_handler, self.bytes)

    def test_init(self):
        self.assertEqual(self.cpa.data, self.trace_handler)
        self.assertEqual(self.cpa.bytes, self.bytes)
        self.assertEqual(self.cpa.tiles, self.trace_handler.tiles)
        self.assertIsInstance(self.cpa.sbox, np.ndarray)
        self.assertIsInstance(self.cpa.keys, np.ndarray)
        self.assertIsInstance(self.cpa.weights, np.ndarray)
        self.assertEqual(self.cpa.count, 0)
        self.assertIsInstance(self.cpa.sample_sum, np.ndarray)
        self.assertIsInstance(self.cpa.sample_sq_sum, np.ndarray)
        self.assertIsInstance(self.cpa.model_sum, np.ndarray)
        self.assertIsInstance(self.cpa.model_sq_sum, np.ndarray)
        self.assertIsInstance(self.cpa.prod_sum, np.ndarray)
        self.assertIsInstance(self.cpa.resultsAnalyzer, MagicMock)
        self.assertEqual(self.cpa.bytes_to_graph, [])
        self.assertIsInstance(self.cpa.results, np.ndarray)

    @patch.object(CPA, 'finalize')
    @patch.object(CPA, 'get_model')
    @patch.object(CPA, 'update')
    @patch.object(CPA, 'data')
    def test_run_1x1(self, mock_data, mock_update, mock_get_model, mock_finalize):
        # Set up mock data
        plaintexts = np.array([[0, 1, 2], [3, 4, 5]])
        samples = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
        mock_data.grab.side_effect = [(plaintexts, samples), ()]

        # Set up mock get_model
        model = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
        mock_get_model.return_value = model

        # Set up mock finalize
        result = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
        mock_finalize.return_value = result

        # Run the method
        self.cpa.run_1x1()

        # Assertions
        mock_data.grab.assert_called_with()
        mock_get_model.assert_called_with(plaintexts[:, 0])
        mock_update.assert_called_with(samples, model)
        mock_finalize.assert_called_with()
        self.assertEqual(self.cpa.results[0, :], result)

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