import numpy as np
import openmdao.api as om
from openmdao.utils import units
from scipy import interpolate
from simupy.block_diagram import DEFAULT_INTEGRATOR_OPTIONS, SimulationMixin
from simupy.systems import DynamicalSystem

from aviary.mission.gasp_based.ode.params import ParamPort


def add_SGM_required_inputs(group: om.Group, inputs_to_add: dict):
    blank_execcomp = om.ExecComp()
    for input, details in inputs_to_add.items():
        blank_execcomp.add_input(input, **details)
    group.add_subsystem('SGM_required_inputs',
                        blank_execcomp,
                        promotes_inputs=list(inputs_to_add.keys()),
                        )


# Subproblem used as a basis for forward in time integration phases.
class SimuPyProblem(SimulationMixin):
    def __init__(
        self,
        ode,
        time_independent=False,
        t_name="t_curr",
        states=None,
        alternate_state_names=None,
        blocked_state_names=None,
        alternate_state_rate_names=None,
        parameters=None,
        outputs=None,
        controls=None,
        include_state_outputs=False,
        rate_suffix="_rate",
        DEBUG=False,
        max_allowable_time=1_000_000,
        adjoint_int_opts=DEFAULT_INTEGRATOR_OPTIONS.copy(),

    ):
        """
        states: a dictionary of the form {state_name:{'units':unit, 'rate':state_rate_name, 'rate_units':state_rate_units}}
        alternate_state_names: a dictionary of the form {state_to_replace:new_state_name}
        blocked_state_names: a list of the form [state_name_to_block]
        alternate_state_rate_names: a dictionary of the form {state_to_modify:new_state_rate_name}
        parameters: a dictionary of the form {parameter_name:unit}
        outputs: a dictionary of the form {output_name:unit}
        controls: a dictionary of the form {control_name:unit}
        include_state_outputs : automatically add the state to the input
        works well for auto-parsed naming, does not check for duplication before adding
        states, parameters, outputs, and controls can also be input as a list of keys for the dictionary 
        """

        default_om_list_args = dict(prom_name=True, val=False,
                                    out_stream=None, units=True)

        self.DEBUG = DEBUG
        self.max_allowable_time = max_allowable_time
        self.adjoint_int_opts = adjoint_int_opts
        self.adjoint_int_opts['nsteps'] = 5000
        self.adjoint_int_opts['name'] = "dop853"

        self.dt = 0.0
        prob = om.Problem()
        prob.model.add_subsystem(
            "ODE_group",
            ode,
            promotes=["*"],
        )
        self.ode = ode

        self.prob = prob
        prob.setup(check=False, force_alloc_complex=True)
        prob.final_setup()

        self.time_independent = time_independent
        if type(states) is list:
            states = {state: {'units': None, 'rate': state+rate_suffix, 'rate_units': None}
                      for state in states}
        if type(parameters) is list:
            parameters = {parameter: None for parameter in parameters}
        if type(outputs) is list:
            outputs = {output: None for output in outputs}
        if controls is None:
            controls = {}
        elif type(controls) is list:
            controls = {control: None for control in controls}
        if alternate_state_names is None:
            alternate_state_names = {}
        if alternate_state_rate_names is None:
            alternate_state_rate_names = {}

        data = prob.model.list_inputs(includes=controls.keys(), **default_om_list_args)
        control_data = {val.pop('prom_name'): val for key, val in data}
        for control_name, control_units in controls.items():
            if control_units is None:
                controls[control_name] = control_data[control_name]["units"]
        self.controls = controls

        if (
            states is None
            or parameters is None
            or outputs is None  # or
            # event_names is None
        ):
            data = prob.model.list_outputs(**default_om_list_args)
            model_outputs = [val['prom_name'] for key, val in data]
            data = prob.model.list_inputs(**default_om_list_args)
            model_inputs = [val['prom_name'] for key, val in data]

        if states is None:
            states = {
                output.replace(rate_suffix, ""): {'units': None, 'rate': output}
                for output in model_outputs
                if output.endswith(rate_suffix)
                or output in alternate_state_names.keys()
            }

            if blocked_state_names is not None:
                for name in blocked_state_names:
                    if name in states.keys():
                        states.pop(name)
            if alternate_state_names:
                for state_key, val in alternate_state_names.items():
                    if state_key in states.keys():  # Used to rename an existing state
                        states[val] = states.pop(state_key)
                    else:  # Used to add a new state
                        states[val]: {'units': None, 'rate': val+rate_suffix}

        if alternate_state_rate_names:
            for state_name, val in alternate_state_rate_names.items():
                if state_name in states.keys():
                    states[state_name]['rate'] = val

        state_rate_names = [val['rate'] for _, val in states.items()]
        data = prob.model.list_outputs(includes=state_rate_names, **default_om_list_args)
        rate_data = {val.pop('prom_name'): val for key, val in data}
        for state_name, state_data in states.items():
            if state_data['rate_units'] is None:
                state_data['rate_units'] = rate_data[state_data['rate']]["units"]
            if state_data['units'] is None:
                rate_unit = rate_data[state_data['rate']]["units"]
                states[state_name]['units'] = units.simplify_unit("s*" + rate_unit)

        if parameters is None:
            parameters = {
                inp: None
                for inp in set(model_inputs)
                if inp not in list(states.keys()) + list(controls.keys()) + [t_name]
            }

        data = prob.model.list_inputs(includes=parameters.keys(), **default_om_list_args)
        parameter_data = {val.pop('prom_name'): val for key, val in data}
        for parameter_name, parameter_units in parameters.items():
            if parameter_units is None:
                parameters[parameter_name] = parameter_data[parameter_name]["units"]

        if outputs is None:
            outputs = {
                outp: None
                for outp in outputs
                if outp not in state_rate_names  # + event_names
            }

        for output_name, output_units in outputs.items():
            if output_units is None:
                outputs[output_name] = next(iter(prob.model.get_io_metadata(
                    includes=[output_name],
                    metadata_keys=["units"]
                ).values()))["units"]

        if include_state_outputs or outputs == {}:  # prevent empty outputs
            outputs = states + outputs

        self.t_name = t_name
        self.states = states

        self.parameters = parameters
        self.outputs = outputs

        self.dim_state = len(states)
        self.dim_output = len(outputs)
        self.dim_input = len(controls)
        self.dim_parameters = len(parameters)
        # TODO: add defensive checks to make sure dimensions match in both setup and
        # calls

        if DEBUG or True:
            om.n2(prob, outfile="n2_simupy_problem.html", show_browser=False)
            with open('input_list_simupy.txt', 'w') as outfile:
                prob.model.list_inputs(out_stream=outfile,)
            print(states)

    @property
    def time(self):
        return self.prob.get_val(self.t_name)[0]

    @time.setter
    def time(self, value):
        if self.time_independent or self.time == value:
            return
        self.prob.set_val(self.t_name, value)

    @property
    def state(self):
        return np.array(
            [
                self.prob.get_val(state_name, units=state_data['units'])[0]
                for state_name, state_data in self.states.items()
            ]
        )

    @state.setter
    def state(self, value):
        if np.all(self.state == value):
            return
        for state_name, elem_val in zip(
            self.states.keys(), value
        ):
            self.prob.set_val(state_name, elem_val,
                              units=self.states[state_name]['units'])

    def compute_along_traj(self, ts, xs):
        self.prob.set_val(self.t_name, ts)
        for state_name, elem_val in zip(self.states.keys(), xs.T):
            self.prob.set_val(state_name, elem_val,
                              units=self.states[state_name]['units'])

        self.prob.run_model()

    @property
    def control(self):
        return np.array(
            [
                self.prob.get_val(control_name, units=unit)[0]
                for control_name, unit in self.controls.items()
            ]
        )

    @control.setter
    def control(self, value):
        if value is None:
            value = np.array([])
        if (self.control.size == value.size) and np.all(self.control == value):
            return
        for control_name, elem_val in zip(
            self.controls, value
        ):
            self.prob.set_val(control_name, elem_val, units=self.controls[control_name])

    @property
    def parameter(self):
        return np.array(
            [
                self.prob.get_val(parameter_name, units=unit)[0]
                for parameter_name, unit in self.parameters.items()
            ]
        )

    @parameter.setter
    def parameter(self, value):
        if np.all(self.parameter == value):
            return
        for parameter_name, elem_val in zip(
            self.paramaters, value
        ):
            self.prob.set_val(parameter_name, elem_val)

    @property
    def state_rate(self):
        return np.array(
            [
                self.prob.get_val(state_data['rate'], units=state_data['rate_units'])[0]
                for _, state_data in self.states.items()
            ]
        )

    @property
    def output(self):
        return np.array(
            [
                self.prob.get_val(output_name, units=unit)[0]
                for output_name, unit in self.outputs.items()
            ]
        )

    @property
    def events(self):
        return np.array(
            [
                self.prob.get_val(event_name, units=unit)[0]
                for event_name, unit in zip(self.event_names, self.event_units)
            ]
        )

    @property
    def compute(self):
        # TODO: create cache -- t, x (optionally u) as hash
        # I guess just set_val for every input and output?
        # for compute_totals, force run_model. anywhere else that needs forced
        # run_model?
        return self.prob.run_model

    @property
    def compute_totals(self):
        return self.prob.compute_totals

    def state_equation_function(self, t, x, u=None):
        self.time = t
        self.state = x
        self.control = u
        self.compute()
        return self.state_rate

    def output_equation_function(self, t, x):
        if self.output_nan:
            return np.ones(self.dim_output) * np.nan
        self.time = t
        self.state = x
        self.compute()
        return self.output

    def prepare_to_integrate(self, t0, x0):
        self.output_nan = False
        # self.time = t0
        # self.state = x0
        # self.prob.run_model()
        return self.output_equation_function(t0, x0)

    def update_equation_function(self, t, x, event_channels=None):
        self.output_nan = True
        return x

    @property
    def get_val(self):
        return self.prob.get_val

    @property
    def set_val(self):
        return self.prob.set_val


