# ____________________________________________________________________________________
#
# 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 platform
import pyomo.environ as pyo
import pyomo.common.unittest as unittest
from pyomo.common.fileutils import find_library
from pyomo.opt import check_available_solvers
from pyomo.core.base.external import nan

flib = find_library("asl_external_demo")
is_pypy = platform.python_implementation().lower().startswith('pypy')


@unittest.skipUnless(flib, 'Could not find the "asl_external_demo.so" library')
class TestAMPLExternalFunction(unittest.TestCase):
    @unittest.skipIf(is_pypy, 'Cannot evaluate external functions under pypy')
    def test_eval_function(self):
        m = pyo.ConcreteModel()
        m.tf = pyo.ExternalFunction(library=flib, function="demo_function")
        self.assertAlmostEqual(m.tf("sum", 1, 2, 3)(), 6, 4)
        self.assertAlmostEqual(m.tf("inv", 1, 2, 3)(), 1.8333333, 4)
        m.cbrt = pyo.ExternalFunction(library=flib, function="safe_cbrt")
        self.assertAlmostEqual(m.cbrt(6)(), 1.81712059, 4)
        self.assertStructuredAlmostEqual(
            m.cbrt.evaluate_fgh([0]), (0, [100951], [-1.121679e13]), reltol=1e-5
        )

    @unittest.skipIf(is_pypy, 'Cannot evaluate external functions under pypy')
    def test_eval_function_fgh(self):
        m = pyo.ConcreteModel()
        m.tf = pyo.ExternalFunction(library=flib, function="demo_function")

        f, g, h = m.tf.evaluate_fgh(("sum", 1, 2, 3))
        self.assertEqual(f, 6)
        self.assertEqual(g, [nan, 1, 1, 1])
        self.assertEqual(h, [nan, nan, 0, nan, 0, 0, nan, 0, 0, 0])

        f, g, h = m.tf.evaluate_fgh(("inv", 1, 2, 3))
        self.assertAlmostEqual(f, 1.8333333, 4)
        self.assertStructuredAlmostEqual(g, [nan, -1, -1 / 4, -1 / 9])
        self.assertStructuredAlmostEqual(
            h, [nan, nan, 2, nan, 0, 1 / 4, nan, 0, 0, 2 / 27]
        )

    @unittest.skipUnless(
        check_available_solvers('ipopt'), "The 'ipopt' solver is not available"
    )
    def test_solve_function(self):
        m = pyo.ConcreteModel()
        m.sum_it = pyo.ExternalFunction(library=flib, function="demo_function")
        m.cbrt = pyo.ExternalFunction(library=flib, function="safe_cbrt")
        m.x = pyo.Var(initialize=4.0)
        m.y = pyo.Var(initialize=0)
        m.y.fix()
        # Note: this also tests passing constant expressions to external
        # functions in the NL writer
        m.c = pyo.Constraint(expr=1.5 == m.sum_it("inv", 3, m.x, 1 / (m.y + 1)))
        m.o = pyo.Objective(expr=m.cbrt(m.x))
        solver = pyo.SolverFactory("ipopt")
        solver.solve(m, tee=True)
        self.assertAlmostEqual(m.x(), 6, 4)
