#  ___________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright (c) 2008-2022
#  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.contrib.pynumero.dependencies import numpy as np, numpy_available
if not numpy_available:
    raise unittest.SkipTest('pynumero MA27 tests require numpy')
import ctypes
from pyomo.contrib.pynumero.linalg.ma27 import MA27Interface


@unittest.skipIf(not MA27Interface.available(), reason='MA27 not available')
class TestMA27Interface(unittest.TestCase):

    def test_get_cntl(self):
        ma27 = MA27Interface()
        self.assertEqual(ma27.get_icntl(1), 6)

        self.assertAlmostEqual(ma27.get_cntl(1), 1e-1) # Numerical pivot threshold
        self.assertAlmostEqual(ma27.get_cntl(3), 0.0) # Null pivot threshold

    def test_set_icntl(self):
        ma27 = MA27Interface()
        ma27.set_icntl(5, 4) # Set output printing to max verbosity
        ma27.set_icntl(8, 1) # Keep factors when we run out of space
                             # (so MA27ED can be used)
        icntl5 = ma27.get_icntl(5)
        icntl8 = ma27.get_icntl(8)
        self.assertEqual(icntl5, 4)
        self.assertEqual(icntl8, 1)

        with self.assertRaisesRegex(TypeError, 'must be an integer'):
            ma27.set_icntl(1.0, 0)
        with self.assertRaisesRegex(IndexError, 'is out of range'):
            ma27.set_icntl(100, 0)
        with self.assertRaises(ctypes.ArgumentError):
            ma27.set_icntl(1, 0.0)

    def test_set_cntl(self):
        ma27 = MA27Interface()
        ma27.set_cntl(1, 1e-8)
        ma27.set_cntl(3, 1e-12)
        self.assertAlmostEqual(ma27.get_cntl(1), 1e-8)
        self.assertAlmostEqual(ma27.get_cntl(3), 1e-12)

    def test_do_symbolic_factorization(self):
        ma27 = MA27Interface()

        n = 5
        ne = 7
        irn = np.array([1,1,2,2,3,3,5], dtype=np.intc)
        icn = np.array([1,2,3,5,3,4,5], dtype=np.intc)
        # These arrays, copied out of HSL docs, contain Fortran indices.
        # Interfaces accept C indices as this is what I typically expect.
        irn = irn - 1
        icn = icn - 1

        bad_icn = np.array([1,2,3,5,3,4], dtype=np.intc)
        # ^No need to update these indices

        ma27.do_symbolic_factorization(n, irn, icn)

        self.assertEqual(ma27.get_info(1), 0)
        self.assertEqual(ma27.get_info(5), 14) # Min required num. integer words
        self.assertEqual(ma27.get_info(6), 20) # Min required num. real words

        with self.assertRaisesRegex(AssertionError, 'Dimension mismatch'):
            ma27.do_symbolic_factorization(n, irn, bad_icn)

    def test_do_numeric_factorization(self):
        ma27 = MA27Interface()

        n = 5
        ne = 7
        irn = np.array([1,1,2,2,3,3,5], dtype=np.intc)
        icn = np.array([1,2,3,5,3,4,5], dtype=np.intc)
        irn = irn - 1
        icn = icn - 1
        ent = np.array([2.,3.,4.,6.,1.,5.,1.], dtype=np.double)
        ma27.do_symbolic_factorization(n, irn, icn)

        status = ma27.do_numeric_factorization(irn, icn, n, ent)
        self.assertEqual(status, 0)

        expected_ent = [2.,3.,4.,6.,1.,5.,1.,]
        for i in range(ne):
            self.assertAlmostEqual(ent[i], expected_ent[i])
    
        self.assertEqual(ma27.get_info(15), 2) # 2 negative eigenvalues
        self.assertEqual(ma27.get_info(14), 1) # 1 2x2 pivot

        # Check that we can successfully perform another numeric factorization
        # with same symbolic factorization
        ent2 = np.array([1.5, 5.4, 1.2, 6.1, 4.2, 3.3, 2.0], dtype=np.double)
        status = ma27.do_numeric_factorization(irn, icn, n, ent2)
        self.assertEqual(ma27.get_info(15), 2)
        self.assertEqual(status, 0)

        with self.assertRaisesRegex(AssertionError, 'Dimension mismatch'):
            ma27.do_numeric_factorization(irn, icn, n+1, ent)

        # Check that we can successfully perform another symbolic and 
        # numeric factorization with the same ma27 struct
        #
        # n is still 5, ne has changed to 8.
        irn = np.array([1,1,2,2,3,3,5,1], dtype=np.intc)
        icn = np.array([1,2,3,5,3,4,5,5], dtype=np.intc)
        irn = irn - 1
        icn = icn - 1
        ent = np.array([2.,3.,4.,6.,1.,5.,1.,3.], dtype=np.double)
        status = ma27.do_symbolic_factorization(n, irn, icn)
        self.assertEqual(status, 0)
        status = ma27.do_numeric_factorization(irn, icn, n, ent)
        self.assertEqual(status, 0)
        self.assertEqual(ma27.get_info(15), 3)

    def test_do_backsolve(self):
        ma27 = MA27Interface()

        n = 5
        ne = 7
        irn = np.array([1,  1,  2,  2,  3,  3,  5], dtype=np.intc)
        icn = np.array([1,  2,  3,  5,  3,  4,  5], dtype=np.intc)
        ent = np.array([2., 3., 4., 6., 1., 5., 1.], dtype=np.double)
        irn = irn - 1
        icn = icn - 1
        rhs = np.array([8.,45.,31.,15.,17.], dtype=np.double)
        status = ma27.do_symbolic_factorization(n, irn, icn)
        status = ma27.do_numeric_factorization(irn, icn, n, ent)
        sol = ma27.do_backsolve(rhs)

        expected_sol = [1,2,3,4,5]
        old_rhs = np.array([8.,45.,31.,15.,17.])
        for i in range(n):
            self.assertAlmostEqual(sol[i], expected_sol[i])
            self.assertEqual(old_rhs[i], rhs[i])

        # Check that we can perform a numeric factorization with different ordering
        irn_mod = np.array([1,  2,  2,  1,  3,  3,  5], dtype=np.intc)
        icn_mod = np.array([2,  3,  5,  1,  3,  4,  5], dtype=np.intc)
        ent_mod = np.array([3., 4., 6., 2., 1., 5., 1.], dtype=np.double)
        irn_mod -= 1
        icn_mod -= 1
        status = ma27.do_numeric_factorization(irn_mod, icn_mod, n, ent_mod)
        sol = ma27.do_backsolve(rhs)
        self.assertTrue(np.allclose(sol, np.array(expected_sol)))

        # Check that we can perform a numeric factorization with differnt lengths
        # due to extra zero entries
        irn_mod = irn_mod[1:]
        icn_mod = icn_mod[1:]
        ent_mod = ent_mod[1:]
        status = ma27.do_numeric_factorization(irn_mod, icn_mod, n, ent_mod)
        sol = ma27.do_backsolve(rhs)
        expected_sol = np.array([4.0, 1.91666666667, 3.0, 4.06666666667, 5.5])
        self.assertTrue(np.allclose(sol, expected_sol))


if __name__ == '__main__':
    unittest.main()
