## 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 sys

import pytest
import scotchpy as SCOTCH
import numpy as np
import scotchpy.common as _common
try:
    import scotchpy.dgraph as PTSCOTCH
except AttributeError:
    pytest.skip('ptscotch not available', allow_module_level=True)
from scotchpy.common import _ptscotch_lib, _libptscotch_mpi4py_lib
import tempfile
import importlib.resources
import os

try:
    from mpi4py import MPI
    if _ptscotch_lib and _libptscotch_mpi4py_lib:
        flag_mpi4py_available = True
    else:
        flag_mpi4py_available = False
except ImportError:
    flag_mpi4py_available = False

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

@pytest.mark.mpi
def testDgraphCheck():

    if not flag_mpi4py_available:
        pytest.xfail("MPI not found")

    comm = MPI.COMM_WORLD
    proclocnum = comm.Get_rank()

    comm.Barrier()
    graf = PTSCOTCH.DGraph(comm)

    if proclocnum == 0:
        with open(os.path.join(_data_dir, "bump_b100000.grf")) as fp:
            graf.load(fp)
    else:
        graf.load(None)

    if graf.check():
        print("Invalid graph")
        sys.exit(1)

    graf.exit()
    comm.Barrier()

@pytest.mark.mpi
def testDgraphBand():

    if not flag_mpi4py_available:
        pytest.xfail("MPI not found")

    comm = MPI.COMM_WORLD
    procglbnbr = comm.Get_size()
    proclocnum = comm.Get_rank()
    comm.Barrier()

    graf = PTSCOTCH.DGraph(comm)
    if proclocnum == 0:
        with open(os.path.join(_data_dir, "bump_b100000.grf")) as fp:
            graf.load(fp)
    else:
        graf.load(None)

    comm.Barrier()
    baseval, vertglbnbr, vertlocnbr, _, _, _, _, _, _, _, _, _, _, _, _ = graf.data()
    fronloctab = np.full(vertlocnbr, -1, _common.proper_int)

    bandgraf = PTSCOTCH.DGraph(comm)

    fronloctab[0] = baseval
    if proclocnum == 1:
        fronlocnbr = 0
    else:
        fronlocnbr = 1

    graf.band(fronlocnbr, fronloctab, 4, bandgraf)

    fronloctab = None

    (
        baseval,
        bandvertglbnbr,
        bandvertlocnbr,
        _,
        _,
        _,
        _,
        _,
        bandvlblloctab,
        _,
        _,
        _,
        _,
        _,
        _,
    ) = bandgraf.data()

    tmpdir = tempfile.TemporaryDirectory()
    tmpfile = os.path.join(tmpdir.name, "test_scotch_dgraph_band_output.txt")
    if proclocnum == 0:
        try:
            with open(tmpfile, "x") as file:
                pass
        except FileExistsError:
            pass

    comm.Barrier()

    for procnum in range(procglbnbr):
        comm.Barrier()

        if procnum == proclocnum:
            with open(tmpfile, "a+") as file:
                if procnum == 0:
                    file.write(f"{bandvertglbnbr}\n")

                for bandvertlocnum in range(bandvertlocnbr):
                    file.write(f"{bandvlblloctab[bandvertlocnum]}\t1\n")

    graf.exit()
    bandgraf.exit()
    tmpdir.cleanup()


@pytest.mark.mpi
def testDgraphCoarsen():

    if not flag_mpi4py_available:
        pytest.xfail("MPI not found")

    comm = MPI.COMM_WORLD
    procglbnbr = comm.Get_size()
    proclocnum = comm.Get_rank()
    comm.Barrier()
    graf = PTSCOTCH.DGraph(comm)

    if proclocnum == 0:
        with open(os.path.join(_data_dir, "bump_b100000.grf")) as fp:
            graf.load(fp)
    else:
        graf.load(None)

    _, vertglbnbr, vertlocnbr, _, _, _, _, _, _, _, _, _, _, _, _ = graf.data()

    coarrat = 0.8

    for i in range(3):
        if i == 0:
            foldval = _common.COARSENNONE
            foldstr = "Plain coarsening"
        elif i == 1:
            foldval = _common.COARSENFOLD
            foldstr = "Folding"
        else:
            foldval = _common.COARSENFOLDDUP
            foldstr = "Folding with duplication"

        if proclocnum == 0:
            print(foldstr)


        coargraf = PTSCOTCH.DGraph(comm)
        o, multloctab = graf.coarsen(0, coarrat, foldval, coargraf)
        coarstr = ""
        if o == 0:
            coarstr = "coarse graph created"
        if o == 2:
            coarstr = "folded graph not created here"
        if o == 3:
            coarstr = "cannot create coarse graph"
        _, coarvertglbnbr, coarvertlocnbr, _, _, _, _, _, _, _, _, _, _, _, _ = (
            coargraf.data()

        )
        for procnum in range(procglbnbr):
            comm.Barrier()
            if procnum == proclocnum:
                if coarvertlocnbr and vertglbnbr :
                   print(
                      f"{procnum}: {coarstr} ({vertlocnbr} / {coarvertlocnbr} / {float(coarvertglbnbr) / float(vertglbnbr)})"
                   )

        coargraf.exit()

    graf.exit()
    comm.Barrier()

