import torch

from style_bert_vits2.models.matcha_flow import (
    DifferentialAttentionV2,
    MatchaFlow,
    configure_matcha_only_training,
)


class DummySynthesizer(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.use_matcha = True
        self.legacy = torch.nn.Linear(4, 4)
        self.matcha = torch.nn.Linear(4, 4)


def test_matcha_only_freezes_legacy_parameters() -> None:
    model = DummySynthesizer()

    trainable = configure_matcha_only_training(model)

    assert trainable == sum(parameter.numel() for parameter in model.matcha.parameters())
    assert all(parameter.requires_grad for parameter in model.matcha.parameters())
    assert not any(parameter.requires_grad for parameter in model.legacy.parameters())


def test_differential_attention_v2_on_cpu() -> None:
    attention = DifferentialAttentionV2(channels=16, num_heads=4, dropout=0.0)
    inputs = torch.randn(2, 7, 16, requires_grad=True)
    valid = torch.tensor(
        [[True, True, True, True, True, True, True],
         [True, True, True, True, True, False, False]]
    )

    output = attention(inputs, valid)
    output.square().mean().backward()

    assert output.shape == inputs.shape
    assert output.device.type == "cpu"
    assert torch.count_nonzero(output[1, 5:]) == 0
    assert attention.lambda_projection.weight.grad is not None


def test_matcha_flow_loss_and_sampling_with_odd_length() -> None:
    model = MatchaFlow(
        latent_channels=8,
        speaker_channels=4,
        channels=16,
        num_heads=2,
        dropout=0.0,
        use_diff_attention=True,
    )
    target = torch.randn(2, 8, 7)
    mu = torch.randn_like(target)
    mask = torch.tensor(
        [[[1, 1, 1, 1, 1, 1, 1]], [[1, 1, 1, 1, 1, 0, 0]]],
        dtype=target.dtype,
    )
    speaker = torch.randn(2, 4, 1)

    loss = model.compute_loss(target, mask, mu, speaker)
    loss.backward()

    assert loss.ndim == 0
    assert torch.isfinite(loss)
    assert any(parameter.grad is not None for parameter in model.parameters())

    sample = model(mu, mask, n_timesteps=2, speaker=speaker)
    assert sample.shape == target.shape
    assert torch.count_nonzero(sample[1, :, 5:]) == 0