import torch
import torch.nn as nn
import torch.nn.functional as F
from .base_class_models import BaseModel

class FiLM(nn.Module):
    def __init__(self, num_features, style_dim):
        super(FiLM, self).__init__()
        self.gamma = nn.Linear(style_dim, num_features)
        self.beta = nn.Linear(style_dim, num_features)

    def forward(self, x, style):
        gamma = self.gamma(style).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
        beta = self.beta(style).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
        return gamma * x + beta

class UNetBlock(nn.Module):
    def __init__(self, in_channels, out_channels, style_dim=None):
        super(UNetBlock, self).__init__()
        self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1)
        self.conv2 = nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1)
        self.norm = nn.BatchNorm3d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.film = FiLM(out_channels, style_dim) if style_dim else None

    def forward(self, x, style=None):
        x = self.relu(self.norm(self.conv1(x)))
        x = self.relu(self.norm(self.conv2(x)))
        if self.film:
            x = self.film(x, style)
        return x

class UNet3D(BaseModel):
    def __init__(self, N: int = 128,
                 in_channels: int = 2,
                 out_channels: int = 1,
                 style_dim: int = 2,
                 device: torch.device = torch.device('cpu')):
        """
        3D U-Net model with optional FiLM layers for style conditioning.
        Parameters:
        - N: Size of the input data: data will have the shape (B, C, N, N, N) with B the batch size, C the number of channels.
        - in_channels: Number of input channels (default is 2).
        - out_channels: Number of output channels (default is 1).
        - style_dim: Dimension of the style vector (default is 2).
        - device: Device to load the model onto (default is CPU).
        This model implements a 3D U-Net architecture with downsampling and upsampling blocks.
        The model uses convolutional layers with ReLU activations and batch normalization.
        The FiLM layers are used to condition the feature maps on style parameters.
        """
        
        super().__init__(N=N,
                     in_channels=in_channels,
                     out_channels=out_channels,
                     style_parameters=style_dim,
                     device=device)
        
        self.enc1 = UNetBlock(in_channels, 32, style_dim)
        self.pool1 = nn.MaxPool3d(2)
        self.enc2 = UNetBlock(32, 64, style_dim)
        self.pool2 = nn.MaxPool3d(2)
        self.bottleneck = UNetBlock(64, 128, style_dim)

        self.up2 = nn.ConvTranspose3d(128, 64, kernel_size=2, stride=2)
        self.dec2 = UNetBlock(128, 64)
        self.up1 = nn.ConvTranspose3d(64, 32, kernel_size=2, stride=2)
        self.dec1 = UNetBlock(64, 32)
        self.final = nn.Conv3d(32, out_channels, kernel_size=1)

    def forward(self, x, style):
        e1 = self.enc1(x, style)
        e2 = self.enc2(self.pool1(e1), style)
        b = self.bottleneck(self.pool2(e2), style)
        d2 = self.dec2(torch.cat([self.up2(b), e2], dim=1))
        d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
        return self.final(d1)