import unittest

import numpy as np
from numpy.testing import assert_almost_equal

import openmdao.api as om
from openmdao.utils.testing_utils import use_tempdirs

import dymos as dm
from dymos.transcriptions.grid_data import GridData
from dymos.transcriptions.pseudospectral.components import CollocationComp
from dymos.transcriptions.pseudospectral.components.state_independents import StateIndependentsComp
from dymos.utils.testing_utils import assert_check_partials

# Modify class so we can run it standalone.
from dymos.utils.misc import CompWrapperConfig
CollocationComp = CompWrapperConfig(CollocationComp)
StateIndependentsComp = CompWrapperConfig(StateIndependentsComp)


@use_tempdirs
class TestCollocationCompSolOpt(unittest.TestCase):

    def setUp(self):
        dm.options['include_check_partials'] = True

    def tearDown(self):
        dm.options['include_check_partials'] = False

    def _make_state_idx_map(self, state_name, options, grid_data, state_idx_map):
        """
        Provides error checking for solve_segments and establishes necessary data structures.

        Parameters
        ----------
        state_name : str
            The name of the state being configured.
        options : StateOptionsDictionary
            The StateOptionsDictionary for the state being configured.
        grid_data : GridData
            The GridData object which holds the grid information for this test.
        state_idx_map : dict
            The state_idx_map to be populated by this method.  This map is equivalent to the
            one generated by PseudospectralBase._configure_solve_segments.
        """
        if state_idx_map is None:
            state_idx_map = {}
        state_idx_map[state_name] = {'solver': None, 'indep': None}

        state_input_idxs = grid_data.subset_node_indices['state_input']
        num_state_input_nodes = grid_data.subset_num_nodes['state_input']
        compressed = options.get('compressed', False)

        # Sanity-checks for solve segments
        # If solve_segments is used at all, we cannot fix the state at both ends of the phase.
        if options['solve_segments']:
            # Forward propagation
            if options['solve_segments'] in {True, 'forward'}:
                if compressed:
                    state_idx_map[state_name]['solver'] = np.arange(1, num_state_input_nodes, dtype=int)
                    state_idx_map[state_name]['indep'] = np.zeros((1,), dtype=int)
                else:
                    left_idxs = grid_data.subset_node_indices['segment_ends'][0::2]
                    state_idx_map[state_name]['solver'] = [i for i in range(num_state_input_nodes)
                                                           if state_input_idxs[i] not in left_idxs]
                    state_idx_map[state_name]['indep'] = [i for i in range(num_state_input_nodes)
                                                          if state_input_idxs[i] in left_idxs]

            # Backward propagation
            elif options['solve_segments'] in {'backward'}:
                if compressed:
                    # The optimizer controls the last state input node, all others are solver-controlled
                    state_idx_map[state_name]['indep'] = np.array([num_state_input_nodes - 1], dtype=int)
                    state_idx_map[state_name]['solver'] = np.arange(num_state_input_nodes - 1, dtype=int)
                else:
                    # The optimizer controls the last state input node in each segment, all others are solver_controlled
                    right_idxs = grid_data.subset_node_indices['segment_ends'][1::2]
                    state_idx_map[state_name]['indep'] = [i for i in range(num_state_input_nodes)
                                                          if state_input_idxs[i] in right_idxs]
                    state_idx_map[state_name]['solver'] = [i for i in range(num_state_input_nodes)
                                                           if state_input_idxs[i] not in right_idxs]
        else:
            # No solver used to solve these nodes.  All state input nodes are the indep nodes.
            state_idx_map[state_name]['solver'] = []
            state_idx_map[state_name]['indep'] = np.arange(len(state_input_idxs), dtype=int)

    def make_prob(self, transcription, n_segs, order, compressed):

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

        gd = GridData(num_segments=n_segs, segment_ends=np.arange(n_segs+1),
                      transcription=transcription, transcription_order=order, compressed=compressed)

        state_options = {'x': {'units': 'm', 'shape': (1, ), 'fix_initial': True,
                               'fix_final': False, 'solve_segments': False,
                               'input_initial': False, 'val': 1.0, 'lower': None,
                               'upper': None},
                         'v': {'units': 'm/s', 'shape': (3, 2), 'fix_initial': False,
                               'fix_final': True, 'solve_segments': True,
                               'input_initial': False, 'val': np.ones((3, 2)),
                               'lower': None, 'upper': None}}

        indep_comp = om.IndepVarComp()
        p.model.add_subsystem('indep', indep_comp, promotes_outputs=['*'])

        indep_comp.add_output(
            'dt_dstau',
            val=np.zeros((gd.subset_num_nodes['col']))
        )

        indep_comp.add_output(
            'f_approx:x',
            val=np.zeros((gd.subset_num_nodes['col'], 1)), units='m')
        indep_comp.add_output(
            'f_computed:x',
            val=np.zeros((gd.subset_num_nodes['col'], 1)), units='m')

        indep_comp.add_output(
            'f_approx:v',
            val=np.zeros((gd.subset_num_nodes['col'], 3, 2)), units='m/s')
        indep_comp.add_output(
            'f_computed:v',
            val=np.zeros((gd.subset_num_nodes['col'], 3, 2)), units='m/s')

        p.model.add_subsystem('defect_comp',
                              subsys=CollocationComp(grid_data=gd, state_options=state_options))

        indep = StateIndependentsComp(grid_data=gd, state_options=state_options)
        p.model.add_subsystem('state_indep', indep, promotes_outputs=['*'])

        p.model.connect('f_approx:x', 'defect_comp.f_approx:x')
        p.model.connect('f_approx:v', 'defect_comp.f_approx:v')
        p.model.connect('f_computed:x', 'defect_comp.f_computed:x')
        p.model.connect('f_computed:v', 'defect_comp.f_computed:v')
        p.model.connect('dt_dstau', 'defect_comp.dt_dstau')
        p.model.connect('defect_comp.defects:v', 'state_indep.defects:v')

        self.state_idx_map = {}
        for state_name, options in state_options.items():
            self._make_state_idx_map(state_name, options, gd, self.state_idx_map)
        p.model.state_indep.configure_io(self.state_idx_map)

        p.setup(force_alloc_complex=True)

        p['dt_dstau'] = np.random.random(gd.subset_num_nodes['col'])

        p['f_approx:x'] = np.random.random((gd.subset_num_nodes['col'], 1))
        p['f_approx:v'] = np.random.random((gd.subset_num_nodes['col'], 3, 2))

        p['f_computed:x'] = np.random.random((gd.subset_num_nodes['col'], 1))
        p['f_computed:v'] = np.random.random((gd.subset_num_nodes['col'], 3, 2))

        p.run_model()
        p.model.run_apply_nonlinear()

        return p

    def test_results(self):
        p = self.make_prob('gauss-lobatto', n_segs=4, order=3, compressed=False)

        dt_dstau = p['dt_dstau']

        assert_almost_equal(p['defect_comp.defects:x'],
                            dt_dstau[:, np.newaxis] * (p['f_approx:x']-p['f_computed:x']))

        # solver_nodes = p.model.state_indep.solver_node_idx[:-1]  # fix_final
        solver_nodes = self.state_idx_map['v']['solver']
        assert_almost_equal(p.model._residuals['state_indep.states:v'][solver_nodes],
                            dt_dstau[:, np.newaxis, np.newaxis] *
                            (p['f_approx:v']-p['f_computed:v']))

        p = self.make_prob('gauss-lobatto', n_segs=4, order=3, compressed=True)

        dt_dstau = p['dt_dstau']

        assert_almost_equal(p['defect_comp.defects:x'],
                            dt_dstau[:, np.newaxis] * (p['f_approx:x']-p['f_computed:x']))

        solver_nodes = self.state_idx_map['v']['solver']
        assert_almost_equal(p.model._residuals['state_indep.states:v'][solver_nodes],
                            dt_dstau[:, np.newaxis, np.newaxis] *
                            (p['f_approx:v']-p['f_computed:v']))

    def test_partials(self):
        np.set_printoptions(linewidth=1024, edgeitems=1e1000)

        p = self.make_prob('radau-ps', n_segs=2, order=5, compressed=False)
        cpd = p.check_partials(compact_print=True, method='fd', out_stream=None)
        del cpd['state_indep']
        assert_check_partials(cpd)

        p = self.make_prob('radau-ps', n_segs=2, order=5, compressed=True)
        cpd = p.check_partials(compact_print=True, method='fd', out_stream=None)
        del cpd['state_indep']
        assert_check_partials(cpd)

        p = self.make_prob('gauss-lobatto', n_segs=3, order=5, compressed=False)
        cpd = p.check_partials(compact_print=True, method='fd', out_stream=None)
        del cpd['state_indep']
        assert_check_partials(cpd)

        p = self.make_prob('gauss-lobatto', n_segs=4, order=3, compressed=True)
        cpd = p.check_partials(compact_print=True, method='fd', out_stream=None)
        del cpd['state_indep']
        assert_check_partials(cpd)


if __name__ == '__main__':  # pragma: no cover
    unittest.main()
