import pickle

import numpy as np
import pytest

from keras.src import backend
from keras.src import dtype_policies
from keras.src import losses as losses_module
from keras.src import ops
from keras.src import testing
from keras.src.losses.loss import Loss
from keras.src.losses.loss import squeeze_or_expand_to_same_rank


class ExampleLoss(Loss):
    def call(self, y_true, y_pred):
        return (y_true - y_pred) ** 2


class LossTest(testing.TestCase):
    def setUp(self):
        self._global_dtype_policy = dtype_policies.dtype_policy.dtype_policy()
        self._floatx = backend.floatx()
        return super().setUp()

    def tearDown(self):
        dtype_policies.dtype_policy.set_dtype_policy(self._global_dtype_policy)
        backend.set_floatx(self._floatx)
        return super().tearDown()

    def test_squeeze_or_expand(self):
        x1 = ops.ones((3,))
        x2 = ops.ones((3, 1))
        x1, x2 = squeeze_or_expand_to_same_rank(x1, x2)
        self.assertEqual(ops.shape(x1), (3, 1))
        self.assertEqual(ops.shape(x2), (3, 1))

        x1 = ops.ones((3, 2))
        x2 = ops.ones((3, 2, 1))
        x1, x2 = squeeze_or_expand_to_same_rank(x1, x2)
        self.assertEqual(ops.shape(x1), (3, 2))
        self.assertEqual(ops.shape(x2), (3, 2))

        x1 = ops.ones((3,))
        x2 = ops.ones((3, 1))
        x2, x1 = squeeze_or_expand_to_same_rank(x2, x1)
        self.assertEqual(ops.shape(x1), (3, 1))
        self.assertEqual(ops.shape(x2), (3, 1))

        x1 = ops.ones((3, 2))
        x2 = ops.ones((3, 2, 1))
        x2, x1 = squeeze_or_expand_to_same_rank(x2, x1)
        self.assertEqual(ops.shape(x1), (3, 2))
        self.assertEqual(ops.shape(x2), (3, 2))

    def test_reduction(self):
        y_true = np.array([1.0, 0.0, 1.0, 0.0])
        y_pred = np.array([0.1, 0.2, 0.3, 0.4])

        # No reduction
        loss_fn = ExampleLoss(reduction=None)
        loss = loss_fn(y_true, y_pred)
        self.assertEqual(backend.standardize_dtype(loss.dtype), "float32")
        self.assertAllClose((y_true - y_pred) ** 2, loss)

        # sum
        loss_fn = ExampleLoss(reduction="sum")
        loss = loss_fn(y_true, y_pred)
        self.assertEqual(backend.standardize_dtype(loss.dtype), "float32")
        self.assertAllClose(np.sum((y_true - y_pred) ** 2), loss)

        # sum_over_batch_size or mean
        loss_fn = ExampleLoss(reduction="sum_over_batch_size")
        loss = loss_fn(y_true, y_pred)
        self.assertEqual(backend.standardize_dtype(loss.dtype), "float32")
        self.assertAllClose(np.sum((y_true - y_pred) ** 2) / 4, loss)

        # bad reduction
        with self.assertRaisesRegex(ValueError, "Invalid value for argument"):
            ExampleLoss(reduction="abc")

    @pytest.mark.skipif(
        backend.backend() == "numpy",
        reason="Numpy backend does not support masking.",
    )
    def test_mask(self):
        mask = np.array([True, False, True, True])
        y_true = np.array([1.0, 0.0, 1.0, 0.0])
        y_pred = np.array([0.1, 0.2, 0.3, 0.4])

        masked_y_true = np.array([1.0, 1.0, 0.0])
        masked_y_pred = np.array([0.1, 0.3, 0.4])

        mask = ops.convert_to_tensor(mask)
        y_true = ops.convert_to_tensor(y_true)
        y_pred = ops.convert_to_tensor(y_pred)
        y_pred._keras_mask = mask

        loss_fn = ExampleLoss()
        loss = loss_fn(y_true, y_pred)
        self.assertEqual(backend.standardize_dtype(loss.dtype), "float32")
        self.assertAllClose(
            np.sum((masked_y_true - masked_y_pred) ** 2) / 3, loss
        )

        # Test edge case where everything is masked.
        mask = np.array([False, False, False, False])
        y_pred._keras_mask = mask
        loss = loss_fn(y_true, y_pred)
        self.assertEqual(backend.standardize_dtype(loss.dtype), "float32")
        self.assertAllClose(loss, 0)  # No NaN.

    def test_sample_weight(self):
        sample_weight = np.array([0.4, 0.3, 0.2, 0.1])
        y_true = np.array([1.0, 0.0, 1.0, 0.0])
        y_pred = np.array([0.1, 0.2, 0.3, 0.4])

        loss_fn = ExampleLoss()
        loss = loss_fn(y_true, y_pred, sample_weight=sample_weight)
        self.assertEqual(backend.standardize_dtype(loss.dtype), "float32")
        self.assertAllClose(
            np.sum(sample_weight * (y_true - y_pred) ** 2) / 4, loss
        )

        # Test edge case where every weight is 0.
        sample_weight = np.array([0.0, 0.0, 0.0, 0.0])
        loss = loss_fn(y_true, y_pred, sample_weight=sample_weight)
        self.assertEqual(backend.standardize_dtype(loss.dtype), "float32")
        self.assertAllClose(loss, 0)  # No NaN.

    @pytest.mark.skipif(
        backend.backend() == "numpy",
        reason="Numpy backend does not support masking.",
    )
    def test_mask_and_sample_weight(self):
        sample_weight = np.array([0.4, 0.3, 0.2, 0.1])
        y_true = np.array([1.0, 0.0, 1.0, 0.0])
        y_pred = np.array([0.1, 0.2, 0.3, 0.4])
        mask = np.array([True, False, True, True])

        masked_sample_weight = np.array([0.4, 0.2, 0.1])
        masked_y_true = np.array([1.0, 1.0, 0.0])
        masked_y_pred = np.array([0.1, 0.3, 0.4])

        mask = ops.convert_to_tensor(mask)
        y_true = ops.convert_to_tensor(y_true)
        y_pred = ops.convert_to_tensor(y_pred)
        y_pred._keras_mask = mask

        loss_fn = ExampleLoss()
        loss = loss_fn(y_true, y_pred, sample_weight=sample_weight)
        self.assertEqual(backend.standardize_dtype(loss.dtype), "float32")
        self.assertAllClose(
            np.sum(masked_sample_weight * (masked_y_true - masked_y_pred) ** 2)
            / 3,
            loss,
        )

    @pytest.mark.skipif(
        backend.backend() == "numpy",
        reason="Numpy backend does not support masking.",
    )
    def test_mask_and_sample_weight_rank2(self):
        # check loss of inputs with duplicate rows doesn't change
        sample_weight = np.array([0.4, 0.3, 0.2, 0.1])
        y_true = np.array([1.0, 0.0, 1.0, 0.0])
        y_pred = np.array([0.1, 0.2, 0.3, 0.4])
        mask = np.array([True, False, True, True])

        mask = ops.convert_to_tensor(mask)
        y_true = ops.convert_to_tensor(y_true)
        y_pred = ops.convert_to_tensor(y_pred)
        y_pred._keras_mask = mask

        loss_fn = ExampleLoss()
        rank1_loss = loss_fn(y_true, y_pred, sample_weight=sample_weight)

        # duplicate rows
        mask = ops.tile(ops.expand_dims(mask, axis=0), (2, 1))
        y_true = ops.tile(ops.expand_dims(y_true, axis=0), (2, 1))
        y_pred = ops.tile(ops.expand_dims(y_pred, axis=0), (2, 1))
        sample_weight = ops.tile(ops.expand_dims(sample_weight, axis=0), (2, 1))
        y_pred._keras_mask = mask
        rank2_loss = loss_fn(y_true, y_pred, sample_weight=sample_weight)
        self.assertAllClose(rank1_loss, rank2_loss)

    # @testing.parametrize(
    #     "uprank", ["mask", "sample_weight", "y_true", "y_pred"])
    # TODO: use parameterization decorator
    @pytest.mark.skipif(
        backend.backend() == "numpy",
        reason="Numpy backend does not support masking.",
    )
    def test_rank_adjustment(self):
        for uprank in ["mask", "sample_weight", "ys"]:
            sample_weight = np.array([0.4, 0.3, 0.2, 0.1])
            y_true = np.array([1.0, 0.0, 1.0, 0.0])
            y_pred = np.array([0.1, 0.2, 0.3, 0.4])
            mask = np.array([True, False, True, True])

            if uprank == "mask":
                mask = np.expand_dims(mask, -1)
            elif uprank == "sample_weight":
                sample_weight = np.expand_dims(sample_weight, -1)
            elif uprank == "ys":
                y_true = np.expand_dims(y_true, -1)
                y_pred = np.expand_dims(y_pred, -1)

            masked_sample_weight = np.array([0.4, 0.2, 0.1])
            masked_y_true = np.array([1.0, 1.0, 0.0])
            masked_y_pred = np.array([0.1, 0.3, 0.4])

            mask = ops.convert_to_tensor(mask)
            y_true = ops.convert_to_tensor(y_true)
            y_pred = ops.convert_to_tensor(y_pred)
            y_pred._keras_mask = mask

            loss_fn = ExampleLoss()
            loss = loss_fn(y_true, y_pred, sample_weight=sample_weight)
            self.assertEqual(backend.standardize_dtype(loss.dtype), "float32")
            self.assertAllClose(
                np.sum(
                    masked_sample_weight * (masked_y_true - masked_y_pred) ** 2
                )
                / 3,
                loss,
            )

    def test_mixed_dtypes(self):
        sample_weight = np.array([0.4, 0.3, 0.2, 0.1], dtype="float64")
        y_true = np.array([1.0, 0.0, 1.0, 0.0], dtype="int32")
        y_pred = np.array([0.1, 0.2, 0.3, 0.4], dtype="float32")

        loss_fn = ExampleLoss()
        loss = loss_fn(y_true, y_pred, sample_weight=sample_weight)
        self.assertEqual(backend.standardize_dtype(loss.dtype), "float32")
        self.assertAllClose(
            np.sum(sample_weight * (y_true - y_pred) ** 2) / 4,
            loss,
        )

    def test_pickle(self):
        loss = losses_module.get("mse")
        loss = pickle.loads(pickle.dumps(loss))
        self.assertEqual(loss, losses_module.mean_squared_error)

    def test_get_method(self):
        loss = losses_module.get("mse")
        self.assertEqual(loss, losses_module.mean_squared_error)

        loss = losses_module.get(None)
        self.assertEqual(loss, None)

        with self.assertRaises(ValueError):
            losses_module.get("typo")

    def test_dtype_arg(self):
        y_true = np.array([1.0, 0.0, 1.0, 0.0], dtype="float32")
        y_pred = np.array([0.1, 0.2, 0.3, 0.4], dtype="float32")

        # Note: we use float16 and not float64 to test this because
        # JAX will map float64 to float32.
        loss_fn = ExampleLoss(dtype="float16")
        loss = loss_fn(y_true, y_pred)
        self.assertDType(loss, "float16")

        # Test DTypePolicy for `dtype` argument
        loss_fn = ExampleLoss(dtype=dtype_policies.DTypePolicy("mixed_float16"))
        loss = loss_fn(y_true, y_pred)
        self.assertDType(loss, "float16")

        # `dtype` setter should raise AttributeError
        with self.assertRaises(AttributeError):
            loss_fn.dtype = "bfloat16"

    def test_default_dtype(self):
        y_true = np.array([1.0, 0.0, 1.0, 0.0], dtype="float32")
        y_pred = np.array([0.1, 0.2, 0.3, 0.4], dtype="float32")

        # Defaults to `keras.config.floatx()` not global `dtype_policy`
        dtype_policies.dtype_policy.set_dtype_policy("mixed_float16")
        loss_fn = ExampleLoss()
        loss = loss_fn(y_true, y_pred)
        self.assertDType(loss, "float32")

        backend.set_floatx("float16")
        loss_fn = ExampleLoss()
        loss = loss_fn(y_true, y_pred)
        self.assertDType(loss, backend.floatx())