## 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 os
import numpy as np
import scotchpy as SCOTCH
import scotchpy.common as _common
from scotchpy.common import proper_int as SCOTCH_num
import tempfile
import importlib.resources

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

def graph1():
    # using kwargs
    return SCOTCH.build_graph(
        baseval=1,
        vertnbr=7,
        verttab=[1, 4, 10, 13, 16, 19, 22, 25],
        vendtab=None,
        velotab=[4, 1, 4, 4, 4, 4, 4],
        vlbltab=None,  # vlbltab
        edgenbr=24,
        edgetab=[
            3,
            2,
            6,
            3,
            4,
            1,
            7,
            6,
            5,
            1,
            2,
            4,
            2,
            7,
            3,
            7,
            2,
            6,
            2,
            1,
            5,
            5,
            2,
            4,
        ],
        edlotab=[
            1,
            1,
            1,
            2,
            2,
            1,
            2,
            3,
            3,
            1,
            2,
            2,
            2,
            1,
            2,
            1,
            3,
            3,
            3,
            1,
            3,
            1,
            2,
            1,
        ],
    )


def graph2():
    # using pos args
    return SCOTCH.build_graph(
        1,
        7,
        [1, 4, 10, 13, 16, 19, 22, 25],
        None,
        [4, 1, 4, 4, 4, 4, 4],
        None,  # vlbltab
        24,
        [3, 2, 6, 3, 4, 1, 7, 6, 5, 1, 2, 4, 2, 7, 3, 7, 2, 6, 2, 1, 5, 5, 2, 4],
        [1, 1, 1, 2, 2, 1, 2, 3, 3, 1, 2, 2, 2, 1, 2, 1, 3, 3, 3, 1, 3, 1, 2, 1],
    )


def graph3():
    # using the least args possible
    return SCOTCH.build_graph(
        verttab=[1, 4, 10, 13, 16, 19, 22, 25],
        edgetab=[
            3,
            2,
            6,
            3,
            4,
            1,
            7,
            6,
            5,
            1,
            2,
            4,
            2,
            7,
            3,
            7,
            2,
            6,
            2,
            1,
            5,
            5,
            2,
            4,
        ],
    )


def graph4():
    # same but with vendtab
    baseval = 1001
    verttab = [1001, 1004, 1010, 1013, 1016, 1019, 1022, 1025]
    vendtab = [1004, 1010, 1013, 1016, 1019, 1022, 1025]
    edgetab = [
        1003,
        1002,
        1006,
        1003,
        1004,
        1001,
        1007,
        1006,
        1005,
        1001,
        1002,
        1004,
        1002,
        1007,
        1003,
        1007,
        1002,
        1006,
        1002,
        1001,
        1005,
        1005,
        1002,
        1004,
    ]
    verttab = _common.properly_format(verttab)
    vendtab = _common.properly_format(vendtab)
    edgetab = _common.properly_format(edgetab)
    return SCOTCH.build_graph(
        baseval=baseval, verttab=verttab, vendtab=vendtab, edgetab=edgetab
    )


def graph5():
    # really non-compact
    return SCOTCH.build_graph(
        baseval=1,
        verttab=[1, 6, 14, 19, 24, 29, 34, 37],
        vendtab=[4, 12, 17, 22, 27, 32, 37],
        edgetab=[
            3,
            2,
            6,
            0,
            0,
            3,
            4,
            1,
            7,
            6,
            5,
            0,
            0,
            1,
            2,
            4,
            0,
            0,
            2,
            7,
            3,
            0,
            0,
            7,
            2,
            6,
            0,
            0,
            2,
            1,
            5,
            0,
            0,
            5,
            2,
            4,
            0,
            0,
        ],
    )


def testBuildKWArgs():
    testGraph = graph1()
    checkk = testGraph.check()
    assert checkk == 0


def testBuildPosArgs():
    testGraph = graph2()
    checkk = testGraph.check()
    assert checkk == 0


def testCompactImplicitArgs():
    testGraph = graph3()
    checkk = testGraph.check()
    assert checkk == 0


