# (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.


import pytest
import numpy as np

from loki import Subroutine
from loki.jit_build import jit_compile_and_run, clean_test
from loki.expression import symbols as sym, Array
from loki.frontend import available_frontends
from loki.ir import (
    nodes as ir, FindNodes, FindVariables, pragmas_attached
)

from loki.transformations.utilities import find_driver_loops
from loki.transformations.loop_blocking import (
    LoopSplittingVariables, split_loop, block_loop_arrays
)
from loki.types import SymbolAttributes, BasicType
from loki.logging import ERROR

"""
Tests for loop splitting and blocking utilities.
"""

# The number of additional variables created during loop splitting,
# e.g. new loop index for outer loop, new loop bounds
LOKI_LOOP_SLIT_VAR_ADDITION = 7


def loop_symbols(node):
    return [
        (loop.variable, loop.bounds.start, loop.bounds.stop, loop.bounds.step)
        for loop in FindNodes(ir.Loop).visit(node)
    ]


def assignment_symbols(node):
    return [(assign.lhs, assign.rhs) for assign in FindNodes(ir.Assignment).visit(node)]


def array_access_symbols(node):
    return [
        (var.name, var.dimensions) for var in FindVariables().visit(node)
        if isinstance(var, Array)
    ]


def test_loop_splitting_vars(caplog):
    loop_var = sym.Variable(name='i', type=SymbolAttributes(BasicType.INTEGER))

    # Test init with int block size
    blk_size = 1
    loop_splitting_vars = LoopSplittingVariables(loop_var, blk_size)
    loop_var = loop_splitting_vars.loop_var
    assert isinstance(loop_var, sym.Scalar) and loop_var.type.dtype == BasicType.INTEGER

    # Test init with scalar
    blk_size = sym.Variable(name='i', type=SymbolAttributes(BasicType.INTEGER))
    loop_splitting_vars = LoopSplittingVariables(loop_var, blk_size)
    loop_var = loop_splitting_vars.loop_var
    assert isinstance(loop_var, sym.Scalar) and loop_var.type.dtype == BasicType.INTEGER

    # Test init with scalar
    blk_size = sym.IntLiteral(1)
    loop_splitting_vars = LoopSplittingVariables(loop_var, blk_size)
    loop_var = loop_splitting_vars.loop_var
    assert isinstance(loop_var, sym.Scalar) and loop_var.type.dtype == BasicType.INTEGER

    # Test init with bad value
    blk_size = 'i'
    with caplog.at_level(ERROR):
        with pytest.raises(ValueError):
            loop_splitting_vars = LoopSplittingVariables(loop_var, blk_size)
        assert len(caplog.records) == 1
        assert ("LoopSplittingVariables: Block size argument must be an integer constant or a " +
                "scalar variable" in caplog.records[0].message)


@pytest.mark.parametrize('frontend', available_frontends())
def test_1d_splitting(frontend):
    """
    Apply loop blocking of simple loops into two loops
    """
    fcode = """
subroutine test_1d_splitting(a, b, n)
  implicit none
  integer, intent(in) :: n
  real(kind=8), intent(inout) :: a(n)
  real(kind=8), intent(inout) :: b(n)
  integer :: i
  !$loki driver-loop
  do i=1,n
    a(i) = real(i, kind=8)
  end do
end subroutine test_1d_splitting
    """
    block_size = 117
    routine = Subroutine.from_source(fcode, frontend=frontend)
    loops = FindNodes(ir.Loop).visit(routine.ir)
    num_loops = len(loops)
    num_vars = len(routine.variable_map)
    with pragmas_attached(routine, ir.Loop):
        loops = find_driver_loops(routine.body, targets=None)
    split_loop(routine, loops[0], block_size)
    loops = FindNodes(ir.Loop).visit(routine.ir)

    assert len(loops) == num_loops + 1, \
        f"Total number of loops transformation is: {len(loops)} but expected {num_loops + 1}"
    assert len(routine.variable_map) == num_vars + LOKI_LOOP_SLIT_VAR_ADDITION, (
        f"Total number of variables after loop splitting is: {len(routine.variable_map)} "
        f"but expected {num_vars + LOKI_LOOP_SLIT_VAR_ADDITION}"
    )

    assert loop_symbols(routine.ir) == [
        ('i_loop_block_idx', 1, 'i_loop_num_blocks', None),
        ('i_loop_local', 1, 'i_loop_block_end - i_loop_block_start + 1', None)
    ]
    # Ensure inner loop bound has been parametrised correctly
    assert routine.variable_map['i_loop_block_size'].type.initial == block_size

    assert assignment_symbols(routine.body) == [
        ('i_loop_num_blocks', '1 + (-1 + n) / i_loop_block_size'),
        ('i_loop_block_start', '(i_loop_block_idx - 1)*i_loop_block_size + 1'),
        ('i_loop_block_end', 'min(i_loop_block_idx*i_loop_block_size, n)'),
        ('i_loop_iter_num', 'i_loop_block_start + i_loop_local - 1'),
        ('i', 'i_loop_iter_num'), ('a(i)', 'real(i, kind=8)')
    ]


