# -*- coding: utf-8 -*-
# GT4Py - GridTools Framework
#
# Copyright (c) 2014-2023, ETH Zurich
# All rights reserved.
#
# This file is part of the GT4Py project and the GridTools framework.
# GT4Py is free software: you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the
# Free Software Foundation, either version 3 of the License, or any later
# version. See the LICENSE.txt file at the top-level directory of this
# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
#
# SPDX-License-Identifier: GPL-3.0-or-later

import numpy as np
import pytest

import gt4py.next as gtx
from gt4py.next import broadcast, float64, int32, int64, max_over, min_over, neighbor_sum, where
from gt4py.next.program_processors.runners import dace_iterator, gtfn_cpu

from next_tests.integration_tests import cases
from next_tests.integration_tests.cases import (
    V2E,
    Edge,
    IDim,
    Ioff,
    JDim,
    Joff,
    KDim,
    V2EDim,
    Vertex,
    cartesian_case,
    unstructured_case,
)
from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import (
    fieldview_backend,
    reduction_setup,
)


def test_maxover_execution(unstructured_case):
    """Testing max_over functionality."""

    if unstructured_case.backend in [
        gtfn_cpu.run_gtfn or fieldview_backend,
        gtfn_cpu.run_gtfn_imperative,
        dace_iterator.run_dace_iterator,
    ]:
        pytest.skip("not yet supported.")

    @gtx.field_operator
    def maxover_fieldoperator(
        inp_field: gtx.Field[[Vertex, V2EDim], int32]
    ) -> gtx.Field[[Vertex], int32]:
        return max_over(inp_field, axis=V2EDim)

    inp = gtx.np_as_located_field(Vertex, V2EDim)(unstructured_case.offset_provider["V2E"].table)

    cases.verify(
        unstructured_case,
        maxover_fieldoperator,
        inp,
        out=cases.allocate(unstructured_case, maxover_fieldoperator, cases.RETURN)(),
        ref=np.max(unstructured_case.offset_provider["V2E"].table, axis=1),
    )


def test_maxover_execution_negatives(unstructured_case):
    """Testing max_over functionality for negative values in array."""

    if unstructured_case.backend in [
        gtfn_cpu.run_gtfn,
        gtfn_cpu.run_gtfn_imperative,
        dace_iterator.run_dace_iterator,
    ]:
        pytest.skip("not yet supported.")

    v2e_table = unstructured_case.offset_provider["V2E"].table

    @gtx.field_operator
    def maxover_negvals(edge_f: cases.EField) -> cases.VField:
        out = max_over(edge_f(V2E), axis=V2EDim)
        return out

    cases.verify_with_default_data(
        unstructured_case, maxover_negvals, ref=lambda edge_f: np.max(edge_f[v2e_table], axis=1)
    )


def test_minover_execution(unstructured_case):
    """Testing max_over functionality for negative values in array."""
    if unstructured_case.backend in [
        gtfn_cpu.run_gtfn,
        gtfn_cpu.run_gtfn_imperative,
        dace_iterator.run_dace_iterator,
    ]:
        pytest.skip("not yet supported.")

    v2e_table = unstructured_case.offset_provider["V2E"].table

    @gtx.field_operator
    def maxover_negvals(edge_f: cases.EField) -> cases.VField:
        out = min_over(edge_f(V2E), axis=V2EDim)
        return out

    cases.verify_with_default_data(
        unstructured_case, maxover_negvals, ref=lambda edge_f: np.min(edge_f[v2e_table], axis=1)
    )


def test_reduction_execution(unstructured_case):
    """Testing a trivial neighbor sum."""

    if unstructured_case.backend == dace_iterator.run_dace_iterator:
        pytest.xfail("Not supported in DaCe backend: reductions")

    @gtx.field_operator
    def reduction(edge_f: cases.EField) -> cases.VField:
        return neighbor_sum(edge_f(V2E), axis=V2EDim)

    @gtx.program
    def fencil(edge_f: cases.EField, out: cases.VField):
        reduction(edge_f, out=out)

    cases.verify_with_default_data(
        unstructured_case,
        fencil,
        ref=lambda edge_f: np.sum(edge_f[unstructured_case.offset_provider["V2E"].table], axis=1),
    )


