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

from importlib.util import find_spec
from packaging.version import Version
from io import StringIO
import sys
import unittest

import numpy as np

import openmdao.api as om
from openmdao.core.driver import Driver
from openmdao.utils.units import convert_units
from openmdao.utils.assert_utils import assert_near_equal, assert_warnings, assert_check_totals, assert_no_warning
from openmdao.utils.general_utils import printoptions, set_pyoptsparse_opt
from openmdao.utils.testing_utils import use_tempdirs
from openmdao.test_suite.components.paraboloid import Paraboloid
from openmdao.test_suite.components.sellar import SellarDerivatives
from openmdao.test_suite.components.simple_comps import DoubleArrayComp, NonSquareArrayComp
from openmdao.utils.om_warnings import OpenMDAOWarning
from openmdao.utils.mpi import MPI
from openmdao.drivers.pyoptsparse_driver import pyOptSparseDriver

try:
    from openmdao.vectors.petsc_vector import PETScVector
except ImportError:
    PETScVector = None


has_pyoptsparse = find_spec('pyoptsparse') is not None


@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['z'][0], 5.0 )

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

        designvars = prob.driver.get_constraint_values()
        self.assertEqual(designvars['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['z'][0], 1.0)
        self.assertEqual(dv['z'][1], -0.5)

        prob.driver.set_design_var('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()['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'][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', 'con1'], wrt=['z'],
                                             return_format='dict')
        assert_near_equal(base[('con1', 'z')][0], derivs['con1']['z'][0], 1e-5)
        assert_near_equal(base[('obj', 'z')][0]*2.0, derivs['obj']['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()

        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)])
        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, ))))
        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()
        prob.final_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),
                         "Driver: Value (['bad_option']) of option 'debug_print' is not one of "
                         "['desvars', 'nl_cons', 'ln_cons', 'objs', 'totals'].")

    def test_debug_print_approx(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.driver = om.ScipyOptimizeDriver()
        prob.driver.options['optimizer'] = 'SLSQP'
        prob.driver.options['tol'] = 1e-9
        prob.driver.options['disp'] = False
        prob.driver.options['maxiter'] = 0

        prob.model.approx_totals()
        prob.driver.options['debug_print'] = ['totals']

        prob.setup()

        stdout = sys.stdout
        strout = StringIO()
        sys.stdout = strout
        try:
            prob.run_driver()
        finally:
            sys.stdout = stdout
        output = strout.getvalue().split('\n')
        self.assertIn("{('obj', 'z'):", output[12])
        self.assertIn("{('con1', 'z'):", output[13])

    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 Version(np.__version__) >= Version("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 Version(np.__version__) >= Version("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 Version(np.__version__) >= Version("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_units_basic(self):
        prob = om.Problem()
        model = prob.model

        ivc = om.IndepVarComp()
        ivc.add_output('x', 35.0, units='degF')

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

        model.add_subsystem('comp2', om.ExecComp('y2 = 3.0*x',
                                                 x={'val': 2.0, 'units': 'degF'},
                                                 y2={'val': 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['x'][0], 3.0 * 5 / 9, 1e-8)

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

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

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

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

        stdout = sys.stdout
        strout = StringIO()
        sys.stdout = strout
        try:
            prob.list_driver_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')

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

        model.add_subsystem('comp2', om.ExecComp('y2 = 3.0*x',
                                                 x={'val': 2.0, 'units': 'degF'},
                                                 y2={'val': 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['x'][0], 35.0, 1e-8)

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

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

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

        meta = model.get_constraints()
        self.assertEqual(meta['y1']['scaler'], None)
        self.assertEqual(meta['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')

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

        model.add_subsystem('comp2', om.ExecComp('y2 = 3.0*x',
                                                 x={'val': 2.0, 'units': 'degF'},
                                                 y2={'val': 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['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['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['y1'][0], ((38.0 * 5 / 9) + 77.0) * 3.5, 1e-8)

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

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

        stdout = sys.stdout
        strout = StringIO()
        sys.stdout = strout
        try:
            prob.list_driver_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])

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

        cr = om.CaseReader(prob.get_outputs_dir() / "cases.sql")
        cases = cr.list_cases('driver', out_stream=None)
        case = cr.get_case(cases[0])

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

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

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

    def test_units_compute_totals(self):
        p = om.Problem()

        p.model.add_subsystem('stuff', om.ExecComp(['y = x', 'cy = x'],
                                                   x={'units': 'inch'},
                                                   y={'units': 'kg'},
                                                   cy={'units': 'kg'}),
                              promotes=['*'])

        p.model.add_design_var('x', units='ft')
        p.model.add_objective('y', units='lbm')
        p.model.add_constraint('cy', units='lbm', lower=0)

        p.setup()

        p['x'] = 1.0
        p.run_model()

        J_driver = p.driver._compute_totals()

        fact = convert_units(1.0, 'kg/inch', 'lbm/ft')
        assert_near_equal(J_driver['y', 'x'][0,0], fact, 1e-5)
        assert_near_equal(J_driver['cy', 'x'][0,0], fact, 1e-5)

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

        ivc = om.IndepVarComp()
        ivc.add_output('x', 35.0, units='degF')

        model.add_subsystem('p', ivc, promotes=['x'])
        model.add_subsystem('comp1', om.ExecComp('y1 = 2.0*x',
                                                 x={'val': 2.0, 'units': 'degF'},
                                                 y1={'val': 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 = "<model> <class Group>: 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')

        model.add_subsystem('p', ivc, promotes=['x'])
        model.add_subsystem('comp1', om.ExecComp('y1 = 2.0*x',
                                                 x={'val': 2.0, 'units': 'degF'},
                                                 y1={'val': 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 = "<model> <class Group>: 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)

        model.add_subsystem('p', ivc, promotes=['x'])
        model.add_subsystem('comp1', om.ExecComp('y1 = 2.0*x',
                                                 x={'val': 2.0},
                                                 y1={'val': 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 = "<model> <class Group>: 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)

        model.add_subsystem('p', ivc, promotes=['x'])
        model.add_subsystem('comp1', om.ExecComp('y1 = 2.0*x',
                                                 x={'val': 2.0},
                                                 y1={'val': 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 = "<model> <class Group>: Target for constraint x has no units, but 'ft' units were specified."
        self.assertEqual(str(context.exception), msg)

    def test_get_desvar_subsystem(self):
        # Test for a bug where design variables in a subsystem were not fully set up.
        prob = om.Problem()
        model = prob.model

        sub = model.add_subsystem('sub', om.Group())
        sub.add_subsystem('comp', Paraboloid(), promotes=['*'])

        prob.set_solver_print(level=0)

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

        sub.add_design_var('x', lower=-50.0, upper=50.0)
        sub.add_design_var('y', lower=-50.0, upper=50.0)
        sub.add_objective('f_xy')
        sub.add_constraint('y', lower=-40.0)

        prob.setup()

        prob.set_val('sub.x', 50.)
        prob.set_val('sub.y', 50.)

        prob.run_driver()

        assert_near_equal(prob['sub.x'], 6.66666667, 1e-6)
        assert_near_equal(prob['sub.y'], -7.3333333, 1e-6)

        prob.set_val('sub.x', 50.)
        prob.set_val('sub.y', 50.)

        prob.run_model()

        totals=prob.check_totals(out_stream=None)
        assert_check_totals(totals)

    def test_get_vars_to_record(self):
        recorder = om.SqliteRecorder("cases.sql")

        prob = om.Problem(name='ABCD')
        prob.add_recorder(recorder)

        prob.model.add_subsystem('mag', om.ExecComp('y=x**2'),
                                 promotes_inputs=['*'], promotes_outputs=['*'])

        prob.model.add_subsystem('sum', om.ExecComp('z=sum(y)'),
                                 promotes_inputs=['*'], promotes_outputs=['*'])


        prob.driver = om.ScipyOptimizeDriver()
        prob.driver.add_recorder(recorder)


        prob.recording_options['record_inputs'] = True
        prob.recording_options['excludes'] = ['*x*', '*aa*']
        prob.recording_options['includes'] = ['*z*', '*bb*']

        prob.driver.recording_options['record_inputs'] = True
        prob.driver.recording_options['excludes'] = ['*y*', '*cc*']
        prob.driver.recording_options['includes'] = ['*x*', '*dd*']

        prob.setup()

        expected_warnings = (
            (OpenMDAOWarning, "Problem ABCD: No matches for pattern '*aa*' in recording_options['excludes']."),
            (OpenMDAOWarning, "Problem ABCD: No matches for pattern '*bb*' in recording_options['includes']."),
            (OpenMDAOWarning, "ScipyOptimizeDriver: No matches for pattern '*cc*' in recording_options['excludes']."),
            (OpenMDAOWarning, "ScipyOptimizeDriver: No matches for pattern '*dd*' in recording_options['includes'].")
        )

        with assert_warnings(expected_warnings):
            prob.final_setup()

    def test_record_residuals_includes_excludes(self):
        import openmdao.api as om
        from openmdao.test_suite.components.sellar import SellarProblem

        prob = SellarProblem()

        recorder = om.SqliteRecorder("rec_resids.sql")
        prob.driver.add_recorder(recorder)

        # just want record residuals
        prob.driver.recording_options['record_inputs'] = False
        prob.driver.recording_options['record_outputs'] = False
        prob.driver.recording_options['record_residuals'] = True
        prob.driver.recording_options['excludes'] = ['*y1*', '*x']   # x is an input, which we are not recording
        prob.driver.recording_options['includes'] = ['*con*', '*z']  # z is an input, which we are not recording

        prob.setup()

        expected_warnings = (
            (om.OpenMDAOWarning, "Driver: No matches for pattern '*x' in recording_options['excludes']."),
            (om.OpenMDAOWarning, "Driver: No matches for pattern '*z' in recording_options['includes']."),
        )

        with assert_warnings(expected_warnings):
            prob.final_setup()

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

        c1 = om.ExecComp()
        c1.add_expr('y = x1**2 - x2**2')
        c1.add_expr('g1 = 2 * x1 + 5')
        c1.add_expr('g2 = -2 * x1 - 5')
        c1.add_expr('g3 = x2 - 5')

        c1.add_objective('y')
        c1.add_constraint('g1', lower=0)
        c1.add_constraint('g2', lower=0)
        c1.add_constraint('g3', lower=0)

        prob.model.add_subsystem('c1', c1, promotes=['*'])
        prob.setup()
        prob.driver = om.ScipyOptimizeDriver()

        with self.assertRaises(RuntimeError) as cm:
            prob.run_driver()

        self.assertEqual(str(cm.exception), 'Problem has no design variables.')


@use_tempdirs
class TestCheckRelevance(unittest.TestCase):
    def setup_problem(self, driver=None):
        prob = om.Problem(driver=driver)
        model = prob.model

        # the customary Paraboloid model
        model.add_subsystem('p1', om.IndepVarComp('x', 50.0), promotes=['*'])
        model.add_subsystem('p2', om.IndepVarComp('y', 50.0), promotes=['*'])
        model.add_subsystem('comp', Paraboloid(), promotes=['*'])
        model.add_subsystem('con', om.ExecComp('c = - x + y'), promotes=['*'])

        model.add_design_var('x', lower=-50.0, upper=50.0)
        model.add_design_var('y', lower=-50.0, upper=50.0)
        model.add_constraint('c', equals=-15.0)

        # add extra comps for testing constraints/objectives
        model.add_subsystem('p3', om.IndepVarComp('z', 50.0), promotes=['*'])
        model.add_subsystem('bad', om.ExecComp('bad = - z'), promotes=['*'])

        return prob

    def test_simple_paraboloid_irrelevant_constraint_run_once(self):
        # setup problem with default driver
        prob = self.setup_problem()

        prob.model.add_objective('f_xy')

        # add constraint that does not depend on a design var
        prob.model.add_constraint('bad', equals=-15.0)

        prob.setup()

        # the default driver does not support optimization, so a constraint that
        # does not depend on a design variable should not trigger a RuntimeError
        prob.run_driver()

    def test_simple_paraboloid_irrelevant_constraint(self):
        # define model with optimizing driver
        prob = self.setup_problem(driver=om.ScipyOptimizeDriver())

        prob.model.add_objective('f_xy')

        # add constraint that does not depend on a design var
        prob.model.add_constraint('bad', equals=-15.0)

        prob.driver.options['singular_jac_behavior'] = 'error'

        prob.setup()

        # since driver is an optimizer, a constraint that does not depend
        # on a design variable should trigger a RuntimeError
        with self.assertRaises(RuntimeError) as err:
            prob.run_driver()

        self.assertTrue("constraint(s) ['bad'] do not depend on any design variables."
                        in str(err.exception))


@unittest.skipUnless(MPI and PETScVector, "MPI and PETSc are required.")
class TestDriverMPI(unittest.TestCase):

    N_PROCS = 2

    def test_distrib_desvar_unsupported(self):

        class MyDriver(Driver):

            def __init__(self, generator=None, **kwargs):
                super().__init__(**kwargs)
                self.supports['distributed_design_vars'] = False
                self.supports._read_only = True

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

        ivc = om.IndepVarComp(distributed=True)
        ivc.add_output('invec')
        model.add_subsystem('p', ivc)
        model.add_design_var('p.invec', lower=0.0, upper=1.0)

        prob.driver = MyDriver()

        prob.setup()

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

        self.assertEqual(str(context.exception),
                         "Distributed design variables are not supported by this driver, "
                         "but the following variables are distributed: [p.invec]")

    def test_distrib_desvar_get_set(self):

        comm = MPI.COMM_WORLD
        size = 3 if comm.rank == 0 else 2

        class DistribComp(om.ExplicitComponent):

            def setup(self):
                self.add_input('w', val=1., distributed=True) # this will connect to a non-distributed IVC
                self.add_input('x', shape=size, distributed=True) # this will connect to a distributed IVC

                self.add_output('y', shape=1, distributed=True) # all-gathered output, duplicated on all procs
                self.add_output('z', shape=size, distributed=True) # distributed output

                self.declare_partials('y', 'x')
                self.declare_partials('y', 'w')
                self.declare_partials('z', 'x')

            def compute(self, inputs, outputs):
                x = inputs['x']
                local_y = np.sum((x-5)**2)
                y_g = np.zeros(self.comm.size)
                self.comm.Allgather(local_y, y_g)

                outputs['y'] = np.sum(y_g) + (inputs['w']-10)**2
                outputs['z'] = x**2

            def compute_partials(self, inputs, J):
                x = inputs['x']
                J['y', 'x'] = 2*(x-5)
                J['y', 'w'] = 2*(inputs['w']-10)
                J['z', 'x'] = np.diag(2*x)

        p = om.Problem()
        model = p.model
        driver = p.driver

        # distributed indep var, 'x'
        d_ivc = model.add_subsystem('d_ivc', om.IndepVarComp(distributed=True),
                                    promotes=['*'])
        d_ivc.add_output('x', 2*np.ones(size))

        # non-distributed indep var, 'w'
        ivc = model.add_subsystem('ivc', om.IndepVarComp(distributed=False),
                                  promotes=['*'])
        # previous version of this test set w to 2 different values depending on rank, which
        # is not valid for a non-distributed variable.  Changed to be the same on all procs.
        ivc.add_output('w', 3.0)

        # distributed component, 'dc'
        model.add_subsystem('dc', DistribComp(), promotes=['*'])

        # distributed design var, 'x'
        model.add_design_var('x', lower=-100, upper=100)
        model.add_objective('y')

        # driver that supports distributed design vars
        driver.supports._read_only = False
        driver.supports['distributed_design_vars'] = True

        p.setup()
        p.run_model()

        # get distributed design var
        driver.get_design_var_values(get_remote=None)

        assert_near_equal(driver.get_design_var_values(get_remote=True)['x'],
                          [2, 2, 2, 2, 2])

        assert_near_equal(driver.get_design_var_values(get_remote=False)['x'],
                          2*np.ones(size))

        # set distributed design var, set_remote=True
        driver.set_design_var('x', [3, 3, 3, 3, 3], set_remote=True)

        assert_near_equal(driver.get_design_var_values(get_remote=True)['x'],
                          [3, 3, 3, 3, 3])

        # set distributed design var, set_remote=False
        if comm.rank == 0:
            driver.set_design_var('x', 5.0*np.ones(size), set_remote=False)
        else:
            driver.set_design_var('x', 9.0*np.ones(size), set_remote=False)

        assert_near_equal(driver.get_design_var_values(get_remote=True)['x'],
                          [5, 5, 5, 9, 9])

        p.run_driver()

        assert_near_equal(p.get_val('dc.y', get_remote=True), [81, 81])
        assert_near_equal(p.get_val('dc.z', get_remote=True), [25, 25, 25, 81, 81])

    def test_distrib_desvar_bug(self):
        class MiniModel(om.Group):
            def setup(self):
                self.add_subsystem('dv', om.IndepVarComp('x', 3.0))
                self.add_subsystem('comp', om.ExecComp('y = (x-2)**2'))
                self.connect('dv.x', 'comp.x')

                self.add_design_var('dv.x', lower=-10, upper=10)

        class ParModel(om.ParallelGroup):
            def setup(self):
                self.add_subsystem('g0', MiniModel())
                self.add_subsystem('g1', MiniModel())

        p = om.Problem()

        p.model.add_subsystem('par_group', ParModel())
        p.model.add_subsystem('obj', om.ExecComp('f = y0 + y1'))

        p.model.connect('par_group.g0.comp.y', 'obj.y0')
        p.model.connect('par_group.g1.comp.y', 'obj.y1')

        p.model.add_objective('obj.f')

        p.driver = om.ScipyOptimizeDriver()

        p.setup()
        p.run_driver()

        p.model.list_outputs()

        dvs = p.driver.get_design_var_values(get_remote=True)

        assert_near_equal(dvs['par_group.g0.dv.x'], 2)
        assert_near_equal(dvs['par_group.g1.dv.x'], 2)

    def test_scalar_scaler_and_adder(self):
        # build the model
        prob = om.Problem()
        dvs = prob.model.add_subsystem('dv', om.IndepVarComp())
        dvs.add_output("x_vec", [3.0, -4.0])
        prob.model.add_subsystem('paraboloid', om.ExecComp('f = (x-3)**2 + x*y + (y+4)**2 - 3'))

        # setup the optimization
        prob.driver = om.ScipyOptimizeDriver()
        prob.driver.options['optimizer'] = 'SLSQP'
        prob.model.connect("dv.x_vec", 'paraboloid.x', src_indices=0)
        prob.model.connect("dv.x_vec", 'paraboloid.y', src_indices=1)

        # Passing scaler as a vector here causes the problem
        prob.model.dv.add_design_var('x_vec', lower=-50, upper=50,
                                     scaler=[0.5, 0.5], adder=[0.0, 0.0])
        prob.model.add_objective('paraboloid.f')

        prob.setup()

        # run the optimization
        prob.run_driver()

        # location of the minimum
        x_star = prob.get_val('paraboloid.x')
        y_star = prob.get_val('paraboloid.y')

        assert_near_equal(x_star, 6.6666, tolerance=1.0E-3)
        assert_near_equal(y_star, -7.3333, tolerance=1.0E-3)


class TestLinearOnlyDVs(unittest.TestCase):
    def build_model(self, driver):

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

        model.add_subsystem('p1', om.IndepVarComp('x', 50.0), promotes=['*'])
        model.add_subsystem('p2', om.IndepVarComp('y', 50.0), promotes=['*'])
        model.add_subsystem('comp', Paraboloid(), promotes=['*'])
        model.add_subsystem('con', om.ExecComp('c = x - y + k'), promotes=['*'])

        prob.set_solver_print(level=0)

        model.add_design_var('x', lower=-50.0, upper=50.0)
        model.add_design_var('y', lower=-50.0, upper=50.0)
        model.add_design_var('k', lower=-1, upper=1)
        model.add_objective('f_xy')
        model.add_constraint('c', lower=15.0, linear=True)

        prob.setup()

        return prob

    @unittest.skipIf(not has_pyoptsparse, "pyoptsparse is required.")
    def test_pyoptsparse(self):
        OPTIMIZER = set_pyoptsparse_opt('SLSQP')[1]

        driver = pyOptSparseDriver(optimizer=OPTIMIZER, print_results=False)
        driver.opt_settings['ACC'] = 1e-9

        prob = self.build_model(driver)
        with assert_no_warning(category=om.DerivativesWarning):
            prob.run_driver()

    def test_scipy_opt(self):
        driver = om.ScipyOptimizeDriver(optimizer='SLSQP', disp=False)
        prob = self.build_model(driver)
        with assert_no_warning(category=om.DerivativesWarning):
            prob.run_driver()


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