import torch
import torch.nn as nn
from einops import rearrange

class SEBlock(nn.Module):
    """Squeeze-and-Excitation Block."""
    def __init__(self, channels, reduction=16):
        super(SEBlock, self).__init__()
        self.fc1 = nn.Linear(channels, channels // reduction, bias=False)
        self.fc2 = nn.Linear(channels // reduction, channels, bias=False)
        self.relu = nn.ReLU(inplace=True)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        b, c, d, h, w = x.shape
        y = x.mean(dim=[2, 3, 4])
        y = self.fc1(y)
        y = self.relu(y)
        y = self.fc2(y)
        y = self.sigmoid(y).view(b, c, 1, 1, 1)
        return x * y


class ResNetBlock(nn.Module):
    """Residual Block with convolutional layers and an SE module."""
    def __init__(self, in_channels, out_channels, dropout_rate=0.1):
        super(ResNetBlock, self).__init__()
        self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1)
        self.in1 = nn.InstanceNorm3d(out_channels, affine=True)
        self.prelu = nn.PReLU()
        self.dropout = nn.Dropout3d(dropout_rate)

        self.conv2 = nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1)
        self.in2 = nn.InstanceNorm3d(out_channels, affine=True)

        self.se = SEBlock(out_channels)
        self.residual = (
            nn.Conv3d(in_channels, out_channels, kernel_size=1)
            if in_channels != out_channels else nn.Identity()
        )

    def forward(self, x):
        residual = self.residual(x)
        x = self.conv1(x)
        x = self.in1(x)
        x = self.prelu(x)
        x = self.dropout(x)
        x = self.conv2(x)
        x = self.in2(x)
        x = self.se(x)
        x += residual
        return self.prelu(x)


class Down(nn.Module):
    """Downsampling block with strided convolution, then a ResNetBlock."""
    def __init__(self, in_channels, out_channels, stride=(2,2,2), dropout_rate=0.1):
        super().__init__()
        self.stride_conv = nn.Conv3d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
        self.res_block = ResNetBlock(out_channels, out_channels, dropout_rate)

    def forward(self, x):
        x = self.stride_conv(x)
        x = self.res_block(x)
        return x


class Up(nn.Module):
    """Upsampling block: skip-connection concat, convolutions, then transposed conv."""
    def __init__(self, in_channels, out_channels, upsample_channels,
                 stride=(2,2,2), dropout_rate=0.1):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.InstanceNorm3d(out_channels, affine=True),
            nn.PReLU(),
            nn.Dropout3d(dropout_rate),
            nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.InstanceNorm3d(out_channels, affine=True),
            nn.PReLU(),
            nn.Dropout3d(dropout_rate),
        )
        self.up = nn.ConvTranspose3d(out_channels, upsample_channels,
                                     kernel_size=stride, stride=stride)

    def forward(self, decoder_x, skip_x):
        x = torch.cat([skip_x, decoder_x], dim=1)
        x = self.double_conv(x)
        x = self.up(x)
        return x


class Bottleneck(nn.Module):
    """Bottleneck block with strided conv, ResNetBlock, and transposed conv."""
    def __init__(self, in_channels, out_channels, stride=(2,2,2), dropout_rate=0.1):
        super().__init__()
        self.stride_conv = nn.Conv3d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
        self.res_block = ResNetBlock(out_channels, out_channels, dropout_rate)
        self.up = nn.ConvTranspose3d(out_channels, out_channels, kernel_size=stride, stride=stride)

    def forward(self, x):
        x = self.stride_conv(x)
        x = self.res_block(x)
        x = self.up(x)
        return x


class Final(nn.Module):
    """Final block that merges the top-level up with the inc output, then outputs num_classes."""
    def __init__(self, num_classes):
        super().__init__()
        self.final_conv = nn.Sequential(
            nn.Conv3d(64, 32, kernel_size=3, padding=1),
            nn.InstanceNorm3d(32, affine=True),
            nn.PReLU(),
            nn.Conv3d(32, 32, kernel_size=3, padding=1),
            nn.InstanceNorm3d(32, affine=True),
            nn.PReLU(),
            nn.Conv3d(32, num_classes, kernel_size=1),
        )

    def forward(self, top_x, inc_x):
        x = torch.cat([inc_x, top_x], dim=1)
        return self.final_conv(x)


class NNUNet3DEncoder(nn.Module):
    """
    3D UNet-like architecture with dynamic depth. 
    If depth=4, it replicates the original structure:
      inc -> down1..down4 -> bottleneck -> up4..up1 -> final
      plus 4 auxiliary outputs.
    """
    def __init__(
        self,
        dropout_rate=0.1,
        depth=5,
        stride_list=[[1, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]],
        channels=[32, 64, 128, 256, 320, 512]
    ):
        super().__init__()
        self.depth = depth
        self.channels = channels

        if depth >= len(channels):
            raise ValueError(f"depth={depth} too large for {channels}")

        self.inc = ResNetBlock(1, channels[0], dropout_rate)

        self.down_blocks = nn.ModuleList()
        for i in range(depth):
            self.down_blocks.append(
                Down(channels[i], channels[i+1],
                     stride=stride_list[i], dropout_rate=dropout_rate)
            )

    def forward(self, x, return_encoded_tokens=False):
        """
        Args:
            x: (batch_size, channels, depth, height, width)
            return_encoded_tokens: (bool) if True, return the encoded tokens
        Returns:
            encoded_tokens: (batch_size, depth, height, width, channels)
        """
        x0 = self.inc(x)

        downs = x0
        for i in range(self.depth):
            downs = self.down_blocks[i](downs)

        downs = rearrange(downs, 'b c d h w -> b d h w c')
        if return_encoded_tokens:
            return downs

        return downs