#!/usr/bin/python
"""Class and helper functions for fitting the EVAC+ model."""

import logging

import numpy as np

from dipy.align.reslice import reslice
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, warning_for_keywords
from dipy.utils.deprecator import deprecated_params
from dipy.utils.optpkg import optional_package

tf, have_tf, _ = optional_package("tensorflow", min_version="2.18.0")
if have_tf:
    from tensorflow.keras.layers import (
        Add,
        Concatenate,
        Conv3D,
        Conv3DTranspose,
        Dropout,
        Layer,
        LayerNormalization,
        ReLU,
        Softmax,
    )
    from tensorflow.keras.models import Model
else:

    class Model:
        pass

    class Layer:
        pass

    logging.warning(
        "This model requires Tensorflow.\
                    Please install these packages using \
                    pip. If using mac, please refer to this \
                    link for installation. \
                    https://github.com/apple/tensorflow_macos"
    )

logging.basicConfig()
logger = logging.getLogger("EVAC+")


def prepare_img(image):
    """Function to prepare image for model input
    Specific to EVAC+

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

    Returns
    -------
    input_data : dict
    """
    input1 = np.moveaxis(image, -1, 0)
    input1 = np.expand_dims(input1, -1)

    input2, _ = reslice(image, np.eye(4), (1, 1, 1), (2, 2, 2))
    input2 = np.moveaxis(input2, -1, 0)
    input2 = np.expand_dims(input2, -1)

    input3, _ = reslice(image, np.eye(4), (1, 1, 1), (4, 4, 4))
    input3 = np.moveaxis(input3, -1, 0)
    input3 = np.expand_dims(input3, -1)

    input4, _ = reslice(image, np.eye(4), (1, 1, 1), (8, 8, 8))
    input4 = np.moveaxis(input4, -1, 0)
    input4 = np.expand_dims(input4, -1)

    input5, _ = reslice(image, np.eye(4), (1, 1, 1), (16, 16, 16))
    input5 = np.moveaxis(input5, -1, 0)
    input5 = np.expand_dims(input5, -1)

    input_data = {
        "input_1": input1,
        "input_2": input2,
        "input_3": input3,
        "input_4": input4,
        "input_5": input5,
    }

    return input_data


class Block(Layer):
    @warning_for_keywords()
    def __init__(
        self,
        out_channels,
        kernel_size,
        strides,
        padding,
        drop_r,
        n_layers,
        *,
        layer_type="down",
    ):
        super(Block, self).__init__()
        self.layer_list = []
        self.layer_list2 = []
        self.n_layers = n_layers
        for _ in range(n_layers):
            self.layer_list.append(
                Conv3D(out_channels, kernel_size, strides=strides, padding=padding)
            )
            self.layer_list.append(Dropout(drop_r))
            self.layer_list.append(LayerNormalization())
            self.layer_list.append(ReLU())
        if layer_type == "down":
            self.layer_list2.append(Conv3D(1, 2, strides=2, padding="same"))
            self.layer_list2.append(ReLU())
        elif layer_type == "up":
            self.layer_list2.append(Conv3DTranspose(1, 2, strides=2, padding="same"))
            self.layer_list2.append(ReLU())

        self.channel_sum = ChannelSum()
        self.add = Add()

    def call(self, input, passed):
        x = input
        for layer in self.layer_list:
            x = layer(x)

        x = self.channel_sum(x)
        fwd = self.add([x, passed])
        x = fwd

        for layer in self.layer_list2:
            x = layer(x)

        return fwd, x


class ChannelSum(Layer):
    def __init__(self):
        super(ChannelSum, self).__init__()

    def call(self, inputs):
        return tf.reduce_sum(inputs, axis=-1, keepdims=True)


