# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import gc
import logging
from collections.abc import Sequence

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


# Set up logging configuration
logger = logging.getLogger(__name__)


def _empty_cuda_cache(save_mem: bool) -> None:
    if torch.cuda.is_available() and save_mem:
        torch.cuda.empty_cache()
    return


def init_from_ckpt(self, path):
        # Load the checkpoint
        try:
            sd = torch.load(path, map_location="cpu", weights_only=False)['state_dict']
        except:
            try: 
                sd = torch.load(path, map_location="cpu", weights_only=False)['model']
            except:
                sd = torch.load(path, map_location="cpu", weights_only=False)

        # Remove '.module' prefix from keys if present
        new_sd = {}
        for k, v in sd.items():
            if "loss" in k: continue
            new_key = k.replace("module.", "") if k.startswith("module.") else k
            new_sd[new_key] = v
        
        # Try to load the state dictionary, but allow for dimension mismatches
        try:
            msg = self.load_state_dict(new_sd, strict=True)
        except RuntimeError as e:
            # If error is due to dimension mismatch (2D vs 3D), print warning and skip loading affected layers
            logger.warning(f"Strict load failed: {str(e)}")
            logger.warning("This might be due to loading 2D weights into 3D model or vice versa.")
            logger.warning("Loading weights with strict=False to skip incompatible layers.")
            msg = self.load_state_dict(new_sd, strict=False)

        # Clean up memory
        torch.cuda.empty_cache()

        logger.info(f"Loading pre-trained {self.__class__.__name__}")
        logger.debug("Missing keys: %s", msg.missing_keys)
        logger.debug("Unexpected keys: %s", msg.unexpected_keys)
        logger.info(f"Restored from {path}")

