""" Unit tests for the problem interface."""

import unittest
import os
import json
import re
import types
import base64
import zlib

import numpy as np

import openmdao.api as om

from openmdao.test_suite.components.sellar import SellarStateConnection
from openmdao.visualization.n2_viewer.n2_viewer import _get_viewer_data, n2
from openmdao.recorders.sqlite_recorder import SqliteRecorder
from openmdao.test_suite.test_examples.test_betz_limit import ActuatorDisc
from openmdao.utils.assert_utils import assert_near_equal
from openmdao.utils.mpi import MPI
from openmdao.utils.shell_proc import check_call
from openmdao.utils.testing_utils import use_tempdirs, set_env_vars_context
from openmdao.test_suite.test_examples.beam_optimization.multipoint_beam_group import MultipointBeamGroup

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

# Whether to pop up a browser window for each N2
DEBUG_BROWSER = False

parent_dir = os.path.dirname(os.path.realpath(__file__))


def extract_compressed_model(filename):
    """
    Load an N2 html, find the compressed data string, uncompress and decode it.
    """
    file = open(filename, 'r', encoding='utf-8')
    for line in file:
        if re.search('var compressedModel', line):
            b64_data = line.replace('var compressedModel = "', '').replace('";', '')
            break

    file.close()
    compressed_data = base64.b64decode(b64_data)
    model_data = json.loads(zlib.decompress(compressed_data).decode("utf-8"))

    return model_data


def save_viewer_data(viewer_data, filename):
    """
    Save viewer data to JSON file for use in future testing.
    """
    from openmdao.utils.testing_utils import _ModelViewerDataTreeEncoder
    with open(os.path.join(parent_dir, filename), 'w') as json_file:
        json.dump(viewer_data['tree'], json_file, cls=_ModelViewerDataTreeEncoder, indent=4)

