import numpy as np
import torch

from speech_bridge_gemma.longcat_acoustic_head import AcousticSample
from speech_bridge_gemma.longcat_ctc_aligned_acoustic_head import PhoneSample, align_samples_with_ctc, ctc_alignment_cache_key


def acoustic_sample(semantic: list[int]) -> AcousticSample:
    frames = len(semantic)
    return AcousticSample(
        sample_id="same",
        text="texto",
        semantic_codes=torch.tensor(semantic, dtype=torch.long),
        acoustic_codes=torch.zeros((3, frames), dtype=torch.long),
        target_audio=np.zeros(160, dtype=np.float32),
        target_sample_rate=16000,
        duration_sec=0.01,
        speaker_id="spk",
    )


def phone_sample(semantic: list[int]) -> PhoneSample:
    return PhoneSample(
        base=acoustic_sample(semantic),
        phone_tokens=["t", "e"],
        phone_target_ids=torch.tensor([1, 2], dtype=torch.long),
    )


def test_ctc_alignment_cache_key_includes_semantic_codes() -> None:
    assert ctc_alignment_cache_key(phone_sample([1, 2, 3])) != ctc_alignment_cache_key(phone_sample([1, 9, 3]))


class DummyCtc:
    blank_id = 0

    def eval(self) -> None:
        return None

    def __call__(self, semantic: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor:
        logits = torch.zeros((semantic.shape[0], semantic.shape[1], 3), dtype=torch.float32)
        logits[:, :, 1] = 10.0
        return logits


def test_align_samples_with_ctc_deduplicates_matching_pending_keys() -> None:
    first = phone_sample([1, 2, 3])
    second = phone_sample([1, 2, 3])

    stats = align_samples_with_ctc(DummyCtc(), [first, second], "cpu", {}, batch_size=2, log_every=0)

    assert stats["computed"] == 1
    assert first.aligned_phone_ids is not None
    assert second.aligned_phone_ids is not None
    assert first.aligned_phone_ids.tolist() == second.aligned_phone_ids.tolist()