def test_reduction_expression(unstructured_case):
    """Test reduction with an expression directly inside the call."""
    if unstructured_case.backend in [gtfn_cpu.run_gtfn, gtfn_cpu.run_gtfn_imperative]:
        pytest.skip("Has a bug.")
    if unstructured_case.backend == dace_iterator.run_dace_iterator:
        pytest.xfail("Not supported in DaCe backend: reductions")

    @gtx.field_operator
    def reduce_expr(edge_f: cases.EField) -> cases.VField:
        tmp_nbh_tup = edge_f(V2E), edge_f(V2E)
        tmp_nbh = tmp_nbh_tup[0]
        return 3 * neighbor_sum(-edge_f(V2E) * tmp_nbh * 2, axis=V2EDim)

    @gtx.program
    def fencil(edge_f: cases.EField, out: cases.VField):
        reduce_expr(edge_f, out=out)

    cases.verify_with_default_data(
        unstructured_case,
        fencil,
        ref=lambda edge_f: 3
        * np.sum(-edge_f[unstructured_case.offset_provider["V2E"].table] ** 2 * 2, axis=1),
    )


def test_reduction_with_common_expression(unstructured_case):
    if unstructured_case.backend == dace_iterator.run_dace_iterator:
        pytest.xfail("Not supported in DaCe backend: reductions")

    @gtx.field_operator
    def testee(flux: cases.EField) -> cases.VField:
        return neighbor_sum(flux(V2E) + flux(V2E), axis=V2EDim)

    cases.verify_with_default_data(
        unstructured_case,
        testee,
        ref=lambda flux: np.sum(flux[unstructured_case.offset_provider["V2E"].table] * 2, axis=1),
    )


def test_conditional_nested_tuple(cartesian_case):
    if cartesian_case.backend == dace_iterator.run_dace_iterator:
        pytest.xfail("Not supported in DaCe backend: tuple returns")

    @gtx.field_operator
    def conditional_nested_tuple(
        mask: cases.IBoolField, a: cases.IFloatField, b: cases.IFloatField
    ) -> tuple[
        tuple[cases.IFloatField, cases.IFloatField],
        tuple[cases.IFloatField, cases.IFloatField],
    ]:
        return where(mask, ((a, b), (b, a)), ((5.0, 7.0), (7.0, 5.0)))

    size = cartesian_case.default_sizes[IDim]
    bool_field = np.random.choice(a=[False, True], size=(size))
    mask = cases.allocate(cartesian_case, conditional_nested_tuple, "mask").strategy(
        cases.ConstInitializer(bool_field)
    )()
    a = cases.allocate(cartesian_case, conditional_nested_tuple, "a")()
    b = cases.allocate(cartesian_case, conditional_nested_tuple, "b")()

    cases.verify(
        cartesian_case,
        conditional_nested_tuple,
        mask,
        a,
        b,
        out=cases.allocate(cartesian_case, conditional_nested_tuple, cases.RETURN)(),
        ref=np.where(
            mask,
            ((a, b), (b, a)),
            ((np.full(size, 5.0), np.full(size, 7.0)), (np.full(size, 7.0), np.full(size, 5.0))),
        ),
    )


def test_broadcast_simple(cartesian_case):
    if cartesian_case.backend == dace_iterator.run_dace_iterator:
        pytest.xfail("Not supported in DaCe backend: broadcast")

    @gtx.field_operator
    def simple_broadcast(inp: cases.IField) -> cases.IJField:
        return broadcast(inp, (IDim, JDim))

    cases.verify_with_default_data(
        cartesian_case, simple_broadcast, ref=lambda inp: inp[:, np.newaxis]
    )


def test_broadcast_scalar(cartesian_case):
    if cartesian_case.backend == dace_iterator.run_dace_iterator:
        pytest.xfail("Not supported in DaCe backend: broadcast")
    size = cartesian_case.default_sizes[IDim]

    @gtx.field_operator
    def scalar_broadcast() -> gtx.Field[[IDim], float64]:
        return broadcast(float(1.0), (IDim,))

    cases.verify_with_default_data(cartesian_case, scalar_broadcast, ref=lambda: np.ones(size))


def test_broadcast_two_fields(cartesian_case):
    if cartesian_case.backend == dace_iterator.run_dace_iterator:
        pytest.xfail("Not supported in DaCe backend: broadcast")

    @gtx.field_operator
    def broadcast_two_fields(inp1: cases.IField, inp2: gtx.Field[[JDim], int32]) -> cases.IJField:
        a = broadcast(inp1, (IDim, JDim))
        b = broadcast(inp2, (IDim, JDim))
        return a + b

    cases.verify_with_default_data(
        cartesian_case, broadcast_two_fields, ref=lambda a, b: a[:, np.newaxis] + b[np.newaxis, :]
    )