@use_tempdirs
class TestViewerData(unittest.TestCase):

    def recursive_tree_compare(self, tree, expected_tree, path=None, val_tol=1.0E-12):
        """
        Compare two tree structures and raise if they differ.

        Parameters
        ----------
        tree : dict
            The tree being tested.
        expected_tree : dict
            The expected tree being tested against.
        path : list[str | int]
            The path through the original tree to the currently tested subtree.
        val_tol : float
            The tolerance for comparing floating point values.

        Raises
        ------
        ValueError
            When both trees have the same item but with a different value.
        KeyError
            When the expected tree has a path that does not exist in tree, or vice-versa.
        """
        _path = [expected_tree['name']] if path is None else path + [expected_tree['name']]

        tol_check_keys = ['val', 'val_min', 'val_max']

        # Test that expected and tree have same keys
        self.assertSetEqual(set(expected_tree.keys()), set(tree.keys()), msg=f'Keys differ at {_path}')

        if 'dtype' in expected_tree and 'dtype' == 'ndarray':
            # Compare val, min_val, and max_val using tolerance
            for key in tol_check_keys:
                if key in expected_tree:
                    expected_val = np.asarray(expected_tree[key])
                    val = np.asarray(tree[key])
                    try:
                        assert_near_equal(expected_val, val, tolerance=val_tol)
                    except Exception:
                        raise AssertionError(f'{_path + [key]} did not match to expected value')

        for i, key in enumerate(expected_tree.keys()):

            expected_val = expected_tree[key]
            val = tree[key]

            if 'dtype' in expected_tree and key in tol_check_keys:
                continue
            elif key == 'children':
                continue

            # Iterables are serialized as lists
            if isinstance(expected_val, list):
                expected_val = tuple(expected_val)
            if isinstance(val, list):
                val = tuple(val)

            self.assertEqual(expected_val, val, msg=f'Tree items differ at {_path + [key]}')

        if 'children' in expected_tree:
            for i, exs in enumerate(expected_tree['children']):
                ts = tree['children'][i]
                self.recursive_tree_compare(ts, exs, path=_path + ['children', i], val_tol=val_tol)

    def check_viewer_data(self, viewer_data, filename, partials=True):
        """
        Check viewer data against expected.
        """
        # check model tree from JSON file
        with open(os.path.join(parent_dir, filename)) as json_file:
            expected_tree = json.load(json_file)

        # np.testing.assert_equal(viewer_data['tree'], expected_tree, err_msg='', verbose=True)
        self.recursive_tree_compare(viewer_data['tree'], expected_tree)

        # check additional items depending on the problem
        if filename.startswith('sellar'):
            expected_pathnames, expected_conns, expected_abs2prom, expected_declare_partials, \
            expected_driver_name, expected_design_vars_names, expected_responses_names = self.expected_sellar()
        elif filename.startswith('betz'):
            expected_pathnames, expected_conns, expected_abs2prom, expected_declare_partials, \
            expected_driver_name, expected_design_vars_names, expected_responses_names = self.expected_betz()

        # check expected system pathnames
        pathnames = viewer_data['sys_pathnames_list']
        self.assertListEqual(sorted(pathnames), expected_pathnames)

        # check expected connections, after mapping cycle_arrows indices back to pathnames
        connections = sorted(viewer_data['connections_list'],
                             key=lambda x: (x['tgt'], x['src']))
        expected_conns = sorted(expected_conns,
                                key=lambda x: (x['tgt'], x['src']))
        self.assertEqual(len(connections), len(expected_conns))
        for conn in connections:
            if 'cycle_arrows' in conn and conn['cycle_arrows']:
                cycle_arrows = []
                for src, tgt in conn['cycle_arrows']:
                    cycle_arrows.append(' '.join([pathnames[src], pathnames[tgt]]))
                conn['cycle_arrows'] = sorted(cycle_arrows)
        for c, ex in zip(connections, expected_conns):
            self.assertEqual(c['src'], ex['src'])
            self.assertEqual(c['tgt'], ex['tgt'])
            self.assertEqual(c.get('cycle_arrows', []), ex.get('cycle_arrows', []))

        # check expected abs2prom map
        self.assertListEqual(sorted(viewer_data['abs2prom']), sorted(expected_abs2prom))

        # check expected declare_partials_list (available after final_setup)
        if partials:
            self.assertListEqual(sorted(viewer_data['declare_partials_list']),
                                 sorted(expected_declare_partials))

        self.assertEqual(viewer_data['driver']['name'], expected_driver_name)

        self.assertListEqual(sorted(dv for dv in viewer_data['design_vars']),
                             sorted(expected_design_vars_names))
        self.assertListEqual(sorted(resp for resp in viewer_data['responses']),
                             sorted(expected_responses_names))

    def expected_sellar(self):
        pathnames = json.loads('["sub.d1", "sub.d2", "sub.state_eq_group.state_eq"]')
        conns = json.loads("""
            [
                {"src": "sub.d1.y1", "tgt": "con_cmp1.y1"},
                {"src": "sub.d2.y2", "tgt": "con_cmp2.y2"},
                {"src": "_auto_ivc.v1", "tgt": "obj_cmp.x"},
                {"src": "sub.d1.y1", "tgt": "obj_cmp.y1"},
                {"src": "sub.d2.y2", "tgt": "obj_cmp.y2"},
                {"src": "_auto_ivc.v0", "tgt": "obj_cmp.z"},
                {"src": "_auto_ivc.v1", "tgt": "sub.d1.x"},
                {"src": "sub.state_eq_group.state_eq.y2_command", "tgt": "sub.d1.y2"},
                {"src": "_auto_ivc.v0", "tgt": "sub.d1.z"},
                {"src": "sub.d1.y1", "tgt": "sub.d2.y1"},
                {"src": "_auto_ivc.v0", "tgt": "sub.d2.z"},
                {"src": "sub.d2.y2", "tgt": "sub.state_eq_group.state_eq.y2_actual", "cycle_arrows": ["sub.d1 sub.d2", "sub.state_eq_group.state_eq sub.d1"]}
            ]
        """)
        abs2prom = json.loads("""
            {
                "input": {
                    "sub.state_eq_group.state_eq.y2_actual": "state_eq.y2_actual",
                    "sub.d1.z": "z",
                    "sub.d1.x": "x",
                    "sub.d1.y2": "d1.y2",
                    "sub.d2.z": "z",
                    "sub.d2.y1": "y1",
                    "obj_cmp.x": "x",
                    "obj_cmp.y1": "y1",
                    "obj_cmp.y2": "obj_cmp.y2",
                    "obj_cmp.z": "z",
                    "con_cmp1.y1": "y1",
                    "con_cmp2.y2": "con_cmp2.y2"
                },
                "output": {
                    "sub.state_eq_group.state_eq.y2_command": "state_eq.y2_command",
                    "sub.d1.y1": "y1",
                    "sub.d2.y2": "d2.y2",
                    "obj_cmp.obj": "obj",
                    "con_cmp1.con1": "con1",
                    "con_cmp2.con2": "con2",
                    "_auto_ivc.v0": "_auto_ivc.v0",
                    "_auto_ivc.v1": "_auto_ivc.v1"
                }
            }
        """)
        declare_partials = json.loads("""
            [
                "sub.state_eq_group.state_eq.y2_command > sub.state_eq_group.state_eq.y2_actual",
                "sub.d1.y1 > sub.d1.z",
                "sub.d1.y1 > sub.d1.x",
                "sub.d1.y1 > sub.d1.y2",
                "sub.d2.y2 > sub.d2.z",
                "sub.d2.y2 > sub.d2.y1",
                "obj_cmp.obj > obj_cmp.x",
                "obj_cmp.obj > obj_cmp.y1",
                "obj_cmp.obj > obj_cmp.y2",
                "obj_cmp.obj > obj_cmp.z",
                "con_cmp1.con1 > con_cmp1.y1",
                "con_cmp2.con2 > con_cmp2.y2"
            ]
        """)

        driver_name = 'Driver'
        design_vars_names = []
        responses_names = []

        return pathnames, conns, abs2prom, declare_partials, driver_name, design_vars_names, responses_names

    def expected_betz(self):
        pathnames = json.loads('[]')
        conns = json.loads("""
            [{"src": "indeps.a", "tgt": "a_disk.a"},
             {"src": "indeps.Area", "tgt": "a_disk.Area"},
             {"src": "indeps.rho", "tgt": "a_disk.rho"},
             {"src": "indeps.Vu", "tgt": "a_disk.Vu"}]
        """)
        abs2prom = json.loads("""
            {"input": {"a_disk.a": "a", "a_disk.Area": "Area", "a_disk.rho": "rho", "a_disk.Vu": "Vu"},
            "output": {"indeps.a": "a", "indeps.Area": "Area", "indeps.rho": "rho", "indeps.Vu": "Vu",
                       "a_disk.Vr": "a_disk.Vr", "a_disk.Vd": "a_disk.Vd", "a_disk.Ct": "a_disk.Ct",
                       "a_disk.thrust": "a_disk.thrust", "a_disk.Cp": "a_disk.Cp", "a_disk.power": "a_disk.power"}}
        """)
        declare_partials = json.loads("""
            ["a_disk.Vr > a_disk.a", "a_disk.Vr > a_disk.Vu", "a_disk.Vd > a_disk.a", "a_disk.Ct > a_disk.a",
             "a_disk.thrust > a_disk.a", "a_disk.thrust > a_disk.Area", "a_disk.thrust > a_disk.rho",
             "a_disk.thrust > a_disk.Vu", "a_disk.Cp > a_disk.a", "a_disk.power > a_disk.a",
             "a_disk.power > a_disk.Area", "a_disk.power > a_disk.rho", "a_disk.power > a_disk.Vu"]
        """)

        driver_name = 'ScipyOptimizeDriver'
        design_vars_names = ['indeps.a',  'indeps.Area']
        responses_names = ['a_disk.Cp',]

        return pathnames, conns, abs2prom, declare_partials, driver_name, design_vars_names, responses_names

    def test_viewer_data(self):
        """
        Verify that the correct model structure data exists when stored as compared
        to the expected structure, using the SellarStateConnection model.

        Note: Use save_viewer_data() to regenerate JSON file if needed when updating test
        """
        filename = "sellarstate_test_viewer_data.sql"

        p = om.Problem(model=SellarStateConnection(), allow_post_setup_reorder=False)
        p.driver.add_recorder(SqliteRecorder(filename))
        p.setup()

        # Uncomment to update regression data
        # save_viewer_data(_get_viewer_data(p), 'sellar_no_values.json')

        # there should be no values when data is generated before final_setup
        self.check_viewer_data(_get_viewer_data(p), 'sellar_no_values.json', partials=False)

        # there should be initial values when data is generated after final_setup
        p.final_setup()

        # Uncomment to update regression data
        # save_viewer_data(_get_viewer_data(p), 'sellar_initial_values.json')

        self.check_viewer_data(_get_viewer_data(p), 'sellar_initial_values.json')

        # recorded viewer data should match
        self.check_viewer_data(_get_viewer_data(p.get_outputs_dir() / filename), 'sellar_initial_values.json')

        # there should be final values when data is generated after run_model
        p.run_model()

        # Uncomment to update regression data
        # save_viewer_data(_get_viewer_data(p), 'sellar_final_values.json')
        # save_viewer_data(_get_viewer_data(p, values=False), 'sellar_no_values_run.json')

        self.check_viewer_data(_get_viewer_data(p), 'sellar_final_values.json')

        # there should be no values when data is generated (after run) with values=False
        self.check_viewer_data(_get_viewer_data(p, values=False), 'sellar_no_values_run.json')

    def test_viewer_data_from_optimization_problem(self):
        """
        Verify that the correct model structure data exists when stored as compared
        to the expected structure, using the ActuatorDisc model.
        """
        prob = om.Problem()
        indeps = prob.model.add_subsystem('indeps', om.IndepVarComp(), promotes=['*'])
        indeps.add_output('a', .5)
        indeps.add_output('Area', 10.0, units='m**2')
        indeps.add_output('rho', 1.225, units='kg/m**3')
        indeps.add_output('Vu', 10.0, units='m/s')

        prob.model.add_subsystem('a_disk', ActuatorDisc(),
                                 promotes_inputs=['a', 'Area', 'rho', 'Vu'])

        # setup the optimization
        prob.driver = om.ScipyOptimizeDriver(optimizer='SLSQP')

        prob.model.add_design_var('a', lower=0., upper=1.)
        prob.model.add_design_var('Area', lower=0., upper=1.)

        # negative one so we maximize the objective
        prob.model.add_objective('a_disk.Cp', scaler=-1)

        prob.setup()
        prob.final_setup()

        # Uncomment to update regression data
        # save_viewer_data(_get_viewer_data(prob), 'betz_tree.json')

        # check expected viewer data
        self.check_viewer_data(_get_viewer_data(prob), 'betz_tree.json')

    def test_viewer_data_from_subgroup(self):
        """
        Test error message when asking for viewer data for a subgroup.
        """
        p = om.Problem(model=SellarStateConnection())
        p.setup()

        msg = "Viewer data is not available for sub-Group 'sub'."
        with self.assertRaises(TypeError) as cm:
            _get_viewer_data(p.model.sub)

        self.assertEqual(str(cm.exception), msg)

    def test_viewer_data_from_None(self):
        """
        Test error message when asking for viewer data for an invalid source.
        """
        msg = "Viewer data is not available for 'None'." + \
              "The source must be a Problem, model or the filename of a recording."

        with self.assertRaises(TypeError) as cm:
            _get_viewer_data(None)

        self.assertEqual(str(cm.exception), msg)

    def test_handle_ndarray_system_option(self):
        class SystemWithNdArrayOption(om.ExplicitComponent):
            def initialize(self):
                self.options.declare('arr', types=(np.ndarray,))

            def setup(self):
                self.add_input('x', val=0.0)
                self.add_output('f_x', val=0.0)

            def compute(self, inputs, outputs):
                x = inputs['x']
                outputs['f_x'] = (x - 3.0) ** 2

        prob = om.Problem()
        prob.model.add_subsystem('comp', SystemWithNdArrayOption(arr=np.ones(2)))
        prob.setup()
        prob.final_setup()

        viewer_data = _get_viewer_data(prob)
        np.testing.assert_equal(viewer_data['tree']['children'][1]['options']['arr'],
                                np.ones(2))

    def test_system_option_too_large(self):
        class SystemWithLargeOption(om.ExplicitComponent):
            def initialize(self):
                self.options.declare('large_option', types=(np.ndarray,))

            def setup(self):
                self.add_input('x', val=0.0)
                self.add_output('f_x', val=0.0)

            def compute(self, inputs, outputs):
                x = inputs['x']
                outputs['f_x'] = (x - 3.0) ** 2

        prob = om.Problem()
        comp = prob.model.add_subsystem('comp', SystemWithLargeOption())
        comp.options['large_option'] = np.zeros(int(1e4))
        prob.setup()
        prob.final_setup()

        viewer_data = _get_viewer_data(prob)
        self.assertEqual(viewer_data['tree']['children'][1]['options']['large_option'],
                         'Too Large to Display')

    def test_viewer_data_with_submodel(self):

        self.maxDiff = None

        def check_viewer_data_with_submodel(sql_filename):
            # create sub problem
            submodel = om.Group()
            submodel.add_subsystem('subcomp', om.ExecComp('y = x1**2 + x2**2 + x3**2'), promotes=['*'])

            subprob = om.Problem(name='subproblem', model=submodel)

            # create top-level problem
            p = om.Problem(name='top')
            p.model.add_subsystem('submodelcomp',
                                  om.SubmodelComp(problem=subprob, inputs=['*'], outputs=['*'], do_coloring=False),
                                  promotes=['*'])
            p.model.add_subsystem('supercomp',
                                  om.ExecComp('z = 3 * y'),
                                  promotes=['*'])

            p.model.add_recorder(om.SqliteRecorder(sql_filename))

            p.setup()
            p.final_setup()

            # extract viewer data from N2 for problem and subproblem
            om.n2(p, title='N2 for Problem', outfile='N2problem.html', show_browser=DEBUG_BROWSER)
            problem_data = extract_compressed_model('N2problem.html')

            # in this particular case, the value of subcomp.y will differ between this case and the
            # later script based case because in the script case, the hook function triggers after final_setup of the
            # subproblem, while in this case, the n2 is generated after final_setup of the
            # top problem, and during that final_setup (in _setup_partials of the submodelcomp), the
            # subproblem is executed, which sets the value of subcomp.y to 3.0 instead of 1.0.
            # So here we'll reset the value of subcomp.y to 1.0 so that we can
            # compare two data dictionaries using assertDictEqual.
            subprob.set_val('y', 1.0)
            om.n2(subprob, title='N2 for SubProblem', outfile='N2subprob.html', show_browser=DEBUG_BROWSER)
            subprob_data = extract_compressed_model('N2subprob.html')

            # check problem data generated from recording against data generated from problem
            check_call(f"openmdao n2 {p.get_outputs_dir() / sql_filename} -o N2recording.html"
                       f"{' --no_browser' if not DEBUG_BROWSER else ''}")
            recording_data = extract_compressed_model('N2recording.html')

            self.assertDictEqual(problem_data, recording_data)

            # source for run script, copied from above (without the recorder)
            src = """if __name__ == '__main__':
            import openmdao.api as om

            # create sub problem
            submodel = om.Group()
            submodel.add_subsystem('subcomp', om.ExecComp('y = x1**2 + x2**2 + x3**2'), promotes=['*'])

            subprob = om.Problem(name='subproblem', model=submodel)

            # create top-level problem
            p = om.Problem(name='top')
            p.model.add_subsystem('submodelcomp',
                                om.SubmodelComp(problem=subprob, inputs=['*'], outputs=['*'], do_coloring=False),
                                promotes=['*'])
            p.model.add_subsystem('supercomp', om.ExecComp('z = 3 * y'), promotes=['*'])

            p.setup()

            p.run_model()
            """
            with open("submodel_script.py", 'w') as f:
                f.write(src)

            # check problem data generated from script against data generated from problem
            check_call("openmdao n2 submodel_script.py -o N2_top.html --problem=top"
                       f"{' --no_browser' if not DEBUG_BROWSER else ''}")
            n2_top_data = extract_compressed_model('N2_top.html')

            self.assertDictEqual(problem_data, n2_top_data)

            # check subproblem data generated from script against data generated from subproblem
            # NOTE: design vars and responses are added in SubmodelComp's setup, which is not executed
            #       when invoking the n2 command on the subproblem, which exits after subproblem setup
            check_call("openmdao n2 submodel_script.py -o N2_subprob.html --problem=subproblem"
                       f"{' --no_browser' if not DEBUG_BROWSER else ''}")
            n2_sub_data = extract_compressed_model('N2_subprob.html')

            subprob_data['design_vars'] = {}
            subprob_data['responses'] = {}

            self.assertDictEqual(subprob_data, n2_sub_data)

        # check that it works correctly regardless of whether the report system is active or not
        # the n2 command should disable the reports system if active
        for val in ('0', '1'):
            with self.subTest(OPENMDAO_REPORTS=val):
                with set_env_vars_context(OPENMDAO_REPORTS=val):
                    from openmdao.core.problem import _clear_problem_names
                    _clear_problem_names()

                    check_viewer_data_with_submodel(f'recording{val}.sql')


