# (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 ctypes as ct
import os
from pathlib import Path
import numpy as np
import pytest

from loki import Subroutine, FortranMaxTransformation, execute, delete
from loki.build import Builder, Obj, Lib, jit_compile, clean_test
from loki.build.max_compiler import (
    compile_all, compile_maxj, compile_max, generate_max,
    get_max_includes, get_max_libs, get_max_libdirs
)
from loki.frontend import available_frontends


def is_maxeler_available():
    """
    Check if Maxeler environment variables are specified.
    """
    maxeler_vars = {'MAXCOMPILERDIR', 'MAXELEROSDIR'}
    return maxeler_vars <= os.environ.keys()


class MaxelerNotAvailableError(RuntimeError):
    pass

# Skip tests in this module if Maxeler environment not present
# pytestmark = pytest.mark.skipif(not check_maxeler(),
#                                 reason='Maxeler compiler not installed')
pytestmark = pytest.mark.xfail(
    condition=not is_maxeler_available(), raises=MaxelerNotAvailableError,
    reason='Maxeler runtime environment not available'
)


@pytest.fixture(scope='module', name='simulator')
def fixture_simulator():
    maxeler_available = is_maxeler_available()

    class MaxCompilerSim:

        def __init__(self):
            name = f'pytest_{os.getpid()}'
            self.base_cmd = ['maxcompilersim', '-n', name]
            if maxeler_available:
                os.environ['SLIC_CONF'] = f'use_simulation={name}'
                self.maxeleros = ct.CDLL(os.environ['MAXELEROSDIR'] + '/lib/libmaxeleros.so')

        def __del__(self):
            if maxeler_available:
                del self.maxeleros
                del os.environ['SLIC_CONF']

        def restart(self):
            if maxeler_available:
                cmd = self.base_cmd + ['-c', 'MAX5C', 'restart']
                execute(cmd)

        def stop(self):
            if maxeler_available:
                cmd = self.base_cmd + ['stop']
                execute(cmd)

        def run(self, target, *args):
            if not maxeler_available:
                raise MaxelerNotAvailableError
            cmd = [str(target)]
            if args is not None:
                cmd += [str(a) for a in args]
            self.restart()
            env = os.environ.copy()
            maxelerosdir = os.environ['MAXELEROSDIR']
            ld_preload = os.environ.get('LD_PRELOAD', '')
            env['LD_PRELOAD'] = f'{maxelerosdir}/lib/libmaxeleros.so:{ld_preload}'
            execute(cmd, env=env)
            self.stop()

        def call(self, fn, *args, **kwargs):
            if not maxeler_available:
                raise MaxelerNotAvailableError
            self.restart()
            ret = fn(*args, **kwargs)
            self.stop()
            return ret

    return MaxCompilerSim()


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


@pytest.fixture(scope='module', name='builder')
def fixture_builder(here):
    include_dirs = get_max_includes() if is_maxeler_available() else []
    yield Builder(source_dirs=here, include_dirs=include_dirs, build_dir=here/'build')
    Obj.clear_cache()


def max_transpile(routine, path, builder, frontend, objects=None, wrap=None):
    builder.clean()

    # Create transformation object and apply
    f2max = FortranMaxTransformation()
    f2max.apply(routine, path=path)

    if not is_maxeler_available():
        raise MaxelerNotAvailableError

    # Generate simulation object file from maxj kernel
    compile_maxj(src=f2max.maxj_kernel_path.parent, build_dir=builder.build_dir)
    max_path = generate_max(manager=f2max.maxj_manager_path.stem, maxj_src=f2max.maxj_src,
                            max_filename=routine.name, build_dir=builder.build_dir,
                            package=routine.name)
    max_obj = compile_max(max_path, f'{max_path.stem}_max.o', build_dir=builder.build_dir)
    max_include = max_obj.parent/(f'{routine.name}_MAX5C_DFE_SIM/results')

    # Build and wrap the cross-compiled library
    objects = (objects or []) + [Obj(source_path=f2max.c_path), Obj(source_path=f2max.wrapperpath)]
    lib = Lib(name=f'fmax_{routine.name}_{frontend}', objs=objects, shared=False)
    lib.build(builder=builder, include_dirs=[max_include], external_objs=[max_obj])

    return lib.wrap(modname=f'mod_{routine.name}_{frontend}', builder=builder,
                    sources=(wrap or []) + [f2max.wrapperpath.name],
                    libs=get_max_libs(), lib_dirs=get_max_libdirs(),
                    kind_map=Path(__file__).parent.parent/'kind_map')


