"""

Note
----

This file is a pytorch adapted version from the
sources of the SynthSeg project - https://github.com/BBillot/SynthSeg.
All weights and model artchitecture are original from the SynthSeg project.
It remains licensed as the rest of SynthSeg
(Apache 2.0 license as of January 2026).

# ## ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
#
#   See LICENSE.txt file distributed along with the SynthSeg package for the
#   copyright and license terms.
#
# ## ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
"""

import numpy as np
from scipy.ndimage import gaussian_filter

from dipy.data import get_fnames
from dipy.nn.utils import (
    normalize,
    recover_img,
    set_logger_level,
    transform_img,
)
from dipy.segment.utils import remove_holes_and_islands
from dipy.testing.decorators import doctest_skip_parser
from dipy.utils.logging import logger
from dipy.utils.optpkg import optional_package

torch, have_torch, _ = optional_package("torch", min_version="2.2.0")
if have_torch:
    from torch.nn import (
        BatchNorm3d,
        Conv3d,
        MaxPool3d,
        Module,
        ModuleList,
        Softmax,
        Upsample,
    )
    import torch.nn.functional as F
else:

    class Module:
        pass


class Conv3dELU(Module):
    """
    Mimics TensorFlow Conv3D + ELU fused behavior.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        padding=0,
        stride=1,
        dilation=1,
        bias=True,
    ):
        super().__init__()
        self.conv = Conv3d(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            padding=padding,
            stride=stride,
            dilation=dilation,
            bias=bias,
        )

    def forward(self, x):
        """Forward pass of the Conv3dELU layer.

        Parameters
        ----------
        x : torch.Tensor
            Input tensor.

        Returns
        -------
        torch.Tensor
            Output tensor after convolution and ELU activation.
        """
        x = self.conv(x)
        x = F.elu(x)
        return x


class Block(Module):
    """Building block for the SynthSeg model.

    Parameters
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    kernel_size : int
        Size of the convolutional kernel.
    padding : int
        Padding of the convolution.
    n_layers : int
        Number of convolutional layers in the block.
    layer_type : str
        Type of the block: 'down', 'down2', or 'up'.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        padding,
        n_layers,
        layer_type,
    ):
        super(Block, self).__init__()
        self.n_layers = n_layers
        self.layer_list = ModuleList()
        self.layer_type = layer_type

        cur_channel = in_channels
        if self.layer_type == "up":
            self.layer_list.append(Upsample(scale_factor=2))
            cur_channel = int(cur_channel * 1.5)

        for _ in range(n_layers):
            self.layer_list.append(
                Conv3dELU(
                    cur_channel,
                    out_channels,
                    kernel_size,
                    padding=padding,
                )
            )
            cur_channel = out_channels

        self.layer_list.append(BatchNorm3d(out_channels, eps=1e-3, momentum=0.01))
        if self.layer_type == "down":
            self.layer_list.append(MaxPool3d(2, stride=2))

    def forward(self, input, passed):
        """Forward pass of the Block.

        Parameters
        ----------
        input : torch.Tensor
            Input tensor.
        passed : torch.Tensor
            Tensor from skipped connection.

        Returns
        -------
        x : torch.Tensor
            Output of the convolutional layers.
        skip : torch.Tensor, optional
            Output of the convolutional layers before downsampling.
            Only returned if layer_type is 'down'.
        """
        x = input
        for l_idx, layer in enumerate(self.layer_list):
            x = layer(x)
            if l_idx == 0 and self.layer_type == "up":
                x = torch.cat([passed, x], dim=1)
            if l_idx == self.n_layers - 1 and self.layer_type == "down":
                skip = x

        if self.layer_type == "down":
            return x, skip
        return x


