## SPDX-License-Identifier: BSD-2-Clause
##
## Copyright 2020-2025 Inria & Université de Bordeaux
##
## Redistribution and use in source and binary forms, with or without
## modification, are permitted provided that the following conditions
## are met:
##
## 1. Redistributions of source code must retain the above copyright
## notice, this list of conditions and the following disclaimer.
##
## 2. Redistributions in binary form must reproduce the above
## copyright notice, this list of conditions and the following
## disclaimer in the documentation and/or other materials provided
## with the distribution.
##
## THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
## CONTRIBUTORS “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES,
## INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
## MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
## DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS
## BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
## EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
## TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
## DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
## ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
## TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF
## THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
## SUCH DAMAGE.
##

import math
import numpy as np
import scotchpy as SCOTCH
import importlib.resources
import os

_data_dir = importlib.resources.files('scotchpy').joinpath("data")


# taken from check/test_scotch_graph_map.c


def testMap():
    ARCHNBR = 4
    STRANBR = 2
    COORD = lambda x, y: y * xdimsiz + x
    filename1 = os.path.join(_data_dir, "m4x4.grf")
    filename2 = "./test_scotchpy_graph_map_temp.txt"
    archtab = [None for _ in range(ARCHNBR)]
    stratab = [None for _ in range(STRANBR)]

    try:
        graf = SCOTCH.Graph()
    except SCOTCH.LibraryError:
        assert False, "cannot initialize graph"
        raise

    try:
        graf.load(filename1)
    except SCOTCH.LibraryError:
        assert False, "cannot load graph"
        raise

    vertnbr, _ = graf.size(as_dict=False)
    xdimsiz = int(math.sqrt(vertnbr))

    assert vertnbr == xdimsiz * xdimsiz, "graph is not a square grid"
    parttab = np.full(vertnbr, -1, SCOTCH.common.proper_int)
    parotab = np.full(vertnbr, -1, SCOTCH.common.proper_int)
    vmlotab = np.full(vertnbr, -1, SCOTCH.common.proper_int)

    for vertnum in range(vertnbr):
        vmlotab[vertnum] = vertnum % 3

    for stranum in range(STRANBR):
        try:
            stratab[stranum] = SCOTCH.Strat()
        except SCOTCH.LibraryError:
            assert False, "cannot initialize strategy"
            raise

    stratab[0].graph_map_build(SCOTCH.STRATRECURSIVE, 4, 0.05)
    stratab[1].graph_map_build(SCOTCH.STRATDEFAULT, 4, 0.05)

    for archnum in range(ARCHNBR):
        try:
            archtab[archnum] = SCOTCH.Arch()
        except SCOTCH.LibraryError:
            assert False, "cannot initialize architecture"
            raise

    archtab[0].cmplt(5)
    archtab[1].mesh2(2, 2)
    archtab[2].mesh2(xdimsiz * 2, xdimsiz * 2)  # oversized architecture
    archtab[3].vhcub()
    mapp = SCOTCH.Mapping()
    mapo = SCOTCH.Mapping()

    for stranum in range(STRANBR):
        for archnum in range(ARCHNBR):
            try:
                graf.map_init(mapp, archtab[archnum], parttab)
            except SCOTCH.LibraryError:
                assert False, "cannot initialize mapping (1)"
                raise

            try:
                graf.map_init(mapo, archtab[archnum], parotab)
            except SCOTCH.LibraryError:
                assert False, "cannot initialize mapping (2)"
                raise

            archsiz = archtab[archnum].size()

            for typenum in range(6):
                parttab = np.full(xdimsiz * xdimsiz, -1, SCOTCH.common.proper_int)

                if archnum < 2:
                    for i in range(xdimsiz - 1):
                        parttab[COORD(0, i)] = 0
                        parttab[COORD(i + 1, 0)] = 1
                        parttab[COORD(xdimsiz - 1, i + 1)] = archsiz - 2
                        parttab[COORD(i, xdimsiz - 1)] = archsiz - 1
                else:
                    for i in range(xdimsiz - 1):
                        parttab[COORD(0, i)] = vertnbr - 2
                        parttab[COORD(i + 1, 0)] = vertnbr - 1
                        parttab[COORD(xdimsiz - 1, i + 1)] = vertnbr
                        parttab[COORD(i, xdimsiz - 1)] = vertnbr + 1

                # print("Strat {}, arch {}, type {}".format(stranum, archnum, typenum))

                vmloptr = vmlotab  # .copy()

                try:
                    if typenum == 0:
                        graf.map_compute(mapp, stratab[stranum])
                        parotab[:vertnbr] = parttab[:vertnbr]
                    elif typenum == 1:
                        graf.map_fixed_compute(mapp, stratab[stranum])
                    elif typenum == 2:
                        vmloptr = None
                    elif typenum == 3:
                        graf.remap_compute(mapp, mapo, 0.2, vmloptr, stratab[stranum])
                    elif typenum == 4:
                        vmloptr = None
                    elif typenum == 5:
                        graf.remap_fixed_compute(
                            mapp, mapo, 0.2, vmloptr, stratab[stranum]
                        )
                except SCOTCH.LibraryError:
                    assert False, "cannot compute mapping"
                    raise
            graf.map_exit(mapo)
            graf.map_exit(mapp)


