"""Focused tests for the direct competing-risk floor-event probe."""

from __future__ import annotations

import unittest

import numpy as np
import torch

from train_easycom_floor_event_probe import (
    DirectFloorEventHead,
    FloorEventDataset,
    ProbeSession,
    metric_rows,
)


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


def probe_session() -> ProbeSession:
    target = np.asarray(
        [
            activity(0, 1),
            activity(0, 0),
            activity(1, 0),
            activity(1, 0),
            activity(1, 1),
            activity(0, 1),
            activity(0, 1),
            activity(0, 0),
            activity(0, 0),
            activity(0, 1),
        ],
        dtype=np.float32,
    )
    steps = len(target)
    return ProbeSession(
        name="Session_12",
        states=torch.randn(steps, 2, 4),
        audio=torch.cat((torch.randn(steps, 3), torch.ones(steps, 1)), dim=1),
        oracle=torch.randn(steps, 6),
        oracle_valid=torch.ones(steps, dtype=torch.bool),
        target=target,
        oracle_condition="aggregate-visible",
        oracle_contract={"config": {"oracle_condition": "aggregate-visible"}},
    )


class FloorEventProbeTests(unittest.TestCase):
    def setUp(self) -> None:
        torch.set_num_threads(1)

    def test_dataset_uses_requested_decision_source(self) -> None:
        session = probe_session()
        pre = FloorEventDataset(
            [session], decision_timing="pre-boundary", midturn_pause_max_steps=10
        )
        post = FloorEventDataset(
            [session], decision_timing="post-boundary", midturn_pause_max_steps=10
        )
        pre_response = next(
            item
            for item in pre.examples
            if item.event_index == 0 and item.horizon_index == 1 and item.label
        )
        post_response = next(
            item
            for item in post.examples
            if item.event_index == 0 and item.horizon_index == 1 and item.label
        )
        self.assertEqual(pre_response.source_row, 0)
        self.assertEqual(post_response.source_row, 1)

    def test_head_outputs_two_events_by_ten_horizons(self) -> None:
        model = DirectFloorEventHead(
            state_shape=(2, 4),
            audio_dim=3,
            oracle_dim=6,
            layer_fusion="concat-mlp",
            local_projection_dim=4,
            person_attention_dim=4,
            person_attention_heads=1,
            dropout=0.0,
        )
        output = model(
            torch.randn(5, 2, 4),
            torch.cat((torch.randn(5, 3), torch.ones(5, 1)), dim=1),
            torch.randn(5, 6),
            torch.ones(5),
        )
        self.assertEqual(output.shape, (5, 2, 10))

    def test_fixed_threshold_is_applied_without_reselection(self) -> None:
        predictions = {
            "score": np.tile(np.asarray([0.9, 0.4]), 20),
            "event": np.repeat(np.arange(2), 20),
            "horizon": np.tile(np.repeat(np.arange(10), 2), 2),
            "label": np.tile(np.asarray([True, False]), 20),
            "short_resume": np.zeros(40, dtype=bool),
        }
        thresholds = {
            (event, horizon): 0.95
            for event in range(2)
            for horizon in range(10)
        }
        rows, _ = metric_rows(
            predictions,
            matched_recall=0.8,
            fixed_thresholds=thresholds,
            threshold_source="Session_12",
        )
        self.assertEqual(len(rows), 20)
        self.assertEqual(rows[0]["threshold"], 0.95)
        self.assertEqual(rows[0]["achieved_positive_recall"], 0.0)


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