class Model(Module):
    """The Model class for the SynthSeg architecture.

    Parameters
    ----------
    model_scale : int, optional
        The scale of the model.
    n_levels : int, optional
        Number of levels in the U-Net.
    output_channels : int, optional
        Number of output channels.
    """

    def __init__(self, *, model_scale=24, n_levels=5, output_channels=33):
        super(Model, self).__init__()

        self.block_list = ModuleList()
        self.channels = [model_scale * (2**i) for i in range(n_levels)]
        self.channels_rev = [
            model_scale * (2 ** (n_levels - i - 1)) for i in range(n_levels)
        ]
        self.channels.insert(0, 1)
        self.channels_rev.append(output_channels)

        self.model_scale = model_scale
        self.n_levels = n_levels

        # Block structure
        for level in range(n_levels - 1):
            block = Block(
                self.channels[level],
                self.channels[level + 1],
                kernel_size=3,
                padding=1,
                n_layers=2,
                layer_type="down",
            )
            self.block_list.append(block)

        level = n_levels - 1
        block = Block(
            self.channels[level],
            self.channels[level + 1],
            kernel_size=3,
            padding=1,
            n_layers=2,
            layer_type="down2",
        )
        self.block_list.append(block)

        for level in range(n_levels - 1):
            block = Block(
                self.channels_rev[level],
                self.channels_rev[level + 1],
                kernel_size=3,
                padding=1,
                n_layers=2,
                layer_type="up",
            )
            self.block_list.append(block)

        self.conv_pred = Conv3d(model_scale, 33, 1, padding=0)
        self.softmax = Softmax(dim=1)

    def forward(self, inputs):
        """Forward pass of the SynthSeg model.

        Parameters
        ----------
        inputs : torch.Tensor
            Input tensor.

        Returns
        -------
        torch.Tensor
            Predicted brain segmentation probability map.
        """
        skip_list = []
        x = inputs
        for block in self.block_list[: self.n_levels - 1]:
            x, skip = block(x, None)
            skip_list.append(skip)
        x = self.block_list[self.n_levels - 1](x, None)
        for idx, block in enumerate(self.block_list[self.n_levels :]):
            passed = skip_list[-(idx + 1)]
            x = block(x, passed)
        x = self.conv_pred(x)
        output = self.softmax(x)

        return output


