# -*- coding: utf-8 -*-
# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
# CRIAQ and ANITI - https://www.deel.ai/
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
from typing import get_args
from typing import Optional

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from tqdm import tqdm

from ..datasets.tf_data_handler import TFDataHandler
from ..types import ItemType
from ..types import List
from ..types import TensorType
from ..types import Tuple
from ..types import Union
from ..utils.tf_operator import sanitize_input
from .feature_extractor import FeatureExtractor


class KerasFeatureExtractor(FeatureExtractor):
    """
    Feature extractor based on "model" to construct a feature space
    on which OOD detection is performed. The features can be the output
    activation values of internal model layers, or the output of the model
    (logits).

    Args:
        model: model to extract the features from
        feature_layers_id: list of str or int that identify features to output.
            If int, the rank of the layer in the layer list
            If str, the name of the layer. Defaults to [].
        head_layer_id (int, str): identifier of the head layer.
            If int, the rank of the layer in the layer list
            If str, the name of the layer.
            Defaults to -1
        input_layer_id: input layer of the feature extractor (to avoid useless forwards
            when working on the feature space without finetuning the bottom of the
            model).
            Defaults to None.
        react_threshold: if not None, penultimate layer activations are clipped under
            this threshold value (useful for ReAct). Defaults to None.
        scale_percentile: if not None, the features are scaled
            following the method of Xu et al., ICLR 2024.
            Defaults to None.
        ash_percentile: if not None, the features are scaled following
            the method of Djurisic et al., ICLR 2023.
        return_penultimate (bool): if True, the penultimate values are returned,
            i.e. the input to the head_layer.
    """

    def __init__(
        self,
        model: tf.keras.Model,
        feature_layers_id: List[Union[int, str]] = [],
        head_layer_id: Optional[Union[int, str]] = -1,
        input_layer_id: Optional[Union[int, str]] = None,
        react_threshold: Optional[float] = None,
        scale_percentile: Optional[float] = None,
        ash_percentile: Optional[float] = None,
        return_penultimate: Optional[bool] = False,
    ):
        if input_layer_id is None:
            input_layer_id = 0

        if return_penultimate:
            if isinstance(head_layer_id, str):
                head_layer_id = self.get_layer_index_by_name(model, head_layer_id)

            feature_layers_id.append(head_layer_id - 1)

        super().__init__(
            model=model,
            feature_layers_id=feature_layers_id,
            input_layer_id=input_layer_id,
            head_layer_id=head_layer_id,
            react_threshold=react_threshold,
            scale_percentile=scale_percentile,
            ash_percentile=ash_percentile,
            return_penultimate=return_penultimate,
        )

        self.backend = "tensorflow"
        self.model.layers[-1].activation = getattr(tf.keras.activations, "linear")
        self._last_logits = None

    @staticmethod
    def find_layer(
        model: tf.keras.Model,
        layer_id: Union[str, int],
        index_offset: int = 0,
        return_id: bool = False,
    ) -> Union[tf.keras.layers.Layer, Tuple[tf.keras.layers.Layer, str]]:
        """Find a layer in a model either by his name or by his index.

        Args:
            model (tf.keras.Model): model whose identified layer will be returned
            layer_id (Union[str, int]): layer identifier
            index_offset (int): index offset to find layers located before (negative
                offset) or after (positive offset) the identified layer
            return_id (bool): if True, the layer will be returned with its id

        Raises:
            ValueError: if the layer is not found

        Returns:
            Union[tf.keras.layers.Layer, Tuple[tf.keras.layers.Layer, str]]:
                the corresponding layer and its id if return_id is True.
        """
        if isinstance(layer_id, str):
            layers_names = [layer.name for layer in model.layers]
            layer_id = layers_names.index(layer_id)
        if isinstance(layer_id, int):
            layer_id += index_offset
            layer = model.get_layer(index=layer_id)
        else:
            raise ValueError(f"Could not find any layer {layer_id}.")

        if return_id:
            return layer, layer_id
        else:
            return layer

    @staticmethod
    def get_layer_index_by_name(model: tf.keras.Model, layer_id: str) -> int:
        """Get the index of a layer by its name.

        Args:
            model (tf.keras.Model): The model containing the layers.
            layer_id (str): The name of the layer.

        Returns:
            int: The index of the layer.

        Raises:
            ValueError: If the layer with the given name is not found.
        """
        layers_names = [layer.name for layer in model.layers]
        if layer_id not in layers_names:
            raise ValueError(f"Layer with name '{layer_id}' not found in the model.")
        return layers_names.index(layer_id)

    # @tf.function
    # TODO check with Thomas about @tf.function
    def prepare_extractor(self) -> tf.keras.models.Model:
        """Constructs the feature extractor model

        Returns:
            tf.keras.models.Model: truncated model (extractor)
        """
        input_layer = self.find_layer(self.model, self.input_layer_id)
        new_input = tf.keras.layers.Input(tensor=input_layer.input)
        output_tensors = [
            self.find_layer(self.model, id).output for id in self.feature_layers_id
        ]

        # === If react method, clip activations from penultimate layer ===
        if self.react_threshold is not None:
            penultimate_layer = self.find_layer(
                self.model, self.head_layer_id, index_offset=-1
            )
            penult_extractor = tf.keras.models.Model(
                new_input, penultimate_layer.output
            )
            last_layer = self.find_layer(self.model, self.head_layer_id)

            # clip penultimate activations
            x = tf.clip_by_value(
                penult_extractor(new_input),
                clip_value_min=tf.float32.min,
                clip_value_max=self.react_threshold,
            )
            # apply ultimate layer on clipped activations
            output_tensors.append(last_layer(x))

        # === If SCALE method, scale activations from penultimate layer ===
        # === If ASH method, scale and prune activations from penultimate layer ===
        elif (self.scale_percentile is not None) or (self.ash_percentile is not None):
            penultimate_layer = self.find_layer(
                self.model, self.head_layer_id, index_offset=-1
            )
            penult_extractor = tf.keras.models.Model(
                new_input, penultimate_layer.output
            )
            last_layer = self.find_layer(self.model, self.head_layer_id)

            # apply scaling on penultimate activations
            penultimate = penult_extractor(new_input)
            if self.scale_percentile is not None:
                output_percentile = tfp.stats.percentile(
                    penultimate, 100 * self.scale_percentile, axis=1
                )
            else:
                output_percentile = tfp.stats.percentile(
                    penultimate, 100 * self.ash_percentile, axis=1
                )

            mask = penultimate > tf.reshape(output_percentile, (-1, 1))
            filtered_penultimate = tf.where(
                mask, penultimate, tf.zeros_like(penultimate)
            )
            s = tf.math.exp(
                tf.reduce_sum(penultimate, axis=1)
                / tf.reduce_sum(filtered_penultimate, axis=1)
            )

            if self.scale_percentile is not None:
                x = penultimate * tf.expand_dims(s, 1)
            else:
                x = filtered_penultimate * tf.expand_dims(s, 1)
            # apply ultimate layer on scaled activations
            output_tensors.append(last_layer(x))

        else:
            output_tensors.append(
                self.find_layer(self.model, self.head_layer_id).output
            )

        extractor = tf.keras.models.Model(new_input, output_tensors)
        return extractor

    @sanitize_input
    def predict_tensor(
        self,
        tensor: TensorType,
        postproc_fns: Optional[List[tf.keras.Model]] = None,
    ) -> Tuple[List[tf.Tensor], tf.Tensor]:
        """Get the projection of tensor in the feature space of self.model

        Args:
            tensor (TensorType): input tensor (or dataset elem)
            postproc_fns (Optional[List[tf.keras.Model]]): postprocessing function
                to apply to each feature immediately after forward. Default to None.

        Returns:
            Tuple[List[tf.Tensor], tf.Tensor]: features, logits
        """
        features = self.forward(tensor)

        if type(features) is not list:
            features = [features]

        # split features and logits
        logits = features.pop()

        if postproc_fns is not None:
            features = [
                postproc_fn(feature)
                for feature, postproc_fn in zip(features, postproc_fns)
            ]

        self._last_logits = logits
        return features, logits

    @tf.function
    def forward(self, tensor: TensorType) -> List[tf.Tensor]:
        return self.extractor(tensor, training=False)

    def predict(
        self,
        dataset: Union[ItemType, tf.data.Dataset],
        postproc_fns: Optional[List[tf.keras.Model]] = None,
        verbose: bool = False,
        numpy_concat: bool = False,
        **kwargs,
    ) -> Tuple[List[tf.Tensor], dict]:
        """Get the projection of the dataset in the feature space of self.model

        Args:
            dataset (Union[ItemType, tf.data.Dataset]): input dataset
            postproc_fns (Optional[tf.keras.Model]): postprocessing function to apply
                to each feature immediately after forward. Default to None.
            verbose (bool): if True, display a progress bar. Defaults to False.
            numpy_concat (bool): if True, each mini-batch is immediately moved
                to CPU and converted to a NumPy array before concatenation.
                That keeps GPU memory constant at one batch, at the cost of a small
                host-device transfer overhead. Defaults to False.
            kwargs (dict): additional arguments not considered for prediction

        Returns:
            List[tf.Tensor], dict: features and extra information (logits, labels) as a
                dictionary.
        """
        labels = None

        if isinstance(dataset, get_args(ItemType)):
            tensor = TFDataHandler.get_input_from_dataset_item(dataset)
            features, logits = self.predict_tensor(tensor, postproc_fns)

            # Get labels if dataset is a tuple/list
            if isinstance(dataset, (list, tuple)):
                labels = TFDataHandler.get_label_from_dataset_item(dataset)
        else:
            # Determine if dataset yields labels
            contains_labels = TFDataHandler.get_item_length(dataset) > 1

            # Buffers for accumulation
            features_per_layer = [[] for _ in self.feature_layers_id]
            logits_list = []
            labels_list = [] if contains_labels else None

            # Iterate through dataset
            for elem in tqdm(dataset, desc="Predicting", disable=not verbose):
                tensor = TFDataHandler.get_input_from_dataset_item(elem)
                feats_batch, logits_batch = self.predict_tensor(tensor, postproc_fns)

                # To host/NumPy if requested
                if numpy_concat:
                    feats_batch = [f.numpy() for f in feats_batch]
                    logits_batch = (
                        logits_batch.numpy() if logits_batch is not None else None
                    )

                # Accumulate per layer
                for i, f in enumerate(feats_batch):
                    features_per_layer[i].append(f)
                if logits_batch is not None:
                    logits_list.append(logits_batch)
                if contains_labels:
                    lbl = TFDataHandler.get_label_from_dataset_item(elem)
                    labels_list.append(lbl)

            # Concatenate
            labels = tf.concat(labels_list, axis=0) if labels_list is not None else None

            if numpy_concat:
                features = [np.concatenate(lst, axis=0) for lst in features_per_layer]
                logits = np.concatenate(logits_list, axis=0) if logits_list else None
                labels = labels.numpy() if labels is not None else None
            else:
                features = [tf.concat(lst, axis=0) for lst in features_per_layer]
                logits = tf.concat(logits_list, axis=0) if logits_list else None

        info = {"labels": labels, "logits": logits}
        return features, info

    def get_weights(self, layer_id: Union[int, str]) -> List[tf.Tensor]:
        """Get the weights of a layer

        Args:
            layer_id (Union[int, str]): layer identifier

        Returns:
            List[tf.Tensor]: weights and biases matrixes
        """
        return self.find_layer(self.model, layer_id).get_weights()

    def _default_postproc_fn(self, feat: tf.Tensor) -> tf.Tensor:
        """
        Default postprocessing function to apply to each feature immediately after
        forward pass.

        This function applies global average pooling if the input tensor has rank 4
        (e.g., [batch, height, width, channels]) or rank 3
        (e.g., [batch, sequence_length, features]). If the tensor is already 2D, it is
        returned as is.

        Args:
            feat (tf.Tensor): Input tensor.

        Returns:
            tf.Tensor: Postprocessed tensor where spatial (or temporal) dimensions have
                been averaged out.

        Raises:
            NotImplementedError: If the input tensor has a rank other than 2, 3, or 4.
        """
        tensor_rank = len(feat.shape)

        if tensor_rank == 4:
            # Assumes input is in channels_last format: [batch, height, width, channels]
            # Applies global average pooling over height and width.
            pooled = tf.keras.layers.GlobalAveragePooling2D()(feat)
            # The resulting tensor has shape [batch, channels]
            return pooled
        elif tensor_rank == 3:
            # Assumes input is in channels_last format: [batch, seq_length, features]
            # Applies global average pooling over the sequence_length dimension.
            pooled = tf.keras.layers.GlobalAveragePooling1D()(feat)
            # The resulting tensor has shape [batch, features]
            return pooled
        elif tensor_rank == 2:
            # If the tensor is already 2D, no further processing is needed.
            return feat
        else:
            raise NotImplementedError(
                "Postprocessing function not implemented for tensors with"
                + " rank {}.".format(tensor_rank)
            )
