#!/usr/bin/env python3
""" S3FD Face detection plugin
https://arxiv.org/abs/1708.05237

Adapted from S3FD Port in FAN:
https://github.com/1adrianb/face-alignment
"""
from __future__ import annotations
import logging
import typing as T

from scipy.special import logsumexp
import numpy as np

from keras.layers import (Concatenate, Conv2D, Input, Layer, Maximum, MaxPooling2D, ZeroPadding2D)
from keras.models import Model
from keras import initializers, ops

from lib.logger import parse_class_init
from lib.utils import get_module_objects
from ._base import BatchType, Detector
from . import s3fd_defaults as cfg

if T.TYPE_CHECKING:
    from keras import KerasTensor

logger = logging.getLogger(__name__)


class Detect(Detector):
    """ S3FD detector for face recognition """
    def __init__(self, **kwargs) -> None:
        git_model_id = 11
        model_filename = "s3fd_keras_v2.h5"
        super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs)
        self.model: S3fd
        self.name = "S3FD"
        self.input_size = 640
        self.vram = 1088  # 1034 in testing
        self.vram_per_batch = 960  # 922 in testing
        self.batchsize = cfg.batch_size()

    def init_model(self) -> None:
        """ Initialize S3FD Model"""
        assert isinstance(self.model_path, str)
        confidence = cfg.confidence() / 100
        self.model = S3fd(self.model_path, self.batchsize, confidence)
        placeholder_shape = (self.batchsize, self.input_size, self.input_size, 3)
        placeholder = np.zeros(placeholder_shape, dtype="float32")
        self.model(placeholder)

    def process_input(self, batch: BatchType) -> None:
        """ Compile the detection image(s) for prediction """
        assert isinstance(self.model, S3fd)
        batch.feed = self.model.prepare_batch(np.array(batch.image))

    def predict(self, feed: np.ndarray) -> np.ndarray:
        """ Run model to get predictions """
        assert isinstance(self.model, S3fd)
        predictions = self.model(feed)
        assert isinstance(predictions, list)
        return self.model.finalize_predictions(predictions)

    def process_output(self, batch) -> None:
        """ Compile found faces for output """
        return


################################################################################
# CUSTOM KERAS LAYERS
################################################################################
class L2Norm(Layer):  # pylint:disable=too-many-ancestors,abstract-method
    """ L2 Normalization layer for S3FD.

    Parameters
    ----------
    n_channels: int
        The number of channels to normalize
    scale: float, optional
        The scaling for initial weights. Default: `1.0`
    """
    def __init__(self, n_channels: int, scale: float = 1.0, **kwargs) -> None:
        super().__init__(**kwargs)
        self._n_channels = n_channels
        self._scale = scale
        self.weight = self.add_weight(name="l2norm",
                                      shape=(self._n_channels, ),
                                      trainable=True,
                                      initializer=initializers.Constant(value=self._scale),
                                      dtype="float32")

    def call(self, inputs: KerasTensor, **kwargs  # pylint:disable=arguments-differ
             ) -> KerasTensor:
        """ Call the L2 Normalization Layer.

        Parameters
        ----------
        inputs: :class:`keras.KerasTensor`
            The input to the L2 Normalization Layer

        Returns
        -------
        :class:`keras.KerasTensor`:
            The output from the L2 Normalization Layer
        """
        norm = ops.sqrt(ops.sum(ops.power(inputs, 2), axis=-1, keepdims=True)) + 1e-10
        var_x = inputs / norm * self.weight
        return var_x

    def get_config(self) -> dict:
        """ Returns the config of the layer.

        Returns
        -------
        dict
            The configuration for the layer
        """
        config = super().get_config()
        config.update({"n_channels": self._n_channels,
                       "scale": self._scale})
        return config