def testBuildIncompact():
    testGraph = graph4()
    checkk = testGraph.check()
    baseval, vertnbr, verttab, vendtab, velotab, vlbltab, _, edgetab, _ = (
        testGraph.data(as_dict=False)
    )
    assert checkk == 0


def testBuildNonCompact():
    testGraph = graph5()
    checkk = testGraph.check()
    assert checkk == 0


def testDiamPV():
    g = graph2()
    pv = g.diam_pv()
    assert pv == 10


def testStat():
    g = graph2()
    stats = g.stat(as_dict=False)
    witness = (
        1,
        4,
        25,
        25 / 7,
        0.7346938775510203,
        3,
        6,
        24 / 7,
        0.7346938775510203,
        1,
        3,
        44,  # 22 or 44 ?
        11 / 6,
        0.6944444444444445,
    )
    validind = range(14)
    for k in validind:
        assert stats[k] == witness[k]


def testDataCompact():
    kwargs = {
        "baseval": 1,
        "vertnbr": 7,
        "verttab": [1, 4, 10, 13, 16, 19, 22, 25],
        "vendtab": None,
        "velotab": [4, 1, 4, 4, 4, 4, 4],
        "vlbltab": None,  # vlbltab
        "edgenbr": 24,
        "edgetab": [
            3,
            2,
            6,
            3,
            4,
            1,
            7,
            6,
            5,
            1,
            2,
            4,
            2,
            7,
            3,
            7,
            2,
            6,
            2,
            1,
            5,
            5,
            2,
            4,
        ],
        "edlotab": [
            1,
            1,
            1,
            2,
            2,
            1,
            2,
            3,
            3,
            1,
            2,
            2,
            2,
            1,
            2,
            1,
            3,
            3,
            3,
            1,
            3,
            1,
            2,
            1,
        ],
    }
    g = SCOTCH.build_graph(**kwargs)
    assert g.check() == 0
    data = g.data(as_dict=True)
    for key, val in kwargs.items():
        if isinstance(val, int):
            assert val == data[key]
        else:
            assert np.array_equal(np.array(val), data[key])


def testDataNonCompact1():
    kwargs = {
        "baseval": 1,
        "vertnbr": 7,
        "verttab": [1, 4, 10, 13, 16, 19, 22, 25],
        "vendtab": [4, 10, 13, 16, 19, 22, 25],
        "velotab": [4, 1, 4, 4, 4, 4, 4],
        "vlbltab": None,  # vlbltab
        "edgenbr": 24,
        "edgetab": [
            3,
            2,
            6,
            3,
            4,
            1,
            7,
            6,
            5,
            1,
            2,
            4,
            2,
            7,
            3,
            7,
            2,
            6,
            2,
            1,
            5,
            5,
            2,
            4,
        ],
        "edlotab": [
            1,
            1,
            1,
            2,
            2,
            1,
            2,
            3,
            3,
            1,
            2,
            2,
            2,
            1,
            2,
            1,
            3,
            3,
            3,
            1,
            3,
            1,
            2,
            1,
        ],
    }
    g = SCOTCH.build_graph(**kwargs)
    assert g.check() == 0
    data = g.data(as_dict=True)
    for key, val in kwargs.items():
        if isinstance(val, int):
            assert val == data[key]
        else:
            assert np.array_equal(np.array(val), data[key])


def testDataNonCompact2():
    kwargs = {
        "baseval": 1,
        "verttab": [3, 8, 16, 21, 26, 31, 36, 39],
        "vendtab": [6, 14, 19, 24, 29, 34, 39],
        "edgetab": [
            0,
            0,
            3,
            2,
            6,
            0,
            0,
            3,
            4,
            1,
            7,
            6,
            5,
            0,
            0,
            1,
            2,
            4,
            0,
            0,
            2,
            7,
            3,
            0,
            0,
            7,
            2,
            6,
            0,
            0,
            2,
            1,
            5,
            0,
            0,
            5,
            2,
            4,
            0,
            0,
        ],
        "velotab": [4, 1, 4, 4, 4, 4, 4],
    }
    g = SCOTCH.build_graph(**kwargs)
    assert g.check() == 0
    data = g.data(as_dict=True)
    for key, val in kwargs.items():
        if isinstance(val, int):
            assert val == data[key]
        else:
            assert np.array_equal(np.array(val), data[key])