@pytest.mark.parametrize('frontend', available_frontends())
def test_1d_splitting_multi_var(frontend):
    """
    Apply loop blocking of simple loops into two loops
    """
    fcode = """
subroutine test_1d_splitting_multi_var(a, b, n)
  implicit none
  integer, intent(in) :: n
  real(kind=8), intent(inout) :: a(n)
  real(kind=8), intent(inout) :: b(n)
  real(kind=8) :: c(n)
  integer :: i
  !$loki driver-loop
  do i=1,n
    c(1) = c(1) + i
    a(i) = real(i)
  end do
end subroutine test_1d_splitting_multi_var
    """
    block_size = 117
    routine = Subroutine.from_source(fcode, frontend=frontend)
    loops = FindNodes(ir.Loop).visit(routine.ir)
    num_loops = len(loops)
    num_vars = len(routine.variable_map)
    with pragmas_attached(routine, ir.Loop):
        loops = find_driver_loops(routine.body, targets=None)
    split_loop(routine, loops[0], block_size)
    loops = FindNodes(ir.Loop).visit(routine.ir)

    assert len(loops) == num_loops + 1, \
        f"Total number of loops transformation is: {len(loops)} but expected {num_loops + 1}"
    assert len(routine.variable_map) == num_vars + LOKI_LOOP_SLIT_VAR_ADDITION, (
        f"Total number of variables after loop splitting is: {len(routine.variable_map)} "
        f"but expected {num_vars + LOKI_LOOP_SLIT_VAR_ADDITION}"
    )

    assert loop_symbols(routine.ir) == [
        ('i_loop_block_idx', 1, 'i_loop_num_blocks', None),
        ('i_loop_local', 1, 'i_loop_block_end - i_loop_block_start + 1', None)
    ]
    # Ensure inner loop bound has been parametrised correctly
    assert routine.variable_map['i_loop_block_size'].type.initial == block_size

    assert assignment_symbols(routine.body) == [
        ('i_loop_num_blocks', '1 + (-1 + n) / i_loop_block_size'),
        ('i_loop_block_start', '(i_loop_block_idx - 1)*i_loop_block_size + 1'),
        ('i_loop_block_end', 'min(i_loop_block_idx*i_loop_block_size, n)'),
        ('i_loop_iter_num', 'i_loop_block_start + i_loop_local - 1'),
        ('i', 'i_loop_iter_num'), ('c(1)', 'c(1) + i'), ('a(i)', 'real(i)')
    ]