class SliceO2K(Layer):  # pylint:disable=too-many-ancestors,abstract-method
    """ Custom Keras Slice layer generated by onnx2keras. """
    def __init__(self,
                 starts: list[int],
                 ends: list[int],
                 axes: list[int] | None = None,
                 steps: list[int] | None = None,
                 **kwargs) -> None:
        self._starts = starts
        self._ends = ends
        self._axes = axes
        self._steps = steps
        super().__init__(**kwargs)

    def _get_slices(self, dimensions: int) -> list[tuple[int, ...]]:
        """ Obtain slices for the given number of dimensions.

        Parameters
        ----------
        dimensions: int
            The number of dimensions to obtain slices for

        Returns
        -------
        list
            The slices for the given number of dimensions
        """
        axes = tuple(range(dimensions)) if self._axes is None else self._axes
        steps = (1,) * len(axes) if self._steps is None else self._steps
        assert len(axes) == len(steps) == len(self._starts) == len(self._ends)
        return list(zip(axes, self._starts, self._ends, steps))

    def compute_output_shape(self, input_shape: tuple[int, ...]  # pylint:disable=arguments-differ
                             ) -> tuple[int, ...]:
        """Computes the output shape of the layer.

        Assumes that the layer will be built to match that input shape provided.

        Parameters
        ----------
        input_shape: tuple or list of tuples
            Shape tuple (tuple of integers) or list of shape tuples (one per output tensor of the
            layer). Shape tuples can include ``None`` for free dimensions, instead of an integer.

        Returns
        -------
        tuple
            An output shape tuple.
        """
        in_shape = list(input_shape)
        for a_x, start, end, steps in self._get_slices(len(in_shape)):
            size = in_shape[a_x]
            if a_x == 0:
                raise AttributeError("Can not slice batch axis.")
            if size is None:
                if start < 0 or end < 0:
                    raise AttributeError("Negative slices not supported on symbolic axes")
                logger.warning("Slicing symbolic axis might lead to problems.")
                in_shape[a_x] = (end - start) // steps
                continue
            if start < 0:
                start = size - start
            if end < 0:
                end = size - end
            in_shape[a_x] = (min(size, end) - start) // steps
        return tuple(in_shape)

    def call(self, inputs, **kwargs):  # pylint:disable=unused-argument,arguments-differ
        """This is where the layer's logic lives.

        Parameters
        ----------
        inputs: Input tensor, or list/tuple of input tensors.
            The input to the layer
        **kwargs: Additional keyword arguments.
            Required for parent class but unused
        Returns
        -------
        A tensor or list/tuple of tensors.
            The layer output
        """
        ax_map = dict((x[0], slice(*x[1:])) for x in self._get_slices(ops.ndim(inputs)))
        shape = inputs.shape
        slices = [(ax_map[a] if a in ax_map else slice(None)) for a in range(len(shape))]
        retval = inputs[tuple(slices)]
        return retval

    def get_config(self) -> dict:
        """ Returns the config of the layer.

        Returns
        -------
        dict
            The configuration for the layer
        """
        config = super().get_config()
        config.update({"starts": self._starts,
                       "ends": self._ends,
                       "axes": self._axes,
                       "steps": self._steps})
        return config


