import torch
import torch.nn as nn
import torch.nn.functional as F

class EnvironmentalModule(nn.Module):
    def __init__(self, input_channels, output_channels):
        super(EnvironmentalModule, self).__init__()
        self.conv = nn.Conv2d(input_channels, output_channels, kernel_size=3, padding=1)

    def forward(self, x):
        return self.conv(x)

class SocialImplicitWithEnvironment(nn.Module):
    def __init__(self, social_implicit_model, environmental_module):
        super(SocialImplicitWithEnvironment, self).__init__()
        self.social_implicit_model = social_implicit_model
        self.environmental_module = environmental_module
        self.fc = nn.Linear(2 + 16, 2)

    def forward(self, social_zones, local_features, global_features, segmentation_maps):
        social_implicit_output = self.social_implicit_model(social_zones, local_features, global_features)

        environmental_features = self.environmental_module(segmentation_maps)
        environmental_features = F.avg_pool2d(environmental_features, kernel_size=environmental_features.size()[2:])
        environmental_features = environmental_features.view(environmental_features.size(0), -1)

        combined_features = torch.cat([social_implicit_output, environmental_features], dim=1)
        prediction = self.fc(combined_features)
        return prediction

# Assuming you have the following tensors:
# social_zones, local_features, global_features, segmentation_maps

# Instantiate your original SocialImplicit model and the new EnvironmentalModule
social_implicit_model = SocialImplicit()
environmental_module = EnvironmentalModule(1, 16)

# Instantiate the SocialImplicitWithEnvironment model and pass the required inputs
model = SocialImplicitWithEnvironment(social_implicit_model, environmental_module)
predictions = model(social_zones, local_features, global_features, segmentation_maps)