@pytest.mark.parametrize('frontend', available_frontends())
def test_2d_splitting(frontend):
    fcode = """
    subroutine test_2d_splitting(a, b, n)
      implicit none
      integer, intent(in) :: n
      real(kind=8), intent(inout) :: a(n)
      real(kind=8), intent(inout) :: b(n,n)
      real(kind=8) :: c(n)
      integer :: i
      !$loki driver-loop
      do i=1,n
        a(i) = i
        c(1) = a(i)
        b(:,i) = a(i)
      end do
    end subroutine test_2d_splitting
        """
    block_size = 117
    routine = Subroutine.from_source(fcode, frontend=frontend)
    loops = FindNodes(ir.Loop).visit(routine.ir)
    num_loops = len(loops)
    num_vars = len(routine.variable_map)
    with pragmas_attached(routine, ir.Loop):
        loops = find_driver_loops(routine.body, targets=None)
    split_loop(routine, loops[0], block_size)
    loops = FindNodes(ir.Loop).visit(routine.ir)

    assert len(loops) == num_loops + 1, \
        f"Total number of loops transformation is: {len(loops)} but expected {num_loops + 1}"
    assert len(routine.variable_map) == num_vars + LOKI_LOOP_SLIT_VAR_ADDITION, (
        f"Total number of variables after loop splitting is: {len(routine.variable_map)} "
        f"but expected {num_vars + LOKI_LOOP_SLIT_VAR_ADDITION}"
    )

    assert loop_symbols(routine.ir) == [
        ('i_loop_block_idx', 1, 'i_loop_num_blocks', None),
        ('i_loop_local', 1, 'i_loop_block_end - i_loop_block_start + 1', None)
    ]
    # Ensure inner loop bound has been parametrised correctly
    assert routine.variable_map['i_loop_block_size'].type.initial == block_size

    assert assignment_symbols(routine.body) == [
        ('i_loop_num_blocks', '1 + (-1 + n) / i_loop_block_size'),
        ('i_loop_block_start', '(i_loop_block_idx - 1)*i_loop_block_size + 1'),
        ('i_loop_block_end', 'min(i_loop_block_idx*i_loop_block_size, n)'),
        ('i_loop_iter_num', 'i_loop_block_start + i_loop_local - 1'),
        ('i', 'i_loop_iter_num'), ('a(i)', 'i'), ('c(1)', 'a(i)'), ('b(:, i)', 'a(i)')
    ]



@pytest.mark.parametrize('frontend', available_frontends())
@pytest.mark.parametrize('block_size', [117])
@pytest.mark.parametrize('n', [500])
def test_3d_splitting(tmp_path, frontend, block_size, n):
    fcode = """
    subroutine test_3d_splitting(a, b, c, n)
      implicit none
      integer, intent(in) :: n
      real(kind=8), intent(inout) :: a(n)
      real(kind=8), intent(inout) :: b(2,n)
      real(kind=8), intent(inout) :: c(2,2,n)
      integer :: i
      !$loki driver-loop
      do i=1,n
        a(i) = i
        b(:,i) = a(i)
        c(:,:,i) = a(i)
      end do
    end subroutine test_3d_splitting
        """
    routine = Subroutine.from_source(fcode, frontend=frontend)
    loops = FindNodes(ir.Loop).visit(routine.ir)
    num_loops = len(loops)
    num_vars = len(routine.variable_map)
    with pragmas_attached(routine, ir.Loop):
        loops = find_driver_loops(routine.body, targets=None)
    split_loop(routine, loops[0], block_size)
    loops = FindNodes(ir.Loop).visit(routine.ir)

    assert len(loops) == num_loops + 1, \
        f"Total number of loops transformation is: {len(loops)} but expected {num_loops + 1}"
    assert len(routine.variable_map) == num_vars + LOKI_LOOP_SLIT_VAR_ADDITION, (
        f"Total number of variables after loop splitting is: {len(routine.variable_map)} "
        f"but expected {num_vars + LOKI_LOOP_SLIT_VAR_ADDITION}"
    )

    filepath = tmp_path / (f'{routine.name}_{frontend}.f90')

    a = np.zeros(n, order='F')
    b = np.zeros((2,n), order='F')
    c = np.zeros((2,2,n), order='F')
    jit_compile_and_run(routine, a, b, c, n, filepath=filepath)
    a_ref = np.linspace(1, n, n)
    b_ref = np.tile(a_ref, (2, 1))
    c_ref = np.tile(a_ref, (2,2,1))
    assert np.array_equal(a, a_ref), "a should be equal to a_ref=(1, 2, ..., n)"
    assert np.array_equal(b, b_ref), "b should equal b_ref"
    assert np.array_equal(c, c_ref), "c should equal c_ref"

    clean_test(filepath)

