# (C) Copyright 2018- ECMWF.
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.

from pathlib import Path
import pytest
import numpy as np

from loki.ir import Loop, FindNodes
from loki.scope import Scope
from loki.sourcefile import Sourcefile
from loki.analyse.util_polyhedron import Polyhedron
from loki.expression import symbols as sym, parse_expr

@pytest.fixture(scope="module", name="here")
def fixture_here():
    return Path(__file__).parent


@pytest.fixture(scope='module', name='testdir')
def fixture_testdir(here):
    return here.parent.parent/'tests'



@pytest.mark.parametrize(
    "variables, lbounds, ubounds, A, b, variable_names",
    [
        # do i=0,5: do j=i,7: ...
        (
            ["i", "j"],
            ["0", "i"],
            ["5", "7"],
            [[-1, 0], [1, 0], [1, -1], [0, 1]],
            [0, 5, 0, 7],
            ["i", "j"],
        ),
        # do i=1,n: do j=0,2*i+1: do k=a,b: ...
        (
            ["i", "j", "k"],
            ["1", "0", "a"],
            ["n", "2*i+1", "b"],
            [
                [-1, 0, 0, 0, 0, 0],
                [1, 0, 0, 0, 0, -1],
                [0, -1, 0, 0, 0, 0],
                [-2, 1, 0, 0, 0, 0],
                [0, 0, -1, 1, 0, 0],
                [0, 0, 1, 0, -1, 0],
            ],
            [-1, 0, 0, 1, 0, 0],
            ["i", "j", "k", "a", "b", "n"],
        ),
        # do jk=1,klev: ...
        (["jk"], ["1"], ["klev"], [[-1, 0], [1, -1]], [-1, 0], ["jk", "klev"]),
        # do JK=1,klev-1: ...
        (["JK"], ["1"], ["klev - 1"], [[-1, 0], [1, -1]], [-1, -1], ["jk", "klev"]),
        # do jk=ncldtop,klev: ...
        (
            ["jk"],
            ["ncldtop"],
            ["klev"],
            [[-1, 0, 1], [1, -1, 0]],
            [0, 0],
            ["jk", "klev", "ncldtop"],
        ),
        # do jk=1,KLEV+1: ...
        (["jk"], ["1"], ["KLEV+1"], [[-1, 0], [1, -1]], [-1, 1], ["jk", "klev"]),
    ],
)
def test_polyhedron_from_loop_ranges(variables, lbounds, ubounds, A, b, variable_names):
    """
    Test converting loop ranges to polyedron representation of iteration space.
    """
    scope = Scope()
    loop_variables = [parse_expr(expr, scope) for expr in variables]
    loop_lbounds = [parse_expr(expr, scope) for expr in lbounds]
    loop_ubounds = [parse_expr(expr, scope) for expr in ubounds]
    loop_ranges = [sym.LoopRange((l, u)) for l, u in zip(loop_lbounds, loop_ubounds)]
    p = Polyhedron.from_loop_ranges(loop_variables, loop_ranges)
    assert np.all(p.A == np.array(A, dtype=np.dtype(int)))
    assert np.all(p.b == np.array(b, dtype=np.dtype(int)))
    assert p.variables == variable_names


def test_polyhedron_from_loop_ranges_failures():
    """
    Test known limitation of the conversion from loop ranges to polyhedron.
    """
    # m*n is non-affine and thus can't be represented
    scope = Scope()
    loop_variable = parse_expr("i", scope)
    lower_bound = parse_expr("1", scope)
    upper_bound = parse_expr("m * n", scope)
    loop_range = sym.LoopRange((lower_bound, upper_bound))
    with pytest.raises(ValueError):
        _ = Polyhedron.from_loop_ranges([loop_variable], [loop_range])

    # no functionality to flatten exponentials, yet
    upper_bound = parse_expr("5**2", scope)
    loop_range = sym.LoopRange((lower_bound, upper_bound))
    with pytest.raises(ValueError):
        _ = Polyhedron.from_loop_ranges([loop_variable], [loop_range])


@pytest.mark.parametrize(
    "A, b, variable_names, lower_bounds, upper_bounds",
    [
        # do i=1,n: ...
        ([[-1, 0], [1, -1]], [-1, 0], ["i", "n"], [["1"], ["i"]], [["n"], []]),
        # do i=1,10: ...
        ([[-1], [1]], [-1, 10], ["i"], [["1"]], [["10"]]),
        # do i=0,5: do j=i,7: ...
        (
            [[-1, 0], [1, 0], [1, -1], [0, 1]],
            [0, 5, 0, 7],
            ["i", "j"],
            [["0"], ["i"]],
            [["5", "j"], ["7"]],
        ),
        # do i=1,n: do j=0,2*i+1: do k=a,b: ...
        (
            [
                [-1, 0, 0, 0, 0, 0],
                [1, 0, 0, 0, 0, -1],
                [0, -1, 0, 0, 0, 0],
                [-2, 1, 0, 0, 0, 0],
                [0, 0, -1, 1, 0, 0],
                [0, 0, 1, 0, -1, 0],
            ],
            [-1, 0, 0, 1, 0, 0],
            ["i", "j", "k", "a", "b", "n"],  # variable names
            [["1", "-1 / 2 + j / 2"], ["0"], ["a"], [], ["k"], ["i"]],  # lower bounds
            [["n"], ["1 + 2*i"], ["b"], ["k"], [], []],
        ),  # upper bounds
    ],
)
def test_polyhedron_bounds(A, b, variable_names, lower_bounds, upper_bounds):
    """
    Test the production of lower and upper bounds.
    """
    scope = Scope()
    variables = [parse_expr(v, scope) for v in variable_names]
    p = Polyhedron(A, b, variables)
    for var, ref_bounds in zip(variables, lower_bounds):
        lbounds = p.lower_bounds(var)
        assert len(lbounds) == len(ref_bounds)
        assert all(str(b1) == b2 for b1, b2 in zip(lbounds, ref_bounds))
    for var, ref_bounds in zip(variables, upper_bounds):
        ubounds = p.upper_bounds(var)
        assert len(ubounds) == len(ref_bounds)
        assert all(str(b1) == b2 for b1, b2 in zip(ubounds, ref_bounds))


