"""
Tests connections with Reconfigurable Model Execution.

Tests for absolute and promoted connections, for different nonlinear solvers.
"""
# FIXME With NonlinearRunOnce run_model() fails with ValueError (P.O.)
# FIXME With Newton solver and NLBGS variable sizes are not updated. (P.O.)

import numpy as np
import unittest

from openmdao.api import Problem, Group, IndepVarComp, ExplicitComponent
from openmdao.solvers.linear.direct import DirectSolver
from openmdao.solvers.nonlinear.newton import NewtonSolver
from openmdao.solvers.nonlinear.nonlinear_block_gs import NonlinearBlockGS
from openmdao.utils.assert_utils import assert_near_equal


class ReconfComp1(ExplicitComponent):

    def initialize(self):
        self.size = 1
        self.counter = 0

    def reconfigure(self):
        self.counter += 1

        if self.counter % 2 == 0:
            self.size += 1
            flag = True
        else:
            flag = False
        return flag

    def setup(self):
        self.add_input('x', val=1.0)
        self.add_output('y', val=np.zeros(self.size))
        # All derivatives are defined.
        self.declare_partials(of='*', wrt='*')

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

    def compute_partials(self, inputs, jacobian):
        jacobian['y', 'x'] = 2 * np.ones((self.size, 1))


class ReconfComp2(ReconfComp1):
    """The size of the y input changes the same as way as in ReconfComp"""

    def setup(self):
        self.add_input('y', val=np.zeros(self.size))
        self.add_output('f', val=np.zeros(self.size))
        # All derivatives are defined.
        self.declare_partials(of='*', wrt='*')

    def compute(self, inputs, outputs):
        outputs['f'] = 2 * inputs['y']

    def compute_partials(self, inputs, jacobian):
        jacobian['f', 'y'] = 2 * np.ones((self.size, 1))


