# Copyright 2015-2024 David Hadka
#
# This file is part of Rhodium, a Python module for robust decision
# making and exploratory modeling.
#
# Rhodium is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Rhodium is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Rhodium.  If not, see <http://www.gnu.org/licenses/>.
import unittest
from ..model import Parameter, Response
from ..optimization import evaluate
from ..openmdao import OpenMDAOModel

class TestOpenMDAOModel(unittest.TestCase):

    def test(self):
        try:
            from openmdao.api import IndepVarComp, Component, Problem, Group
        except ImportError as e:
            self.skipTest(f"OpenMDAO not available: {e}")

        # the following example is taken from OpenMDAO's Paraboloid tutorial
        class Paraboloid(Component):
            """ Evaluates the equation f(x,y) = (x-3)^2 + xy + (y+4)^2 - 3 """

            def __init__(self):
                super(Paraboloid, self).__init__()

                self.add_param('x', val=0.0)
                self.add_param('y', val=0.0)

                self.add_output('f_xy', val=0.0)

            def solve_nonlinear(self, params, unknowns, resids):
                """f(x,y) = (x-3)^2 + xy + (y+4)^2 - 3
                Optimal solution (minimum): x = 6.6667; y = -7.3333
                """

                x = params['x']
                y = params['y']

                unknowns['f_xy'] = (x-3.0)**2 + x*y + (y+4.0)**2 - 3.0

            def linearize(self, params, unknowns, resids):
                """ Jacobian for our paraboloid."""

                x = params['x']
                y = params['y']
                J = {}

                J['f_xy', 'x'] = 2.0*x - 6.0 + y
                J['f_xy', 'y'] = 2.0*y + 8.0 + x
                return J

        top = Problem()
        root = top.root = Group()

        root.add('p1', IndepVarComp('x', 3.0))
        root.add('p2', IndepVarComp('y', -4.0))
        root.add('p', Paraboloid())

        root.connect('p1.x', 'p.x')
        root.connect('p2.y', 'p.y')

        top.setup()

        # here is where Rhodium begins
        model = OpenMDAOModel(top)
        model.parameters = [Parameter("x", connect="p1.x"),
                            Parameter("y", connect="p2.y")]
        model.responses = [Response("f_xy", connect="p.f_xy")]

        output = evaluate(model, {"x": 3.0, "y": -4.0})
        self.assertEqual(-15.0, output["f_xy"])