class S3fd():
    """ Keras Network

    Parameters
    ----------
    weights_path: str
        Full path to the S3FD weights file
    batch_size: int
        The batch size to feed the model
    confidence: float
        The confidence level to accept detections at
    """
    def __init__(self, weights_path: str, batch_size: int, confidence: float) -> None:
        logger.debug(parse_class_init(locals()))
        self._batch_size = batch_size
        self._model = self._load_model(weights_path)
        self.confidence = confidence
        self.average_img = np.array([104.0, 117.0, 123.0])
        logger.debug("Initialized: %s", self.__class__.__name__)

    @classmethod
    def conv_block(cls,
                   inputs: KerasTensor,
                   filters: int,
                   idx: int,
                   recursions: int) -> KerasTensor:
        """ First round convolutions with zero padding added.

        Parameters
        ----------
        inputs: :class:`keras.KerasTensor`
            The input tensor to the convolution block
        filters: int
            The number of filters
        idx: int
            The layer index for naming
        recursions: int
            The number of recursions of the block to perform

        Returns
        -------
        :class:`keras.KerasTensor`
            The output tensor from the convolution block
        """
        name = f"conv{idx}"
        var_x = inputs
        for i in range(1, recursions + 1):
            rec_name = f"{name}_{i}"
            var_x = ZeroPadding2D(1, name=f"{rec_name}.zeropad")(var_x)
            var_x = Conv2D(filters,
                           kernel_size=3,
                           strides=1,
                           activation="relu",
                           name=rec_name)(var_x)
        return var_x

    @classmethod
    def conv_up(cls, inputs: KerasTensor, filters: int, idx: int) -> KerasTensor:
        """ Convolution up filter blocks with zero padding added.

        Parameters
        ----------
        inputs: :class:`keras.KerasTensor`
            The input tensor to the convolution block
        filters: int
            The initial number of filters
        idx: int
            The layer index for naming

        Returns
        -------
        :class:`keras.KerasTensor`
            The output tensor from the convolution block
        """
        name = f"conv{idx}"
        var_x = inputs
        for i in range(1, 3):
            rec_name = f"{name}_{i}"
            size = 1 if i == 1 else 3
            if i == 2:
                var_x = ZeroPadding2D(1, name=f"{rec_name}.zeropad")(var_x)
            var_x = Conv2D(filters * i,
                           kernel_size=size,
                           strides=i,
                           activation="relu",
                           name=rec_name)(var_x)
        return var_x

    def _load_model(self, weights_path: str) -> Model:
        """ Keras S3FD Model Definition, adapted from FAN pytorch implementation.

        Parameters
        ----------
        weights_path: str
            Full path to the model's weights

        Returns
        -------
        :class:`keras.models.Model`
            The S3FD model
        """
        input_ = Input(shape=(640, 640, 3))
        var_x = self.conv_block(input_, 64, 1, 2)
        var_x = MaxPooling2D(pool_size=2, strides=2)(var_x)

        var_x = self.conv_block(var_x, 128, 2, 2)
        var_x = MaxPooling2D(pool_size=2, strides=2)(var_x)

        var_x = self.conv_block(var_x, 256, 3, 3)
        f3_3 = var_x
        var_x = MaxPooling2D(pool_size=2, strides=2)(var_x)

        var_x = self.conv_block(var_x, 512, 4, 3)
        f4_3 = var_x
        var_x = MaxPooling2D(pool_size=2, strides=2)(var_x)

        var_x = self.conv_block(var_x, 512, 5, 3)
        f5_3 = var_x
        var_x = MaxPooling2D(pool_size=2, strides=2)(var_x)

        var_x = ZeroPadding2D(3)(var_x)
        var_x = Conv2D(1024, kernel_size=3, strides=1, activation="relu", name="fc6")(var_x)
        var_x = Conv2D(1024, kernel_size=1, strides=1, activation="relu", name="fc7")(var_x)
        ffc7 = var_x

        f6_2 = self.conv_up(var_x, 256, 6)
        f7_2 = self.conv_up(f6_2, 128, 7)

        f3_3 = L2Norm(256, scale=10, name="conv3_3_norm")(f3_3)
        f4_3 = L2Norm(512, scale=8, name="conv4_3_norm")(f4_3)
        f5_3 = L2Norm(512, scale=5, name="conv5_3_norm")(f5_3)

        classes = []
        regs = []

        f3_3 = ZeroPadding2D(1)(f3_3)
        classes.append(Conv2D(4, kernel_size=3, strides=1, name="conv3_3_norm_mbox_conf")(f3_3))
        regs.append(Conv2D(4, kernel_size=3, strides=1, name="conv3_3_norm_mbox_loc")(f3_3))

        f4_3 = ZeroPadding2D(1)(f4_3)
        classes.append(Conv2D(2, kernel_size=3, strides=1, name="conv4_3_norm_mbox_conf")(f4_3))
        regs.append(Conv2D(4, kernel_size=3, strides=1, name="conv4_3_norm_mbox_loc")(f4_3))

        f5_3 = ZeroPadding2D(1)(f5_3)
        classes.append(Conv2D(2, kernel_size=3, strides=1, name="conv5_3_norm_mbox_conf")(f5_3))
        regs.append(Conv2D(4, kernel_size=3, strides=1, name="conv5_3_norm_mbox_loc")(f5_3))

        ffc7 = ZeroPadding2D(1)(ffc7)
        classes.append(Conv2D(2, kernel_size=3, strides=1, name="fc7_mbox_conf")(ffc7))
        regs.append(Conv2D(4, kernel_size=3, strides=1, name="fc7_mbox_loc")(ffc7))

        f6_2 = ZeroPadding2D(1)(f6_2)
        classes.append(Conv2D(2, kernel_size=3, strides=1, name="conv6_2_mbox_conf")(f6_2))
        regs.append(Conv2D(4, kernel_size=3, strides=1, name="conv6_2_mbox_loc")(f6_2))

        f7_2 = ZeroPadding2D(1)(f7_2)
        classes.append(Conv2D(2, kernel_size=3, strides=1, name="conv7_2_mbox_conf")(f7_2))
        regs.append(Conv2D(4, kernel_size=3, strides=1, name="conv7_2_mbox_loc")(f7_2))

        # max-out background label
        chunks = [SliceO2K(starts=[0], ends=[1], axes=[3], steps=None)(classes[0]),
                  SliceO2K(starts=[1], ends=[2], axes=[3], steps=None)(classes[0]),
                  SliceO2K(starts=[2], ends=[3], axes=[3], steps=None)(classes[0]),
                  SliceO2K(starts=[3], ends=[4], axes=[3], steps=None)(classes[0])]

        bmax = Maximum()([chunks[0], chunks[1], chunks[2]])
        classes[0] = Concatenate()([bmax, chunks[3]])

        retval = Model(input_,
                       [classes[0],
                        regs[0],
                        classes[1],
                        regs[1],
                        classes[2],
                        regs[2],
                        classes[3],
                        regs[3],
                        classes[4],
                        regs[4],
                        classes[5],
                        regs[5]])
        retval.load_weights(weights_path)
        retval.make_predict_function()
        return retval

    def prepare_batch(self, batch: np.ndarray) -> np.ndarray:
        """ Prepare a batch for prediction.

        Normalizes the feed images.

        Parameters
        ----------
        batch: class:`numpy.ndarray`
            The batch to be fed to the model

        Returns
        -------
        class:`numpy.ndarray`
            The normalized images for feeding to the model
        """
        batch = batch - self.average_img
        return batch

    def finalize_predictions(self, bounding_boxes_scales: list[np.ndarray]) -> np.ndarray:
        """ Process the output from the model to obtain faces

        Parameters
        ----------
        bounding_boxes_scales: list
            The output predictions from the S3FD model
        """
        ret = []
        batch_size = range(bounding_boxes_scales[0].shape[0])
        for img in batch_size:
            bboxlist = [scale[img:img+1] for scale in bounding_boxes_scales]
            boxes = self._post_process(bboxlist)
            finallist = self._nms(boxes, 0.5)
            ret.append(finallist)
        return np.array(ret, dtype="object")

    def _process_bbox(self,
                      ocls: np.ndarray,
                      oreg: np.ndarray,
                      stride: int) -> list[list[np.ndarray]]:
        """ Process a bounding box """
        retval = []
        for pos in zip(*np.where(ocls[:, :, :, 1] > 0.05)):
            a_c = stride / 2 + pos[2] * stride, stride / 2 + pos[1] * stride
            score = ocls[0, pos[1], pos[2], 1]
            if score >= self.confidence:
                loc = np.ascontiguousarray(oreg[0, pos[1], pos[2], :]).reshape((1, 4))
                priors = np.array([[a_c[0] / 1.0,
                                    a_c[1] / 1.0,
                                    stride * 4 / 1.0,
                                    stride * 4 / 1.0]])
                box = self.decode(loc, priors)
                x_1, y_1, x_2, y_2 = box[0] * 1.0
                retval.append([x_1, y_1, x_2, y_2, score])
        return retval

    def _post_process(self, bboxlist: list[np.ndarray]) -> np.ndarray:
        """ Perform post processing on output
            TODO: do this on the batch.
        """
        retval = []
        for i in range(len(bboxlist) // 2):
            bboxlist[i * 2] = self.softmax(bboxlist[i * 2], axis=3)
        for i in range(len(bboxlist) // 2):
            ocls, oreg = bboxlist[i * 2], bboxlist[i * 2 + 1]
            stride = 2 ** (i + 2)    # 4,8,16,32,64,128
            retval.extend(self._process_bbox(ocls, oreg, stride))

        return_numpy = np.array(retval) if len(retval) != 0 else np.zeros((1, 5))
        return return_numpy

    @staticmethod
    def softmax(inp, axis: int) -> np.ndarray:
        """Compute softmax values for each sets of scores in x."""
        return np.exp(inp - logsumexp(inp, axis=axis, keepdims=True))

    @staticmethod
    def decode(location: np.ndarray, priors: np.ndarray) -> np.ndarray:
        """Decode locations from predictions using priors to undo the encoding we did for offset
        regression at train time.

        Parameters
        ----------
        location: tensor
            location predictions for location layers,
        priors: tensor
            Prior boxes in center-offset form.

        Returns
        -------
        :class:`numpy.ndarray`
            decoded bounding box predictions
        """
        variances = [0.1, 0.2]
        boxes = np.concatenate((priors[:, :2] + location[:, :2] * variances[0] * priors[:, 2:],
                                priors[:, 2:] * np.exp(location[:, 2:] * variances[1])), axis=1)
        boxes[:, :2] -= boxes[:, 2:] / 2
        boxes[:, 2:] += boxes[:, :2]
        return boxes

    @staticmethod
    def _nms(boxes: np.ndarray, threshold: float) -> np.ndarray:
        """ Perform Non-Maximum Suppression """
        retained_box_indices = []

        areas = (boxes[:, 2] - boxes[:, 0] + 1) * (boxes[:, 3] - boxes[:, 1] + 1)
        ranked_indices = boxes[:, 4].argsort()[::-1]
        while ranked_indices.size > 0:
            best_rest = ranked_indices[0], ranked_indices[1:]

            max_of_xy = np.maximum(boxes[best_rest[0], :2], boxes[best_rest[1], :2])
            min_of_xy = np.minimum(boxes[best_rest[0], 2:4], boxes[best_rest[1], 2:4])
            width_height = np.maximum(0, min_of_xy - max_of_xy + 1)
            intersection_areas = width_height[:, 0] * width_height[:, 1]
            iou = intersection_areas / (areas[best_rest[0]] +
                                        areas[best_rest[1]] - intersection_areas)

            overlapping_boxes = (iou > threshold).nonzero()[0]
            if len(overlapping_boxes) != 0:
                overlap_set = ranked_indices[overlapping_boxes + 1]
                vote = np.average(boxes[overlap_set, :4], axis=0, weights=boxes[overlap_set, 4])
                boxes[best_rest[0], :4] = vote
            retained_box_indices.append(best_rest[0])

            non_overlapping_boxes = (iou <= threshold).nonzero()[0]
            ranked_indices = ranked_indices[non_overlapping_boxes + 1]
        return boxes[retained_box_indices]

    def __call__(self, inputs: np.ndarray) -> np.ndarray:
        """ Get predictions from the S3FD model

        Parameters
        ----------
        inputs: :class:`numpy.ndarray`
            The input to S3FD

        Returns
        -------
        :class:`numpy.ndarray`
            The output from S3FD
        """
        return self._model.predict(inputs, verbose=0, batch_size=self._batch_size)


__all__ = get_module_objects(__name__)