class SGMTrajBase(om.ExplicitComponent):
    def initialize(self):
        # needs to get passed to each ODE
        # TODO: param_dict
        self.options.declare("param_dict",
                             default=ParamPort.param_data)
        self.DEBUG = False
        self.max_allowable_time = 1_000_000
        self.adjoint_int_opts = DEFAULT_INTEGRATOR_OPTIONS.copy()
        self.adjoint_int_opts['nsteps'] = 5000
        self.adjoint_int_opts['name'] = "dop853"

    def setup_params(
            self,
            ODEs,
            promote_all_auto_ivc=False,
            traj_final_state_output=None,
            traj_promote_final_output=None,
            traj_promote_initial_input=None,
            traj_initial_state_input=None,
            traj_event_trigger_input=None,
    ):
        """
        API requirements:
            pass ODE's,
            next_problem = f(current_problem, current_result)
            initial_state/time/etc 
            next_state from last state/output/event information

            pass in terminal and integrand output functions with derivatives (components)
            -- anything special for final state, final time?
            declare initial state(s) as parameters to take derivative wrt
            assume all other inputs are parameters for deriv? 
        """
        if traj_final_state_output is None:
            traj_final_state_output = []
        if traj_promote_final_output is None:
            traj_promote_final_output = []
        if traj_promote_initial_input is None:
            traj_promote_initial_input = {}
        if traj_initial_state_input is None:
            traj_initial_state_input = []
        if traj_event_trigger_input is None:
            traj_event_trigger_input = []

        if promote_all_auto_ivc:
            for ode in ODEs:
                ode: SimuPyProblem
                ode.auto_ivc_vars = ode.prob.model.list_inputs(
                    is_indep_var=True, units=True, out_stream=None)
                for abs_name, data in ode.auto_ivc_vars:
                    prom_name = data.pop('prom_name')
                    if '.' in prom_name:
                        continue
                    traj_promote_initial_input[prom_name] = data

        self.traj_promote_initial_input = {
            **self.options["param_dict"], **traj_promote_initial_input}
        for name, kwargs in self.traj_promote_initial_input.items():
            self.add_input(name, **kwargs)

        final_suffix = "_final"
        self.traj_final_state_output = {
            final_state_output: {
                **dict(
                    name=final_state_output+final_suffix,
                    state_name=final_state_output,
                ),
                **self.add_output(
                    final_state_output+final_suffix,
                    units=ODEs[-1].states[final_state_output]['units'],
                )
            }
            for final_state_output in traj_final_state_output
        }
        self.traj_promote_final_output = {
            promoted_final_output: {
                **dict(
                    name=promoted_final_output+final_suffix,
                    output_name=promoted_final_output,
                ),
                **self.add_output(
                    promoted_final_output+final_suffix,
                    units=ODEs[-1].outputs[promoted_final_output],
                ),
            }
            for promoted_final_output in traj_promote_final_output
        }
        self.all_traj_outputs = {
            **self.traj_final_state_output,
            **self.traj_promote_final_output,
        }
        initial_suffix = "_initial"
        self.traj_initial_state_input = {
            initial_state_input: {
                **dict(name=initial_state_input+initial_suffix),
                **self.add_input(
                    initial_state_input+initial_suffix,
                    units=ODEs[0].states[initial_state_input]['units'],
                )
            }
            for initial_state_input in traj_initial_state_input
        }

        # TODO: assumes state, not output
        trigger_suffix = "trigger"
        self.traj_event_trigger_input = {
            event_trigger_input: {
                **dict(name="_".join([
                    event_trigger_input[0].__class__.__name__,
                    event_trigger_input[1],
                    trigger_suffix
                ])),
                **self.add_input(
                    "_".join([
                        event_trigger_input[0].__class__.__name__,
                        event_trigger_input[1],
                        trigger_suffix
                    ]),
                    units=event_trigger_input[0].states[event_trigger_input[1]]['units'],
                )
            }
            for event_trigger_input in traj_event_trigger_input
        }
        self.ODEs = ODEs
        self.declare_partials(["*"], ["*"],)

    def compute_params(self, inputs):
        for input in self.traj_promote_initial_input.keys():
            for ode in self.ODEs:
                try:
                    ode.set_val(input, inputs[input])
                except KeyError:
                    if self.DEBUG:
                        print(
                            "*** Input not found:",
                            ode,
                            input
                        )
                    pass

    def compute_traj_loop(self, first_problem, inputs, outputs, t0=0., state0=None):
        if self.DEBUG:
            print("initializing compute_traj_loop")
        sim_results = []
        sim_problems = [first_problem]
        t = t0
        if state0 is not None:
            state = state0
        else:
            state = np.array([
                inputs[state_name+"_initial"].squeeze()
                if state_name in self.traj_initial_state_input
                else 0.
                for state_name in first_problem.states.keys()
            ]).squeeze()

        while True:
            current_problem = sim_problems[-1]
            current_problem.initial_condition = state

            sim_result = current_problem.simulate(
                (t, self.max_allowable_time),
            )
            if sim_result.t.shape[0] == 2:
                print("\n"*3, "IMMEDIATE PHASE TERMINATION", current_problem, "\n"*2)
            sim_results.append(sim_result)

            t = sim_result.t[-1]
            x = sim_result.x[-1, :]

            # TODO: is there a better way to do this? Perhaps don't use for loop -- use
            # while True ?Ij
            try:
                try_next_problem = (yield current_problem, sim_result)
            except GeneratorExit:
                if self.DEBUG:
                    print("stop iteration 1")
                break

            if try_next_problem is not None:
                next_problem = try_next_problem
            else:
                try:
                    next_problem = (yield current_problem, sim_result)
                except GeneratorExit:
                    if self.DEBUG:
                        print("stop iteration 2")
                    break

                if self.DEBUG:
                    print(" was on problem:", current_problem,
                          "\n got back:", next_problem)
            # compute the output at the final condition to make sure all outputs are current
            current_problem.output_equation_function(t, x)
            state = np.array(
                [
                    current_problem.prob.get_val(state_name, units=state_data['units'])
                    for state_name, state_data in next_problem.states.items()
                ]
            ).squeeze()
            sim_problems.append(next_problem)

        if self.DEBUG:
            print("ended loop")

        # wrap main loop
        self.sim_results = sim_results
        self.sim_problems = sim_problems

        # trajectory-specific outputs
        for output in self.traj_final_state_output:
            output_name = self.traj_final_state_output[output]["name"]
            state_name = self.traj_final_state_output[output]["state_name"]

            outputs[output_name] = sim_results[-1].x[
                -1,
                list(sim_problems[-1].states.keys()).index(state_name)
            ]

        for output in self.traj_promote_final_output:
            promoted_name = self.traj_promote_final_output[output]["name"]
            output_name = self.traj_promote_final_output[output]["output_name"]

            outputs[promoted_name] = sim_results[-1].y[
                -1,
                list(sim_problems[-1].outputs.keys()).index(output_name)
            ]

        self.last_inputs = np.array(list(inputs.values()))

    def compute_partials(self, inputs, J):
        self.compute_params(inputs)
        # defensive check -- should really make sure ALL inputs are the same, need a
        # deep copy
        # just calling compute_params doesn't fix it -- possibly different trajectory!

        if np.any(self.last_inputs != np.array(list(inputs.values()))):
            raise ValueError(
                "Attempting to run compute_partials when the last compute"
                " inputs did not match",
            )

        param_dict = self.options["param_dict"]

        # assume the first problem has the most states?
        costate_reses = {output: [] for output in self.all_traj_outputs}
        tf_total = self.sim_results[-1].t[-1]

        next_res = self.sim_results[-1]
        next_prob = self.sim_problems[-1]

        df_dxs = []
        df_dparams = []
        dg_dxs = []
        f_minuses = []
        f_pluses = [np.zeros(self.sim_problems[-1].dim_state)]
        state_updates = [next_res.x[-1, :]]
        dh_dxs = [np.eye(next_prob.dim_state)]

        dh_dparams = [np.zeros((next_prob.dim_state, len(param_dict)))]

        # keep directionality of forward for plant, backward for adjoint by caching
        # everything in order to ensure fewest number of ODE calls with smallest
        # step-size sum(abs(x[i]-x[i-1]), )
        costate_ics = []
        param_derivs = []

        for output in self.all_traj_outputs:
            output_name = self.all_traj_outputs[output]["name"]
            costate = np.zeros(next_prob.dim_state)
            param_deriv = np.zeros(len(param_dict))

            if output in self.traj_final_state_output:
                costate[list(next_prob.states.keys()).index(output)] = 1.
            else:  # in self.traj_promote_final_output

                next_prob.state_equation_function(next_res.t[-1], next_res.x[-1, :])
                costate[:] = next_prob.compute_totals(
                    output,
                    next_prob.states.keys(),
                    return_format='array'
                ).squeeze()

                param_deriv[:] = next_prob.compute_totals(
                    output,
                    list(param_dict.keys()),
                    return_format='array'
                ).squeeze()

            param_derivs.append(param_deriv)
            costate_ics.append(costate)

        # pre-compute data for adjoint
        for phase_idx, res, prob in zip(
            range(len(self.sim_results), 0, -1),
            self.sim_results[::-1],
            self.sim_problems[::-1],
        ):
            # build time-varying co-state matrix
            df_dx_data = np.empty(res.x.shape + (res.x.shape[-1],))

            if param_dict:
                df_dparam_data = np.empty(res.x.shape + (len(param_dict),))

            last_res_idx = res.t.shape[0] - 1

            num_active_event_channels = 0

            state_rate = prob.state_equation_function(res.t[-1], res.x[-1, :])
            if (prob is not self.sim_problems[0]):
                f_minuses.append(state_rate)

            for channel_idx, channel_name in enumerate(prob.event_channel_names):

                if np.argmin(np.abs(res.e[-1, :])) not in [channel_idx]:
                    continue

                num_active_event_channels += 1
                dg_dx = np.zeros((1, prob.dim_state))

                if channel_name in prob.states.keys():
                    dg_dx[0, list(prob.states.keys()).index(channel_name)] = 1.
                else:
                    dg_dx[0, :] = prob.compute_totals(
                        [channel_name],
                        prob.states.keys(),
                        return_format='array'
                    )

                # TODO: actually save dg_dts
                if channel_name == prob.t_name:
                    dg_dt = 1.
                else:
                    dg_dt = 0.

            dg_dxs.append(dg_dx)

            if num_active_event_channels != 1:
                raise ValueError("Somehow ended a phase without an event? assume "
                                 "time in the future?? but currently no time-based "
                                 "events are used")

            for idx, (t, x) in enumerate(zip(res.t[::-1], res.x[::-1, :])):
                state_rate = prob.state_equation_function(t, x)

                if (idx == last_res_idx) and (prob is not self.sim_problems[0]):
                    next_prob = self.sim_problems[self.sim_problems.index(prob)-1]

                    f_plus = np.zeros(next_prob.dim_state)
                    plus_rate = state_rate

                    # NOTE / TODO: should enforce that all states in all ODEs exist
                    # in eachother (even if only as output). Don't like assuming
                    # zero
                    # state_update = np.zeros(next_prob.dim_state)
                    state_update = np.ones(next_prob.dim_state)*np.inf
                    dh_dx = np.zeros((next_prob.dim_state,)*2)
                    dh_dparam = np.zeros((next_prob.dim_state, len(param_dict)))

                    # here and co-state assume number of states is only decreasing
                    # forward in time
                    for state_name in next_prob.states.keys():
                        state_idx = list(next_prob.states.keys()).index(state_name)

                        if state_name in prob.states.keys():
                            f_plus[
                                state_idx
                            ] = plus_rate[list(prob.states.keys()).index(state_name)]

                            # state_update[
                            #    list(next_prob.states.keys()).index(state_name)
                            # ] = x[list(prob.states.keys()).index(state_name)]

                            # TODO: make sure index multiplying next_pronb costate
                            # lines up -- since costate is pre-filled to next_prob's
                            # order, the continuous terms should be right
                            # column should map to
                            dh_dx[state_idx, state_idx] = 1.

                        elif state_name in prob.outputs.keys():
                            state_update[
                                state_idx
                            ] = res.y[-1, list(prob.outputs.keys()).index(state_name)]

                            dh_j_dx = prob.compute_totals(
                                [state_name],
                                prob.states.keys(),
                                return_format='array').squeeze()

                            dh_dparam[state_idx, :] = prob.compute_totals(
                                [state_name],
                                list(param_dict.keys()),
                                return_format='array'
                            ).squeeze()

                            for state_name_2 in prob.states.keys():
                                # I'm actually computing dh_dx.T
                                # dh_dx rows are new state, columns are old state
                                # now, dh_dx.T rows are old state, columns are new
                                # so I think this is right
                                dh_dx[
                                    list(next_prob.states.keys()).index(state_name_2),
                                    state_idx,
                                ] = dh_j_dx[list(prob.states.keys()).index(state_name_2)]

                        else:
                            state_update[
                                state_idx
                            ] = 0.

                    f_pluses.append(f_plus)
                    state_updates.append(state_update)
                    dh_dxs.append(dh_dx)
                    dh_dparams.append(dh_dparam)

                state_rate_names = [val['rate'] for _, val in prob.states.items()]
                df_dx_data[idx, :, :] = prob.compute_totals(state_rate_names,
                                                            prob.states.keys(),
                                                            return_format='array').T
                if param_dict:
                    df_dparam_data[idx, ...] = prob.compute_totals(
                        state_rate_names,
                        list(param_dict.keys()),
                        return_format='array'
                    )

            k = min(3, res.t.shape[0]-1)
            skip_interp = (k == 1) and np.isclose(res.t[0], res.t[1])

            # TODO: why is this failing?
            if skip_interp:
                mean_df_dx = np.mean(df_dx_data, axis=0)

                df_dxs.append(lambda t: mean_df_dx)
            else:
                try:
                    df_dxs.append(interpolate.make_interp_spline(
                        tf_total - res.t[::-1],
                        df_dx_data,
                        k=k
                    ))
                except Exception as e:

                    print("EXCEPTION!", k, res.t, df_dx_data, sep='\n')
                    print("check: ", (k == 1) and (res.t[0] == res.t[1]))
                    breakpoint()
                    raise e

            if param_dict:
                if skip_interp:
                    mean_df_dparam = np.mean(df_dparam_data, axis=0)
                    df_dparams.append(lambda t: mean_df_dparam)
                else:
                    df_dparams.append(interpolate.make_interp_spline(
                        tf_total - res.t[::-1],
                        df_dparam_data,
                        k=k
                    ))
            else:
                df_dparams.append(None)

        if self.DEBUG:
            print("data....")
            print("dgs", dg_dxs)
            print("f-", f_minuses)
            print("f+", f_pluses)

            print("size check:", len(self.sim_problems), len(dg_dxs), len(f_minuses),
                  len(f_pluses), )

        # main loop
        for output, costate_ic, param_deriv in zip(self.all_traj_outputs, costate_ics,
                                                   param_derivs):

            output_name = self.all_traj_outputs[output]["name"]
            next_prob = self.sim_problems[-1]
            costate = costate_ic
            lamda_dot_plus = np.zeros_like(costate)

            # self.sim_results[-1].x[-1, next_prob.state_names.index(output)]
            if self.DEBUG:
                print("\nstarting partial for %s" % output, costate)

            dg_dt = 0.

            for (
                phase_idx,
                res,
                prob,
                df_dx,
                df_dparam,
                dg_dx,
                f_minus,
                f_plus,
                state_update,
                dh_dx,
                dh_dparam,
            ) in zip(
                range(len(self.sim_results), 0, -1),
                self.sim_results[::-1],
                self.sim_problems[::-1],
                df_dxs,
                df_dparams,
                dg_dxs,
                f_minuses,
                f_pluses,
                state_updates,
                dh_dxs,
                dh_dparams,
            ):

                t0, tf = tf_total - res.t[[-1, 0]]

                # assumes only 1 of time, state, or output dependence
                # assume no discontinuous state update, would need an API for that in
                # compute as well --
                # but assume some form of event has happened

                # already checked that event_channel_names was well-defined in the
                # pre-compute, so will just assign the co-state just once
                for channel_idx, channel_name in enumerate(prob.event_channel_names):
                    if np.argmin(np.abs(res.e[-1, :])) not in [channel_idx]:
                        continue

                    state_disc = res.x[-1] - state_update
                    state_disc[np.where(np.isinf(state_update))] = 0.

                    if channel_name != prob.t_name:
                        lamda_dot = df_dx(res.t[-1]) @ costate
                        # lamda_dot_plus = lamda_dot
                        if self.DEBUG:
                            if np.any(state_disc):
                                print("update is non-zero!", prob, prob.states.keys(),
                                      state_disc, costate, lamda_dot)
                                print(
                                    "inner product becomes...",
                                    state_disc[None,
                                               :] @ dh_dx @ lamda_dot_plus[:, None],
                                    state_disc[None,
                                               :] @ dh_dx.T @ lamda_dot_plus[:, None]
                                )
                            print("dh_dx for", prob, prob.states.keys(), "\n",  dh_dx)
                            print("costate", costate)
                        costate_update_terms = [
                            dh_dx.T @ costate[:, None],
                            # costate[:, None],
                            # TODO: should this be f_plus? probably not
                            (dg_dx.T @ (f_plus - f_minus)
                             [None, :] @ costate[:, None]) / (dg_dx@f_minus),
                            # don't believe in lamda_dot terms anymore
                            # -(dg_dx.T @ state_disc[None, :] @ dh_dx.T @ lamda_dot_plus[:, None]) / (dg_dx@f_minus),

                        ]

                        # TODO: is this wrong?
                        costate[:] = np.sum(costate_update_terms, axis=0).squeeze()

                    if (
                        (event_key := (prob, channel_name, channel_idx))
                        in self.traj_event_trigger_input
                    ):
                        event_trigger_name = self.traj_event_trigger_input[event_key]["name"]
                        if self.DEBUG:
                            print("setting event trigger data", event_trigger_name)
                        J[output_name, event_trigger_name] = (
                            + costate[None, :] @ (f_minus - f_plus) /
                            (dg_dt + dg_dx@f_minus)
                            # +(lamda_dot_plus[None, :] @ dh_dx @ state_disc[None, :])/(dg_dt + dg_dx@f_minus)
                        )

                    # how to account for terminal event? through costate IC.
                    # TODO: Is this wrong?
                    param_deriv += (costate[None, :] @ dh_dparam).squeeze()

                # build co-state systems

                def co_state_rate(t, costate, *args):
                    return df_dx(t) @ costate

                if self.DEBUG:
                    print('dim_state:', prob.dim_state, "ic:", costate)

                costate_sys = DynamicalSystem(state_equation_function=co_state_rate,
                                              dim_state=prob.dim_state)
                costate_sys.initial_condition = costate

                # simulate co-state system
                co_res = costate_sys.simulate(
                    (t0, tf), integrator_options=self.adjoint_int_opts)
                costate_reses[output].append(co_res)

                if param_dict:
                    df_dparam_val = df_dparam(co_res.t)
                    param_deriv_integrand_data = np.matmul(
                        co_res.x[:, None, :],
                        df_dparam_val
                    ).squeeze()
                    try:
                        param_deriv_integrand = interpolate.make_interp_spline(
                            co_res.t,
                            np.atleast_1d(param_deriv_integrand_data),
                            # k=df_dparam.k
                            k=min(3, co_res.t.shape[0]-1)
                        )
                    except ValueError as e:
                        print(
                            "HIT VALUE ERROR!",
                            output,
                            prob,
                            co_res.t.shape,
                            co_res.x.shape,
                            df_dparam_val.shape,
                            df_dparam.k,
                            "final_results:\n\n",
                            t0, tf,
                            co_res.t,
                            co_res.x,
                        )
                        raise e
                    param_deriv_integrand_antideriv = param_deriv_integrand.antiderivative()

                    # TODO: is the sign wrong here?
                    param_deriv -= (
                        param_deriv_integrand_antideriv(t0)
                        - param_deriv_integrand_antideriv(tf)
                    )

                # consume initial condition
                if prob is not self.sim_problems[0]:
                    next_prob = self.sim_problems[self.sim_problems.index(prob)-1]
                else:
                    break
                costate = np.zeros(next_prob.dim_state)
                lamda_dot_plus = np.zeros_like(costate)
                lamda_dot_plus_rate = co_state_rate(co_res.t[-1], co_res.x[-1])

                # TODO: do co-states need unit changes? probably not...
                for state_name in prob.state_names:
                    costate[list(next_prob.states.keys()).index(
                        state_name)] = co_res.x[-1, list(prob.states.keys()).index(state_name)]
                    lamda_dot_plus[
                        list(next_prob.states.keys()).index(state_name)
                    ] = lamda_dot_plus_rate[list(prob.states.keys()).index(state_name)]

            for state_to_deriv, metadata in self.traj_initial_state_input.items():
                param_name = metadata["name"]
                J[output_name, param_name] = costate_reses[output][-1].x[
                    -1,
                    list(prob.states.keys()).index(state_to_deriv)
                ]
            for param_deriv_val, param_deriv_name in zip(param_deriv, param_dict):
                J[output_name, param_deriv_name] = param_deriv_val
        self.costate_reses = costate_reses
