""" Unit tests for the Driver base class."""


from distutils.version import LooseVersion
from io import StringIO
import sys
import unittest

import numpy as np

import openmdao.api as om
from openmdao.utils.assert_utils import assert_near_equal, assert_warning, assert_check_partials
from openmdao.utils.general_utils import printoptions
from openmdao.utils.testing_utils import use_tempdirs
from openmdao.test_suite.components.sellar import SellarDerivatives
from openmdao.test_suite.components.simple_comps import DoubleArrayComp, NonSquareArrayComp


@use_tempdirs
class TestDriver(unittest.TestCase):

    def test_basic_get(self):

        prob = om.Problem()
        prob.model = model = SellarDerivatives()

        model.add_design_var('z')
        model.add_objective('obj')
        model.add_constraint('con1', lower=0)
        prob.set_solver_print(level=0)

        prob.setup()
        prob.run_driver()

        designvars = prob.driver.get_design_var_values()
        self.assertEqual(designvars['pz.z'][0], 5.0 )

        designvars = prob.driver.get_objective_values()
        self.assertEqual(designvars['obj_cmp.obj'], prob['obj'] )

        designvars = prob.driver.get_constraint_values()
        self.assertEqual(designvars['con_cmp1.con1'], prob['con1']  )

    def test_scaled_design_vars(self):

        prob = om.Problem()
        prob.model = model = SellarDerivatives()

        model.add_design_var('z', ref=5.0, ref0=3.0)
        model.add_objective('obj')
        model.add_constraint('con1', lower=0)
        prob.set_solver_print(level=0)

        prob.setup()

        # Conclude setup but don't run model.
        prob.final_setup()

        dv = prob.driver.get_design_var_values()
        self.assertEqual(dv['pz.z'][0], 1.0)
        self.assertEqual(dv['pz.z'][1], -0.5)

        prob.driver.set_design_var('pz.z', np.array((2.0, -2.0)))
        self.assertEqual(prob['z'][0], 7.0)
        self.assertEqual(prob['z'][1], -1.0)

    def test_scaled_constraints(self):

        prob = om.Problem()
        prob.model = model = SellarDerivatives()

        model.add_design_var('z')
        model.add_objective('obj')
        model.add_constraint('con1', lower=0, ref=2.0, ref0=3.0)
        prob.set_solver_print(level=0)

        prob.setup()
        prob.run_model()

        cv = prob.driver.get_constraint_values()['con_cmp1.con1'][0]
        base = prob['con1']
        self.assertEqual((base-3.0)/(2.0-3.0), cv)

    def test_scaled_objectves(self):

        prob = om.Problem()
        prob.model = model = SellarDerivatives()

        model.add_design_var('z')
        model.add_objective('obj', ref=2.0, ref0=3.0)
        model.add_constraint('con1', lower=0)
        prob.set_solver_print(level=0)

        prob.setup()
        prob.run_model()

        cv = prob.driver.get_objective_values()['obj_cmp.obj'][0]
        base = prob['obj']
        self.assertEqual((base-3.0)/(2.0-3.0), cv)

    def test_scaled_derivs(self):

        prob = om.Problem()
        prob.model = model = SellarDerivatives()

        model.add_design_var('z')
        model.add_objective('obj')
        model.add_constraint('con1')
        prob.set_solver_print(level=0)

        prob.setup()
        prob.run_model()

        base = prob.compute_totals(of=['obj', 'con1'], wrt=['z'])

        prob = om.Problem()
        prob.model = model = SellarDerivatives()

        model.add_design_var('z', ref=2.0, ref0=0.0)
        model.add_objective('obj', ref=1.0, ref0=0.0)
        model.add_constraint('con1', lower=0, ref=2.0, ref0=0.0)
        prob.set_solver_print(level=0)

        prob.setup()
        prob.run_model()

        derivs = prob.driver._compute_totals(of=['obj_cmp.obj', 'con_cmp1.con1'], wrt=['pz.z'],
                                             return_format='dict')
        assert_near_equal(base[('con1', 'z')][0], derivs['con_cmp1.con1']['pz.z'][0], 1e-5)
        assert_near_equal(base[('obj', 'z')][0]*2.0, derivs['obj_cmp.obj']['pz.z'][0], 1e-5)

    def test_vector_scaled_derivs(self):

        prob = om.Problem()
        model = prob.model

        model.add_subsystem('px', om.IndepVarComp(name="x", val=np.ones((2, ))))
        comp = model.add_subsystem('comp', DoubleArrayComp())
        model.connect('px.x', 'comp.x1')

        model.add_design_var('px.x', ref=np.array([2.0, 3.0]), ref0=np.array([0.5, 1.5]))
        model.add_objective('comp.y1', ref=np.array([[7.0, 11.0]]), ref0=np.array([5.2, 6.3]))
        model.add_constraint('comp.y2', lower=0.0, upper=1.0, ref=np.array([[2.0, 4.0]]), ref0=np.array([1.2, 2.3]))

        prob.setup()
        prob.run_driver()

        with assert_warning(DeprecationWarning, "'global_names' is deprecated in calls to _compute_totals. Use 'use_abs_names' instead."):
            derivs = prob.driver._compute_totals(of=['comp.y1'], wrt=['px.x'], global_names=True,
                                                 return_format='dict')

        oscale = np.array([1.0/(7.0-5.2), 1.0/(11.0-6.3)])
        iscale = np.array([2.0-0.5, 3.0-1.5])
        J = comp.JJ[0:2, 0:2]

        # doing this manually so that I don't inadvertantly make an error in the vector math in both the code and test.
        J[0, 0] *= oscale[0]*iscale[0]
        J[0, 1] *= oscale[0]*iscale[1]
        J[1, 0] *= oscale[1]*iscale[0]
        J[1, 1] *= oscale[1]*iscale[1]
        assert_near_equal(J, derivs['comp.y1']['px.x'], 1.0e-3)

        obj = prob.driver.get_objective_values()
        obj_base = np.array([ (prob['comp.y1'][0]-5.2)/(7.0-5.2), (prob['comp.y1'][1]-6.3)/(11.0-6.3) ])
        assert_near_equal(obj['comp.y1'], obj_base, 1.0e-3)

        con = prob.driver.get_constraint_values()
        con_base = np.array([ (prob['comp.y2'][0]-1.2)/(2.0-1.2), (prob['comp.y2'][1]-2.3)/(4.0-2.3) ])
        assert_near_equal(con['comp.y2'], con_base, 1.0e-3)

    def test_vector_bounds_inf(self):

        # make sure no overflow when there is no specified upper/lower bound and significatn scaling
        prob = om.Problem()
        model = prob.model

        model.add_subsystem('px', om.IndepVarComp(name="x", val=np.ones((2, ))))
        comp = model.add_subsystem('comp', DoubleArrayComp())
        model.connect('px.x', 'comp.x1')

        model.add_design_var('px.x', ref=np.array([.1, 1e-6]))
        model.add_constraint('comp.y2', ref=np.array([.2, 2e-6]))

        prob.setup()

        desvars = model.get_design_vars()

        self.assertFalse(np.any(np.isinf(desvars['px.x']['upper'])))
        self.assertFalse(np.any(np.isinf(-desvars['px.x']['lower'])))

        responses = prob.model.get_responses()

        self.assertFalse(np.any(np.isinf(responses['comp.y2']['upper'])))
        self.assertFalse(np.any(np.isinf(-responses['comp.y2']['lower'])))

    def test_vector_scaled_derivs_diff_sizes(self):

        prob = om.Problem()
        model = prob.model

        model.add_subsystem('px', om.IndepVarComp(name="x", val=np.ones((2, ))))
        comp = model.add_subsystem('comp', NonSquareArrayComp())
        model.connect('px.x', 'comp.x1')

        model.add_design_var('px.x', ref=np.array([2.0, 3.0]), ref0=np.array([0.5, 1.5]))
        model.add_objective('comp.y1', ref=np.array([[7.0, 11.0, 2.0]]), ref0=np.array([5.2, 6.3, 1.2]))
        model.add_constraint('comp.y2', lower=0.0, upper=1.0, ref=np.array([[2.0]]), ref0=np.array([1.2]))

        prob.setup()
        prob.run_driver()

        derivs = prob.driver._compute_totals(of=['comp.y1'], wrt=['px.x'],
                                                   return_format='dict')

        oscale = np.array([1.0/(7.0-5.2), 1.0/(11.0-6.3), 1.0/(2.0-1.2)])
        iscale = np.array([2.0-0.5, 3.0-1.5])
        J = comp.JJ[0:3, 0:2]

        # doing this manually so that I don't inadvertantly make an error in the vector math in both the code and test.
        J[0, 0] *= oscale[0]*iscale[0]
        J[0, 1] *= oscale[0]*iscale[1]
        J[1, 0] *= oscale[1]*iscale[0]
        J[1, 1] *= oscale[1]*iscale[1]
        J[2, 0] *= oscale[2]*iscale[0]
        J[2, 1] *= oscale[2]*iscale[1]
        assert_near_equal(J, derivs['comp.y1']['px.x'], 1.0e-3)

        obj = prob.driver.get_objective_values()
        obj_base = np.array([ (prob['comp.y1'][0]-5.2)/(7.0-5.2), (prob['comp.y1'][1]-6.3)/(11.0-6.3), (prob['comp.y1'][2]-1.2)/(2.0-1.2) ])
        assert_near_equal(obj['comp.y1'], obj_base, 1.0e-3)

        con = prob.driver.get_constraint_values()
        con_base = np.array([ (prob['comp.y2'][0]-1.2)/(2.0-1.2)])
        assert_near_equal(con['comp.y2'], con_base, 1.0e-3)

    def test_debug_print_option(self):

        prob = om.Problem()
        prob.model = model = SellarDerivatives()

        model.add_design_var('z')
        model.add_objective('obj')
        model.add_constraint('con1', lower=0)
        model.add_constraint('con2', lower=0, linear=True)
        prob.set_solver_print(level=0)

        prob.setup()

        # Make sure nothing prints if debug_print is the default of empty list
        stdout = sys.stdout
        strout = StringIO()
        sys.stdout = strout
        try:
            prob.run_driver()
        finally:
            sys.stdout = stdout
        output = strout.getvalue().split('\n')
        self.assertEqual(output, [''])

        # Make sure everything prints when all options are on
        prob.driver.options['debug_print'] = ['desvars','ln_cons','nl_cons','objs']
        stdout = sys.stdout
        strout = StringIO()
        sys.stdout = strout
        try:
            prob.run_driver(reset_iter_counts=False)
        finally:
            sys.stdout = stdout
        output = strout.getvalue().split('\n')
        self.assertEqual(output.count("Driver debug print for iter coord: rank0:Driver|1"), 1)
        self.assertEqual(output.count("Design Vars"), 1)
        self.assertEqual(output.count("Nonlinear constraints"), 1)
        self.assertEqual(output.count("Linear constraints"), 1)
        self.assertEqual(output.count("Objectives"), 1)

        # Check to make sure an invalid debug_print option raises an exception
        with self.assertRaises(ValueError) as context:
            prob.driver.options['debug_print'] = ['bad_option']
        self.assertEqual(str(context.exception),
                         "Option 'debug_print' contains value 'bad_option' which is not one of ['desvars', 'nl_cons', 'ln_cons', 'objs', 'totals'].")

    def test_debug_print_desvar_physical_with_indices(self):
        prob = om.Problem()
        model = prob.model

        size = 3
        model.add_subsystem('p1', om.IndepVarComp('x', np.array([50.0] * size)))
        model.add_subsystem('p2', om.IndepVarComp('y', np.array([50.0] * size)))
        model.add_subsystem('comp', om.ExecComp('f_xy = (x-3.0)**2 + x*y + (y+4.0)**2 - 3.0',
                                                x=np.zeros(size), y=np.zeros(size),
                                                f_xy=np.zeros(size)))
        model.add_subsystem('con', om.ExecComp('c = - x + y',
                                               c=np.zeros(size), x=np.zeros(size),
                                               y=np.zeros(size)))

        model.connect('p1.x', 'comp.x')
        model.connect('p2.y', 'comp.y')
        model.connect('p1.x', 'con.x')
        model.connect('p2.y', 'con.y')

        prob.set_solver_print(level=0)

        prob.driver = om.ScipyOptimizeDriver()
        prob.driver.options['optimizer'] = 'SLSQP'
        prob.driver.options['tol'] = 1e-9
        prob.driver.options['disp'] = False

        model.add_design_var('p1.x', indices=[1], lower=-50.0, upper=50.0, ref=[5.0,])
        model.add_design_var('p2.y', indices=[1], lower=-50.0, upper=50.0)
        model.add_objective('comp.f_xy', index=1)
        model.add_constraint('con.c', indices=[1], upper=-15.0)

        prob.setup()

        prob.driver.options['debug_print'] = ['desvars',]
        stdout = sys.stdout
        strout = StringIO()
        sys.stdout = strout

        try:
            # formatting has changed in numpy 1.14 and beyond.
            if LooseVersion(np.__version__) >= LooseVersion("1.14"):
                with printoptions(precision=2, legacy="1.13"):
                    prob.run_driver()
            else:
                with printoptions(precision=2):
                    prob.run_driver()
        finally:
            sys.stdout = stdout
        output = strout.getvalue().split('\n')
        # should see unscaled (physical) and the full arrays, not just what is indicated by indices
        self.assertEqual(output[3], "{'p1.x': array([ 50.,  50.,  50.]), 'p2.y': array([ 50.,  50.,  50.])}")

    def test_debug_print_response_physical(self):
        prob = om.Problem()
        model = prob.model

        size = 3
        model.add_subsystem('p1', om.IndepVarComp('x', np.array([50.0] * size)))
        model.add_subsystem('p2', om.IndepVarComp('y', np.array([50.0] * size)))
        model.add_subsystem('comp', om.ExecComp('f_xy = (x-3.0)**2 + x*y + (y+4.0)**2 - 3.0',
                                                x=np.zeros(size), y=np.zeros(size),
                                                f_xy=np.zeros(size)))
        model.add_subsystem('con', om.ExecComp('c = - x + y + 1',
                                               c=np.zeros(size), x=np.zeros(size),
                                               y=np.zeros(size)))

        model.connect('p1.x', 'comp.x')
        model.connect('p2.y', 'comp.y')
        model.connect('p1.x', 'con.x')
        model.connect('p2.y', 'con.y')

        prob.set_solver_print(level=0)

        prob.driver = om.ScipyOptimizeDriver()
        prob.driver.options['optimizer'] = 'SLSQP'
        prob.driver.options['tol'] = 1e-9
        prob.driver.options['disp'] = False

        model.add_design_var('p1.x', indices=[1], lower=-50.0, upper=50.0)
        model.add_design_var('p2.y', indices=[1], lower=-50.0, upper=50.0)
        model.add_objective('comp.f_xy', index=1, ref=1.5)
        model.add_constraint('con.c', indices=[1], upper=-15.0, ref=1.02)

        prob.setup()

        prob.driver.options['debug_print'] = ['objs', 'nl_cons']
        stdout = sys.stdout
        strout = StringIO()
        sys.stdout = strout

        try:
            # formatting has changed in numpy 1.14 and beyond.
            if LooseVersion(np.__version__) >= LooseVersion("1.14"):
                with printoptions(precision=2, legacy="1.13"):
                    prob.run_driver()
            else:
                with printoptions(precision=2):
                    prob.run_driver()
        finally:
            sys.stdout = stdout
        output = strout.getvalue().split('\n')
        # should see unscaled (physical) and the full arrays, not just what is indicated by indices
        self.assertEqual(output[3], "{'con.c': array([ 1.])}")
        self.assertEqual(output[6], "{'comp.f_xy': array([ 7622.])}")

    def test_debug_desvar_shape(self):
        # Desvar should always be printed un-flattened.

        prob = om.Problem()
        model = prob.model

        model.add_subsystem('p', om.IndepVarComp('x', val=np.array([[1.0, 3, 4], [7, 2, 5]])))

        model.add_design_var('p.x', np.array([[1.0, 3, 4], [7, 2, 5]]))
        prob.driver.options['debug_print'] = ['desvars']

        prob.setup()

        stdout = sys.stdout
        strout = StringIO()
        sys.stdout = strout

        try:
            # formatting has changed in numpy 1.14 and beyond.
            if LooseVersion(np.__version__) >= LooseVersion("1.14"):
                with printoptions(precision=2, legacy="1.13"):
                    prob.run_driver()
            else:
                with printoptions(precision=2):
                    prob.run_driver()
        finally:
            sys.stdout = stdout

        output = strout.getvalue().split('\n')

        self.assertEqual(output[3], "{'p.x': array([[ 1.,  3.,  4.],")
        self.assertEqual(output[4], '       [ 7.,  2.,  5.]])}')

    def test_unsupported_discrete_desvar(self):
        prob = om.Problem()

        indep = om.IndepVarComp()
        indep.add_discrete_output('xI', val=0)
        prob.model.add_subsystem('p', indep)

        prob.model.add_design_var('p.xI')

        prob.driver = om.ScipyOptimizeDriver()

        prob.setup()

        with self.assertRaises(RuntimeError) as context:
            prob.final_setup()

        msg = "Discrete design variables are not supported by this driver: p.xI"
        self.assertEqual(str(context.exception), msg)

    def test_driver_recording_options_deprecated(self):
        prob = om.Problem()
        msg = "The recording option, record_model_metadata, on Driver is deprecated. " \
              "Recording of model metadata will always be done"
        with assert_warning(DeprecationWarning, msg):
            prob.driver.recording_options['record_model_metadata'] = True

    def test_units_basic(self):
        prob = om.Problem()
        model = prob.model

        ivc = om.IndepVarComp()
        ivc.add_output('x', 35.0, units='degF', lower=32.0, upper=212.0)

        model.add_subsystem('p', ivc, promotes=['x'])
        model.add_subsystem('comp1', om.ExecComp('y1 = 2.0*x',
                                                 x={'value': 2.0, 'units': 'degF'},
                                                 y1={'value': 2.0, 'units': 'degF'}),
                            promotes=['x', 'y1'])

        model.add_subsystem('comp2', om.ExecComp('y2 = 3.0*x',
                                                 x={'value': 2.0, 'units': 'degF'},
                                                 y2={'value': 2.0, 'units': 'degF'}),
                            promotes=['x', 'y2'])

        model.add_design_var('x', units='degC', lower=0.0, upper=100.0)
        model.add_constraint('y1', units='degC', lower=0.0, upper=100.0)
        model.add_objective('y2', units='degC')

        prob.setup()

        prob.run_driver()

        dv = prob.driver.get_design_var_values()
        assert_near_equal(dv['p.x'][0], 3.0 * 5 / 9, 1e-8)

        obj = prob.driver.get_objective_values(driver_scaling=True)
        assert_near_equal(obj['comp2.y2'][0], 73.0 * 5 / 9, 1e-8)

        con = prob.driver.get_constraint_values(driver_scaling=True)
        assert_near_equal(con['comp1.y1'][0], 38.0 * 5 / 9, 1e-8)

        meta = model.get_design_vars()
        assert_near_equal(meta['p.x']['lower'], 0.0, 1e-7)
        assert_near_equal(meta['p.x']['upper'], 100.0, 1e-7)

        meta = model.get_constraints()
        assert_near_equal(meta['comp1.y1']['lower'], 0.0, 1e-7)
        assert_near_equal(meta['comp1.y1']['upper'], 100.0, 1e-7)

        stdout = sys.stdout
        strout = StringIO()
        sys.stdout = strout
        try:
            prob.list_problem_vars(desvar_opts=['units'], objs_opts=['units'], cons_opts=['units'])
        finally:
            sys.stdout = stdout
        output = strout.getvalue().split('\n')

        self.assertTrue('1.666' in output[5])
        self.assertTrue('21.111' in output[12])
        self.assertTrue('40.555' in output[19])
        self.assertTrue('degC' in output[5])
        self.assertTrue('degC' in output[12])
        self.assertTrue('degC' in output[19])

    def test_units_equal(self):
        prob = om.Problem()
        model = prob.model

        ivc = om.IndepVarComp()
        ivc.add_output('x', 35.0, units='degF', lower=32.0, upper=212.0)

        model.add_subsystem('p', ivc, promotes=['x'])
        model.add_subsystem('comp1', om.ExecComp('y1 = 2.0*x',
                                                 x={'value': 2.0, 'units': 'degF'},
                                                 y1={'value': 2.0, 'units': 'degF'}),
                            promotes=['x', 'y1'])

        model.add_subsystem('comp2', om.ExecComp('y2 = 3.0*x',
                                                 x={'value': 2.0, 'units': 'degF'},
                                                 y2={'value': 2.0, 'units': 'degF'}),
                            promotes=['x', 'y2'])

        model.add_design_var('x', units='degF', lower=32.0, upper=212.0)
        model.add_constraint('y1', units='degF', lower=32.0, upper=212.0)
        model.add_objective('y2', units='degF')

        prob.setup()

        prob.run_driver()

        dv = prob.driver.get_design_var_values()
        assert_near_equal(dv['p.x'][0], 35.0, 1e-8)

        obj = prob.driver.get_objective_values(driver_scaling=True)
        assert_near_equal(obj['comp2.y2'][0], 105.0, 1e-8)

        con = prob.driver.get_constraint_values(driver_scaling=True)
        assert_near_equal(con['comp1.y1'][0], 70.0, 1e-8)

        meta = model.get_design_vars()
        self.assertEqual(meta['p.x']['scaler'], None)
        self.assertEqual(meta['p.x']['adder'], None)

        meta = model.get_constraints()
        self.assertEqual(meta['comp1.y1']['scaler'], None)
        self.assertEqual(meta['comp1.y1']['adder'], None)

    def test_units_with_scaling(self):
        prob = om.Problem()
        model = prob.model

        ivc = om.IndepVarComp()
        ivc.add_output('x', 35.0, units='degF', lower=32.0, upper=212.0)

        model.add_subsystem('p', ivc, promotes=['x'])
        model.add_subsystem('comp1', om.ExecComp('y1 = 2.0*x',
                                                 x={'value': 2.0, 'units': 'degF'},
                                                 y1={'value': 2.0, 'units': 'degF'}),
                            promotes=['x', 'y1'])

        model.add_subsystem('comp2', om.ExecComp('y2 = 3.0*x',
                                                 x={'value': 2.0, 'units': 'degF'},
                                                 y2={'value': 2.0, 'units': 'degF'}),
                            promotes=['x', 'y2'])

        model.add_design_var('x', units='degC', lower=0.0, upper=100.0, scaler=3.5, adder=77.0)
        model.add_constraint('y1', units='degC', lower=0.0, upper=100.0, scaler=3.5, adder=77.0)
        model.add_objective('y2', units='degC', scaler=3.5, adder=77.0)

        recorder = om.SqliteRecorder('cases.sql')
        prob.driver.add_recorder(recorder)

        prob.driver.recording_options['record_objectives'] = True
        prob.driver.recording_options['record_constraints'] = True
        prob.driver.recording_options['record_desvars'] = True

        prob.setup()

        prob.run_driver()

        dv = prob.driver.get_design_var_values()
        assert_near_equal(dv['p.x'][0], ((3.0 * 5 / 9) + 77.0) * 3.5, 1e-8)

        obj = prob.driver.get_objective_values(driver_scaling=True)
        assert_near_equal(obj['comp2.y2'][0], ((73.0 * 5 / 9) + 77.0) * 3.5, 1e-8)

        con = prob.driver.get_constraint_values(driver_scaling=True)
        assert_near_equal(con['comp1.y1'][0], ((38.0 * 5 / 9) + 77.0) * 3.5, 1e-8)

        meta = model.get_design_vars()
        assert_near_equal(meta['p.x']['lower'], ((0.0) + 77.0) * 3.5, 1e-7)
        assert_near_equal(meta['p.x']['upper'], ((100.0) + 77.0) * 3.5, 1e-7)

        meta = model.get_constraints()
        assert_near_equal(meta['comp1.y1']['lower'], ((0.0) + 77.0) * 3.5, 1e-7)
        assert_near_equal(meta['comp1.y1']['upper'], ((100.0) + 77.0) * 3.5, 1e-7)

        stdout = sys.stdout
        strout = StringIO()
        sys.stdout = strout
        try:
            prob.list_problem_vars(desvar_opts=['units'], objs_opts=['units'], cons_opts=['units'])
        finally:
            sys.stdout = stdout
        output = strout.getvalue().split('\n')

        self.assertTrue('275.33' in output[5])
        self.assertTrue('343.3888' in output[12])
        self.assertTrue('411.444' in output[19])
        self.assertTrue('degC' in output[5])
        self.assertTrue('degC' in output[12])
        self.assertTrue('degC' in output[19])

        totals = prob.check_totals(out_stream=None, driver_scaling=True)

        for key, val in totals.items():
            assert_near_equal(val['rel error'][0], 0.0, 1e-6)

        cr = om.CaseReader("cases.sql")
        cases = cr.list_cases('driver')
        case = cr.get_case(cases[0])

        dv = case.get_design_vars()
        assert_near_equal(dv['p.x'][0], ((3.0 * 5 / 9) + 77.0) * 3.5, 1e-8)

        obj = case.get_objectives()
        assert_near_equal(obj['comp2.y2'][0], ((73.0 * 5 / 9) + 77.0) * 3.5, 1e-8)

        con = case.get_constraints()
        assert_near_equal(con['comp1.y1'][0], ((38.0 * 5 / 9) + 77.0) * 3.5, 1e-8)

    def test_units_error_messages(self):
        prob = om.Problem()
        model = prob.model

        ivc = om.IndepVarComp()
        ivc.add_output('x', 35.0, units='degF', lower=32.0, upper=212.0)

        model.add_subsystem('p', ivc, promotes=['x'])
        model.add_subsystem('comp1', om.ExecComp('y1 = 2.0*x',
                                                 x={'value': 2.0, 'units': 'degF'},
                                                 y1={'value': 2.0, 'units': 'degF'}),
                            promotes=['x', 'y1'])

        model.add_design_var('x', units='ft', lower=0.0, upper=100.0, scaler=3.5, adder=77.0)
        prob.setup()

        with self.assertRaises(RuntimeError) as context:
            prob.final_setup()

        msg = "Group (<model>): Target for design variable x has 'degF' units, but 'ft' units were specified."
        self.assertEqual(str(context.exception), msg)

        prob = om.Problem()
        model = prob.model

        ivc = om.IndepVarComp()
        ivc.add_output('x', 35.0, units='degF', lower=32.0, upper=212.0)

        model.add_subsystem('p', ivc, promotes=['x'])
        model.add_subsystem('comp1', om.ExecComp('y1 = 2.0*x',
                                                 x={'value': 2.0, 'units': 'degF'},
                                                 y1={'value': 2.0, 'units': 'degF'}),
                            promotes=['x', 'y1'])

        model.add_constraint('x', units='ft', lower=0.0, upper=100.0)
        prob.setup()

        with self.assertRaises(RuntimeError) as context:
            prob.final_setup()

        msg = "Group (<model>): Target for constraint x has 'degF' units, but 'ft' units were specified."
        self.assertEqual(str(context.exception), msg)

        prob = om.Problem()
        model = prob.model

        ivc = om.IndepVarComp()
        ivc.add_output('x', 35.0, units=None, lower=32.0, upper=212.0)

        model.add_subsystem('p', ivc, promotes=['x'])
        model.add_subsystem('comp1', om.ExecComp('y1 = 2.0*x',
                                                 x={'value': 2.0},
                                                 y1={'value': 2.0, 'units': 'degF'}),
                            promotes=['x', 'y1'])

        model.add_design_var('x', units='ft', lower=0.0, upper=100.0, scaler=3.5, adder=77.0)
        prob.setup()

        with self.assertRaises(RuntimeError) as context:
            prob.final_setup()

        msg = "Group (<model>): Target for design variable x has no units, but 'ft' units were specified."
        self.assertEqual(str(context.exception), msg)

        prob = om.Problem()
        model = prob.model

        ivc = om.IndepVarComp()
        ivc.add_output('x', 35.0, units=None, lower=32.0, upper=212.0)

        model.add_subsystem('p', ivc, promotes=['x'])
        model.add_subsystem('comp1', om.ExecComp('y1 = 2.0*x',
                                                 x={'value': 2.0},
                                                 y1={'value': 2.0, 'units': 'degF'}),
                            promotes=['x', 'y1'])

        model.add_constraint('x', units='ft', lower=0.0, upper=100.0)
        prob.setup()

        with self.assertRaises(RuntimeError) as context:
            prob.final_setup()

        msg = "Group (<model>): Target for constraint x has no units, but 'ft' units were specified."
        self.assertEqual(str(context.exception), msg)


