import torch
import torch.nn as nn

from .duration_predictor import TTSDurationModel


class DPNetwork(TTSDurationModel):
    """Inheritance keeps state_dict keys (`sentence_encoder.*` / `predictor.*`) aligned with ONNX/checkpoints."""

    def __init__(
        self,
        vocab_size: int = 37,
        latent_channels: int = 144,
        style_dp: int = 8,
        style_dim: int = 16,
        sentence_encoder_cfg: dict = None,
        style_encoder_cfg: dict = None,
        predictor_cfg: dict = None,
    ):
        super().__init__(
            vocab_size=vocab_size,
            style_dp=style_dp,
            style_dim=style_dim,
            ref_in_channels=latent_channels,
            sentence_encoder_cfg=sentence_encoder_cfg,
            style_encoder_cfg=style_encoder_cfg,
            predictor_cfg=predictor_cfg,
        )

    def forward(
        self,
        text_ids: torch.Tensor,
        z_ref: torch.Tensor | None = None,
        text_mask: torch.Tensor | None = None,
        ref_mask: torch.Tensor | None = None,
        style_dp: torch.Tensor | None = None,
        return_log: bool = False,
    ) -> torch.Tensor:
        if text_mask is not None and text_mask.dtype != torch.float32:
            text_mask = text_mask.float()

        if ref_mask is not None and ref_mask.dtype != torch.float32:
            ref_mask = ref_mask.float()
        elif ref_mask is None and z_ref is not None:
            B, C, T_ref = z_ref.shape
            ref_mask = torch.ones(B, 1, T_ref, device=z_ref.device, dtype=torch.float32)

        return super().forward(
            text_ids,
            z_ref=z_ref,
            text_mask=text_mask,
            ref_mask=ref_mask,
            style_dp=style_dp,
            return_log=return_log,
        )