import sys
import types

import torch


def maximum_path(score: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    path = torch.zeros_like(score)
    for batch in range(score.shape[0]):
        mel_length = int(mask[batch].any(-1).sum())
        text_length = int(mask[batch].any(0).sum())
        for frame in range(mel_length):
            token = min(frame * text_length // mel_length, text_length - 1)
            path[batch, frame, token] = 1
    return path


alignment_stub = types.ModuleType(
    "style_bert_vits2.models.monotonic_alignment"
)
alignment_stub.maximum_path = maximum_path
sys.modules["style_bert_vits2.models.monotonic_alignment"] = alignment_stub

from style_bert_vits2.models.mel_synthesizer import (
    JointMelSynthesizer,
    MelAcousticModel,
    MelVocoder,
)


class DummyEncoder(torch.nn.Module):
    def __init__(self, hidden: int, n_mels: int) -> None:
        super().__init__()
        self.embedding = torch.nn.Embedding(16, hidden)
        self.projection = torch.nn.Conv1d(hidden, n_mels * 2, 1)

    def forward(self, x, x_lengths, *_args, g=None):
        hidden = self.embedding(x).transpose(1, 2)
        mask = (
            torch.arange(x.shape[1], device=x.device)[None, :] < x_lengths[:, None]
        ).unsqueeze(1).to(hidden)
        mu, logs = self.projection(hidden * mask).chunk(2, dim=1)
        return hidden, mu * mask, logs.tanh() * mask, mask


class DummyDurationPredictor(torch.nn.Module):
    def __init__(self, hidden: int) -> None:
        super().__init__()
        self.projection = torch.nn.Conv1d(hidden, 1, 1)

    def forward(self, x, mask, g=None):
        return self.projection(x) * mask


class DummyStochasticDurationPredictor(DummyDurationPredictor):
    def forward(self, x, mask, w=None, g=None, reverse=False, noise_scale=1.0):
        prediction = super().forward(x, mask, g)
        if reverse:
            return prediction
        return ((prediction - torch.log(w + 1e-6)) ** 2 * mask).sum()


class DummyVocoder(torch.nn.Module):
    def __init__(self, n_mels: int) -> None:
        super().__init__()
        self.projection = torch.nn.Conv1d(n_mels, 1, 1)

    def forward(self, mel, g=None):
        return self.projection(mel)


def make_acoustic() -> MelAcousticModel:
    hidden, n_mels, speaker_channels = 8, 4, 2
    return MelAcousticModel(
        DummyEncoder(hidden, n_mels),
        DummyStochasticDurationPredictor(hidden),
        DummyDurationPredictor(hidden),
        n_mel_channels=n_mels,
        n_speakers=2,
        gin_channels=speaker_channels,
        hidden_channels=hidden,
        matcha_channels=8,
        matcha_num_heads=2,
        matcha_dropout=0.0,
        matcha_sigma_min=1e-4,
        matcha_n_timesteps=2,
        matcha_use_diff_attention=False,
    )


def test_acoustic_model_outputs_mel_and_losses() -> None:
    model = make_acoustic()
    x = torch.tensor([[1, 2, 3], [1, 2, 0]])
    x_lengths = torch.tensor([3, 2])
    mel = torch.randn(2, 4, 7)
    mel_lengths = torch.tensor([7, 5])
    sid = torch.tensor([0, 1])

    output = model(
        x, x_lengths, mel, mel_lengths, sid, out_size=4, generate=True
    )
    loss = (
        output["duration_loss"] + output["prior_loss"] + output["flow_loss"]
        + output["generated_mel"].abs().mean()
    )
    loss.backward()

    assert output["generated_mel"].shape == (2, 4, 4)
    assert output["target_mel"].shape == (2, 4, 4)
    assert all(torch.isfinite(output[name]) for name in (
        "duration_loss", "prior_loss", "flow_loss"
    ))
    assert model.enc_p.embedding.weight.grad is not None
    assert any(p.grad is not None for p in model.matcha.parameters())


def test_joint_model_backpropagates_through_vocoder_to_acoustic() -> None:
    acoustic = make_acoustic()
    vocoder = MelVocoder(DummyVocoder(4), n_speakers=2, gin_channels=2)
    model = JointMelSynthesizer(acoustic, vocoder)
    waveform, output = model(
        torch.tensor([[1, 2, 3]]),
        torch.tensor([3]),
        torch.randn(1, 4, 6),
        torch.tensor([6]),
        torch.tensor([0]),
        out_size=4,
        n_timesteps=2,
    )

    waveform.square().mean().backward()

    assert waveform.shape == (1, 1, 4)
    assert model.vocoder.generator.projection.weight.grad is not None
    assert any(p.grad is not None for p in acoustic.matcha.parameters())
    assert output["generated_mel"].grad_fn is not None