class TestReconfConnections(unittest.TestCase):

    @unittest.expectedFailure
    def test_promoted_connections(self):
        p = Problem()

        p.model = model = Group()
        model.add_subsystem('c1', IndepVarComp('x', 1.0), promotes_outputs=['x'])
        model.add_subsystem('c2', ReconfComp1(), promotes_inputs=['x'], promotes_outputs=['y'])
        model.add_subsystem('c3', ReconfComp2(), promotes_inputs=['y'],
                            promotes_outputs=['f'])

        p.setup()
        p['x'] = 3.

        self.assertEqual(len(p['y']), 1)
        # First run the model once; counter = 1, size of y = 1
        p.run_model()

        totals = p.compute_totals(wrt=['x'], of=['y'])
        assert_near_equal(p['x'], 3.0)
        assert_near_equal(p['y'], 6.0)
        assert_near_equal(totals['y', 'x'], [[2.0]])

        # Run the model again, which will trigger reconfiguration; counter = 2, size of y = 2
        p.run_model()  # Fails with ValueError
        self.assertEqual(len(p['y']), 2)

    @unittest.expectedFailure
    def test_abs_connections(self):
        p = Problem()

        p.model = model = Group()
        model.add_subsystem('c1', IndepVarComp('x', 1.0), promotes_outputs=['x'])
        model.add_subsystem('c2', ReconfComp1(), promotes_inputs=['x'])
        model.add_subsystem('c3', ReconfComp2(), promotes_outputs=['f'])
        model.connect('c2.y', 'c3.y')
        p.setup()
        p['x'] = 3.

        self.assertEqual(len(p['c2.y']), 1)
        self.assertEqual(len(p['c3.y']), 1)
        # Run the model again, which will trigger reconfiguration; counter = 2, size of y = 2
        p.run_model()

        # Run the model again, which will trigger reconfiguration; counter = 2, size of y = 2
        p.run_model()  # Fails with ValueError
        self.assertEqual(len(p['c2.y']), 2)
        self.assertEqual(len(p['c3.y']), 2)

    @unittest.expectedFailure
    def test_reconf_comp_connections_newton_solver(self):
        p = Problem()

        p.model = model = Group()
        model.linear_solver = DirectSolver()
        model.nonlinear_solver = NewtonSolver(solve_subsystems=False)
        model.add_subsystem('c1', IndepVarComp('x', 1.0), promotes_outputs=['x'])
        model.add_subsystem('c2', ReconfComp1(), promotes_inputs=['x'])
        model.add_subsystem('c3', ReconfComp2(), promotes_outputs=['f'])
        model.connect('c2.y', 'c3.y')
        p.setup()
        p['x'] = 3.

        # First run the model once; counter = 1, size of y = 1
        p.run_model()

        self.assertEqual(len(p['c2.y']), 1)
        self.assertEqual(len(p['c3.y']), 1)
        # Run the model again, which will trigger reconfiguration; counter = 2, size of y = 2
        p.run_model()

        self.assertEqual(len(p['c2.y']), 2)
        self.assertEqual(len(p['c3.y']), 2)
        assert_near_equal(p['c3.y'], [6., 6.])

    @unittest.expectedFailure
    def test_reconf_comp_connections_nlbgs_solver(self):
        p = Problem()

        p.model = model = Group()
        model.linear_solver = DirectSolver()
        model.nonlinear_solver = NonlinearBlockGS()
        model.add_subsystem('c1', IndepVarComp('x', 1.0), promotes_outputs=['x'])
        model.add_subsystem('c2', ReconfComp1(), promotes_inputs=['x'])
        model.add_subsystem('c3', ReconfComp2(), promotes_outputs=['f'])
        model.connect('c2.y', 'c3.y')
        p.setup()
        p['x'] = 3.

        # First run the model once; counter = 1, size of y = 1
        p.run_model()

        self.assertEqual(len(p['c2.y']), 1)
        self.assertEqual(len(p['c3.y']), 1)
        # Run the model again, which will trigger reconfiguration; counter = 2, size of y = 2
        p.run_model()

        self.assertEqual(len(p['c2.y']), 2)
        self.assertEqual(len(p['c3.y']), 2)
        assert_near_equal(p['c3.y'], [6., 6.])

    @unittest.expectedFailure
    def test_promoted_connections_newton_solver(self):
        p = Problem()

        p.model = model = Group()
        model.linear_solver = DirectSolver()
        model.nonlinear_solver = NewtonSolver(solve_subsystems=False)
        model.add_subsystem('c1', IndepVarComp('x', 1.0), promotes_outputs=['x'])
        model.add_subsystem('c2', ReconfComp1(), promotes_inputs=['x'], promotes_outputs=['y'])
        model.add_subsystem('c3', ReconfComp2(), promotes_inputs=['y'], promotes_outputs=['f'])
        p.setup()
        p['x'] = 3.

        # First run the model once; counter = 1, size of y = 1
        p.run_model()
        self.assertEqual(len(p['y']), 1)

        # Run the model again, which will trigger reconfiguration; counter = 2, size of y = 2
        p.run_model()
        self.assertEqual(len(p['y']), 2)
        assert_near_equal(p['y'], [6., 6.])

    @unittest.expectedFailure
    def test_test_promoted_connections_nlbgs_solver(self):
        p = Problem()

        p.model = model = Group()
        model.linear_solver = DirectSolver()
        model.nonlinear_solver = NonlinearBlockGS()
        model.nonlinear_solver.options['reraise_child_analysiserror'] = True
        model.add_subsystem('c1', IndepVarComp('x', 1.0), promotes_outputs=['x'])
        model.add_subsystem('c2', ReconfComp1(), promotes_inputs=['x'], promotes_outputs=['y'])
        model.add_subsystem('c3', ReconfComp2(), promotes_inputs=['y'], promotes_outputs=['f'])
        p.setup()
        p['x'] = 3.

        # First run the model once; counter = 1, size of y = 1
        p.run_model()
        self.assertEqual(len(p['y']), 1)

        # Run the model again, which will trigger reconfiguration; counter = 2, size of y = 2
        p.run_model()
        self.assertEqual(len(p['y']), 2)
        assert_near_equal(p['y'], [6., 6.])

    def test_reconf_comp_not_connected(self):
        p = Problem()

        p.model = model = Group()
        model.add_subsystem('c1', IndepVarComp('x', 1.0), promotes_outputs=['x'])
        model.add_subsystem('c2', ReconfComp1(), promotes_inputs=['x'])
        model.add_subsystem('c3', ReconfComp2(), promotes_outputs=['f'])
        # c2.y not connected to c3.y
        p.setup()
        p['x'] = 3.

        # First run the model once; counter = 1, size of y = 1
        p.run_model()

        self.assertEqual(len(p['c2.y']), 1)
        self.assertEqual(len(p['c3.y']), 1)

        # Run the model again, which will trigger reconfiguration; counter = 2, size of y = 2
        p.run_model()
        self.assertEqual(len(p['c3.y']), 2)
        self.assertEqual(len(p['c2.y']), 2)


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