@warning_for_keywords()
def init_model(*, model_scale=16):
    """Function to create model for EVAC+

    Parameters
    ----------
    model_scale : int, optional
        The scale of the model
        Should match the saved weights from fetcher
        Default is 16

    Returns
    -------
    model : tf.keras.Model
    """
    inputs = tf.keras.Input(shape=(128, 128, 128, 1), name="input_1")
    raw_input_2 = tf.keras.Input(shape=(64, 64, 64, 1), name="input_2")
    raw_input_3 = tf.keras.Input(shape=(32, 32, 32, 1), name="input_3")
    raw_input_4 = tf.keras.Input(shape=(16, 16, 16, 1), name="input_4")
    raw_input_5 = tf.keras.Input(shape=(8, 8, 8, 1), name="input_5")
    # Encode
    fwd1, x = Block(
        model_scale, kernel_size=5, strides=1, padding="same", drop_r=0.2, n_layers=1
    )(inputs, inputs)

    x = Concatenate()([x, raw_input_2])

    fwd2, x = Block(
        model_scale * 2,
        kernel_size=5,
        strides=1,
        padding="same",
        drop_r=0.5,
        n_layers=2,
    )(x, x)

    x = Concatenate()([x, raw_input_3])

    fwd3, x = Block(
        model_scale * 4,
        kernel_size=5,
        strides=1,
        padding="same",
        drop_r=0.5,
        n_layers=3,
    )(x, x)

    x = Concatenate()([x, raw_input_4])

    fwd4, x = Block(
        model_scale * 8,
        kernel_size=5,
        strides=1,
        padding="same",
        drop_r=0.5,
        n_layers=3,
    )(x, x)

    x = Concatenate()([x, raw_input_5])

    _, up = Block(
        model_scale * 16,
        kernel_size=5,
        strides=1,
        padding="same",
        drop_r=0.5,
        n_layers=3,
        layer_type="up",
    )(x, x)

    x = Concatenate()([fwd4, up])

    _, up = Block(
        model_scale * 8,
        kernel_size=5,
        strides=1,
        padding="same",
        drop_r=0.5,
        n_layers=3,
        layer_type="up",
    )(x, up)

    x = Concatenate()([fwd3, up])

    _, up = Block(
        model_scale * 4,
        kernel_size=5,
        strides=1,
        padding="same",
        drop_r=0.5,
        n_layers=3,
        layer_type="up",
    )(x, up)

    x = Concatenate()([fwd2, up])

    _, up = Block(
        model_scale * 2,
        kernel_size=5,
        strides=1,
        padding="same",
        drop_r=0.5,
        n_layers=2,
        layer_type="up",
    )(x, up)

    x = Concatenate()([fwd1, up])

    _, pred = Block(
        model_scale,
        kernel_size=5,
        strides=1,
        padding="same",
        drop_r=0.5,
        n_layers=1,
        layer_type="none",
    )(x, up)

    pred = Conv3D(2, 1, padding="same")(pred)
    output = Softmax(axis=-1)(pred)

    model = Model(
        {
            "input_1": inputs,
            "input_2": raw_input_2,
            "input_3": raw_input_3,
            "input_4": raw_input_4,
            "input_5": raw_input_5,
        },
        output[..., 0],
    )
    return model