@pytest.mark.skipif(not is_maxeler_available(), reason='Maxeler runtime environment not available')
def test_max_simulator(simulator):
    """
    Starts and stops the Maxeler Simulator.
    """
    simulator.restart()
    simulator.stop()
    assert True


@pytest.mark.skipif(not is_maxeler_available(), reason='Maxeler runtime environment not available')
def test_max_passthrough(simulator, here):
    """
    A simple test streaming data to the DFE and back to CPU.
    """
    build_dir = here/'build'
    compile_all(c_src=here/'passthrough', maxj_src=here/'passthrough', build_dir=build_dir,
                target='PassThrough', manager='PassThroughMAX5CManager', package='passthrough')
    simulator.run(build_dir/'PassThrough')


@pytest.mark.skipif(not is_maxeler_available(), reason='Maxeler runtime environment not available')
def test_max_passthrough_ctypes(simulator, here):
    """
    A simple test streaming data to the DFE and back to CPU, called via ctypes
    """
    # First, build shared library
    build_dir = here/'build'
    compile_all(c_src=here/'passthrough', maxj_src=here/'passthrough', build_dir=build_dir,
                target='libPassThrough.so', manager='PassThroughMAX5CManager', package='passthrough')
    lib = ct.CDLL(build_dir/'libPassThrough.so')

    # Extract function interfaces for CPU and DFE version
    func_cpu = lib.PassThroughCPU
    func_cpu.restype = None
    func_cpu.argtypes = [ct.c_int, ct.POINTER(ct.c_uint32), ct.POINTER(ct.c_uint32)]

    func_dfe = lib.passthrough
    func_dfe.restype = None
    func_dfe.argtypes = [ct.c_uint64, ct.c_void_p, ct.c_size_t, ct.c_void_p, ct.c_size_t]

    # Create input/output data structures
    size = 1024
    data_in = [i+1 for i in range(size)]
    data_out = size * [0]

    array_type = ct.c_uint32 * size
    size_bytes = ct.c_size_t(ct.sizeof(ct.c_uint32) * size)
    data_in = array_type(*data_in)
    expected_out = array_type(*data_out)
    data_out = array_type(*expected_out)

    # Run CPU function
    func_cpu(ct.c_int(size), data_in, expected_out)
    assert list(data_in) == list(expected_out)

    # Run DFE function
    simulator.call(func_dfe, ct.c_uint64(size), data_in, size_bytes, data_out, size_bytes)
    assert list(data_in) == list(data_out)


@pytest.mark.parametrize('frontend', available_frontends())
def test_max_routine_axpy_scalar(here, builder, simulator, frontend):

    fcode = """
subroutine routine_axpy_scalar(a, x, y)
  ! A simple standard routine that computes x = a * x + y for
  ! scalar arguments
  use iso_fortran_env, only: real64
  implicit none
  real(kind=real64), intent(in) :: a, y
  real(kind=real64), intent(inout) :: x

  x = a * x + y
end subroutine routine_axpy_scalar
    """
    routine = Subroutine.from_source(fcode, frontend=frontend)
    filepath = here/(f'routine_axpy_scalar_{frontend}.f90')
    function = jit_compile(routine, filepath=filepath, objname='routine_axpy_scalar')

    # Test the reference solution
    a = -3.
    x = np.zeros(shape=(1,), order='F') + 2.
    y = np.zeros(shape=(1,), order='F') + 10.
    function(a=a, x=x, y=y)
    assert np.all(a * 2. + y == x)

    simulator.restart()

    # TODO: For some reason we have to generate and run the kernel twice in the same instance of
    # the simulator to actually get any results other than 0. Probably doing something wrong with
    # the Maxeler language...
    for _ in range(2):
        # Generate the transpiled kernel
        max_kernel = max_transpile(routine, here, builder, frontend)

        # Test the transpiled kernel
        a = -3.
        x = np.zeros(shape=(1,), order='F') + 2.
        y = np.zeros(shape=(1,), order='F') + 10.

        max_kernel.routine_axpy_scalar_c_fc_mod.routine_axpy_scalar_c_fc(ticks=1, a=a, x=x, y=y)

    simulator.stop()
