from __future__ import annotations

import unittest

import numpy as np

from evaluate_annotated_floor_events import AnnotatedOpportunity
from evaluate_egospeak_event_proxy import (
    decision_source_rows,
    evaluate_family,
    horizon_indices,
    score_activity_proxy,
)
from evaluate_floor_state_changes import NUM_CLASSES, NUM_HORIZONS


def opportunity(
    uid: str,
    *,
    family: str = "entry",
    positive: bool = True,
    start: int = 20,
    end: int = 21,
    anchor: float = 4.0,
    region_start: float = 4.0,
    region_end: float = 4.0,
    source: str = "WEARER_ONSET",
) -> AnnotatedOpportunity:
    return AnnotatedOpportunity(
        uid=uid,
        session="Session_4",
        family=family,  # type: ignore[arg-type]
        positive=positive,
        subtype=("WEARER_ENTRY" if family == "entry" else "INT_YIELD")
        if positive
        else ("ENTRY_OTHER_TO_OTHER" if family == "entry" else "INT_OTHER_BACKCHANNEL"),
        start_row=start,
        end_row_exclusive=end,
        anchor_seconds=anchor,
        region_start_seconds=region_start,
        region_end_seconds=region_end,
        candidate_source=source,
    )


def arrays(length: int = 50) -> tuple[np.ndarray, np.ndarray]:
    return (
        np.zeros((length, NUM_CLASSES, NUM_HORIZONS), dtype=np.float32),
        np.ones((length, NUM_HORIZONS), dtype=bool),
    )


class EgoSpeakEventProxyTests(unittest.TestCase):
    def test_positive_uses_three_pre_anchor_source_rows(self) -> None:
        item = opportunity("positive", anchor=4.0)
        self.assertEqual(
            decision_source_rows(item, window_rows=3),
            (17, 20, "pre_behavior_anchor"),
        )

    def test_independent_hard_negative_uses_only_equal_width_region_prefix(self) -> None:
        item = opportunity(
            "negative",
            positive=False,
            start=10,
            end=40,
            anchor=2.0,
            region_start=2.0,
            region_end=8.0,
            source="OTHER_TO_OTHER",
        )
        self.assertEqual(
            decision_source_rows(item, window_rows=3),
            (10, 13, "hard_negative_opportunity_onset"),
        )

    def test_short_hard_negative_is_not_silently_padded(self) -> None:
        item = opportunity(
            "short",
            positive=False,
            start=10,
            end=12,
            source="OTHER_TO_OTHER",
        )
        self.assertIsNone(decision_source_rows(item, window_rows=3))

    def test_relabelled_raw_boundary_uses_pre_anchor_rows(self) -> None:
        item = opportunity(
            "backchannel",
            positive=False,
            start=5,
            end=35,
            anchor=4.0,
            region_start=1.0,
            region_end=7.0,
            source="WEARER_ONSET",
        )
        self.assertEqual(
            decision_source_rows(item, window_rows=3),
            (17, 20, "pre_behavior_anchor"),
        )

    def test_trajectory_follows_target_aligned_diagonal(self) -> None:
        item = opportunity("positive", anchor=4.0)
        prediction, valid = arrays()
        selected = horizon_indices([0.6, 0.8, 1.0])
        # Best decision row is 18. Target rows are source+horizon_index+1.
        prediction[21, 1, 2] = 0.6
        prediction[22, 1, 3] = 0.9
        prediction[23, 1, 4] = 0.3
        scored = score_activity_proxy(
            item,
            {"Session_4": (prediction, valid)},
            window_rows=3,
            selected_horizons=selected,
        )
        self.assertIsNotNone(scored)
        score, details = scored  # type: ignore[misc]
        self.assertAlmostEqual(score, 0.6)
        self.assertEqual(details["best_source_row"], 18)
        self.assertEqual(len(details["best_wearer_trajectory"]), 3)

    def test_yield_inverts_sustained_wearer_activity(self) -> None:
        item = opportunity("yield", family="int", anchor=4.0, source="WEARER_OFFSET")
        prediction, valid = arrays()
        selected = horizon_indices([0.6, 0.8])
        for source_row in range(17, 20):
            prediction[source_row + 3, 1, 2] = 0.2
            prediction[source_row + 4, 1, 3] = 0.2
        scored = score_activity_proxy(
            item,
            {"Session_4": (prediction, valid)},
            window_rows=3,
            selected_horizons=selected,
        )
        self.assertIsNotNone(scored)
        self.assertAlmostEqual(scored[0], 0.8)  # type: ignore[index]

    def test_each_opportunity_contributes_one_metric_row(self) -> None:
        positive = opportunity("p", anchor=4.0)
        negative = opportunity(
            "n",
            positive=False,
            start=25,
            end=30,
            anchor=5.0,
            region_start=5.0,
            region_end=6.0,
            source="OTHER_TO_OTHER",
        )
        prediction, valid = arrays()
        selected = horizon_indices([0.6])
        for source_row in range(17, 20):
            prediction[source_row + 3, 1, 2] = 0.9
        metric, rows = evaluate_family(
            [positive, negative],
            {"Session_4": (prediction, valid)},
            "entry",
            [],
            window_rows=3,
            selected_horizons=selected,
        )
        self.assertEqual(len(rows), 2)
        self.assertEqual(metric["num_positive"], 1)
        self.assertEqual(metric["num_hard_negative"], 1)
        self.assertAlmostEqual(metric["average_precision"], 1.0)


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