import torch
import torch.nn as nn
from torchvision import models


class BoneAgeModel(nn.Module):
    def __init__(self):
        super(BoneAgeModel, self).__init__()
        self.backbone = models.efficientnet_b4(weights=None)
        num_features = self.backbone.classifier[1].in_features
        self.backbone.classifier = nn.Identity()
        self.gender_fc = nn.Linear(1, 32)
        self.fc = nn.Sequential(
            nn.Linear(num_features + 32, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 128),
            nn.ReLU(),
            nn.Linear(128, 1)
        )

    def forward(self, image, gender):
        img_features = self.backbone(image)
        gender_features = torch.relu(self.gender_fc(gender.unsqueeze(1)))
        combined = torch.cat([img_features, gender_features], dim=1)
        return self.fc(combined).squeeze(1)