def testLoadOpen():
    # Using some test file found in the check/ directory
    g = SCOTCH.Graph()
    with open(os.path.join(_data_dir, "m4x4.grf")) as fp:
        g.load(fp)
    assert g.check() == 0


def testLoadString():
    g = SCOTCH.Graph()
    g.load(os.path.join(_data_dir, "m4x4.grf"))
    assert g.check() == 0


def testLoadBinaryString():
    g = SCOTCH.Graph()
    g.load(os.path.join(_data_dir, "m4x4.grf"))
    assert g.check() == 0


def testLoadGraph():
    g = SCOTCH.load_graph(os.path.join(_data_dir, "m4x4.grf"))
    assert g.check() == 0


def testSaveOpen():
    filename = "./test_scotchpy_graph_save_output.txt"
    g = graph1()
    with open(filename, "w+") as fp1:
        g.save(fp1)
    with open(filename) as fp2:
        h = SCOTCH.load_graph(fp2)
    datag = g.data()
    datah = h.data()
    for valg, valh in zip(datag, datah):
        if isinstance(valg, int):
            assert valg == valh
        else:
            assert np.array_equal(valg, valh)
    os.remove(filename)


def testLibraryError():
    try:
        testGraph = SCOTCH.build_graph(
            baseval=1,
            vertnbr=7,
            verttab=[1, 5, 10, 13, 16, 19, 22, 25],
            vendtab=None,
            velotab=[4, 1, 4, 4, 4, 4, 4],
            vlbltab=None,  # vlbltab
            edgenbr=24,
            edgetab=[
                3,
                2,
                6,
                3,
                4,
                1,
                7,
                6,
                5,
                1,
                2,
                4,
                2,
                7,
                3,
                7,
                2,
                6,
                2,
                1,
                5,
                5,
                2,
                4,
            ],
            edlotab=[
                1,
                1,
                1,
                2,
                2,
                1,
                2,
                3,
                3,
                1,
                2,
                2,
                2,
                1,
                2,
                1,
                3,
                3,
                3,
                1,
                3,
                1,
                2,
                1,
            ],
        )
        testGraph.check()
    except SCOTCH.LibraryError as e:
        assert "Inside Python ERROR" in "\n".join(e.args)
    else:
        assert False, "A LibraryException was supposed to be raised"


# taken from check/test_scotch_graph_part_ovl.c


def testPartOvl():
    partnbr = 4
    filename1 = os.path.join(_data_dir, "m4x4.grf")
    tmpdir = tempfile.TemporaryDirectory()
    filename2 = os.path.join(tmpdir.name, "test_scotchpy_graph_part_ovl_output.txt")
    string = None

    assert partnbr >= 1

    try:
        strat = SCOTCH.Strat()
    except SCOTCH.LibraryError:
        assert False, "cannot initialize strategy"
        raise

    if string:
        try:
            strat.graph_part_ovl(string)
        except SCOTCH.LibraryError:
            assert False, "invalid user-provided strategy"
            raise

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

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

    baseval, vertnbr, verttab, vendtab, velotab, vlbltab, _, edgetab, _ = graf.data(
        as_dict=False
    )
    vendtab = vendtab or verttab[1:]
    parttab = [0 for _ in range(vertnbr)]

    try:
        graf.part_ovl(partnbr, strat, parttab)
    except SCOTCH.LibraryError:
        assert False, "cannot compute mapping"
        raise

    if baseval:
        edgetab = np.concatenate(([None for _ in range(baseval)], edgetab))
        parttab = np.concatenate(([None for _ in range(baseval)], parttab))
    loadtab = [0 for _ in range(partnbr)]
    flagtab = [~0 for _ in range(partnbr)]

    for vertnum in range(vertnbr):
        veloval = velotab[vertnum] if velotab else 1
        partval = parttab[vertnum + baseval]
        if partval >= 0:
            loadtab[partval] += veloval
        else:
            for edgenum in range(verttab[vertnum], vendtab[vertnum]):
                vertend = edgetab[edgenum]
                partend = parttab[vertend]
                if partend < 0:
                    continue
                if flagtab[partend] == vertnum:
                    continue
                loadtab[partend] += veloval
                flagtab[partend] = vertnum

    loadsum = loadmax = 0
    loadmin = float("inf")  # SCOTCH_NUMMAX, but this is better
    for partnum in range(partnbr):
        loadsum += loadtab[partnum]
        if loadtab[partnum] > loadmax:
            loadmax = loadtab[partnum]
        if loadtab[partnum] < loadmin:
            loadmin = loadtab[partnum]
        # assert  False ,"M\tCompload[{}]\t{}\n".format(partnum, loadtab[partnum]))

    loadavg = loadsum / partnbr
    # assert  False ,"M\tCompLoadAvg\t{}\n".format(loadavg))
    # assert  False ,"M\tCompLoadMax/Avg\t{}\n".format(loadmax/loadavg))

    with open(filename2, "w+") as file2:
        file2.write("{}\n".format(vertnbr))
        for vertnum in range(vertnbr):
            file2.write(
                "{}\t{}\n".format(
                    vlbltab[vertnum] if vlbltab else vertnum, parttab[vertnum + baseval]
                )
            )
    tmpdir.cleanup()