@pytest.mark.mpi
def testDgraphGrow():

    if not flag_mpi4py_available:
        pytest.xfail("MPI not found")

    comm = MPI.COMM_WORLD
    procglbnbr = comm.Get_size()
    proclocnum = comm.Get_rank()
    comm.Barrier()
    graf = PTSCOTCH.DGraph(comm)

    if proclocnum == 0:
        with open(os.path.join(_data_dir, "bump_b100000.grf")) as fp:
            graf.load(fp)
    else:
        graf.load(None)

    comm.Barrier()
    graf.ghst()

    baseval, vertglbnbr, vertlocnbr, _, vertgstnbr, _, _, _, _, _, _, _, _, _, _ = (
        graf.data()
    )
    seedloctab = np.full(vertlocnbr, -1, _common.proper_int)
    partgsttab = np.full(vertgstnbr, ~0, _common.proper_int)

    _common.random_reset()
    seedloctab[0] = baseval + _common.random_val(vertlocnbr)
    seedloctab[1] = baseval + _common.random_val(vertlocnbr)
    seedloctab[2] = baseval + _common.random_val(vertlocnbr)

    partgsttab[seedloctab[0] - baseval] = 0
    partgsttab[seedloctab[1] - baseval] = 1
    partgsttab[seedloctab[2] - baseval] = 2

    graf.grow(3, seedloctab, 4, partgsttab)
    seedloctab = None
    tmpdir = tempfile.TemporaryDirectory()
    tmpfile = os.path.join(tmpdir.name, "test_scotchpy_dgraph_grow_output.txt")
    if proclocnum == 0:
        try:
            with open(tmpfile, "x") as file:
                pass
        except FileExistsError:
            pass
    comm.Barrier()

    statdat = MPI.Status()

    for procnum in range(procglbnbr):
        comm.Barrier()

        if procnum == proclocnum:
            with open(tmpfile, "a+") as file:
                if procnum == 0:
                    file.write(f"{vertglbnbr}\n")
                    vertlocadj = baseval
                else:
                    vertlocadj = comm.recv(source=procnum - 1, tag=0, status=statdat)

                for vertlocnum in range(vertlocnbr):
                    file.write(f"{vertlocadj + vertlocnum}\t{partgsttab[vertlocnum]}\n")

            if procnum < (procglbnbr - 1):
                vertlocadj += vertlocnbr
                comm.send(vertlocadj, dest=procnum + 1, tag=0)
    graf.exit()
    tmpdir.cleanup()

@pytest.mark.mpi
def testDgraphInduce():

    if not flag_mpi4py_available:
        pytest.xfail("MPI not found")

    comm = MPI.COMM_WORLD
    proclocnum = comm.Get_rank()
    comm.Barrier()
    orggraf = PTSCOTCH.DGraph(comm)

    if proclocnum == 0:
        with open(os.path.join(_data_dir, "bump_b100000.grf")) as fp:
            orggraf.load(fp)
    else:
        orggraf.load(None)

    comm.Barrier()
    baseval, _, orgvertlocnbr, _, _, _, _, _, _, _, _, _, _, _, _ = orggraf.data()
    orgpartloctab = np.full(orgvertlocnbr, -1, dtype=int)
    indlistloctab = np.full(orgvertlocnbr, -1, dtype=int)
    _common.random_reset()
    orgpartloctab = np.zeros(orgvertlocnbr, dtype=int)
    indlistloctab = np.arange(baseval, baseval + orgvertlocnbr, dtype=int)

    np.random.shuffle(indlistloctab)

    indvertlocnbr = (orgvertlocnbr + 1) // 2

    orgpartloctab[indlistloctab[:indvertlocnbr] - baseval] = 1

    indgraf = PTSCOTCH.DGraph(comm)
    indgraf.init(comm)
    orggraf.induce_part(orgpartloctab, 1, indvertlocnbr, indgraf)
    indgraf.check()
    orggraf.exit()
    indgraf.exit()


@pytest.mark.mpi
def testDgraphRedist():

    if not flag_mpi4py_available:
        pytest.xfail("MPI not found")

    comm = MPI.COMM_WORLD
    procglbnbr = comm.Get_size()
    proclocnum = comm.Get_rank()

    comm.Barrier()

    srcgraf = PTSCOTCH.DGraph(comm)
    dstgraf = PTSCOTCH.DGraph(comm)
    dstgraf.init(comm)

    if proclocnum == 0:
        with open(os.path.join(_data_dir, "bump_b100000.grf")) as fp:
            srcgraf.load(fp)
    else:
        srcgraf.load(None)

    srcgraf.check()
    comm.Barrier()

    _, vertglbnbr, vertlocnbr, _, _, _, _, _, _, _, _, _, _, _, _ = srcgraf.data()

    partloctab = np.zeros(vertlocnbr, dtype=int)
    for vertlocnum in range(vertlocnbr):
        partloctab[vertlocnum] = (vertlocnum // 3) % procglbnbr

    srcgraf.redist(partloctab, None, -1, -1, dstgraf)

    dstgraf.exit()
    srcgraf.exit()
    comm.Barrier()