#!/usr/bin/env python3

import torch
import tempfile
import sys
from pathlib import Path

# Add parent directory to path for imports
sys.path.append(str(Path(__file__).parent.parent))

from utils.bct import BlockType, Block, BlockSequence


def test_block_sequence_serialization():
    # Create test data
    text_block_type = BlockType(name="text", is_causal=True)
    semantic_block_type = BlockType(name="semantic", is_causal=True)

    text_block = Block(spec=text_block_type, inputs={"text_input": torch.randn(10, 256)}, targets={})

    semantic_block = Block(
        spec=semantic_block_type,
        inputs={"semantic_input": torch.randn(8, 512)},
        targets={"semantic_output": torch.randn(8, 512)},
    )

    original_seq = BlockSequence([text_block, semantic_block])

    # Test save/load
    with tempfile.TemporaryDirectory() as tmpdir:
        path = Path(tmpdir) / "test_sequence.pt"

        # Save
        original_seq.save(path)
        print(f"✓ Saved BlockSequence to {path}")

        # Load
        loaded_seq = BlockSequence.load(path)
        print(f"✓ Loaded BlockSequence from {path}")

        # Verify
        assert len(loaded_seq) == len(original_seq)
        assert loaded_seq.n_tokens == original_seq.n_tokens

        for orig_block, loaded_block in zip(original_seq, loaded_seq):
            assert orig_block.spec.name == loaded_block.spec.name
            assert orig_block.spec.is_causal == loaded_block.spec.is_causal

            for key in orig_block.inputs:
                assert torch.allclose(orig_block.inputs[key], loaded_block.inputs[key])

            for key in orig_block.targets:
                assert torch.allclose(orig_block.targets[key], loaded_block.targets[key])

        print("✓ All tensors match after serialization")


if __name__ == "__main__":
    print("Testing BlockSequence serialization...")
    test_block_sequence_serialization()
    print()
    print("🎉 Test passed!")