# taken from check/test_scotch_graph_color.c


def testColor():
    filename = os.path.join(_data_dir, "m4x4.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)
    colotab = [-1 for _ in range(vertnbr)]
    cnbrtab = [0 for _ in range(vertnbr)]

    try:
        graf.color(colotab, 0)
    except SCOTCH.LibraryError:
        assert False, "cannot color graph"
        raise
    colonbr = max(colotab)

    # supplementary test
    assert -1 not in colotab, "incorrect graph coloration"
    return  # to avoid unnecessary assert False

    assert False, "Number of colors: {}".format(colonbr)

    for vertnum in range(vertnbr):
        cnbrtab[colotab[vertnum]] += 1

    for colonum in range(colonbr):
        assert False, "Color {}: {}".format(colonum, cnbrtab[colonum])


# taken from check/test_scotch_graph_diam.c
def testDiam():
    filename = os.path.join(_data_dir, "m4x4.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

    try:
        diamval = graf.diam_pv()
    except SCOTCH.LibraryError:
        assert False, "cannot compute graph pseudo-diameter"
        raise

    assert diamval >= 0

    # assert  False ,"Graph pseudo-diameter: {}".format(diamval))


def testInduce():
    # taken from check/test_scotch_graph_induce.c

    filename = os.path.join(_data_dir, "m4x4.grf")

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

    try:
        orggraf.load(filename)
    except SCOTCH.LibraryError:
        assert False, "cannot load graph"
        raise
    baseval, orgvertnbr, _, _, _, _, _, _, _ = orggraf.data(as_dict=False)

    orgparttab = [0 for _ in range(orgvertnbr)]

    indlisttab = [0 for _ in range(orgvertnbr)]

    _common.random_reset()

    for orgvertnum in range(orgvertnbr):
        orgparttab[orgvertnum] = orgvertnum + baseval

    for orgvertnum in range(orgvertnbr - 1):
        randval = _common.random_val(orgvertnbr - orgvertnum)
        verttmp = orgparttab[orgvertnum + randval]
        indlisttab[orgvertnum + randval] = indlisttab[orgvertnum]
        indlisttab[orgvertnum] = verttmp

    indvertnbr = (orgvertnbr + 1) // 2

    orgparttab = [0 for _ in range(orgvertnbr)]
    for indvertnum in range(indvertnbr):
        orgparttab[indlisttab[indvertnum] - baseval] = 1

    indgraf = SCOTCH.Graph()
    try:
        orggraf.induce_list(indvertnbr, indlisttab, indgraf)
    except SCOTCH.LibraryError:
        assert False, "cannot induce graph"
        raise
    indgraf.check()
    indgraf.exit()
    indgraf.init()
    try:
        orggraf.induce_part(indvertnbr, orgparttab, 1, indgraf)
    except SCOTCH.LibraryError:
        assert False, "cannot induce graph"
        raise
    indgraf.check()
    indgraf.exit()
    orggraf.exit()
