## 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 scotchpy as SCOTCH
import scotchpy.common as common
import importlib.resources
import os

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

# taken from check/test_scotch_graph_coarsen.c


def checkCoarMult(coarmulttab, coarvertnbr, finevertnbr, baseval):

    finevertnnd = finevertnbr + baseval
    coarverttmp = finevertnbr
    for coarvertnum in range(coarvertnbr):
        finevertnum0 = coarmulttab[coarvertnum * 2]
        finevertnum1 = coarmulttab[coarvertnum * 2 + 1]
        if (
            finevertnum0 < baseval
            or finevertnum0 >= finevertnnd
            or finevertnum1 < baseval
            or finevertnum1 >= finevertnnd
        ):
            raise ValueError("Invalid multinode array (1)")
        if finevertnum0 != finevertnum1:
            coarverttmp -= 1
    if coarverttmp != coarvertnbr:
        raise ValueError("Invalid multinode array (2)")


def checkFineCoar(finematetab, coarvertnbr, finevertnbr, baseval):

    finevertnnd = finevertnbr + baseval
    coarverttmp = finevertnbr
    for finevertnum in range(baseval, finevertnnd):
        finevertend = finematetab[finevertnum]

        if (finevertnnd <= finevertend) or (finevertend < baseval):
            assert False, "invalid mate array (1)"
        if finevertend > finevertnum:
            coarverttmp -= 1
            if finematetab[finevertend] != finevertnum:
                assert False, "invalid mate array (2)"
    if coarverttmp != coarvertnbr:
        assert False, "invalid mate array (3)"


def testCoarsen():

    for filename in (
        os.path.join(_data_dir, "bump.grf"),
        os.path.join(_data_dir, "bump_b100000.grf"),
    ):
        try:
            finegraf = SCOTCH.Graph(init=True)
        except SCOTCH.LibraryError:
            assert False, "cannot initialize graph"
            raise
        coargraf = SCOTCH.Graph(init=False)

        with open(filename) as fp:
            try:
                finegraf.load(fp, -1, 0)
            except SCOTCH.LibraryError:
                assert False, "cannot load {} graph".format(filename)
                raise
        baseval, finevertnbr, fineverttax, finevendtax, _, _, _, fineedgetax, _ = (
            finegraf.data(as_dict=False)
        )
        print("baseval = {}".format(baseval))
        print("finevertnbr = {}".format(finevertnbr))
        print("fineverttax = {}".format(fineverttax))
        print("finevendtax = {}".format(finevendtax))
        print("fineedgetax = {}".format(fineedgetax))
        if finegraf.check() != 0:
            assert False, "invalid graph"
            raise
        baseval, finevertnbr, fineverttax, finevendtax, _, _, _, fineedgetax, _ = (
            finegraf.data(as_dict=False)
        )
        print("Fine graph baseval = {}".format(baseval))
        if fineverttax is not None:
            fineverttax = [-1 for _ in range(baseval)] + list(fineverttax)
        if finevendtax is not None:
            finevendtax = [-1 for _ in range(baseval)] + list(finevendtax)
        if fineedgetax is not None:
            fineedgetax = [-1 for _ in range(baseval)] + list(fineedgetax)
        finematetax = [-1 for _ in range(finevertnbr)]
        coarvertnbr = 0

        try:
            coarvertnbr = finegraf.coarsen_match(
                coarvertnbr, 1.0, common.COARSENNOMERGE, finematetax
            )
        except SCOTCH.LibraryError:
            assert False, "cannot compute matching"
            raise

        finematetax = [-1 for _ in range(baseval)] + finematetax

        checkFineCoar(finematetax, coarvertnbr, finevertnbr, baseval)

        print("Mate array has {} vertices".format(coarvertnbr))
        print("Graph mated with a ratio of {}".format(coarvertnbr / finevertnbr))

        coarmulttab = [-1 for _ in range(coarvertnbr * 2)]

        try:
            finegraf.coarsen_build(
                coarvertnbr, finematetax[baseval:], coargraf, coarmulttab
            )
        except SCOTCH.LibraryError:
            assert False, "cannot compute coarse graph (1)"
            raise

        coarvertnbr, coaredgenbr = coargraf.size()

        print(
            "Coarse graph has {} vertices and {} edges".format(coarvertnbr, coaredgenbr)
        )
        print("Graph coarsened with a ratio of {}".format(coarvertnbr / finevertnbr))

        checkCoarMult(coarmulttab, coarvertnbr, finevertnbr, baseval)

        coargraf.exit()
        del coarmulttab

        finematetax[baseval:] = [~0 for _ in range(finevertnbr)]

        # match build
        finevertnnd = finevertnbr + baseval
        coarvertnbr = 0
        for finevertnum in range(baseval, finevertnnd):
            if finematetax[finevertnum] >= 0:
                continue

            coarvertnbr += 1

            finematenum = finevertnum
            if fineverttax is not None and finevendtax is not None:
                if fineverttax[finevertnum] == finevendtax[finevertnum]:
                    while True:
                        finevertnum += 1
                        if finevertnum >= finevertnnd:
                            finematetax[finematenum] = finematenum
                            brk = True
                            break
                        if finematetax[finevertnum] < 0:
                            break
                    if brk:
                        break
                else:
                    for fineedgenum in range(
                        fineverttax[finevertnum], finevendtax[finevertnum]
                    ):
                        finevertend = fineedgetax[fineedgenum]
                        if finematetax[finevertend] < 0:
                            finematenum = finevertend
                            break

            finematetax[finevertnum] = finematenum
            finematetax[finematenum] = finevertnum

        print("Mate array has {} vertices".format(coarvertnbr))
        print("Graph mated with a ratio of {}".format(coarvertnbr / finevertnbr))

        coarmulttab = [-1 for _ in range(2 * coarvertnbr)]

        try:
            finegraf.coarsen_build(
                coarvertnbr, finematetax[baseval:], coargraf, coarmulttab
            )
        except SCOTCH.LibraryError:
            assert False, "cannot compute coarse graph (2)"
            raise

        coarvertnbr, coaredgenbr = coargraf.size()

        print(
            "Coarse graph has {} vertices and {} edges".format(coarvertnbr, coaredgenbr)
        )
        print("Graph coarsened with a ratio of {}".format(coarvertnbr / finevertnbr))

        coargraf.exit()
        del coarmulttab
        del finematetax

        coarmulttab = [-1 for _ in range(2 * finevertnbr)]

        try:
            finegraf.coarsen(1, 1.0, common.COARSENNONE, coargraf, coarmulttab)
        except SCOTCH.LibraryError:
            assert False, "cannot coarsen graph"
            raise

        coarvertnbr, _ = coargraf.size()
        print(
            "Coarse graph has {} vertices and {} edges".format(coarvertnbr, coaredgenbr)
        )
        print("Graph coarsened with a ratio of {}".format(coarvertnbr / finevertnbr))
        checkCoarMult(coarmulttab, coarvertnbr, finevertnbr, baseval)
