import tempfile
import unittest
from pathlib import Path

import numpy as np
import torch

from egospeak_omni_residual import (
    EgoSpeakOmniResidual,
    EgoSpeakOmniWindowDataset,
    MiniROADABase,
    align_source_predictions,
    load_easycom_session_arrays,
    load_frozen_miniroada,
    miniroad_normalized_loss,
)


class EgoSpeakOmniResidualTest(unittest.TestCase):
    def test_raw_miniroada_state_dict_loads_strictly_and_freezes(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            checkpoint = Path(directory) / "base.pth"
            source = MiniROADABase(
                rgb_dim=4,
                audio_dim=6,
                embedding_dim=8,
                hidden_dim=5,
                num_horizons=10,
            )
            torch.save(source.state_dict(), checkpoint)
            loaded = load_frozen_miniroada(
                checkpoint,
                rgb_dim=4,
                audio_dim=6,
                embedding_dim=8,
                hidden_dim=5,
                num_horizons=10,
            )
            self.assertFalse(loaded.training)
            self.assertTrue(
                all(not parameter.requires_grad for parameter in loaded.parameters())
            )
            for key, value in source.state_dict().items():
                torch.testing.assert_close(loaded.state_dict()[key], value)

    def test_zero_residual_exactly_recovers_frozen_base(self) -> None:
        torch.manual_seed(7)
        base = MiniROADABase(
            rgb_dim=4,
            audio_dim=6,
            embedding_dim=8,
            hidden_dim=5,
            num_horizons=10,
            dropout=0.4,
        )
        model = EgoSpeakOmniResidual(
            base,
            omni_dim=7,
            fusion_dim=4,
            tail_embedding_dim=3,
            dropout=0.3,
        )
        model.train()
        self.assertFalse(model.base.training)
        self.assertTrue(
            all(not parameter.requires_grad for parameter in model.base.parameters())
        )

        rgb = torch.randn(2, 9, 4)
        audio = torch.randn(2, 9, 6)
        omni = torch.randn(2, 7)
        tail = torch.tensor([0, 9], dtype=torch.long)
        fused, frozen = model.forward_last(rgb, audio, omni, tail)
        torch.testing.assert_close(fused, frozen, rtol=0.0, atol=0.0)

        targets = torch.zeros_like(fused)
        targets[..., 1] = 1
        optimizer = torch.optim.SGD(model.trainable_parameters(), lr=0.1)
        loss = miniroad_normalized_loss(fused, targets)
        loss.backward()
        self.assertGreater(float(model.residual_output.weight.grad.abs().sum()), 0.0)
        optimizer.step()
        updated, frozen_again = model.forward_last(rgb, audio, omni, tail)
        self.assertGreater(
            float((updated - frozen_again).abs().max().detach()),
            0.0,
        )

    def test_feature_contract_and_window_alignment(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            ego = root / "ego"
            omni = root / "omni"
            for subdir in (
                "rgb_kinetics_resnet50",
                "audio_wav2vec2",
                "target_perframe",
            ):
                (ego / subdir).mkdir(parents=True, exist_ok=True)
            omni.mkdir()
            steps = 15
            rgb = np.arange(steps * 4, dtype=np.float32).reshape(steps, 4)
            audio = np.arange(steps * 6, dtype=np.float32).reshape(steps, 6)
            target = np.zeros((steps, 3), dtype=np.float32)
            target[:, 0] = 1
            np.save(ego / "rgb_kinetics_resnet50" / "Session_4.npy", rgb)
            np.save(ego / "audio_wav2vec2" / "Session_4.npy", audio)
            np.save(ego / "target_perframe" / "Session_4.npy", target)
            torch.save(
                {
                    "mode": "open-stream-replay",
                    "states": torch.arange(steps * 7, dtype=torch.float32).view(
                        steps, 7
                    ),
                    "times": (torch.arange(steps, dtype=torch.float64) + 1) / 5,
                    "label_indices": torch.arange(steps, dtype=torch.long),
                    "start": 0.0,
                    "config": {
                        "step_seconds": 0.2,
                        "native_block_seconds": 2.0,
                    },
                },
                omni / "Session_4.pt",
            )
            session = load_easycom_session_arrays(
                ego,
                omni,
                "Session_4",
                rgb_dim=4,
                audio_dim=6,
            )
            self.assertEqual(
                session.tail_bins.tolist(), list(range(10)) + list(range(5))
            )
            dataset = EgoSpeakOmniWindowDataset(
                [session],
                window_size=4,
                stride=1,
                num_horizons=3,
                seed=0,
            )
            first_index = dataset.indices.index((0, 0))
            rgb_window, audio_window, state, tail, future, _, source = dataset[
                first_index
            ]
            self.assertEqual(int(source), 0)
            self.assertEqual(int(tail), 0)
            torch.testing.assert_close(rgb_window[-1], torch.from_numpy(rgb[0]))
            self.assertEqual(int(torch.count_nonzero(rgb_window[:-1])), 0)
            torch.testing.assert_close(audio_window[-1], torch.from_numpy(audio[0]))
            torch.testing.assert_close(state, session.omni[0])
            torch.testing.assert_close(future, torch.from_numpy(target[1:4]))

    def test_source_prediction_alignment(self) -> None:
        target_steps = 14
        raw = np.arange(
            (target_steps - 10) * 10 * 3,
            dtype=np.float32,
        ).reshape(target_steps - 10, 10, 3)
        pred, valid = align_source_predictions(raw, target_steps)
        np.testing.assert_array_equal(pred[1, :, 0], raw[0, 0, :])
        np.testing.assert_array_equal(pred[13, :, 9], raw[3, 9, :])
        self.assertTrue(valid[1, 0])
        self.assertFalse(valid[0, 0])
        self.assertTrue(valid[13, 9])


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