"""
--------------------------------------------------------------------------------
Blocking tests

Tests that variables are correctly blocked, and that blocked loops produce
the correct output.
--------------------------------------------------------------------------------
"""

@pytest.mark.parametrize('frontend', available_frontends())
def test_1d_blocking(frontend):
    """
    Apply loop blocking of simple loops into two loops
    """
    fcode = """
subroutine test_1d_blocking(a, b, n)
  implicit none
  integer, intent(in) :: n
  real(kind=8), intent(inout) :: a(n)
  real(kind=8), intent(inout) :: b(n)
  integer :: i
  !$loki driver-loop
  do i=1,n
    a(i) = real(i)
  end do
end subroutine test_1d_blocking
    """
    block_size = 117
    routine = Subroutine.from_source(fcode, frontend=frontend)
    loops = FindNodes(ir.Loop).visit(routine.ir)
    with pragmas_attached(routine, ir.Loop):
        loops = find_driver_loops(routine.body,
                                  targets=None)

    num_loops = len(loops)
    num_vars = len(routine.variable_map)

    splitting_vars, inner_loop, outer_loop = split_loop(routine, loops[0], block_size)
    loops = FindNodes(ir.Loop).visit(routine.ir)
    assert len(loops) == num_loops + 1, \
        f"Total number of loops transformation is: {len(loops)} but expected {num_loops + 1}"
    assert len(routine.variable_map) == num_vars + LOKI_LOOP_SLIT_VAR_ADDITION, (
        f"Total number of variables after loop splitting is: {len(routine.variable_map)} "
        f"but expected {num_vars + LOKI_LOOP_SLIT_VAR_ADDITION}"
    )
    num_vars = len(routine.variable_map)

    blocking_indices = ['i']
    block_loop_arrays(routine, splitting_vars, inner_loop, outer_loop, blocking_indices)
    for var in FindVariables().visit(inner_loop.body):
        if isinstance(var, Array):
            for idx in blocking_indices:
                assert idx not in var.dimensions, "The variable should be blocked and the local variable used"

    assert len(routine.variable_map) == num_vars+1, "Expected 1 loop blocking to be added"

    assert loop_symbols(routine.ir) == [
        ('i_loop_block_idx', 1, 'i_loop_num_blocks', None),
        ('i_loop_local', 1, 'i_loop_block_end - i_loop_block_start + 1', None)
    ]
    # Ensure inner loop bound has been parametrised correctly
    assert routine.variable_map['i_loop_block_size'].type.initial == block_size

    assert assignment_symbols(routine.body) == [
        ('i_loop_num_blocks', '1 + (-1 + n) / i_loop_block_size'),
        ('i_loop_block_start', '(i_loop_block_idx - 1)*i_loop_block_size + 1'),
        ('i_loop_block_end', 'min(i_loop_block_idx*i_loop_block_size, n)'),
        ('a_block(1:i_loop_block_end - i_loop_block_start + 1)', 'a(i_loop_block_start:i_loop_block_end)'),
        ('i_loop_iter_num', 'i_loop_block_start + i_loop_local - 1'),
        ('i', 'i_loop_iter_num'), ('a_block(i_loop_local)', 'real(i)'),
        ('a(i_loop_block_start:i_loop_block_end)', 'a_block(1:i_loop_block_end - i_loop_block_start + 1)')
    ]
    assert array_access_symbols(inner_loop.body) == [('a_block', ('i_loop_local',))]


