FEATURE_NAMES = ["accelerometer_0",
        "accelerometer_1",
        "accelerometer_2",
        "goal_compass_1",
        "goal_compass_2",
        "goal_dist",
        "gyro_0",
        "gyro_1",
        "gyro_2",
        "hazards_lidar_0",
        "hazards_lidar_1",
        "hazards_lidar_2",
        "hazards_lidar_3",
        "hazards_lidar_4",
        "hazards_lidar_5",
        "hazards_lidar_6",
        "hazards_lidar_7",
        "hazards_lidar_8",
        "hazards_lidar_9",
        "velocimeter_0",
        "velocimeter_1",
        "velocimeter_2",
        "left_wheel", "right_wheel"]

class FeatureExtractor:
    def __init__(self, restart_name, n_burn_in=0, n_lookahead=1):
        """
        Parameters
        ----------
        restart_name : str
            The name of the 0/1 column indicating restarts in the time series.
        """
        self.restart_name = restart_name
        self.n_burn_in = n_burn_in
        self.n_lookahead = n_lookahead

    def transform(self, X_df):
        """Transform time series into list of states.
        We use the observables at time t as the state

        Be careful not to use any information from the future (X_ds[t + 1:])
        when constructing X_df[t].
        Parameters
        ----------
        X_df_raw : xarray.Dataset
            The raw time series.
        Return
        ------
        X_df : pandas Dataframe

        """
        return X_df[FEATURE_NAMES]
