# original implementation is from treex https://github.com/cgarciae/treex
# which is itself based on Tensorflow Keras:
# https://github.com/tensorflow/tensorflow/blob/v2.2.0/tensorflow/python/keras/losses.py#L44-L201

import typing as tp
from abc import abstractmethod
from enum import Enum

import jax.numpy as jnp
import numpy as np
from flax import struct

from ..utils import _get_name

IndexLike = tp.Union[str, int, tp.Sequence[tp.Union[str, int]]]
PathLike = tp.Tuple[IndexLike, ...]
ScalarLike = tp.Union[float, np.ndarray, jnp.ndarray]


class Reduction(Enum):
    """
    Types of loss reduction.

    Contains the following values:
    * `NONE`: Weighted losses with one dimension reduced (axis=-1, or axis
        specified by loss function). When this reduction type used with built-in
        Keras training loops like `fit`/`evaluate`, the unreduced vector loss is
        passed to the optimizer but the reported loss will be a scalar value.
    * `SUM`: Scalar sum of weighted losses.
    * `SUM_OVER_BATCH_SIZE`: Scalar `SUM` divided by number of elements in losses.
    """

    # AUTO = "auto"
    NONE = "none"
    SUM = "sum"
    SUM_OVER_BATCH_SIZE = "sum_over_batch_size"

    @classmethod
    def all(cls):
        return (
            # cls.AUTO,
            cls.NONE,
            cls.SUM,
            cls.SUM_OVER_BATCH_SIZE,
        )

    @classmethod
    def validate(cls, key):
        if key not in cls.all():
            raise ValueError("Invalid Reduction Key %s." % key)


class Loss:
    """
    Loss base class.

    To be implemented by subclasses:

    * `call()`: Contains the logic for loss calculation.

    Example subclass implementation:

    ```python
    class MeanSquaredError(Loss):
        def call(self, target, preds):
            return jnp.mean(jnp.square(preds - target), axis=-1)
    ```

    Please see the [Modules, Losses, and Metrics Guide]
    (https://poets-ai.github.io/elegy/guides/modules-losses-metrics/#losses) for more
    details on this.
    """

    def __init__(
        self,
        reduction: tp.Optional[Reduction] = None,
        weight: tp.Optional[ScalarLike] = None,
        on: tp.Optional[IndexLike] = None,
        name: tp.Optional[str] = None,
    ):
        """
        Initializes `Loss` class.

        Arguments:
            reduction: (Optional) Type of `tx.losses.Reduction` to apply to
                loss. Default value is `SUM_OVER_BATCH_SIZE`. For almost all cases
                this defaults to `SUM_OVER_BATCH_SIZE`.
            weight: Optional weight contribution for the total loss. Defaults to `1`.
            on: A string or integer, or iterable of string or integers, that
                indicate how to index/filter the `target` and `preds`
                arguments before passing them to `call`. For example if `on = "a"` then
                `target = target["a"]`. If `on` is an iterable
                the structures will be indexed iteratively, for example if `on = ["a", 0, "b"]`
                then `target = target["a"][0]["b"]`, same for `preds`. For more information
                check out [Keras-like behavior](https://poets-ai.github.io/elegy/guides/modules-losses-metrics/#keras-like-behavior).
            name: Optional name for the instance, if not provided lower snake_case version
                of the name of the class is used instead.
        """
        self.name = name if name is not None else _get_name(self)
        self.weight = (
            jnp.asarray(weight, dtype=jnp.float32)
            if weight is not None
            else jnp.array(1.0, dtype=jnp.float32)
        )
        self._reduction = reduction if reduction is not None else Reduction.NONE
        self._labels_filter = (on,) if isinstance(on, (str, int)) else on
        self._signature_f = self.call

    def __call__(
        self,
        **kwargs,
    ) -> jnp.ndarray:

        if self._labels_filter is not None:
            if "target" in kwargs and kwargs["target"] is not None:
                for index in self._labels_filter:
                    kwargs["target"] = kwargs["target"][index]

            if "preds" in kwargs and kwargs["preds"] is not None:
                for index in self._labels_filter:
                    kwargs["preds"] = kwargs["preds"][index]

        sample_weight: tp.Optional[jnp.ndarray] = kwargs.pop("sample_weight", None)

        values = self.call(**kwargs)

        return reduce_loss(values, sample_weight, self.weight, self._reduction)

    def __hash__(self):
        return hash(self.name)

    @abstractmethod
    def call(self, *args, **kwargs) -> jnp.ndarray:
        ...


def reduce_loss(
    values: jnp.ndarray, sample_weight: tp.Optional[jnp.ndarray], weight, reduction
) -> jnp.ndarray:

    values = jnp.asarray(values)

    if sample_weight is not None:
        # expand `sample_weight` dimensions until it has the same rank as `values`
        while sample_weight.ndim < values.ndim:
            sample_weight = sample_weight[..., jnp.newaxis]

        values *= sample_weight

    if reduction == Reduction.NONE:
        loss = values
    elif reduction == Reduction.SUM:
        loss = jnp.sum(values)
    elif reduction == Reduction.SUM_OVER_BATCH_SIZE:
        loss = jnp.sum(values) / jnp.prod(jnp.array(values.shape))
    else:
        raise ValueError(f"Invalid reduction '{reduction}'")

    return loss * weight


class LossLog(struct.PyTreeNode):
    loss_fn: Loss = struct.field(pytree_node=False)
    weight: float = 1.0
    cnt: jnp.ndarray = jnp.array(0.0)
    sum: jnp.ndarray = jnp.array(0.0)

    def update(self, **kwargs):
        loss = self.loss_fn(**kwargs) * self.weight
        new_log = self.replace(cnt=self.cnt + 1, sum=self.sum + loss)
        return loss, new_log

    def compute(self):
        return self.sum / self.cnt