@use_tempdirs
class TestN2(unittest.TestCase):

    def test_n2_from_problem(self):
        """
        Test that an n2 html file is generated from a Problem.
        """
        p = om.Problem(model=SellarStateConnection())
        p.setup()

        html_filename = "problem_n2.html"

        for values in (None, True, False):
            with self.subTest(values=values):
                if values is not None:
                    n2(p, outfile=html_filename, show_browser=DEBUG_BROWSER, values=values)
                else:
                    n2(p, outfile=html_filename, show_browser=DEBUG_BROWSER)

                # Check that the html file has been created and has something in it.
                self.assertTrue(os.path.isfile(html_filename), f"{html_filename} is not a valid file.")
                self.assertGreater(os.path.getsize(html_filename), 100)

    def test_n2_from_model(self):
        """
        Test that an n2 html file is generated from a model.
        """
        p = om.Problem(model=SellarStateConnection())
        p.setup()

        html_filename = "model_n2.html"

        for values in (None, True, False):
            with self.subTest(values=values):
                if values is not None:
                    n2(p.model, outfile=html_filename, show_browser=DEBUG_BROWSER, values=values)
                else:
                    n2(p.model, outfile=html_filename, show_browser=DEBUG_BROWSER)

                # Check that the html file has been created and has something in it.
                self.assertTrue(os.path.isfile(html_filename), f"{html_filename} is not a valid file.")
                self.assertGreater(os.path.getsize(html_filename), 100)

    def test_n2_from_sql(self):
        """
        Test that an n2 html file is generated from a sqlite file.
        """
        sql_filename = "sellarstate.sql"

        p = om.Problem(model=SellarStateConnection())
        p.driver.add_recorder(SqliteRecorder(sql_filename))

        p.setup()
        p.final_setup()
        p.cleanup()

        for values in (None, True, False):
            with self.subTest(values=values):
                n2_from_prob_html = f"n2_from_prob_{values}.html"
                n2_from_file_html = f"n2_from_file_{values}.html"

                _sql_filename = p.get_outputs_dir() / sql_filename

                if values is not None:
                    n2(p, outfile=n2_from_prob_html, show_browser=DEBUG_BROWSER, values=values)
                    n2(_sql_filename, outfile=n2_from_file_html, show_browser=DEBUG_BROWSER, values=values)
                else:
                    n2(p, outfile=n2_from_prob_html, show_browser=DEBUG_BROWSER)
                    n2(_sql_filename, outfile=n2_from_file_html, show_browser=DEBUG_BROWSER)

                # Compare models from the files generated from the Problem and the recording
                model_data_from_prob = extract_compressed_model(n2_from_prob_html)
                model_data_from_file = extract_compressed_model(n2_from_file_html)

                self.assertDictEqual(model_data_from_prob, model_data_from_file)

                # also check data generated using n2 command
                n2_from_cmd_html = f"n2_from_cmd_{values}.html"
                if values is not False:
                    cmd = f"openmdao n2 --no_browser -o {n2_from_cmd_html} {_sql_filename}"
                else:
                    cmd = f"openmdao n2 --no_values --no_browser  -o {n2_from_cmd_html} {_sql_filename}"
                check_call(cmd)

                model_data_from_cmd = extract_compressed_model(n2_from_cmd_html)

                self.assertDictEqual(model_data_from_prob, model_data_from_cmd)

    def test_n2_command(self):
        """
        Check that there are no errors when running from the command line with a script.
        """
        from openmdao.test_suite.scripts import sellar
        filename = os.path.abspath(sellar.__file__).replace('.pyc', '.py')  # PY2

        browser_arg = '' if DEBUG_BROWSER else '--no_browser'

        for values_arg in ('', '--no_values'):
            with self.subTest(values_arg=values_arg):
                check_call(f'openmdao n2 {values_arg} {browser_arg} {filename}')

    def test_n2_set_title(self):
        """
        Test that an n2 html file is generated with the correct title.
        """
        p = om.Problem(model=SellarStateConnection())
        p.setup()

        html_filename = "title_n2.html"

        n2(p, outfile=html_filename, show_browser=DEBUG_BROWSER,
           title="Sellar State Connection")

        # Check that the html file has been created and has the correct title.
        self.assertTrue(os.path.isfile(html_filename),
                        f"{html_filename} is not a valid file.")

        self.assertTrue('OpenMDAO Model Hierarchy and N2 diagram: Sellar State Connection'
                        in open(html_filename, 'r', encoding='utf-8').read())

    def test_n2_connection_error(self):
        """
        Test that an n2 html file is generated from a Problem even if it has connection errors.
        """
        from openmdao.test_suite.scripts.bad_connection import BadConnectionModel

        p = om.Problem(BadConnectionModel(), name='n2_connection_error')

        msg = "\nCollected errors for problem 'n2_connection_error':" \
              "\n   'sub' <class Group>: Attempted to connect from 'tgt.x' to 'cmp.x', but " \
              "'tgt.x' is an input. All connections must be from an output to an input."

        with self.assertRaises(Exception) as cm:
            p.setup()
            p.final_setup()

        self.assertEqual(cm.exception.args[0], msg)

        html_filename = "conn_n2.html"

        n2(p, outfile=html_filename, show_browser=DEBUG_BROWSER,
           title="Bad Connection")

        # Check that the html file has been created and has something in it.
        self.assertTrue(os.path.isfile(html_filename),
                        f"{html_filename} is not a valid file.")
        self.assertTrue('OpenMDAO Model Hierarchy and N2 diagram: Bad Connection'
                        in open(html_filename, 'r', encoding='utf-8').read())