class SynthSeg:
    """This class is intended for the SynthSeg model.

    The SynthSeg model :footcite:p:`Billot2023` is a deep learning neural network for
    brain segmentation. It uses a U-net architecture and was trained on synthetic
    images generated from label maps. The model is robust to variations in
    contrast and resolution, making it suitable for segmenting a wide range of
    brain scans.

    Note that we are not saving any PVE maps here, only the hard segmentation
    due to the size of the output probability maps.

    References
    ----------
    .. footbibliography::
    """

    @doctest_skip_parser
    def __init__(self, *, verbose=False, use_cuda=False):
        """Model initialization

        Parameters
        ----------
        verbose : bool, optional
            Whether to show information about the processing.
        use_cuda : bool, optional
            Whether to use GPU for processing.
            If False or no CUDA is detected, CPU will be used.
        """
        if not have_torch:
            raise torch()

        log_level = "INFO" if verbose else "CRITICAL"
        set_logger_level(log_level, logger)

        self.model = self.init_model()
        self.model.eval()
        if use_cuda:
            if torch.cuda.is_available():
                self.device = torch.device("cuda")
            else:
                logger.warning("CUDA requested but not found, switching to CPU")
                self.device = torch.device("cpu")
        else:
            self.device = torch.device("cpu")
        self.model = self.model.to(self.device)
        self.fetch_default_weights()
        self.labels_segmentation = np.array(
            [
                0,
                2,
                3,
                4,
                5,
                7,
                8,
                10,
                11,
                12,
                13,
                14,
                15,
                16,
                17,
                18,
                24,
                26,
                28,
                41,
                42,
                43,
                44,
                46,
                47,
                49,
                50,
                51,
                52,
                53,
                54,
                58,
                60,
            ]
        )
        self.topological_classes = np.array(
            [
                0,
                4,
                4,
                4,
                4,
                5,
                5,
                6,
                7,
                8,
                9,
                1,
                2,
                3,
                10,
                11,
                0,
                12,
                13,
                14,
                14,
                14,
                14,
                15,
                15,
                16,
                17,
                18,
                19,
                20,
                21,
                22,
                23,
            ]
        )
        self.flip_indices = np.array(
            [
                0,
                19,
                20,
                21,
                22,
                23,
                24,
                25,
                26,
                27,
                28,
                11,
                12,
                13,
                29,
                30,
                16,
                31,
                32,
                1,
                2,
                3,
                4,
                5,
                6,
                7,
                8,
                9,
                10,
                14,
                15,
                17,
                18,
            ]
        )
        self.label_dict = {
            0: "background",
            2: "left cerebral white matter",
            3: "left cerebral cortex",
            4: "left lateral ventricle",
            5: "left inferior lateral ventricle",
            7: "left cerebellum white matter",
            8: "left cerebellum cortex",
            10: "left thalamus",
            11: "left caudate",
            12: "left putamen",
            13: "left pallidum",
            14: "3rd ventricle",
            15: "4th ventricle",
            16: "brain-stem",
            17: "left hippocampus",
            18: "left amygdala",
            24: "CSF",
            26: "left accumbens area",
            28: "left ventral DC",
            41: "right cerebral white matter",
            42: "right cerebral cortex",
            43: "right lateral ventricle",
            44: "right inferior lateral ventricle",
            46: "right cerebellum white matter",
            47: "right cerebellum cortex",
            49: "right thalamus",
            50: "right caudate",
            51: "right putamen",
            52: "right pallidum",
            53: "right hippocampus",
            54: "right amygdala",
            58: "right accumbens area",
            60: "right ventral DC",
        }

    def init_model(self, model_scale=24, n_levels=5, output_channels=33):
        """Initialize the SynthSeg model.

        Parameters
        ----------
        model_scale : int, optional
            The scale of the model.
        n_levels : int, optional
            Number of levels in the U-Net.
        output_channels : int, optional
            Number of output channels.

        Returns
        -------
        Model
            Initialized SynthSeg model.
        """
        return Model(
            model_scale=model_scale, n_levels=n_levels, output_channels=output_channels
        )

    def fetch_default_weights(self):
        """Load the model pre-training weights to use for the fitting.

        While the user can load different weights, the function
        is mainly intended for the class function 'predict'.
        """
        fetch_model_weights_path = get_fnames(name="synthseg_torch_weights")
        self.load_model_weights(fetch_model_weights_path)

    def load_model_weights(self, weights_path):
        """Load the custom pre-training weights to use for the fitting.

        Parameters
        ----------
        weights_path : str
            Path to the file containing the weights (pth, saved by Pytorch)
        """
        try:
            self.model.load_state_dict(
                torch.load(
                    weights_path,
                    weights_only=True,
                    map_location=self.device,
                )
            )
            self.model.eval()
        except ValueError as e:
            raise ValueError(
                "Expected input for the provided model weights \
                             do not match the declared model"
            ) from e

    def _flip_img_indices(self, img):
        """Flip the label indices for a flipped image.

        Parameters
        ----------
        img : np.ndarray
            The image with label indices.

        Returns
        -------
        np.ndarray
            The image with flipped label indices.
        """
        new_img = np.zeros_like(img)
        for l_idx, flip_lbl in enumerate(self.flip_indices):
            if l_idx != 0:
                new_img = np.where(img == l_idx, flip_lbl, new_img)
        return new_img

    def _prepare_img(self, img):
        """Prepare the image for model input.

        Parameters
        ----------
        img : np.ndarray
            Input image.

        Returns
        -------
        torch.Tensor
            Image tensor ready for the model.
        """
        img = np.expand_dims(img, 1)  # add channel dimension
        img_tensor = torch.tensor(img, dtype=torch.float32).to(self.device)
        return img_tensor

    def __predict(self, x_test):
        """Internal prediction function

        Parameters
        ----------
        x_test : list of np.ndarray
            Image should match the required shape of the model.

        Returns
        -------
        np.ndarray (batch, ...)
            Predicted brain labels
        """
        return self.model(x_test).detach().numpy()

    def predict(
        self,
        T1,
        affine,
        *,
        batch_size=None,
        return_prob=False,
    ):
        """Wrapper function to facilitate prediction of larger dataset.

        Parameters
        ----------
        T1 : np.ndarray or list of np.ndarray
            For a single image, input should be a 3D array.
            If multiple images, it should be a a list or tuple.
        affine : np.ndarray (4, 4) or (batch, 4, 4)
            or list of np.ndarrays with len of batch
            Affine matrix for the T1 image. Should have
            batch dimension if T1 has one.
        batch_size : int, optional
            Number of images per prediction pass. Only available if data
            is provided with a batch dimension.
            Consider lowering it if you get an out of memory error.
            Increase it if you want it to be faster and have a lot of data.
            If None, batch_size will be set to 1 if the provided image
            has a batch dimension.
        return_prob : bool, optional
            Whether to return the probability map instead of a
            label map. Useful for testing.

        Returns
        -------
        pred_output : np.ndarray (...) or (batch, ...)
            Predicted brain labels. If return_prob is True, it will be a
            probability map instead of a label map.

        label_dict : dict
            Dictionary mapping label indices to anatomical structure names.
            Only if return_prob is False.

        mask : np.ndarray (...) or (batch, ...)
            Predicted brain mask.
            Only if return_prob is False.
        """
        affine = np.array(affine)

        if isinstance(T1, (list, tuple)):
            dim = 4
            T1 = np.array(T1)
        elif len(T1.shape) == 3:
            dim = 3
            if batch_size is not None:
                logger.warning(
                    "Batch size specified, but not used",
                    "due to the input not having \
                                a batch dimension",
                )

            T1 = np.expand_dims(T1, 0)
            affine = np.expand_dims(affine, 0)
        else:
            raise ValueError(
                "Input data should be a np.ndarray of dimension 3 or a list/tuple of it"
            )
        if batch_size is None:
            batch_size = 1

        input_data = np.zeros((len(T1), 192, 192, 192))
        params_list = []

        # Normalize the data.
        ori_shape = T1.shape[1:]
        for i, T1_img in enumerate(T1):
            t_img, params = transform_img(
                T1_img,
                affine[i],
                target_voxsize=(1.0, 1.0, 1.0),
                final_size=(192, 192, 192),
                order=3,
            )
            min_v, max_v = np.percentile(t_img, (0.5, 99.5))
            t_img = normalize(t_img, min_v=min_v, max_v=max_v, new_min=0, new_max=1)
            input_data[i] = t_img
            params_list.append(params)
        # Prediction stage
        prediction = np.zeros((len(T1), 192, 192, 192, 33), dtype=np.float32)
        for batch_idx in range(batch_size, len(T1) + 1, batch_size):
            batch = input_data[batch_idx - batch_size : batch_idx]
            batch_start = batch_idx - batch_size
            temp_input = self._prepare_img(batch)
            temp_pred = self.__predict(temp_input)
            temp_pred = gaussian_filter(temp_pred, (0, 0.5, 0.5, 0.5, 0))
            prediction[batch_start:batch_idx] = np.moveaxis(temp_pred, 1, -1)
            temp_input = self._prepare_img(np.flip(batch, axis=1).copy())
            temp_pred = self.__predict(temp_input)
            temp_pred = gaussian_filter(temp_pred, (0, 0.5, 0.5, 0.5, 0))
            temp_pred = np.flip(temp_pred, axis=2)
            temp_pred = np.stack(
                [
                    temp_pred[:, self.flip_indices[i]]
                    for i in range(len(self.flip_indices))
                ],
                axis=1,
            )
            prediction[batch_start:batch_idx] += np.moveaxis(temp_pred, 1, -1)
            prediction[batch_start:batch_idx] /= 2
        remainder = np.mod(len(T1), batch_size)
        if remainder != 0:
            batch = input_data[-remainder:]
            temp_input = self._prepare_img(batch)
            temp_pred = self.__predict(temp_input)
            temp_pred = gaussian_filter(temp_pred, (0, 0.5, 0.5, 0.5, 0))
            prediction[-remainder:] = np.moveaxis(temp_pred, 1, -1)
            temp_input = self._prepare_img(np.flip(batch, axis=1).copy())
            temp_pred = self.__predict(temp_input)
            temp_pred = gaussian_filter(temp_pred, (0, 0.5, 0.5, 0.5, 0))
            temp_pred = np.flip(temp_pred, axis=1)
            temp_pred = np.stack(
                [
                    temp_pred[:, self.flip_indices[i]]
                    for i in range(len(self.flip_indices))
                ],
                axis=1,
            )
            prediction[-remainder:] += np.moveaxis(temp_pred, 1, -1)
            prediction[-remainder:] /= 2

        if return_prob:
            labels = np.zeros((len(T1),) + (192, 192, 192, 33)).astype(np.float32)
        else:
            labels = np.zeros((len(T1),) + ori_shape).astype(np.int32)
            masks = np.zeros((len(T1),) + ori_shape)
        for i in range(len(T1)):
            output = prediction[i]

            tmp_post_patch_seg = output[..., 1:]
            post_patch_seg_mask = np.sum(tmp_post_patch_seg, axis=-1) > 0.25
            post_patch_seg_mask = remove_holes_and_islands(
                post_patch_seg_mask, remove_holes=False
            ).astype(bool)
            output[..., 1:] = output[..., 1:] * np.stack(
                [post_patch_seg_mask] * 32, axis=-1
            )

            post_patch_seg_mask = output > 0.25
            for topology_class in np.unique(self.topological_classes)[1:]:
                tmp_topology_indices = np.where(
                    self.topological_classes == topology_class
                )[0]
                tmp_mask = np.any(
                    post_patch_seg_mask[..., tmp_topology_indices], axis=-1
                )
                tmp_mask = remove_holes_and_islands(
                    tmp_mask, remove_holes=False
                ).astype(bool)
                for idx in tmp_topology_indices:
                    output[..., idx] *= tmp_mask

            output /= np.sum(output, axis=-1)[..., np.newaxis]
            if return_prob:
                labels[i] = output.astype("float32")
                continue
            else:
                temp = self.labels_segmentation[
                    output.argmax(-1).astype("int32")
                ].astype("int32")
                temp = recover_img(temp, params_list[i], order=0)
                labels[i] = np.round(temp).astype(np.int32)
            masks[i] = (labels[i] > 0).astype(np.int32)

        if dim == 3:
            labels = labels[0]
            if not return_prob:
                masks = masks[0]

        if return_prob:
            return labels
        return labels, self.label_dict, masks
