"""
infrastructure/model/vgtlnet.py
────────────────────────────────
VGTL-Net Model Architecture.
Strictly defines the PyTorch BP prediction architecture (SRP). No signal preprocessing.
"""
from __future__ import annotations


def build_bp_mlp(in_features: int):
    """
    MLP head for SBP or DBP prediction.
    """
    import torch.nn as nn
    return nn.Sequential(
        nn.Linear(in_features, 1024),
        nn.BatchNorm1d(1024),
        nn.ReLU(inplace=True),
        nn.Dropout(0.3),
        nn.Linear(1024, 512),
        nn.BatchNorm1d(512),
        nn.ReLU(inplace=True),
        nn.Dropout(0.2),
        nn.Linear(512, 1),
    )


def build_convnextv2_bp_model(pretrained: bool = False):
    """
    Build ConvNeXtV2BPModel (VGTL-Net backbone + dual MLP head).
    """
    try:
        import timm
        import torch.nn as nn

        class ConvNeXtV2BPModel(nn.Module):
            """VGTL-Net: ConvNeXt V2 Tiny + Dual MLP Head for SBP/DBP."""

            def __init__(self, pretrained: bool = False):
                super().__init__()
                self.feature_extractor = timm.create_model(
                    "convnextv2_tiny.fcmae_ft_in22k_in1k",
                    pretrained=pretrained,
                    num_classes=0,
                    global_pool="avg",
                )
                feat_dim = self.feature_extractor.num_features  # 768
                self.mlp_sbp = build_bp_mlp(feat_dim)
                self.mlp_dbp = build_bp_mlp(feat_dim)

            def forward(self, x):
                """
                Args:
                    x: (B, 3, 224, 224) visibility graph image tensor
                Returns:
                    Tuple (sbp_pred, dbp_pred)
                """
                feat = self.feature_extractor(x)
                return self.mlp_sbp(feat).squeeze(-1), self.mlp_dbp(feat).squeeze(-1)

        return ConvNeXtV2BPModel(pretrained=pretrained)

    except ImportError as e:
        raise RuntimeError(
            f"Dependencies for VGTL-Net model are missing: {e}. "
            "Please run: pip install timm"
        ) from e