#    simulator.call(max_kernel.routine_axpy_scalar_fmax_mod.routine_axpy_scalar_fmax, a, x, y)
    assert np.all(a * 2. + y == x)

    clean_test(filepath)
    delete(here/routine.name, force=True)  # Delete MaxJ sources


@pytest.mark.parametrize('frontend', available_frontends())
def test_max_routine_copy_scalar(here, builder, simulator, frontend):

    fcode = """
subroutine routine_copy_scalar(x, y)
  ! A simple routine that copies the value of x to y
  use iso_fortran_env, only: real64
  implicit none
  ! integer, parameter :: jprb = selected_real_kind(13,300)
  real(kind=real64), intent(in) :: x
  real(kind=real64), intent(out) :: y

  y = x
end subroutine routine_copy_scalar
    """
    routine = Subroutine.from_source(fcode, frontend=frontend)
    filepath = here/(f'routine_copy_scalar_{frontend}.f90')
    function = jit_compile(routine, filepath=filepath, objname='routine_copy_scalar')

    # Test the reference solution
    x = np.zeros(1) + 2.
    y = function(x=x)
    assert np.all(y == x)

    simulator.restart()
    # TODO: For some reason we have to generate and run the kernel twice in the same instance of
    # the simulator to actually get any results other than 0. Probably doing something wrong with
    # the Maxeler language...
    for _ in range(2):
        # Generate the transpiled kernel
        max_kernel = max_transpile(routine, here, builder, frontend)

        # Test the transpiled kernel
        x = np.zeros(1) + 2.
        y = max_kernel.routine_copy_scalar_c_fc_mod.routine_copy_scalar_c_fc(ticks=1, x=x)
        print(y)
    simulator.stop()
    assert np.all(y == x)

    clean_test(filepath)
    delete(here/routine.name, force=True)  # Delete MaxJ sources


@pytest.mark.parametrize('frontend', available_frontends())
def test_max_routine_fixed_loop(here, builder, simulator, frontend):

    fcode = """
subroutine routine_fixed_loop(scalar, vector, vector_out, tensor, tensor_out)
  use iso_fortran_env, only: real64
  implicit none
  integer :: n=6, m=4
  real(kind=real64), intent(in) :: scalar
  real(kind=real64), intent(in) :: tensor(6, 4), vector(6)
  real(kind=real64), intent(out) :: tensor_out(4, 6), vector_out(6)
  integer :: i, j

  ! For testing, the operation is:
  do i=1, n
     vector_out(i) = vector(i) + tensor(i, 1) + 1.0
  end do

  do j=1, n
     do i=1, m
        tensor_out(i, j) = tensor(j, i)
     end do
  end do
end subroutine routine_fixed_loop
    """
    routine = Subroutine.from_source(fcode, frontend=frontend)
    filepath = here/(f'routine_fixed_loop_{frontend}.f90')
    function = jit_compile(routine, filepath=filepath, objname='routine_fixed_loop')

    # Test the reference solution
    n, m = 6, 4
    scalar = 2.0
    vector = np.zeros(shape=(n,), order='F') + 3.
    tensor = np.array([list(range(i, i+m)) for i in range(n)], order='F', dtype=np.float64)
    tensor_out = np.zeros(shape=(m, n), order='F')
    ref_vector = vector + np.array(list(range(n)), dtype=np.float64) + 1.
    ref_tensor = np.transpose(tensor)
    function(scalar=scalar, vector=vector, vector_out=vector, tensor=tensor, tensor_out=tensor_out)
    assert np.all(vector == ref_vector)
    assert np.all(tensor_out == ref_tensor)

    # Generate the transpiled kernel
    max_kernel = max_transpile(routine, here, builder, frontend)

    # Test the transpiled kernel
    n, m = 6, 4
    scalar = 2.0
    vector = np.zeros(shape=(n,), order='F') + 3.
    tensor = np.zeros(shape=(n, m), order='F') + 4.
    tensor = np.array([list(range(i, i+m)) for i in range(n)], order='F', dtype=np.float64)
    tensor_out = np.zeros(shape=(m, n), order='F')
    function = max_kernel.routine_fixed_loop_c_fc_mod.routine_fixed_loop_c_fc
    simulator.call(function, ticks=1, scalar=scalar, vector=vector, vector_size=n * 8,
                   vector_out=vector, vector_out_size=n * 8, tensor=tensor, tensor_size=n * m * 8,
                   tensor_out=tensor_out, tensor_out_size=n * m * 8)
    assert np.all(vector == ref_vector)
    assert np.all(tensor_out == ref_tensor)

    clean_test(filepath)
    delete(here/routine.name, force=True)  # Delete MaxJ sources