# taken from check/test_scotch_graph_map_copy.c


def testMapCopy():
    STRANBR = 3
    stratab = [None for _ in range(STRANBR)]
    mapp = SCOTCH.Mapping()
    mapo = SCOTCH.Mapping()
    filename = os.path.join(_data_dir, "bump.grf")

    try:
        graf = SCOTCH.Graph()
    except SCOTCH.LibraryError:
        assert False, "cannot initialize graph"
        raise

    try:
        graf.load(filename)
    except SCOTCH.LibraryError:
        assert False, "cannot load graph"
        raise

    vertnbr, _ = graf.size(as_dict=False)

    parttab = np.full(vertnbr, -1, SCOTCH.common.proper_int)
    parotab = np.full(vertnbr, -1, SCOTCH.common.proper_int)

    for stranum in range(STRANBR):
        try:
            stratab[stranum] = SCOTCH.Strat()
        except SCOTCH.LibraryError:
            assert False, "cannot initialize strategy"
            raise

    stratab[0].graph_map("cf{move=10000,pass=-1,bal=0.05}")
    stratab[1].graph_map(
        "m{vert=120,low=cf{move=10000,pass=-1,bal=0.05},asc=b{bnd=f{move=10000,pass=-1,bal=0.05},org=f{move=10000,pass=-1,bal=0.05}}}"
    )

    try:
        arch = SCOTCH.Arch()
    except SCOTCH.LibraryError:
        assert False, "cannot initialize architecture"
        raise

    arch.cmplt(5)

    for stranum in range(STRANBR - 1):
        for typenum in range(2):
            # print("Strat {}, type {}".format(stranum, typenum))

            if typenum == 0:
                try:
                    graf.map_init(mapp, arch, parttab)
                except SCOTCH.LibraryError:
                    assert False, "cannot initialize mapping (1)"
                    raise
                try:
                    # graf.map_compute(mapp, stratab[STRANBR-1])
                    mapp.compute(stratab[STRANBR - 1])  # canonically equivalent
                except SCOTCH.LibraryError:
                    assert False, "cannot compute mapping (1)"
                    raise
                parotab[:vertnbr] = parttab[:vertnbr]
            elif typenum == 1:
                try:
                    graf.map_init(mapo, arch, parotab)
                except SCOTCH.LibraryError:
                    assert False, "cannot initialize mapping (2)"
                    raise
                try:
                    graf.remap_compute(mapp, mapo, 0.0, None, stratab[stranum])
                except SCOTCH.ScotchException:
                    assert False, "cannot compute mapping (2)"
                    raise
        graf.map_exit(mapo)
        graf.map_exit(mapp)
