import json
import copy


def parse_input(mux_input):
    """Takes an input to MachUpX and converts it to a scene dictionary and lists of airplanes, states, and control states.
    This is really just used to simplify things in the unit tests.

    Parameters
    ----------
    mux_input : dict or str
        Scene input to MachUpX

    Returns
    -------
    dict
        Scene input stripped of airplanes

    list
        Airplane names

    list
        Airplane dictionaries

    list
        State dictionaries

    list
        Control state dictionaries
    """

    # Load input
    if isinstance(mux_input, str):
        with open(mux_input, 'r') as input_handle:
            scene_dict = json.load(input_handle)

    else:
        scene_dict = copy.deepcopy(mux_input)

    # Get aircraft
    airplane_names = []
    airplanes = []
    states = []
    control_states = []
    for key, value in scene_dict["scene"].pop("aircraft", {}).items():
        airplane_names.append(key)

        # Load airplane object
        with open(value["file"], 'r') as airplane_file:
            airplanes.append(json.load(airplane_file))

        states.append(value.get("state", {}))
        control_states.append(value.get("control_state", {}))

    if len(airplane_names) == 1:
        return scene_dict, airplane_names[0], airplanes[0], states[0], control_states[0]
    else:
        return scene_dict, airplane_names, airplanes, states, control_states