import tempfile
import unittest
from pathlib import Path
import numpy as np
from encoder_inputs import frozen_pn_drive, requests_from_forced, write_npz


class EncoderInputTests(unittest.TestCase):
    def test_exact_frozen_encoder_ast_on_synthetic_static_mapping(self):
        import pandas as pd
        local = Path(__file__).resolve().parents[1] / 'models/male-cns-lif/phase4l.py'
        source = local if local.exists() else Path('/data/digital-drosophila/models/male-cns-lif/phase4l.py')
        pin = 'e97d58d8c1f44920d0e953f8f5831ddc86fa0198784e86460e315de4116ac553'
        function, ast_sha = frozen_pn_drive(source, pin)
        self.assertEqual(len(ast_sha), 64)
        mapping = pd.DataFrame(dict(glomerulus=['G2', 'G2', 'G1', 'G1'],
            mapping_confidence=['exact', 'alias-normalized', 'exact', 'unknown'], neuron_idx=[3, 1, 7, 999]))
        cfg = {'odor_encoder': {'max_total_rate_per_glomerulus_hz': 400}}
        target, rates, _ = function('synthetic', np.array([.5, .25], np.float32), mapping,
                                     ['G2', 'G1'], 'matched_energy', 1.5, cfg)
        np.testing.assert_array_equal(target, np.array([3, 1, 7], np.int32))
        np.testing.assert_array_equal(rates, np.array([200, 200, 200], np.float32))
        with self.assertRaises(AssertionError):
            frozen_pn_drive(source, '0' * 64)

    def test_rng_column_order_preserved_until_request_sort(self):
        targets = np.array([10, 1, 3], np.int32)
        forced = np.array([[True, True, False], [False, True, True]], np.bool_)
        actual = requests_from_forced(forced, targets)
        np.testing.assert_array_equal(actual, np.array([[0, 1], [0, 10], [1, 1], [1, 3]], np.int32))
        np.testing.assert_array_equal(targets, np.array([10, 1, 3], np.int32))

    def test_duplicate_targets_rejected_not_merged(self):
        with self.assertRaises(AssertionError):
            requests_from_forced(np.ones((2, 2), np.bool_), np.array([1, 1], np.int32))

    def test_empty_sparse_schedule_has_exact_shape_and_dtype(self):
        result = requests_from_forced(np.zeros((3, 2), np.bool_), np.array([2, 1], np.int32))
        self.assertEqual(result.shape, (0, 2))
        self.assertEqual(result.dtype, np.int32)

    def test_npz_creation_is_exclusive(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / 'schedule.npz'
            write_npz(path, tick=np.array([1], np.int32))
            with np.load(path, allow_pickle=False) as stored:
                np.testing.assert_array_equal(stored['tick'], np.array([1], np.int32))
            with self.assertRaises(FileExistsError):
                write_npz(path, tick=np.array([2], np.int32))


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