from __future__ import annotations

import tempfile
import unittest
from pathlib import Path

import numpy as np

from evaluate_floor_state_changes import (
    NUM_CLASSES,
    NUM_HORIZONS,
    evaluate_horizon,
    load_arrays,
    parse_runs,
    validate_recalls,
)


def activity(wearer: int, others: int) -> list[int]:
    return [int(not wearer and not others), wearer, others]


class FloorStateMetricTests(unittest.TestCase):
    def test_horizon_one_proxies_and_session_boundary(self) -> None:
        # t0->t1: EOT-response; t3->t4: HOLD; t4->t5: early entry;
        # t6->t7: other-entry yield; t8->t9: other-entry continue.
        target = np.asarray(
            [
                activity(0, 1),
                activity(1, 0),
                activity(0, 0),
                activity(0, 1),
                activity(0, 1),
                activity(1, 1),
                activity(1, 0),
                activity(0, 1),
                activity(1, 0),
                activity(1, 1),
            ],
            dtype=np.float32,
        )
        prediction = np.zeros((len(target), NUM_CLASSES, NUM_HORIZONS), dtype=np.float32)
        prediction[:, 1, :] = 0.5
        prediction[1, 1, 0] = 0.9
        prediction[4, 1, 0] = 0.1
        prediction[5, 1, 0] = 0.8
        prediction[7, 1, 0] = 0.1
        prediction[9, 1, 0] = 0.9
        valid = np.ones((len(target), NUM_HORIZONS), dtype=bool)
        session_ids = np.zeros(len(target), dtype=np.int32)

        result = evaluate_horizon(
            target, prediction, valid, session_ids, 0, [0.8]
        )
        floor = result["eot_response_vs_other_hold_proxy"]
        self.assertEqual(floor["num_positive"], 2)
        self.assertEqual(floor["num_negative"], 1)
        self.assertEqual(floor["num_early_entry_proxy"], 1)
        self.assertAlmostEqual(floor["average_precision"], 1.0)
        self.assertAlmostEqual(
            floor["operating_points"][0]["negative_false_positive_rate"], 0.0
        )
        yield_metric = result["other_entry_yield_vs_continue_proxy"]
        self.assertEqual(yield_metric["num_positive"], 1)
        self.assertEqual(yield_metric["num_negative"], 1)
        self.assertAlmostEqual(yield_metric["average_precision"], 1.0)

        # The apparent t4->t5 early entry is removed if t5 begins a new session.
        session_ids[5:] = 1
        bounded = evaluate_horizon(
            target, prediction, valid, session_ids, 0, [0.8]
        )
        self.assertEqual(
            bounded["eot_response_vs_other_hold_proxy"]["num_early_entry_proxy"],
            0,
        )

    def test_load_arrays_validates_layout(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            target_root = root / "target"
            run_root = root / "run"
            (run_root / "pred").mkdir(parents=True)
            (run_root / "valid").mkdir(parents=True)
            target_root.mkdir()
            target = np.asarray([activity(0, 0), activity(1, 0)], dtype=np.float32)
            np.save(target_root / "Session_1.npy", target)
            np.save(
                run_root / "pred/Session_1.npy",
                np.zeros((2, NUM_CLASSES, NUM_HORIZONS), dtype=np.float32),
            )
            np.save(
                run_root / "valid/Session_1.npy",
                np.ones((2, NUM_HORIZONS), dtype=bool),
            )
            loaded = load_arrays(target_root, run_root, ["Session_1"])
            self.assertEqual(loaded[0].shape, (2, NUM_CLASSES))
            self.assertEqual(loaded[1].shape, (2, NUM_CLASSES, NUM_HORIZONS))

    def test_argument_validation(self) -> None:
        self.assertEqual(parse_runs(["G=results/g"]), [("G", Path("results/g"))])
        with self.assertRaises(ValueError):
            parse_runs(["results/g"])
        with self.assertRaises(ValueError):
            parse_runs(["G=a", "G=b"])
        self.assertEqual(validate_recalls([0.9, 0.7]), [0.7, 0.9])
        with self.assertRaises(ValueError):
            validate_recalls([0.0])


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