# ____________________________________________________________________________________
#
# 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
from pyomo.common.dependencies import attempt_import

np, numpy_available = attempt_import(
    'numpy', 'Interior point requires numpy', minimum_version='1.13.0'
)
scipy, scipy_available = attempt_import('scipy', 'Interior point requires scipy')
mumps, mumps_available = attempt_import('mumps')
if not (numpy_available and scipy_available):
    raise unittest.SkipTest('Interior point tests require numpy and scipy')
from scipy.sparse import coo_matrix
import pyomo.contrib.interior_point as ip

if mumps_available:
    from pyomo.contrib.interior_point.linalg.mumps_interface import MumpsInterface
from pyomo.contrib.pynumero.linalg.base import LinearSolverStatus


@unittest.skipIf(not mumps_available, 'mumps is not available')
class TestReallocation(unittest.TestCase):
    def test_reallocate_memory_mumps(self):
        # Create a tri-diagonal matrix with small entries on the diagonal
        n = 10000
        small_val = 1e-7
        big_val = 1e2
        irn = []
        jcn = []
        ent = []
        for i in range(n - 1):
            irn.extend([i + 1, i, i])
            jcn.extend([i, i, i + 1])
            ent.extend([big_val, small_val, big_val])
        irn.append(n - 1)
        jcn.append(n - 1)
        ent.append(small_val)
        irn = np.array(irn)
        jcn = np.array(jcn)
        ent = np.array(ent)

        matrix = coo_matrix((ent, (irn, jcn)), shape=(n, n))

        linear_solver = MumpsInterface()
        linear_solver.do_symbolic_factorization(matrix)

        predicted = linear_solver.get_infog(16)

        # We predict that factorization will take 2 MB
        self.assertEqual(predicted, 2)

        # Explicitly set maximum memory to less than the predicted
        # requirement.
        linear_solver.set_icntl(23, 1)

        res = linear_solver.do_numeric_factorization(matrix, raise_on_error=False)
        self.assertEqual(res.status, LinearSolverStatus.not_enough_memory)

        linear_solver.do_symbolic_factorization(matrix)

        factor = 2
        linear_solver.increase_memory_allocation(factor)

        res = linear_solver.do_numeric_factorization(matrix)
        self.assertEqual(res.status, LinearSolverStatus.successful)

        # Expected memory allocation (MB)
        self.assertEqual(linear_solver._prev_allocation, 2 * predicted)

        actual = linear_solver.get_infog(18)

        # Sanity checks:
        # Make sure actual memory usage is greater than initial guess
        self.assertTrue(predicted < actual)
        # Make sure memory allocation is at least as much as was used
        self.assertTrue(actual <= linear_solver._prev_allocation)


if __name__ == '__main__':
    test_realloc = TestReallocation()
    test_realloc.test_reallocate_memory_mumps()