class EVACPlus:
    """This class is intended for the EVAC+ model.

    The EVAC+ model :footcite:p:`Park2024` is a deep learning neural network for
    brain extraction. It uses a V-net architecture combined with
    multi-resolution input data, an additional conditional random field (CRF)
    recurrent layer and supplementary Dice loss term for this recurrent layer.

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

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

        The model was pre-trained for usage on
        brain extraction of T1 images.

        This model is designed to take as input
        a T1 weighted image.

        Parameters
        ----------
        verbose : bool, optional
            Whether to show information about the processing.
        """
        if not have_tf:
            raise tf()

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

        # EVAC+ network load

        self.model = init_model()
        self.fetch_default_weights()

    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="evac_default_tf_weights")
        print(f"fetched {fetch_model_weights_path}")
        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 (hdf5, saved by tensorflow)
        """
        try:
            self.model.load_weights(weights_path)
        except ValueError as e:
            raise ValueError(
                "Expected input for the provided model weights \
                             do not match the declared model"
            ) from e

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

        Parameters
        ----------
        x_test : np.ndarray (batch, 128, 128, 128, 1)
            Image should match the required shape of the model.

        Returns
        -------
        np.ndarray (batch, ...)
            Predicted brain mask
        """
        return self.model.predict(x_test)

    @deprecated_params(
        "largest_area", new_name="finalize_mask", since="1.10", until="1.12"
    )
    def predict(
        self,
        T1,
        affine,
        *,
        voxsize=(1, 1, 1),
        batch_size=None,
        return_affine=False,
        return_prob=False,
        finalize_mask=True,
    ):
        """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.
            or list of np.ndarrays with len of batch_size
        affine : np.ndarray (4, 4) or (batch, 4, 4)
            Affine matrix for the T1 image. Should have
            batch dimension if T1 has one.
        voxsize : np.ndarray or list or tuple, optional
            (3,) or (batch, 3)
            voxel size of the T1 image.
        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_affine : bool, optional
            Whether to return the affine matrix. Useful if the input was a
            file path.
        return_prob : bool, optional
            Whether to return the probability map instead of a
            binary mask. Useful for testing.
        finalize_mask : bool, optional
            Whether to remove potential holes or islands.
            Useful for solving minor errors.

        Returns
        -------
        pred_output : np.ndarray (...) or (batch, ...)
            Predicted brain mask

        affine : np.ndarray (...) or (batch, ...)
            affine matrix of mask
            only if return_affine is True
        """
        voxsize = np.array(voxsize)
        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)
            voxsize = np.expand_dims(voxsize, 0)
        else:
            raise ValueError(
                "T1 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((128, 128, 128, len(T1)))
        affines = np.zeros((len(T1), 4, 4))
        mid_shapes = np.zeros((len(T1), 3)).astype(int)
        offset_arrays = np.zeros((len(T1), 4, 4)).astype(int)
        scales = np.zeros(len(T1))
        crop_vss = np.zeros((len(T1), 3, 2))
        pad_vss = np.zeros((len(T1), 3, 2))

        # Normalize the data.
        n_T1 = np.zeros(T1.shape)
        for i, T1_img in enumerate(T1):
            n_T1[i] = normalize(T1_img, new_min=0, new_max=1)
            t_img, t_affine, mid_shape, offset_array, scale, crop_vs, pad_vs = (
                transform_img(n_T1[i], affine[i], voxsize=voxsize[i])
            )
            input_data[..., i] = t_img
            affines[i] = t_affine
            mid_shapes[i] = mid_shape
            offset_arrays[i] = offset_array
            scales[i] = scale
            crop_vss[i] = crop_vs
            pad_vss[i] = pad_vs

        # Prediction stage
        prediction = np.zeros((len(T1), 128, 128, 128), dtype=np.float32)
        for batch_idx in range(batch_size, len(T1) + 1, batch_size):
            batch = input_data[..., batch_idx - batch_size : batch_idx]
            temp_input = prepare_img(batch)
            temp_pred = self.__predict(temp_input)
            prediction[:batch_idx] = temp_pred
        remainder = np.mod(len(T1), batch_size)
        if remainder != 0:
            temp_input = prepare_img(input_data[..., -remainder:])
            temp_pred = self.__predict(temp_input)
            prediction[-remainder:] = temp_pred

        output_mask = []
        for i in range(len(T1)):
            output = recover_img(
                prediction[i],
                affines[i],
                mid_shapes[i],
                n_T1[i].shape,
                offset_arrays[i],
                voxsize=voxsize[i],
                scale=scales[i],
                crop_vs=crop_vss[i],
                pad_vs=pad_vss[i],
            )
            if not return_prob:
                output = np.where(output >= 0.5, 1, 0)
                if finalize_mask:
                    output = remove_holes_and_islands(output, slice_wise=True)
            output_mask.append(output)

        if dim == 3:
            output_mask = output_mask[0]
            affine = affine[0]

        output_mask = np.array(output_mask)
        affine = np.array(affine)
        if return_affine:
            return output_mask, affine
        else:
            return output_mask