class TestDriverFeature(unittest.TestCase):

    def test_specify_units(self):
        import openmdao.api as om

        prob = om.Problem()
        model = prob.model

        ivc = om.IndepVarComp()
        ivc.add_output('x', 35.0, units='degF', lower=32.0, upper=212.0)

        model.add_subsystem('p', ivc, promotes=['x'])
        model.add_subsystem('comp1', om.ExecComp('y1 = 2.0*x',
                                                 x={'value': 2.0, 'units': 'degF'},
                                                 y1={'value': 2.0, 'units': 'degF'}),
                            promotes=['x', 'y1'])

        model.add_subsystem('comp2', om.ExecComp('y2 = 3.0*x',
                                                 x={'value': 2.0, 'units': 'degF'},
                                                 y2={'value': 2.0, 'units': 'degF'}),
                            promotes=['x', 'y2'])

        model.add_design_var('x', units='degC', lower=0.0, upper=100.0)
        model.add_constraint('y1', units='degC', lower=0.0, upper=100.0)
        model.add_objective('y2', units='degC')

        prob.setup()
        prob.run_driver()

        print('Model variables')
        assert_near_equal(prob['p.x'][0], 35.0, 1e-8)
        assert_near_equal(prob['comp2.y2'][0], 105.0, 1e-8)
        assert_near_equal(prob['comp1.y1'][0], 70.0, 1e-8)

        print('')
        print('Driver variables')
        dv = prob.driver.get_design_var_values()
        assert_near_equal(dv['p.x'][0], 3.0 * 5 / 9, 1e-8)

        obj = prob.driver.get_objective_values(driver_scaling=True)
        assert_near_equal(obj['comp2.y2'][0], 73.0 * 5 / 9, 1e-8)

        con = prob.driver.get_constraint_values(driver_scaling=True)
        assert_near_equal(con['comp1.y1'][0], 38.0 * 5 / 9, 1e-8)


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