@pytest.mark.parametrize('frontend', available_frontends())
def test_max_routine_copy_stream(here, builder, simulator, frontend):

    fcode = """
subroutine routine_copy_stream(length, scalar, vector_in, vector_out)
  implicit none
  ! A simple standard looking routine to test argument declarations
  ! and generator toolchain
  integer, intent(in) :: length, scalar, vector_in(length)
  integer, intent(out) :: vector_out(length)
  integer :: i

  !$loki dataflow
  do i=1, length
    vector_out(i) = vector_in(i) + scalar
  end do
end subroutine routine_copy_stream
    """
    routine = Subroutine.from_source(fcode, frontend=frontend)
    filepath = here/(f'routine_copy_stream_{frontend}.f90')
    function = jit_compile(routine, filepath=filepath, objname='routine_copy_stream')

    # Test the reference solution
    length = 32
    scalar = 7
    vector_in = np.array(range(length), order='F', dtype=np.intc)
    vector_out = np.zeros(length, order='F', dtype=np.intc)
    function(length=length, scalar=scalar, vector_in=vector_in, vector_out=vector_out)
    assert np.all(vector_out == np.array(range(length)) + scalar)

    # Generate the transpiled kernel
    max_kernel = max_transpile(routine, here, builder, frontend)

    vec_in = np.array(range(length), order='F', dtype=np.intc)
    vec_out = np.zeros(length, order='F', dtype=np.intc)
    function = max_kernel.routine_copy_stream_c_fc_mod.routine_copy_stream_c_fc
    simulator.call(function, ticks=length, length=length, scalar=scalar, vector_in=vec_in,
                   vector_in_size=length * 4, vector_out=vec_out, vector_out_size=length * 4)
    assert np.all(vec_out == np.array(range(length)) + scalar)

    clean_test(filepath)
    delete(here/routine.name, force=True)  # Delete MaxJ sources


@pytest.mark.parametrize('frontend', available_frontends())
def test_max_routine_moving_average(here, builder, simulator, frontend):

    fcode = """
subroutine routine_moving_average(length, data_in, data_out)
  use iso_fortran_env, only: real64
  implicit none
  integer, intent(in) :: length
  real(kind=real64), intent(in) :: data_in(length)
  real(kind=real64), intent(out) :: data_out(length)
  integer :: i
  real(kind=real64) :: prev, next, divisor

  !$loki dataflow
  do i=1, length
    divisor = 1.0
    if (i > 1) then
      prev = data_in(i-1)
      divisor = divisor + 1.0
    else
      prev = 0
    end if
    if (i < length) then
      next = data_in(i+1)
      divisor = divisor + 1.0
    else
      next = 0
    end if
    data_out(i) = (prev + data_in(i) + next) / divisor
  end do
end subroutine routine_moving_average
    """
    routine = Subroutine.from_source(fcode, frontend=frontend)
    filepath = here/(f'routine_moving_average_{frontend}.f90')
    function = jit_compile(routine, filepath=filepath, objname='routine_moving_average')

    # Create random input data
    n = 32
    data_in = np.array(np.random.rand(n), order='F')

    # Compute reference solution
    expected = np.zeros(shape=(n,), order='F')
    expected[0] = (data_in[0] + data_in[1]) / 2.
    expected[1:-1] = (data_in[:-2] + data_in[1:-1] + data_in[2:]) / 3.
    expected[-1] = (data_in[-2] + data_in[-1]) / 2.

    # Test the Fortran kernel
    data_out = np.zeros(shape=(n,), order='F')
    function(n, data_in, data_out)
    assert np.all(data_out == expected)

    # Generate and test the transpiled kernel
    max_kernel = max_transpile(routine, here, builder, frontend)

    data_out = np.zeros(shape=(n,), order='F')
    function = max_kernel.routine_moving_average_c_fc_mod.routine_moving_average_c_fc
    simulator.call(function, ticks=n, length=n, data_in=data_in, data_in_size=n * 8,
                   data_out_size=n * 8, data_out=data_out)
    assert np.all(data_out == expected)

    clean_test(filepath)
    delete(here/routine.name, force=True)  # Delete MaxJ sources