@pytest.mark.parametrize('frontend', available_frontends())
@pytest.mark.parametrize('block_size', [117])
@pytest.mark.parametrize('n', [500])
def test_1d_blocking_multi_intent(tmp_path, frontend, block_size, n):
    """
    Apply loop blocking of simple loops into two loops
    """
    fcode = """
subroutine test_1d_blocking_multi_intent(a, b, n)
  implicit none
  integer, intent(in) :: n
  real(kind=8), intent(in) :: a(n)
  real(kind=8), intent(inout) :: b(n)
  integer :: i
  !$loki driver-loop
  do i=1,n
    b(i) = b(i) + a(i)*a(i)
  end do
end subroutine test_1d_blocking_multi_intent
    """
    routine = Subroutine.from_source(fcode, frontend=frontend)
    loops = FindNodes(ir.Loop).visit(routine.ir)
    with pragmas_attached(routine, ir.Loop):
        loops = find_driver_loops(routine.body,
                                  targets=None)

    num_loops = len(loops)
    num_vars = len(routine.variable_map)
    splitting_vars, inner_loop, outer_loop = split_loop(routine, loops[0], block_size)
    loops = FindNodes(ir.Loop).visit(routine.ir)
    assert len(loops) == num_loops + 1, \
        f"Total number of loops transformation is: {len(loops)} but expected {num_loops + 1}"
    assert len(routine.variable_map) == num_vars + LOKI_LOOP_SLIT_VAR_ADDITION, (
        f"Total number of variables after loop splitting is: {len(routine.variable_map)} "
        f"but expected {num_vars + LOKI_LOOP_SLIT_VAR_ADDITION}"
    )

    num_vars = len(routine.variable_map)
    blocking_indices = ['i']
    block_loop_arrays(routine, splitting_vars, inner_loop, outer_loop, blocking_indices)

    assert len(routine.variable_map) == num_vars+2, "Expected 2 loop blocking to be added"
    for var in FindVariables().visit(inner_loop.body):
        if isinstance(var, Array):
            for idx in blocking_indices:
                assert idx not in var.dimensions, "The variable should be blocked and the local variable used"
    block_loop_arrays(routine, splitting_vars, inner_loop, outer_loop, blocking_indices=['i'])


    filepath = tmp_path / (f'{routine.name}_{frontend}.f90')

    a = np.linspace(1, n, n)
    b = np.ones(n, order='F')
    a_ref = np.linspace(1, n, n)
    b_ref = b + a*a
    jit_compile_and_run(routine, a, b, n, filepath=filepath, isolated=False)
    assert np.array_equal(a, a_ref), "a should be equal to a_ref=(1, 2, ..., n)"
    assert np.array_equal(b, b_ref), "b should equal to (2, 5, ..., 1 + n^2)"
    clean_test(filepath)


@pytest.mark.parametrize('frontend', available_frontends())
@pytest.mark.parametrize('block_size', [117])
@pytest.mark.parametrize('n', [500])
def test_2d_blocking(tmp_path, frontend, block_size, n):
    fcode = """
    subroutine test_2d_blocking(a, b, n)
      implicit none
      integer, intent(in) :: n
      real(kind=8), intent(inout) :: a(n)
      real(kind=8), intent(inout) :: b(n,n)
      real(kind=8) :: c(n)
      integer :: i
      !$loki driver-loop
      do i=1,n
        a(i) = i
        c(1) = a(i)
        b(:,i) = a(i)
      end do
    end subroutine test_2d_blocking
        """
    routine = Subroutine.from_source(fcode, frontend=frontend)
    loops = FindNodes(ir.Loop).visit(routine.ir)
    num_loops = len(loops)
    num_vars = len(routine.variable_map)
    with pragmas_attached(routine, ir.Loop):
        loops = find_driver_loops(routine.body, targets=None)
    splitting_vars, inner_loop, outer_loop = split_loop(routine, loops[0], block_size)
    loops = FindNodes(ir.Loop).visit(routine.ir)

    assert len(loops) == num_loops + 1, \
        f"Total number of loops transformation is: {len(loops)} but expected {num_loops + 1}"
    assert len(routine.variable_map) == num_vars + LOKI_LOOP_SLIT_VAR_ADDITION, (
        f"Total number of variables after loop splitting is: {len(routine.variable_map)} "
        f"but expected {num_vars + LOKI_LOOP_SLIT_VAR_ADDITION}"
    )

    num_vars = len(routine.variable_map)
    blocking_indices = ['i']
    block_loop_arrays(routine, splitting_vars, inner_loop, outer_loop, blocking_indices)

    assert len(routine.variable_map) == num_vars + 2, "Expected 2 loop blocking to be added"
    for var in FindVariables().visit(inner_loop.body):
        if isinstance(var, Array):
            for idx in blocking_indices:
                assert idx not in var.dimensions, "The variable should be blocked and the local variable used"

    filepath = tmp_path / (f'{routine.name}_{frontend}.f90')

    a = np.zeros(n, order='F')
    b = np.zeros((n,n), order='F')
    jit_compile_and_run(routine, a, b, n, filepath=filepath, isolated=False)
    a_ref = np.linspace(1, n, n)
    b_ref = np.tile(a_ref, (n,1))
    assert np.array_equal(a, a_ref), "a should be equal to a_ref=(1, 2, ..., n)"
    assert np.array_equal(b, b_ref), "b[:,1] should equal a and a_ref"

    clean_test(filepath)