class MaisiGroupNorm3D(nn.GroupNorm):
    """
    Dimension-agnostic GroupNorm that supports both 2D (N,C,H,W) and 3D (N,C,D,H,W).

    Args:
        num_groups: Number of groups.
        num_channels: Number of channels.
        eps: Epsilon for numerical stability.
        affine: Whether to use learnable weight & bias.
        norm_float16: Cast normalized output to fp16.
        print_info: Verbose logging.
        save_mem: Whether to clear CUDA cache to reduce memory.
    """

    def __init__(
        self,
        num_groups: int,
        num_channels: int,
        eps: float = 1e-5,
        affine: bool = True,
        norm_float16: bool = False,
        print_info: bool = False,
        save_mem: bool = True,
    ):
        super().__init__(num_groups, num_channels, eps, affine)
        self.norm_float16 = norm_float16
        self.print_info = print_info
        self.save_mem = save_mem

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if self.print_info:
            logger.info(f"MaisiGroupNorm with input size: {input.size()}")

        if input.ndim not in (4, 5):
            raise ValueError("Expected a 4D tensor (2D data) or 5D tensor (3D data).")

        N, C, *spatial = input.shape
        S = len(spatial)  # S = 2 (2D) or S = 3 (3D)

        # reshape into groups
        input = input.view(N, self.num_groups, C // self.num_groups, *spatial)

        # dimensions to compute mean/std across (all except batch/group)
        # Example:
        # 2D -> mean/std over [2,3,4] since shape is [N, G, Cg, H, W]
        # 3D -> mean/std over [2,3,4,5] since shape is [N, G, Cg, D, H, W]
        reduce_dims = list(range(2, 3 + S))

        outputs = []
        for g in range(self.num_groups):
            arr = input[:, g:g+1, ...].to(torch.float32)
            mean = arr.mean(reduce_dims, keepdim=True)
            std = arr.var(reduce_dims, unbiased=False, keepdim=True).add_(self.eps).sqrt_()

            norm = (arr - mean) / std
            if self.norm_float16:
                norm = norm.to(torch.float16)

            outputs.append(norm)

        del input
        _empty_cuda_cache(self.save_mem)

        # large concatenation helper
        out = (
            torch.cat(outputs, dim=1)
            if max(outputs[0].shape) < 500
            else self._cat_inputs(outputs)
        )

        # reshape back
        out = out.view(N, C, *spatial)

        # affine parameters
        if self.affine:
            # expand affine params to match either 4D or 5D
            shape = [1, C] + [1] * S
            out = out * self.weight.view(*shape) + self.bias.view(*shape)

        if self.print_info:
            logger.info(f"MaisiGroupNorm output size: {out.size()}")

        return out

    def _cat_inputs(self, inputs):
        device_type = inputs[0].device.type
        out = inputs[0].clone().cpu() if device_type == "cuda" else inputs[0].clone()
        inputs[0] = 0
        _empty_cuda_cache(self.save_mem)

        for k in range(len(inputs) - 1):
            out = torch.cat((out, inputs[k + 1].cpu()), dim=1)
            inputs[k + 1] = 0
            _empty_cuda_cache(self.save_mem)
            gc.collect()

            if self.print_info:
                logger.info(f"MaisiGroupNorm concat progress: {k+1}/{len(inputs)-1}")

        return out.to("cuda", non_blocking=True) if device_type == "cuda" else out


class MaisiConvolution(nn.Module):
    """
    Convolutional layer with optional print_info output and custom splitting mechanism.

    Args:
        spatial_dims: Number of spatial dimensions (1D, 2D, 3D).
        in_channels: Number of input channels.
        out_channels: Number of output channels.
        num_splits: Number of splits for the input tensor.
        dim_split: Dimension of splitting for the input tensor.
        print_info: Whether to print information.
        save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`.
        Additional arguments for the convolution operation.
        https://docs.monai.io/en/stable/networks.html#convolution
    """

    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        out_channels: int,
        num_splits: int,
        dim_split: int,
        print_info: bool,
        save_mem: bool = True,
        strides: Sequence[int] | int = 1,
        kernel_size: Sequence[int] | int = 3,
        adn_ordering: str = "NDA",
        act: tuple | str | None = "PRELU",
        norm: tuple | str | None = "INSTANCE",
        dropout: tuple | str | float | None = None,
        dropout_dim: int = 1,
        dilation: Sequence[int] | int = 1,
        groups: int = 1,
        bias: bool = True,
        conv_only: bool = False,
        is_transposed: bool = False,
        padding: Sequence[int] | int | None = None,
        output_padding: Sequence[int] | int | None = None,
    ) -> None:
        super().__init__()
        from monai.networks.blocks import Convolution
        self.conv = Convolution(
            spatial_dims=spatial_dims,
            in_channels=in_channels,
            out_channels=out_channels,
            strides=strides,
            kernel_size=kernel_size,
            adn_ordering=adn_ordering,
            act=act,
            norm=norm,
            dropout=dropout,
            dropout_dim=dropout_dim,
            dilation=dilation,
            groups=groups,
            bias=bias,
            conv_only=conv_only,
            is_transposed=is_transposed,
            padding=padding,
            output_padding=output_padding,
        )

        self.dim_split = dim_split
        self.stride = strides[self.dim_split] if isinstance(strides, list) else strides
        self.num_splits = num_splits
        self.print_info = print_info
        self.save_mem = save_mem

    def _split_tensor(self, x: torch.Tensor, split_size: int, padding: int) -> list[torch.Tensor]:
        overlaps = [0] + [padding] * (self.num_splits - 1)
        last_padding = x.size(self.dim_split + 2) % split_size

        slices = [slice(None)] * len(x.shape)
        splits: list[torch.Tensor] = []
        for i in range(self.num_splits):
            slices[self.dim_split + 2] = slice(
                i * split_size - overlaps[i],
                (i + 1) * split_size + (padding if i != self.num_splits - 1 else last_padding),
            )
            splits.append(x[tuple(slices)])

        if self.print_info:
            for j in range(len(splits)):
                logger.info(f"Split {j + 1}/{len(splits)} size: {splits[j].size()}")

        return splits

    def _concatenate_tensors(self, outputs: list[torch.Tensor], split_size: int, padding: int) -> torch.Tensor:
        slices = [slice(None)] * len(outputs[0].shape)
        for i in range(self.num_splits):
            slices[self.dim_split + 2] = slice(None, split_size) if i == 0 else slice(padding, padding + split_size)
            outputs[i] = outputs[i][tuple(slices)]

        if self.print_info:
            for i in range(self.num_splits):
                logger.info(f"Output {i + 1}/{len(outputs)} size after: {outputs[i].size()}")

        if max(outputs[0].size()) < 500:
            x = torch.cat(outputs, dim=self.dim_split + 2)
        else:
            x = outputs[0].clone().to("cpu", non_blocking=True)
            outputs[0] = torch.tensor(0)
            _empty_cuda_cache(self.save_mem)
            for k in range(len(outputs) - 1):
                x = torch.cat((x, outputs[k + 1].cpu()), dim=self.dim_split + 2)
                outputs[k + 1] = torch.tensor(0)
                _empty_cuda_cache(self.save_mem)
                gc.collect()
                if self.print_info:
                    logger.info(f"MaisiConvolution concat progress: {k + 1}/{len(outputs) - 1}.")

            x = x.to("cuda", non_blocking=True)
        return x

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.print_info:
            logger.info(f"Number of splits: {self.num_splits}")

        # number of spatial dims: 2D=2, 3D=3
        spatial_dims = x.ndim - 2

        # requested split dimension must map into the spatial dimensions
        # clamp so it works for both 2D and 3D
        effective_dim = 2 + min(self.dim_split, spatial_dims - 1)

        # compute size of splits
        l = x.size(effective_dim)
        split_size = l // self.num_splits

        # update padding length if necessary
        padding = 3
        if padding % self.stride > 0:
            padding = (padding // self.stride + 1) * self.stride
        if self.print_info:
            logger.info(f"Padding size: {padding}")

        # split tensor into a list
        splits = self._split_tensor(x, split_size, padding)

        del x
        _empty_cuda_cache(self.save_mem)

        # convolution
        outputs = [self.conv(split) for split in splits]
        if self.print_info:
            for j, out in enumerate(outputs):
                logger.info(f"Output {j + 1}/{len(outputs)} size before: {out.size()}")

        # update size of splits and padding length for output
        split_size_out = split_size
        padding_s = padding

        # same adjustment as your original code but with general dims
        # non_dim_split selects a different spatial dimension
        non_dim_split = (self.dim_split + 1) if self.dim_split < (spatial_dims - 1) else 0
        non_dim_split += 2  # shift to correct axis index


        # detect up/downsampling effects
        if outputs[0].size(non_dim_split) // splits[0].size(non_dim_split) == 2:
            split_size_out *= 2
            padding_s *= 2
        elif splits[0].size(non_dim_split) // outputs[0].size(non_dim_split) == 2:
            split_size_out //= 2
            padding_s //= 2

        # concatenate pieces back together
        x = self._concatenate_tensors(outputs, split_size_out, padding_s)

        del outputs
        _empty_cuda_cache(self.save_mem)

        return x


class MaisiUpsample(nn.Module):
    """
    Convolution-based upsampling layer.

    Args:
        spatial_dims: Number of spatial dimensions (1D, 2D, 3D).
        in_channels: Number of input channels to the layer.
        use_convtranspose: If True, use ConvTranspose to upsample feature maps in decoder.
        num_splits: Number of splits for the input tensor.
        dim_split: Dimension of splitting for the input tensor.
        print_info: Whether to print information.
        save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`.
    """

    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        use_convtranspose: bool,
        num_splits: int,
        dim_split: int,
        print_info: bool,
        save_mem: bool = True,
    ) -> None:
        super().__init__()
        self.conv = MaisiConvolution(
            spatial_dims=spatial_dims,
            in_channels=in_channels,
            out_channels=in_channels,
            strides=2 if use_convtranspose else 1,
            kernel_size=3,
            padding=1,
            conv_only=True,
            is_transposed=use_convtranspose,
            num_splits=num_splits,
            dim_split=dim_split,
            print_info=print_info,
            save_mem=save_mem,
        )
        self.use_convtranspose = use_convtranspose
        self.save_mem = save_mem

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        from monai.utils.type_conversion import convert_to_tensor
        if self.use_convtranspose:
            x = self.conv(x)
            x_tensor: torch.Tensor = convert_to_tensor(x)
            return x_tensor

        x = F.interpolate(x, scale_factor=2.0, mode="trilinear" if x.ndim == 5 else "bilinear")
        _empty_cuda_cache(self.save_mem)
        x = self.conv(x)
        _empty_cuda_cache(self.save_mem)

        out_tensor: torch.Tensor = convert_to_tensor(x)
        return out_tensor