def test_broadcast_shifted(cartesian_case):
    if cartesian_case.backend == dace_iterator.run_dace_iterator:
        pytest.xfail("Not supported in DaCe backend: broadcast")

    @gtx.field_operator
    def simple_broadcast(inp: cases.IField) -> cases.IJField:
        bcasted = broadcast(inp, (IDim, JDim))
        return bcasted(Joff[1])

    cases.verify_with_default_data(
        cartesian_case, simple_broadcast, ref=lambda inp: inp[:, np.newaxis]
    )


def test_conditional(cartesian_case):
    @gtx.field_operator
    def conditional(
        mask: cases.IBoolField, a: cases.IFloatField, b: cases.IFloatField
    ) -> cases.IFloatField:
        return where(mask, a, b)

    size = cartesian_case.default_sizes[IDim]
    bool_field = np.random.choice(a=[False, True], size=(size))
    mask = cases.allocate(cartesian_case, conditional, "mask").strategy(
        cases.ConstInitializer(bool_field)
    )()
    a = cases.allocate(cartesian_case, conditional, "a")()
    b = cases.allocate(cartesian_case, conditional, "b")()
    out = cases.allocate(cartesian_case, conditional, cases.RETURN)()

    cases.verify(cartesian_case, conditional, mask, a, b, out=out, ref=np.where(mask, a, b))


def test_conditional_promotion(cartesian_case):
    @gtx.field_operator
    def conditional_promotion(mask: cases.IBoolField, a: cases.IFloatField) -> cases.IFloatField:
        return where(mask, a, 10.0)

    size = cartesian_case.default_sizes[IDim]
    bool_field = np.random.choice(a=[False, True], size=(size))
    mask = cases.allocate(cartesian_case, conditional_promotion, "mask").strategy(
        cases.ConstInitializer(bool_field)
    )()
    a = cases.allocate(cartesian_case, conditional_promotion, "a")()
    out = cases.allocate(cartesian_case, conditional_promotion, cases.RETURN)()

    cases.verify(
        cartesian_case, conditional_promotion, mask, a, out=out, ref=np.where(mask, a, 10.0)
    )


def test_conditional_compareop(cartesian_case):
    @gtx.field_operator
    def conditional_promotion(a: cases.IFloatField) -> cases.IFloatField:
        return where(a != a, a, 10.0)

    cases.verify_with_default_data(
        cartesian_case, conditional_promotion, ref=lambda a: np.where(a != a, a, 10.0)
    )


def test_conditional_shifted(cartesian_case):
    @gtx.field_operator
    def conditional_shifted(
        mask: cases.IBoolField, a: cases.IFloatField, b: cases.IFloatField
    ) -> gtx.Field[[IDim], float64]:
        tmp = where(mask, a, b)
        return tmp(Ioff[1])

    @gtx.program
    def conditional_program(
        mask: cases.IBoolField,
        a: cases.IFloatField,
        b: cases.IFloatField,
        out: cases.IFloatField,
    ):
        conditional_shifted(mask, a, b, out=out)

    size = cartesian_case.default_sizes[IDim] + 1
    mask = gtx.np_as_located_field(IDim)(np.random.choice(a=[False, True], size=(size)))
    a = cases.allocate(cartesian_case, conditional_program, "a").extend({IDim: (0, 1)})()
    b = cases.allocate(cartesian_case, conditional_program, "b").extend({IDim: (0, 1)})()
    out = cases.allocate(cartesian_case, conditional_shifted, cases.RETURN)()

    cases.verify(
        cartesian_case,
        conditional_program,
        mask,
        a,
        b,
        out,
        inout=out,
        ref=np.where(mask, a, b)[1:],
    )


def test_promotion(unstructured_case):
    @gtx.field_operator
    def promotion(
        inp1: gtx.Field[[Edge, KDim], float64], inp2: gtx.Field[[KDim], float64]
    ) -> gtx.Field[[Edge, KDim], float64]:
        return inp1 / inp2

    cases.verify_with_default_data(unstructured_case, promotion, ref=lambda inp1, inp2: inp1 / inp2)
