# ____________________________________________________________________________________
#
# Pyomo: Python Optimization Modeling Objects
# Copyright (c) 2008-2026 National Technology and Engineering Solutions of Sandia, LLC
# Under the terms of Contract DE-NA0003525 with National Technology and Engineering
# Solutions of Sandia, LLC, the U.S. Government retains certain rights in this
# software.  This software is distributed under the 3-clause BSD License.
# ____________________________________________________________________________________
import pyomo.common.unittest as unittest

try:
    import numpy as np
    from scipy.sparse import coo_matrix, tril
except ImportError:
    raise unittest.SkipTest("Pynumero needs scipy and numpy to run linear solver tests")

from pyomo.contrib.pynumero.linalg.mumps_interface import (
    mumps_available,
    MumpsCentralizedAssembledLinearSolver,
)

from pyomo.contrib.pynumero.sparse import BlockMatrix, BlockVector


@unittest.skipIf(
    not mumps_available, "Pynumero needs pymumps to run linear solver tests"
)
class TestMumpsLinearSolver(unittest.TestCase):
    def test_mumps_linear_solver(self):
        A = np.array([[1, 7, 3], [7, 4, -5], [3, -5, 6]], dtype=np.double)
        A = coo_matrix(A)
        A_lower = tril(A)
        x1 = np.arange(3) + 1
        b1 = A * x1
        x2 = np.array(list(reversed(x1)))
        b2 = A * x2

        solver = MumpsCentralizedAssembledLinearSolver()
        solver.do_symbolic_factorization(A)
        solver.do_numeric_factorization(A)
        x, res = solver.do_back_solve(b1)
        self.assertTrue(np.allclose(x, x1))
        x, res = solver.do_back_solve(b2)
        self.assertTrue(np.allclose(x, x2))

        solver = MumpsCentralizedAssembledLinearSolver(sym=2)
        x, res = solver.solve(A_lower, b1)
        self.assertTrue(np.allclose(x, x1))

        block_A = BlockMatrix(2, 2)
        block_A.set_row_size(0, 2)
        block_A.set_row_size(1, 1)
        block_A.set_col_size(0, 2)
        block_A.set_col_size(1, 1)
        block_A.copyfrom(A)

        block_b1 = BlockVector(2)
        block_b1.set_block(0, b1[0:2])
        block_b1.set_block(1, b1[2:])

        block_b2 = BlockVector(2)
        block_b2.set_block(0, b2[0:2])
        block_b2.set_block(1, b2[2:])

        solver = MumpsCentralizedAssembledLinearSolver(
            icntl_options={10: -3}, cntl_options={2: 1e-16}
        )
        solver.do_symbolic_factorization(block_A)
        solver.do_numeric_factorization(block_A)
        x, res = solver.do_back_solve(block_b1)
        self.assertTrue(np.allclose(x, x1))
        x, res = solver.do_back_solve(block_b2)
        self.assertTrue(np.allclose(x, x2))
        self.assertEqual(solver.get_infog(15), 3)