@pytest.mark.parametrize('frontend', available_frontends())
@pytest.mark.parametrize('block_size', [117])
@pytest.mark.parametrize('n', [500])
def test_3d_blocking(tmp_path, frontend, block_size, n):
    fcode = """
    subroutine test_3d_blocking(a, b, c, n)
      implicit none
      integer, intent(in) :: n
      real(kind=8), intent(inout) :: a(n)
      real(kind=8), intent(inout) :: b(2,n)
      real(kind=8), intent(inout) :: c(2,2,n)
      integer :: i
      !$loki driver-loop
      do i=1,n
        a(i) = i
        b(:,i) = a(i)
        c(:,:,i) = a(i)
      end do
    end subroutine test_3d_blocking
        """
    routine = Subroutine.from_source(fcode, frontend=frontend)
    loops = FindNodes(ir.Loop).visit(routine.ir)
    num_loops = len(loops)
    num_vars = len(routine.variable_map)
    with pragmas_attached(routine, ir.Loop):
        loops = find_driver_loops(routine.body, targets=None)
    splitting_vars, inner_loop, outer_loop = split_loop(routine, loops[0], block_size)
    loops = FindNodes(ir.Loop).visit(routine.ir)

    assert len(loops) == num_loops + 1, \
        f"Total number of loops transformation is: {len(loops)} but expected {num_loops + 1}"
    assert len(routine.variable_map) == num_vars + LOKI_LOOP_SLIT_VAR_ADDITION, (
        f"Total number of variables after loop splitting is: {len(routine.variable_map)} "
        f"but expected {num_vars + LOKI_LOOP_SLIT_VAR_ADDITION}"
    )

    num_vars = len(routine.variable_map)
    blocking_indices = ['i']
    block_loop_arrays(routine, splitting_vars, inner_loop, outer_loop, blocking_indices)

    assert len(routine.variable_map) == num_vars + 3, "Expected 3 loop blocking to be added"
    for var in FindVariables().visit(inner_loop.body):
        if isinstance(var, Array):
            for idx in blocking_indices:
                assert idx not in var.dimensions, "The variable should be blocked and the local variable used"

    filepath = tmp_path / (f'{routine.name}_{frontend}.f90')


    a = np.zeros(n, order='F')
    b = np.zeros((2,n), order='F')
    c = np.zeros((2,2,n), order='F')
    jit_compile_and_run(routine, a, b, c, n, filepath=filepath)
    a_ref = np.linspace(1, n, n)
    b_ref = np.tile(a_ref, (2, 1))
    c_ref = np.tile(a_ref, (2,2,1))
    assert np.array_equal(a, a_ref), "a should be equal to a_ref=(1, 2, ..., n)"
    assert np.array_equal(b, b_ref), "b should equal b_ref"
    assert np.array_equal(c, c_ref), "c should equal c_ref"

    clean_test(filepath)
