import unittest

import numpy as np
import pandas as pd
from scipy.sparse import coo_matrix

from connectome_interpreter.utils import modify_coo_matrix


class TestModifyCooMatrix(unittest.TestCase):

    def setUp(self):
        # Create a small sparse matrix with float data type and values between 0 and 1
        self.example_coo = coo_matrix(
            np.array([[0.1, 0.2, 0], [0, 0, 0.3], [0.4, 0, 0.5]], dtype=np.float32)
        )  # Ensure dtype is float32

        # DataFrame for batch updates with float values
        self.updates_df = pd.DataFrame(
            {
                "input_idx": [0, 2],
                "output_idx": [2, 1],
                "value": [0.25, 0.75],  # Use values between 0 and 1
            }
        )

    def test_single_update(self):
        # Test updating a single value with a float value between 0 and 1
        result = modify_coo_matrix(
            self.example_coo, input_idx=1, output_idx=2, value=0.99, re_normalize=False
        )
        self.assertAlmostEqual(result.toarray()[1, 2], 0.99)

    def test_batch_update(self):
        # Test updating using a DataFrame of updates with float values
        result = modify_coo_matrix(
            self.example_coo, updates_df=self.updates_df, re_normalize=False
        )
        self.assertAlmostEqual(result.toarray()[0, 2], 0.25)
        self.assertAlmostEqual(result.toarray()[2, 1], 0.75)

    def test_renormalize(self):
        # Test the re-normalization functionality with float values
        result = modify_coo_matrix(
            self.example_coo, updates_df=self.updates_df, re_normalize=True
        )
        updated_cols = self.updates_df["output_idx"].unique()
        # Check if columns sum up to 1 or close to it
        column_sums = np.array(result.sum(axis=0))[0]
        for col_sum in column_sums[updated_cols]:
            # Assuming a little margin for floating-point arithmetic
            self.assertTrue(0.99 <= col_sum <= 1.01)

    def test_handle_zero_colsums(self):
        # Test how the function handles columns with zero sums after update with float values
        zero_updates_df = pd.DataFrame(
            {
                "input_idx": [1, 0],
                "output_idx": [2, 1],
                "value": [0.0, 0.0],  # Use float zero to keep consistent data type
            }
        )
        result = modify_coo_matrix(
            self.example_coo, updates_df=zero_updates_df, re_normalize=True
        )
        column_sums = np.array(result.sum(axis=0))[0]
        self.assertAlmostEqual(column_sums[1], 0.0)