@pytest.mark.parametrize(
    "polyhedron,is_empty,will_fail",
    [
        # totaly empty polyhedron
        (Polyhedron.from_nested_loops([]), True, False),
        # full matrix --> non trivial problem
        (Polyhedron([[1]], [1]), None, True),
        # empty matrix, full and fullfiled b --> non empty polyhedron
        (Polyhedron([[]], [1]), False, False),
        # empty matrix, full b but not fullfiled b --> empty polyhedron
        (Polyhedron([[]], [-1]), True, False),
    ],
)
def test_check_empty_polyhedron(polyhedron, is_empty, will_fail):
    if will_fail:
        with pytest.raises(RuntimeError):
            _ = polyhedron.is_empty()
    else:
        assert polyhedron.is_empty() == is_empty


def simple_loop_extractor(start_node):
    """Find all loops in the AST and structure them depending on their nesting level"""
    start_loops = FindNodes(Loop, greedy=True).visit(start_node)
    return [FindNodes(Loop).visit(node) for node in start_loops]


def assert_equal_polyhedron(poly_A, poly_B):
    assert poly_A.variables == poly_B.variables
    assert (poly_A.A == poly_B.A).all()
    assert (poly_A.b == poly_B.b).all()


@pytest.mark.parametrize(
    "filename, loop_extractor, polyhedrons_per_subroutine",
    [
        (
            "sources/data_dependency_detection/loop_carried_dependencies.f90",
            simple_loop_extractor,
            {
                "SimpleDependency": [
                    Polyhedron(
                        [[-1, 0], [1, -1]], [-1, 0], [sym.Scalar("i"), sym.Scalar("n")]
                    ),
                ],
                "NestedDependency": [
                    Polyhedron(
                        [[-1, 0, 0], [1, 0, -1], [0, -1, 0], [-1, 1, 0]],
                        [-2, 0, -1, -1],
                        [sym.Scalar("i"), sym.Scalar("j"), sym.Scalar("n")],
                    ),
                ],
                "ConditionalDependency": [
                    Polyhedron(
                        [[-1, 0], [1, -1]],
                        [-2, 0],
                        [sym.Scalar("i"), sym.Scalar("n")],
                    ),
                ],
                "NoDependency": [
                    Polyhedron(
                        [[-1], [1]],
                        [-1, 10],
                        [sym.Scalar("i")],
                    ),
                    Polyhedron(
                        [[-1], [1]],
                        [-1, 5],
                        [sym.Scalar("i")],
                    ),
                ],
            },
        ),
        (
            "sources/data_dependency_detection/various_loops.f90",
            simple_loop_extractor,
            {
                "single_loop": [
                    Polyhedron(
                        [[-1, 0], [1, -1]],
                        [-1, 0],
                        [sym.Scalar("i"), sym.Scalar("n")],
                    ),
                ],
                "single_loop_split_access": [
                    Polyhedron(
                        [[-1, 0], [1, -1]],
                        [-1, 0],
                        [sym.Scalar("i"), sym.Scalar("nhalf")],
                    ),
                ],
                "single_loop_arithmetic_operations_for_access": [
                    Polyhedron(
                        [[-1, 0], [1, -1]],
                        [-1, 0],
                        [sym.Scalar("i"), sym.Scalar("n")],
                    ),
                ],
                "nested_loop_single_dimensions_access": [
                    Polyhedron(
                        [[-1, 0, 0], [1, 0, -1], [0, -1, 0], [0, 1, -1]],
                        [-1, 0, -1, 0],
                        [sym.Scalar("i"), sym.Scalar("j"), sym.Scalar("nhalf")],
                    ),
                ],
                "nested_loop_partially_used": [
                    Polyhedron(
                        [[-1, 0, 0], [1, 0, -1], [0, -1, 0], [0, 1, -1]],
                        [-1, 0, -1, 0],
                        [sym.Scalar("i"), sym.Scalar("j"), sym.Scalar("nfourth")],
                    ),
                ],
                "partially_used_array": [
                    Polyhedron(
                        [[-1, 0], [1, -1]],
                        [-2, 0],
                        [sym.Scalar("i"), sym.Scalar("nhalf")],
                    ),
                ],
            },
        ),
    ],
)
def test_polyhedron_construction_from_nested_loops(
    testdir, filename, loop_extractor, polyhedrons_per_subroutine
):
    source = Sourcefile.from_file(testdir / filename)

    for subroutine in source.all_subroutines:
        expected_polyhedrons = polyhedrons_per_subroutine[subroutine.name]

        list_of_loops = loop_extractor(subroutine.body)

        polyhedrons = [Polyhedron.from_nested_loops(loops) for loops in list_of_loops]

        for polyhedron, expected_polyhedron in zip(polyhedrons, expected_polyhedrons):
            assert_equal_polyhedron(polyhedron, expected_polyhedron)
