import unittest
from collections.abc import Iterable

import numpy as np

import openmdao.api as om
from openmdao.utils.mpi import MPI

try:
    from parameterized import parameterized
except ImportError:
    from openmdao.utils.assert_utils import SkipParameterized as parameterized

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


def _test_func_name(func, num, param):
    args = []
    for p in param.args:
        if isinstance(p, str):
            p = {p.replace('.', '_')}
        elif not isinstance(p, Iterable):
            p = {p}
        for item in p:
            try:
                arg = item.__name__
            except Exception:
                arg = str(item)
            args.append(arg)
    return func.__name__ + '_' + '_'.join(args)


class PathCompEx(om.ExplicitComponent):
    def __init__(self, s=''):
        super().__init__()
        self.s = s

    def setup(self):
        self.add_discrete_input('x', val=self.s)
        self.add_discrete_output('y', val=self.s)

    def compute(self, inputs, outputs, discrete_inputs, discrete_outputs):
        discrete_outputs['y'] = discrete_inputs['x'] + self.pathname + '/'


class SerialTests(unittest.TestCase):
    @parameterized.expand([(3, 'par.C1.x', True),
                           (3, 'par.C2.x', True),
                           (3, 'indeps.x', False)],
                          name_func=_test_func_name)
    def test_fan_out(self, size, toset, auto):
        p = om.Problem()
        model = p.model
        if not auto:
            ivc = model.add_subsystem('indeps', om.IndepVarComp())
            ivc.add_output('x', np.ones(size))

        par = model.add_subsystem('par', om.ParallelGroup())
        par.add_subsystem('C1', om.ExecComp('y = 3 * x', x=np.ones(size), y=np.ones(size)), promotes_inputs=['x'])
        par.add_subsystem('C2', om.ExecComp('y = 5 * x', x=np.ones(size), y=np.ones(size)), promotes_inputs=['x'])

        if not auto:
            model.connect('indeps.x', 'par.x')

        p.setup()

        inval = np.arange(size) + 1.0

        p[toset] = inval
        p.run_model()

        np.testing.assert_allclose(p.get_val('par.C1.y', get_remote=True), inval * 3.)
        np.testing.assert_allclose(p.get_val('par.C2.y', get_remote=True), inval * 5.)

        of = ['par.C1.y', 'par.C2.y']
        if auto:
            wrt = ['par.x']
        else:
            wrt = ['indeps.x']

        J = p.compute_totals(of=of, wrt=wrt, return_format='flat_dict')
        np.testing.assert_allclose(J['par.C1.y', wrt[0]], 3. * np.eye(size))
        np.testing.assert_allclose(J['par.C2.y', wrt[0]], 5. * np.eye(size))

    @parameterized.expand([(3, 'par.C1.x', True, True), (3, 'par.C1.x', True, False),
                           (3, 'par.C2.x', True, True), (3, 'par.C2.x', True, False),
                           (3, 'C3.x', True, True), (3, 'C3.x', True, False),
                           (3, 'indeps.x', False, True), (3, 'indeps.x', False, False)],
                          name_func=_test_func_name)
    def test_fan_out_with_dup(self, size, toset, auto, before):
        # this connects an auto_ivc to 3 variables.  2 are under a parallel group and 1 is
        # duplicated in all procs
        p = om.Problem()
        model = p.model

        if not auto:
            ivc = model.add_subsystem('indeps', om.IndepVarComp())
            ivc.add_output('x', np.ones(size))

        c3 = om.ExecComp('y = 4. * x', x=np.ones(size), y=np.ones(size))

        if before:
            model.add_subsystem('C3', c3, promotes_inputs=['x'])

        par = model.add_subsystem('par', om.ParallelGroup(), promotes_inputs=['x'])
        par.add_subsystem('C1', om.ExecComp('y = 3 * x', x=np.ones(size), y=np.ones(size)), promotes_inputs=['x'])
        par.add_subsystem('C2', om.ExecComp('y = 5 * x', x=np.ones(size), y=np.ones(size)), promotes_inputs=['x'])

        if not before:
            model.add_subsystem('C3', c3, promotes_inputs=['x'])

        if not auto:
            model.connect('indeps.x', 'x')

        p.setup()

        inval = np.arange(size) + 1.0

        p[toset] = inval
        p.run_model()

        np.testing.assert_allclose(p.get_val('par.C1.y', get_remote=True), inval * 3.)
        np.testing.assert_allclose(p.get_val('par.C2.y', get_remote=True), inval * 5.)
        np.testing.assert_allclose(p.get_val('C3.y', get_remote=True), inval * 4.)

        of = ['par.C1.y', 'par.C2.y', 'C3.y']

        if auto:
            wrt = ['x']
        else:
            wrt = ['indeps.x']

        J = p.compute_totals(of=of, wrt=wrt, return_format='flat_dict')
        np.testing.assert_allclose(J['par.C1.y', wrt[0]], 3. * np.eye(size))
        np.testing.assert_allclose(J['par.C2.y', wrt[0]], 5. * np.eye(size))
        np.testing.assert_allclose(J['C3.y', wrt[0]], 4. * np.eye(size))

        # try with absolute names
        if not auto:
            J = p.compute_totals(of=of, wrt=['indeps.x'], return_format='flat_dict')
            np.testing.assert_allclose(J['par.C1.y', 'indeps.x'], 3. * np.eye(size))
            np.testing.assert_allclose(J['par.C2.y', 'indeps.x'], 5. * np.eye(size))
            np.testing.assert_allclose(J['C3.y', 'indeps.x'], 4. * np.eye(size))

    @parameterized.expand([(3, 'par.C1.x', 'par.C2.x', True),
                           (3, 'par.C1.x', 'par.C2.x', True),
                           (3, 'indeps1.x', 'indeps2.x', False)],
                          name_func=_test_func_name)
    def test_fan_in(self, size, toset1, toset2, auto):
        p = om.Problem()
        model = p.model
        if not auto:
            ivc = model.add_subsystem('indeps1', om.IndepVarComp())
            ivc.add_output('x', np.ones(size))
            ivc = model.add_subsystem('indeps2', om.IndepVarComp())
            ivc.add_output('x', np.ones(size))

        par = model.add_subsystem('par', om.ParallelGroup())
        par.add_subsystem('C1', om.ExecComp('y = 3 * x', x=np.ones(size), y=np.ones(size)))
        par.add_subsystem('C2', om.ExecComp('y = 5 * x', x=np.ones(size), y=np.ones(size)))

        model.add_subsystem('sum', om.ExecComp('z = x + y', x=np.ones(size), y=np.ones(size), z=np.ones(size)))

        if not auto:
            model.connect('indeps1.x', 'par.C1.x')
            model.connect('indeps2.x', 'par.C2.x')

        model.connect('par.C1.y', 'sum.x')
        model.connect('par.C2.y', 'sum.y')

        p.setup()

        inval1 = np.arange(size) + 1.0
        inval2 = (np.arange(size) + 1.0)[::-1]

        p[toset1] = inval1
        p[toset2] = inval2

        p.run_model()
        np.testing.assert_allclose(p['sum.z'], inval1 * 3. + inval2 * 5.)

    def test_discrete_fan_out1(self):
        p = om.Problem()
        model = p.model
        par = model.add_subsystem('par', om.ParallelGroup(), promotes=['x'])
        par.add_subsystem('C1', PathCompEx(), promotes=['x'])
        par.add_subsystem('C2', PathCompEx(), promotes=['x'])
        p.setup()
        p.run_model()
        self.assertEqual(p.get_val('par.C1.y', get_remote=True), 'par.C1/')
        self.assertEqual(p.get_val('par.C2.y', get_remote=True), 'par.C2/')

    def test_discrete_fan_out2(self):
        p = om.Problem(name='discrete_fan_out2')
        model = p.model
        par = model.add_subsystem('par', om.ParallelGroup(), promotes=['x'])
        par.add_subsystem('C1', PathCompEx('foo'), promotes=['x'])
        par.add_subsystem('C2', PathCompEx('bar'), promotes=['x'])

        try:
            p.setup()
        except Exception as err:
            self.assertTrue(
              "\nCollected errors for problem 'discrete_fan_out2':"
              "\n   <model> <class Group>: The following inputs, ['par.C1.x', 'par.C2.x'], promoted "
              "to 'x', are connected but their metadata entries ['val'] differ. Call "
              "<group>.set_input_defaults('x', val=?), where <group> is the Group named 'par' to "
              "remove the ambiguity." in str(err))
        else:
            self.fail("Exception expected.")

    def test_obj_using_input_name(self):
        class Phase(om.Group):
            def setup(self):
                self.add_subsystem('C1', om.ExecComp('y=.5*x'))
                self.add_subsystem('C2', om.ExecComp('y=g*x'))
                self.add_subsystem('C3', om.ExecComp('y=-x'))

                # this is the culprit.  Bug when objective is added in group using input name
                self.add_objective('C2.g')
                self.add_design_var('C2.g')

                self.connect('C1.y', 'C2.x')
                self.connect('C2.y', 'C3.x')

        p = om.Problem()
        indep = p.model.add_subsystem('indep', om.IndepVarComp('x'))
        indep.add_output('g')

        p.model.add_subsystem('phase0', Phase())

        p.model.connect('indep.x', 'phase0.C1.x')

        p.model.add_design_var('indep.x')
        p.model.add_constraint('phase0.C3.y', equals=0.0)

        p.setup(force_alloc_complex=True)
        p['indep.x'] = [9.9]
        p['indep.g'] = 9.80665
        p.run_model()
        totals = p.check_totals(compact_print=True, method='cs', out_stream=None)
        for key, meta in totals.items():
            np.testing.assert_allclose(meta['abs error'][0], 0.)

    def test_set_val_auto_ivc(self):
        p = om.Problem()
        p.model.add_subsystem('exec',
                              om.ExecComp('z = a + b**2 + c**3',
                                          z={'shape': (100,)},
                                          a={'shape': (100,)},
                                          b={'shape': (100,)},
                                          c={'shape': (100,)}))

        p.setup()

        p.set_val('exec.a', np.linspace(5, 10, 100))
        p.set_val('exec.b', np.linspace(10, 20, 15))
        p.set_val('exec.c', np.linspace(1, 3, 23))

        with self.assertRaises(ValueError) as cm:
            p.final_setup()

        # can not broadcast input array from shape (15,) into shape (100,) for exec.b
        self.assertTrue(str(cm.exception).startswith(
                        "<model> <class Group>: Failed to set value of 'exec.b': "))


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

    N_PROCS = 2