@pytest.mark.parametrize('frontend', available_frontends())
def test_max_routine_laplace(here, builder, simulator, frontend):
    fcode = """
subroutine routine_laplace(h, data_in, data_out)
  use iso_fortran_env, only: real64
  implicit none
  integer :: m = 32, n = 32
!  real(kind=real64), intent(in) :: h, rhs(m*n), data_in(m*n)
!  real(kind=real64), intent(out) :: data_out(m*n)
  real(kind=real64), intent(in) :: h, data_in(32*32)
  real(kind=real64), intent(out) :: data_out(32*32)
  integer :: i, i_mod_n
  real(kind=real64) :: north, south, east, west

  !$loki dataflow
  do i=1, m*n
    i_mod_n = mod(i, n)
    if (i_mod_n /= 0) then
        north = data_in(i+1)
    else
        north = 0
    endif
    if (i_mod_n /= 1) then
        south = data_in(i-1)
    else
        south = 0
    end if
    if (i > n) then
        west = data_in(i-n)
    else
        west = 0
    end if
    if (i <= (m-1)*n) then
        east = data_in(i+n)
    else
        east = 0
    end if
    data_out(i) = (north + south + east + west - 4 * data_in(i)) / (h * h)
  end do
end subroutine routine_laplace
    """
    routine = Subroutine.from_source(fcode, frontend=frontend)
    filepath = here/(f'routine_laplace_{frontend}.f90')
    function = jit_compile(routine, filepath=filepath, objname='routine_laplace')

    # Create random input data
    m, n = 32, 32
    h, length = 1./m, m * n
    data_in = np.array(np.random.rand(length), order='F')

    # Compute reference solution
    expected = -4. * data_in

    expected[0:n] += data_in[n:2*n]
    expected[1:n] += data_in[0:n-1]
    expected[0:n-1] += data_in[1:n]

    for i in range(1, m-1):
        idx = i*n
        expected[idx+0:idx+n] += data_in[idx-n:idx] + data_in[idx+n:idx+2*n]
        expected[idx+1:idx+n] += data_in[idx+0:idx+n-1]
        expected[idx+0:idx+n-1] += data_in[idx+1:idx+n]

    idx = (m-1)*n
    expected[idx+0:idx+n] += data_in[idx-n:idx]
    expected[idx+1:idx+n] += data_in[idx+0:idx+n-1]
    expected[idx+0:idx+n-1] += data_in[idx+1:idx+n]
    expected /= h*h

    # Test the Fortran kernel
    data_out = np.zeros(shape=(length,), order='F')
    function(h, data_in, data_out)
    assert np.all(abs(data_out - expected) < 1e-12)

    # Generate the transpiled kernel
    max_kernel = max_transpile(routine, here, builder, frontend)

    data_out = np.zeros(shape=(length,), order='F')
    function = max_kernel.routine_laplace_c_fc_mod.routine_laplace_c_fc
    simulator.call(function, ticks=length, h=h, data_in=data_in, data_in_size=length * 8,
                   data_out=data_out, data_out_size=length * 8)
    assert np.all(abs(data_out - expected) < 1e-12)

    clean_test(filepath)
    delete(here/routine.name, force=True)  # Delete MaxJ sources