class MaisiDownsample(nn.Module):
    """
    Convolution-based downsampling layer.

    Args:
        spatial_dims: Number of spatial dimensions (1D, 2D, 3D).
        in_channels: Number of input channels.
        num_splits: Number of splits for the input tensor.
        dim_split: Dimension of splitting for the input tensor.
        print_info: Whether to print information.
        save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`.
    """

    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        num_splits: int,
        dim_split: int,
        print_info: bool,
        save_mem: bool = True,
    ) -> None:
        super().__init__()
        self.pad = (0, 1) * spatial_dims
        self.conv = MaisiConvolution(
            spatial_dims=spatial_dims,
            in_channels=in_channels,
            out_channels=in_channels,
            strides=2,
            kernel_size=3,
            padding=0,
            conv_only=True,
            num_splits=num_splits,
            dim_split=dim_split,
            print_info=print_info,
            save_mem=save_mem,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = F.pad(x, self.pad, mode="constant", value=0.0)
        x = self.conv(x)
        return x


class MaisiResBlock(nn.Module):
    """
    Residual block consisting of a cascade of 2 convolutions + activation + normalisation block, and a
    residual connection between input and output.

    Args:
        spatial_dims: Number of spatial dimensions (1D, 2D, 3D).
        in_channels: Input channels to the layer.
        norm_num_groups: Number of groups for the group norm layer.
        norm_eps: Epsilon for the normalization.
        out_channels: Number of output channels.
        num_splits: Number of splits for the input tensor.
        dim_split: Dimension of splitting for the input tensor.
        norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`.
        print_info: Whether to print information, default to `False`.
        save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`.
    """

    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        norm_num_groups: int,
        norm_eps: float,
        out_channels: int,
        num_splits: int,
        dim_split: int,
        norm_float16: bool = False,
        print_info: bool = False,
        save_mem: bool = True,
    ) -> None:
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = in_channels if out_channels is None else out_channels
        self.save_mem = save_mem

        self.norm1 = MaisiGroupNorm3D(
            num_groups=norm_num_groups,
            num_channels=in_channels,
            eps=norm_eps,
            affine=True,
            norm_float16=norm_float16,
            print_info=print_info,
            save_mem=save_mem,
        )
        self.conv1 = MaisiConvolution(
            spatial_dims=spatial_dims,
            in_channels=self.in_channels,
            out_channels=self.out_channels,
            strides=1,
            kernel_size=3,
            padding=1,
            conv_only=True,
            num_splits=num_splits,
            dim_split=dim_split,
            print_info=print_info,
            save_mem=save_mem,
        )
        self.norm2 = MaisiGroupNorm3D(
            num_groups=norm_num_groups,
            num_channels=out_channels,
            eps=norm_eps,
            affine=True,
            norm_float16=norm_float16,
            print_info=print_info,
            save_mem=save_mem,
        )
        self.conv2 = MaisiConvolution(
            spatial_dims=spatial_dims,
            in_channels=self.out_channels,
            out_channels=self.out_channels,
            strides=1,
            kernel_size=3,
            padding=1,
            conv_only=True,
            num_splits=num_splits,
            dim_split=dim_split,
            print_info=print_info,
            save_mem=save_mem,
        )

        self.nin_shortcut = (
            MaisiConvolution(
                spatial_dims=spatial_dims,
                in_channels=self.in_channels,
                out_channels=self.out_channels,
                strides=1,
                kernel_size=1,
                padding=0,
                conv_only=True,
                num_splits=num_splits,
                dim_split=dim_split,
                print_info=print_info,
                save_mem=save_mem,
            )
            if self.in_channels != self.out_channels
            else nn.Identity()
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        from monai.utils.type_conversion import convert_to_tensor
        h = self.norm1(x)
        _empty_cuda_cache(self.save_mem)

        h = F.silu(h)
        _empty_cuda_cache(self.save_mem)
        h = self.conv1(h)
        _empty_cuda_cache(self.save_mem)

        h = self.norm2(h)
        _empty_cuda_cache(self.save_mem)

        h = F.silu(h)
        _empty_cuda_cache(self.save_mem)
        h = self.conv2(h)
        _empty_cuda_cache(self.save_mem)

        if self.in_channels != self.out_channels:
            x = self.nin_shortcut(x)
            _empty_cuda_cache(self.save_mem)

        out = x + h
        out_tensor: torch.Tensor = convert_to_tensor(out)
        return out_tensor


class MaisiEncoder(nn.Module):
    """
    Convolutional cascade that downsamples the image into a spatial latent space.

    Args:
        spatial_dims: Number of spatial dimensions (1D, 2D, 3D).
        in_channels: Number of input channels.
        num_channels: Sequence of block output channels.
        out_channels: Number of channels in the bottom layer (latent space) of the autoencoder.
        num_res_blocks: Number of residual blocks (see AEKLResBlock) per level.
        norm_num_groups: Number of groups for the group norm layers.
        norm_eps: Epsilon for the normalization.
        attention_levels: Indicate which level from num_channels contain an attention block.
        with_nonlocal_attn: If True, use non-local attention block.
        include_fc: whether to include the final linear layer in the attention block. Default to False.
        use_combined_linear: whether to use a single linear layer for qkv projection in the attention block, default to False.
        use_flash_attention: If True, use flash attention for a memory efficient attention mechanism.
        num_splits: Number of splits for the input tensor.
        dim_split: Dimension of splitting for the input tensor.
        norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`.
        print_info: Whether to print information, default to `False`.
        save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`.
    """

    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        num_channels: Sequence[int],
        out_channels: int,
        num_res_blocks: Sequence[int],
        norm_num_groups: int,
        norm_eps: float,
        attention_levels: Sequence[bool],
        num_splits: int,
        dim_split: int,
        norm_float16: bool = False,
        double_z: bool = True,
        print_info: bool = False,
        save_mem: bool = True,
        with_nonlocal_attn: bool = True,
        include_fc: bool = False,
        use_combined_linear: bool = False,
        use_flash_attention: bool = False,
    ) -> None:
        super().__init__()
        from monai.networks.blocks.spatialattention import SpatialAttentionBlock
        from monai.networks.nets.autoencoderkl import AEKLResBlock

        # Check if attention_levels and num_channels have the same size
        if len(attention_levels) != len(num_channels):
            raise ValueError("attention_levels and num_channels must have the same size")

        # Check if num_res_blocks and num_channels have the same size
        if len(num_res_blocks) != len(num_channels):
            raise ValueError("num_res_blocks and num_channels must have the same size")

        self.vae_stride = 2 ** (len(num_channels) - 1)
        self.dims = spatial_dims

        self.save_mem = save_mem

        self.z_channels = out_channels
        self.double_z = double_z

        blocks: list[nn.Module] = []

        blocks.append(
            MaisiConvolution(
                spatial_dims=spatial_dims,
                in_channels=in_channels,
                out_channels=num_channels[0],
                strides=1,
                kernel_size=3,
                padding=1,
                conv_only=True,
                num_splits=num_splits,
                dim_split=dim_split,
                print_info=print_info,
                save_mem=save_mem,
            )
        )

        output_channel = num_channels[0]
        for i in range(len(num_channels)):
            input_channel = output_channel
            output_channel = num_channels[i]
            is_final_block = i == len(num_channels) - 1

            for _ in range(num_res_blocks[i]):
                blocks.append(
                    MaisiResBlock(
                        spatial_dims=spatial_dims,
                        in_channels=input_channel,
                        norm_num_groups=norm_num_groups,
                        norm_eps=norm_eps,
                        out_channels=output_channel,
                        num_splits=num_splits,
                        dim_split=dim_split,
                        norm_float16=norm_float16,
                        print_info=print_info,
                        save_mem=save_mem,
                    )
                )
                input_channel = output_channel
                if attention_levels[i]:
                    blocks.append(
                        SpatialAttentionBlock(
                            spatial_dims=spatial_dims,
                            num_channels=input_channel,
                            norm_num_groups=norm_num_groups,
                            norm_eps=norm_eps,
                            include_fc=include_fc,
                            use_combined_linear=use_combined_linear,
                            use_flash_attention=use_flash_attention,
                        )
                    )

            if not is_final_block:
                blocks.append(
                    MaisiDownsample(
                        spatial_dims=spatial_dims,
                        in_channels=input_channel,
                        num_splits=num_splits,
                        dim_split=dim_split,
                        print_info=print_info,
                        save_mem=save_mem,
                    )
                )

        if with_nonlocal_attn:
            blocks.append(
                AEKLResBlock(
                    spatial_dims=spatial_dims,
                    in_channels=num_channels[-1],
                    norm_num_groups=norm_num_groups,
                    norm_eps=norm_eps,
                    out_channels=num_channels[-1],
                )
            )

            blocks.append(
                SpatialAttentionBlock(
                    spatial_dims=spatial_dims,
                    num_channels=num_channels[-1],
                    norm_num_groups=norm_num_groups,
                    norm_eps=norm_eps,
                    include_fc=include_fc,
                    use_combined_linear=use_combined_linear,
                    use_flash_attention=use_flash_attention,
                )
            )
            blocks.append(
                AEKLResBlock(
                    spatial_dims=spatial_dims,
                    in_channels=num_channels[-1],
                    norm_num_groups=norm_num_groups,
                    norm_eps=norm_eps,
                    out_channels=num_channels[-1],
                )
            )

        blocks.append(
            MaisiGroupNorm3D(
                num_groups=norm_num_groups,
                num_channels=num_channels[-1],
                eps=norm_eps,
                affine=True,
                norm_float16=norm_float16,
                print_info=print_info,
                save_mem=save_mem,
            )
        )
        blocks.append(
            MaisiConvolution(
                spatial_dims=spatial_dims,
                in_channels=num_channels[-1],
                out_channels=out_channels * 2 if self.double_z else out_channels,
                strides=1,
                kernel_size=3,
                padding=1,
                conv_only=True,
                num_splits=num_splits,
                dim_split=dim_split,
                print_info=print_info,
                save_mem=save_mem,
            )
        )

        self.blocks = nn.ModuleList(blocks)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for block in self.blocks:
            x = block(x)
            _empty_cuda_cache(self.save_mem)
        return x


