#  ___________________________________________________________________________
#
#  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 ctypes
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')
from pyomo.contrib.pynumero.linalg.ma57 import MA57Interface


@unittest.skipIf(not MA57Interface.available(), reason='MA57 not available')
class TestMA57Interface(unittest.TestCase):

    def test_get_cntl(self):
        ma57 = MA57Interface()
        self.assertEqual(ma57.get_icntl(1), 6)
        self.assertEqual(ma57.get_icntl(7), 1)

        self.assertAlmostEqual(ma57.get_cntl(1), 1e-2) # Numerical pivot threshold
        self.assertAlmostEqual(ma57.get_cntl(2), 1e-20) # Null pivot threshold

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

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

    def test_set_cntl(self):
        ma57 = MA57Interface()
        ma57.set_cntl(1, 1e-8)
        ma57.set_cntl(2, 1e-12)
        self.assertAlmostEqual(ma57.get_cntl(1), 1e-8)
        self.assertAlmostEqual(ma57.get_cntl(2), 1e-12)

    def test_do_symbolic_factorization(self):
        ma57 = MA57Interface()

        n = 5
        ne = 7
        irn = np.array([1,1,2,2,3,3,5], dtype=np.intc)
        jcn = np.array([1,2,3,5,3,4,5], dtype=np.intc)
        # Copied these Fortran-style indices from HSL docs. 
        # Interface expects C-style indices, as is typical in Python.
        irn = irn - 1
        jcn = jcn - 1

        bad_jcn = np.array([1,2,3,5,3,4], dtype=np.intc)

        ma57.do_symbolic_factorization(n, irn, jcn)

        self.assertEqual(ma57.get_info(1), 0)
        self.assertEqual(ma57.get_info(4), 0)
        self.assertEqual(ma57.get_info(9), 48) # Min required length of FACT
        self.assertEqual(ma57.get_info(10), 53) # Min required length of IFACT
        self.assertEqual(ma57.get_info(14), 0) # Should not yet be set

        with self.assertRaisesRegex(AssertionError, 'Dimension mismatch'):
            ma57.do_symbolic_factorization(n, irn, bad_jcn)

    def test_do_numeric_factorization(self):
        ma57 = MA57Interface()

        n = 5
        ne = 7
        irn = np.array([1,1,2,2,3,3,5], dtype=np.intc)
        jcn = np.array([1,2,3,5,3,4,5], dtype=np.intc)
        irn = irn - 1
        jcn = jcn - 1
        ent = np.array([2.,3.,4.,6.,1.,5.,1.], dtype=np.double)
        ma57.do_symbolic_factorization(n, irn, jcn)
        ma57.fact_factor = 1.5
        ma57.ifact_factor = 1.5
        # ^ No way to check whether these are handled properly... Would have to
        # access the struct to get LFACT, LIFACT

        status = ma57.do_numeric_factorization(n, ent)
        self.assertEqual(status, 0)
    
        self.assertEqual(ma57.get_info(14), 12) # 12 entries in factors
        self.assertEqual(ma57.get_info(24), 2) # 2 negative eigenvalues
        self.assertEqual(ma57.get_info(22), 1) # 1 2x2 pivot
        self.assertEqual(ma57.get_info(23), 0) # 0 delayed pivots

        ent2 = np.array([1.,5.,1.,6.,4.,3.,2.], dtype=np.double)
        ma57.do_numeric_factorization(n, ent2)
        self.assertEqual(status, 0)

        bad_ent = np.array([2.,3.,4.,6.,1.,5.], dtype=np.double)
        with self.assertRaisesRegex(AssertionError, 'Wrong number of entries'):
            ma57.do_numeric_factorization(n, bad_ent)
        with self.assertRaisesRegex(AssertionError, 'Dimension mismatch'):
            ma57.do_numeric_factorization(n+1, ent)

        n = 5
        ne = 8
        irn = np.array([1,1,2,2,3,3,5,5], dtype=np.intc)
        jcn = np.array([1,2,3,5,3,4,5,1], dtype=np.intc)
        irn = irn - 1
        jcn = jcn - 1
        ent = np.array([2.,3.,4.,6.,1.,5.,1.,-1.3], dtype=np.double)
        status = ma57.do_symbolic_factorization(n, irn, jcn)
        self.assertEqual(status, 0)
        status = ma57.do_numeric_factorization(n, ent)
        self.assertEqual(status, 0)
        self.assertEqual(ma57.get_info(24), 2)
        self.assertEqual(ma57.get_info(23), 0)


    def test_do_backsolve(self):
        ma57 = MA57Interface()

        n = 5
        ne = 7
        irn = np.array([1,1,2,2,3,3,5], dtype=np.intc)
        jcn = np.array([1,2,3,5,3,4,5], dtype=np.intc)
        irn = irn - 1
        jcn = jcn - 1
        ent = np.array([2.,3.,4.,6.,1.,5.,1.], dtype=np.double)
        rhs = np.array([8.,45.,31.,15.,17.], dtype=np.double)
        status = ma57.do_symbolic_factorization(n, irn, jcn)
        status = ma57.do_numeric_factorization(n, ent)
        sol = ma57.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])

        #rhs2 = np.array([[8., 17.],
        #                 [45., 15.],
        #                 [31., 31.],
        #                 [15., 45.],
        #                 [17., 8.]], dtype=np.double)
        #sol = ma57.do_backsolve(rhs2)
        # FIXME
        # This gives unexpected (incorrect) results. 
        # Need to investigate further.


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