@use_tempdirs
class TestUnderMPI(unittest.TestCase):
    N_PROCS = 2

    def test_non_recordable(self):
        dummyModule = types.ModuleType('dummyModule', 'The dummyModule module')

        class myComp(om.ExplicitComponent):
            def initialize(self):
                self.options.declare('foo', recordable=False)

            def setup(self):
                self.add_input('x2', distributed=True)
                self.add_output('x3', distributed=True)

            def compute(self, inputs, outputs):
                outputs['x3'] = inputs['x2'] + 1

        p = om.Problem()
        ivc = p.model.add_subsystem('ivc', om.IndepVarComp())
        ivc.add_output('x1')
        p.model.add_subsystem('myComp', myComp(foo=dummyModule))

        p.model.connect('ivc.x1', 'myComp.x2')
        p.setup()

        # Test for bug where assembling the options metadata under MPI caused a lockup when
        # they were gathered.
        n2(p, show_browser=False)

    @unittest.skipUnless(MPI and PETScVector, "MPI and PETSc are required.")
    def test_initial_value(self):
        E = 1.
        L = 1.
        b = 0.1
        volume = 0.01

        num_cp = 5
        num_elements = 50
        num_load_cases = 2

        prob = om.Problem(model=MultipointBeamGroup(E=E, L=L, b=b, volume=volume,
                                                    num_elements=num_elements, num_cp=num_cp,
                                                    num_load_cases=num_load_cases))

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

        h = prob['interp.h']
        expected = np.array([ 0.14122705,  0.14130706,  0.14154096,  0.1419107,   0.14238706,  0.14293095,
                              0.14349514,  0.14402636,  0.1444677,   0.14476123,  0.14485062,  0.14468388,
                              0.14421589,  0.1434107,   0.14224356,  0.14070252,  0.13878952,  0.13652104,
                              0.13392808,  0.13105565,  0.1279617,   0.12471547,  0.1213954,   0.11808665,
                              0.11487828,  0.11185599,  0.10900669,  0.10621949,  0.10338308,  0.10039485,
                              0.09716531,  0.09362202,  0.08971275,  0.08540785,  0.08070168,  0.07561313,
                              0.0701851,   0.06448311,  0.05859294,  0.05261756,  0.0466733,   0.04088557,
                              0.03538417,  0.03029845,  0.02575245,  0.02186027,  0.01872173,  0.01641869,
                              0.0150119,   0.01453876])

        assert np.linalg.norm(h - expected) < 1e-4

        def check_initial_value(subsys, parallel=False):
            """
            check that 'initial_value' is indicated for variables under a parallel group
            """
            if subsys['type'] == 'subsystem':
                # Group or Component, recurse to children
                parallel = parallel or subsys['class'] == 'openmdao.core.parallel_group:ParallelGroup'
                for child in subsys['children']:
                    check_initial_value(child, parallel)
            else:
                # input or output, check for 'initial_value' flag
                if parallel:
                    assert('initial_value' in subsys and subsys['initial_value'] is True)
                else:
                    assert('initial_value' not in subsys)

        model_data = _get_viewer_data(prob)
        for subsys in model_data['tree']['children']:
            check_initial_value(subsys)


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