class MaisiDecoder(nn.Module):
    """
    Convolutional cascade upsampling from a spatial latent space into an image space.

    Args:
        spatial_dims: Number of spatial dimensions (1D, 2D, 3D).
        num_channels: Sequence of block output channels.
        in_channels: Number of channels in the bottom layer (latent space) of the autoencoder.
        out_channels: Number of output channels.
        num_res_blocks: Number of residual blocks (see AEKLResBlock) per level.
        norm_num_groups: Number of groups for the group norm layers.
        norm_eps: Epsilon for the normalization.
        attention_levels: Indicate which level from num_channels contain an attention block.
        with_nonlocal_attn: If True, use non-local attention block.
        include_fc: whether to include the final linear layer in the attention block. Default to False.
        use_combined_linear: whether to use a single linear layer for qkv projection in the attention block, default to False.
        use_flash_attention: If True, use flash attention for a memory efficient attention mechanism.
        use_convtranspose: If True, use ConvTranspose to upsample feature maps in decoder.
        num_splits: Number of splits for the input tensor.
        dim_split: Dimension of splitting for the input tensor.
        norm_float16: If True, convert output of MaisiGroupNorm3D to float16 format, default to `False`.
        print_info: Whether to print information, default to `False`.
        save_mem: Whether to clean CUDA cache in order to save GPU memory, default to `True`.
    """

    def __init__(
        self,
        spatial_dims: int,
        num_channels: Sequence[int],
        in_channels: int,
        out_channels: int,
        num_res_blocks: Sequence[int],
        norm_num_groups: int,
        norm_eps: float,
        attention_levels: Sequence[bool],
        num_splits: int,
        dim_split: int,
        norm_float16: bool = False,
        double_z: bool = True,
        print_info: bool = False,
        save_mem: bool = True,
        with_nonlocal_attn: bool = True,
        include_fc: bool = False,
        use_combined_linear: bool = False,
        use_flash_attention: bool = False,
        use_convtranspose: bool = False,
    ) -> None:
        super().__init__()
        from monai.networks.blocks.spatialattention import SpatialAttentionBlock
        from monai.networks.nets.autoencoderkl import AEKLResBlock
        self.print_info = print_info
        self.save_mem = save_mem
        self.double_z = double_z

        reversed_block_out_channels = list(reversed(num_channels))

        blocks: list[nn.Module] = []

        blocks.append(
            MaisiConvolution(
                spatial_dims=spatial_dims,
                in_channels=in_channels,
                out_channels=reversed_block_out_channels[0],
                strides=1,
                kernel_size=3,
                padding=1,
                conv_only=True,
                num_splits=num_splits,
                dim_split=dim_split,
                print_info=print_info,
                save_mem=save_mem,
            )
        )

        if with_nonlocal_attn:
            blocks.append(
                AEKLResBlock(
                    spatial_dims=spatial_dims,
                    in_channels=reversed_block_out_channels[0],
                    norm_num_groups=norm_num_groups,
                    norm_eps=norm_eps,
                    out_channels=reversed_block_out_channels[0],
                )
            )
            blocks.append(
                SpatialAttentionBlock(
                    spatial_dims=spatial_dims,
                    num_channels=reversed_block_out_channels[0],
                    norm_num_groups=norm_num_groups,
                    norm_eps=norm_eps,
                    include_fc=include_fc,
                    use_combined_linear=use_combined_linear,
                    use_flash_attention=use_flash_attention,
                )
            )
            blocks.append(
                AEKLResBlock(
                    spatial_dims=spatial_dims,
                    in_channels=reversed_block_out_channels[0],
                    norm_num_groups=norm_num_groups,
                    norm_eps=norm_eps,
                    out_channels=reversed_block_out_channels[0],
                )
            )

        reversed_attention_levels = list(reversed(attention_levels))
        reversed_num_res_blocks = list(reversed(num_res_blocks))
        block_out_ch = reversed_block_out_channels[0]
        for i in range(len(reversed_block_out_channels)):
            block_in_ch = block_out_ch
            block_out_ch = reversed_block_out_channels[i]
            is_final_block = i == len(num_channels) - 1

            for _ in range(reversed_num_res_blocks[i]):
                blocks.append(
                    MaisiResBlock(
                        spatial_dims=spatial_dims,
                        in_channels=block_in_ch,
                        norm_num_groups=norm_num_groups,
                        norm_eps=norm_eps,
                        out_channels=block_out_ch,
                        num_splits=num_splits,
                        dim_split=dim_split,
                        norm_float16=norm_float16,
                        print_info=print_info,
                        save_mem=save_mem,
                    )
                )
                block_in_ch = block_out_ch

                if reversed_attention_levels[i]:
                    blocks.append(
                        SpatialAttentionBlock(
                            spatial_dims=spatial_dims,
                            num_channels=block_in_ch,
                            norm_num_groups=norm_num_groups,
                            norm_eps=norm_eps,
                            include_fc=include_fc,
                            use_combined_linear=use_combined_linear,
                            use_flash_attention=use_flash_attention,
                        )
                    )

            if not is_final_block:
                blocks.append(
                    MaisiUpsample(
                        spatial_dims=spatial_dims,
                        in_channels=block_in_ch,
                        use_convtranspose=use_convtranspose,
                        num_splits=num_splits,
                        dim_split=dim_split,
                        print_info=print_info,
                        save_mem=save_mem,
                    )
                )

        blocks.append(
            MaisiGroupNorm3D(
                num_groups=norm_num_groups,
                num_channels=block_in_ch,
                eps=norm_eps,
                affine=True,
                norm_float16=norm_float16,
                print_info=print_info,
                save_mem=save_mem,
            )
        )
        blocks.append(
            MaisiConvolution(
                spatial_dims=spatial_dims,
                in_channels=block_in_ch,
                out_channels=out_channels,
                strides=1,
                kernel_size=3,
                padding=1,
                conv_only=True,
                num_splits=num_splits,
                dim_split=dim_split,
                print_info=print_info,
                save_mem=save_mem,
            )
        )

        self.blocks = nn.ModuleList(blocks)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for block in self.blocks:
            x = block(x)
            _empty_cuda_cache(self.save_mem)
        return x