import unittest
import numpy as np
import pandas as pd
from scope import contact_scope, edge_ids_for_pairs, blocked_weights


class ScopeTests(unittest.TestCase):
    def test_roi_and_side_not_guessed(self):
        apl = pd.DataFrame({'source_neuron_idx': [0, 0, 1], 'target_neuron_idx': [3, 3, 4],
                            'primary_post': ['CA(R)', 'CentralBrain-unspecified', 'CA(L)']})
        pn = pd.DataFrame({'source_neuron_idx': [3, 3, 3, 4, 4], 'target_neuron_idx': [5]*5,
                           'primary_post': ['CA(R)', 'CA(L)', None, 'CA(R)', 'CA(L)']})
        pairs, allowed, keep = contact_scope(apl, pn, ['CA(L)', 'CA(R)'])
        self.assertEqual(keep.tolist(), [True, False, False, False, True])
        self.assertEqual(pairs.blocked_contact_count.tolist(), [1, 1])
        self.assertEqual(pairs.retained_contact_count.tolist(), [2, 1])
        self.assertEqual(len(allowed), 2)

    def test_zero_identity_and_nonmutation(self):
        w = np.array([.275, -3.14159, 0, -0.0], np.float32)
        raw = np.array([1, 43, 10, 1], np.int64)
        old = w.tobytes(), raw.tobytes()
        result = blocked_weights(w, raw, np.zeros(4, np.int64))
        self.assertEqual(w.tobytes(), result.tobytes())
        self.assertEqual((w.tobytes(), raw.tobytes()), old)

    def test_partial_and_full(self):
        w = np.array([2., -4., 0., 1.], np.float32)
        raw = np.array([4, 4, 3, 7], np.int64)
        result = blocked_weights(w, raw, np.array([1, 4, 2, 0], np.int64))
        np.testing.assert_array_equal(result, np.array([1.5, -0., 0., 1.], np.float32))
        np.testing.assert_array_equal(w, np.array([2., -4., 0., 1.], np.float32))

    def test_invalid_partition_rejected(self):
        for bad in [-1, 3]:
            with self.assertRaises(ValueError):
                blocked_weights(np.array([1.], np.float32), np.array([2], np.int64), np.array([bad], np.int64))

    def test_existing_edge_only(self):
        p = pd.DataFrame({'source_neuron_idx': [0, 1], 'target_neuron_idx': [2, 0]})
        ptr = np.array([0, 2, 3, 3], np.int64); dst = np.array([1, 2, 0], np.int32)
        np.testing.assert_array_equal(edge_ids_for_pairs(p, ptr, dst), [1, 2])
        p.loc[1, 'target_neuron_idx'] = 2
        with self.assertRaises(ValueError): edge_ids_for_pairs(p, ptr, dst)

    def test_duplicates_rejected(self):
        p = pd.DataFrame({'source_neuron_idx': [0], 'target_neuron_idx': [1]})
        with self.assertRaises(ValueError):
            edge_ids_for_pairs(p, np.array([0, 2, 2], np.int64), np.array([1, 1], np.int32))


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