"""Focused tests for the causal EasyCom oracle-information experiment."""

from __future__ import annotations

from pathlib import Path
from types import SimpleNamespace
import tempfile
import unittest

import numpy as np
import torch

from easycom_head import TurnPredictionHead
from local_turn_features import ORACLE_FORMAT, load_frame_set_features, load_local_features
from prepare_easycom_oracle_features import (
    causal_history,
    oracle_person_rows,
    participant_activity,
)


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

    def test_transcript_intervals_are_rasterized_into_half_open_bins(self) -> None:
        rows = [
            SimpleNamespace(
                uid="a",
                participant_id=3,
                dense_start_frame=3,
                dense_end_frame_exclusive=5,
            ),
            SimpleNamespace(
                uid="b",
                participant_id=2,
                dense_start_frame=8,
                dense_end_frame_exclusive=12,
            ),
        ]
        activity, wearer, audit = participant_activity(
            rows, num_steps=4, wearer_id=2
        )
        self.assertEqual(activity[3].tolist(), [True, True, False, False])
        self.assertEqual(wearer.tolist(), [False, False, True, False])
        self.assertEqual(audit["discarded_terminal_transcript_frames"], 0)

    def test_history_at_row_does_not_read_later_activity(self) -> None:
        first = torch.tensor([0, 1, 0, 0, 0], dtype=torch.bool)
        second = first.clone()
        second[3:] = True
        observed = causal_history(first, row=2, history_steps=4)
        changed_future = causal_history(second, row=2, history_steps=4)
        self.assertTrue(torch.equal(observed, changed_future))
        self.assertEqual(
            observed.tolist(),
            [0, 0, 1, 0, 0, 1, 1, 1],
        )
        delayed = causal_history(second, row=2, history_steps=4, delay_steps=2)
        self.assertEqual(
            delayed.tolist(),
            [0, 0, 0, 0, 0, 0, 0, 1],
        )
        visibility = torch.tensor([1, 0, 1, 1, 1], dtype=torch.bool)
        visible_only = causal_history(
            second,
            row=2,
            history_steps=4,
            observed_sequence=visibility,
        )
        self.assertEqual(
            visible_only.tolist(),
            [0, 0, 0, 0, 0, 1, 0, 1],
        )

    def test_shuffled_id_preserves_every_lag_count_but_breaks_histories(self) -> None:
        activity = {
            3: torch.tensor([1, 1, 0, 0, 1, 1], dtype=torch.bool),
            4: torch.tensor([0, 0, 1, 1, 0, 0], dtype=torch.bool),
            5: torch.tensor([0, 1, 0, 1, 0, 1], dtype=torch.bool),
        }
        visible = [[3, 4, 5] for _ in range(6)]
        visibility = {
            person_id: torch.ones(6, dtype=torch.bool)
            for person_id in activity
        }
        records = [
            [
                {"id": str(person_id), "source_clip_index": 0, "source_time": row / 5}
                for person_id in person_ids
            ]
            for row, person_ids in enumerate(visible)
        ]
        original, offsets, _, _ = oracle_person_rows(
            activity,
            visibility,
            visible,
            records,
            condition="per-person",
            history_steps=4,
            delay_steps=5,
            seed=9,
        )
        shuffled, shuffled_offsets, _, audit = oracle_person_rows(
            activity,
            visibility,
            visible,
            records,
            condition="shuffled-id",
            history_steps=4,
            delay_steps=5,
            seed=9,
        )
        self.assertTrue(torch.equal(offsets, shuffled_offsets))
        for row in range(6):
            start, stop = int(offsets[row]), int(offsets[row + 1])
            self.assertTrue(
                torch.equal(
                    original[start:stop, :4].sum(dim=0),
                    shuffled[start:stop, :4].sum(dim=0),
                )
            )
        self.assertGreater(audit["shuffled_values_changed"], 0)

    def test_oracle_format_requires_explicit_loader_and_head_mode(self) -> None:
        n = 4
        times = (torch.arange(n, dtype=torch.float64) + 1) / 5
        reference = {
            "states": torch.randn(n, 8),
            "label_indices": torch.arange(n),
            "times": times,
            "mode": "open-stream-replay",
            "start": 0,
        }
        sidecar = {
            "format": ORACLE_FORMAT,
            "dataset": "easycom",
            "sample_id": "Session_4",
            "times": times.clone(),
            "label_indices": torch.arange(n),
            "partial": False,
            "visual": torch.randn(n, 6),
            "audio": torch.randn(n, 5),
            "visual_valid": torch.ones(n, dtype=torch.bool),
            "audio_valid": torch.ones(n, dtype=torch.bool),
            "persons": {
                "features": torch.randn(n, 6),
                "frame_times": torch.arange(n, dtype=torch.float64) / 5,
                "frame_offsets": torch.arange(n + 1, dtype=torch.long),
                "records": [{"id": "3"} for _ in range(n)],
            },
            "config": {
                "oracle_analysis": True,
                "uses_speech_labels": True,
                "causal_policy": "fully-observed-half-open-5hz-bins",
                "oracle_condition": "per-person",
                "visual_dim": 6,
                "audio_dim": 5,
            },
        }
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "Session_4.pt"
            torch.save(sidecar, path)
            with self.assertRaisesRegex(ValueError, "format"):
                load_local_features(
                    path,
                    reference,
                    sample_id="Session_4",
                    dataset="easycom",
                    modalities="av",
                )
            audio, people, contract = load_frame_set_features(
                path,
                reference,
                sample_id="Session_4",
                dataset="easycom",
                oracle_analysis=True,
            )
        self.assertEqual(audio.shape, (n, 6))
        self.assertEqual(people.dimension, 6)
        self.assertEqual(contract["format"], ORACLE_FORMAT)
        head = TurnPredictionHead(
            hidden_size=8,
            num_horizons=2,
            num_classes=3,
            local_visual_dim=6,
            local_audio_dim=5,
            local_projection_dim=4,
            local_fusion="oracle-frame-set-attention",
            person_attention_dim=4,
            person_attention_heads=1,
        )
        padded, mask = people.padded()
        logits = head(reference["states"], audio, padded, mask)
        self.assertEqual(logits.shape, (n, 2, 3))


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