import pandas as pd


def select_worstcase_snapshots(network):
    """
    Select two worst-case snapshot from time series

    Two time steps in a time series represent worst-case snapshots. These are

    1. Maximum residual load: refers to the point in the time series where the
        (load - generation) achieves its maximum
    2. Minimum residual load: refers to the point in the time series where the
        (load - generation) achieves its minimum

    These to points are identified based on the generation and load time series.

    Parameters
    ----------
    network : :class:`~.grid.network.Network`
        The eDisGo overall container

    Returns
    -------
    type
        Timestamp of snapshot maximum residual load
    type
        Timestamp of snapshot minimum residual load
    """

    grids = [network.mv_grid] + list(network.mv_grid.lv_grids)

    peak_generation = pd.concat(
        [_.peak_generation_per_technology for _ in grids], axis=1).fillna(
        0).sum(axis=1)

    non_solar_wind = [_ for _ in list(peak_generation.index)
                      if _ not in ['wind', 'solar']]
    peak_generation['other'] = peak_generation[non_solar_wind].sum()
    peak_generation.drop(non_solar_wind, inplace=True)

    peak_load = pd.concat(
        [_.consumption for _ in grids], axis=1).fillna(
        0).sum(axis=1)

    residual_load = (
    (network.timeseries.load * peak_load).sum(axis=1) - (
        network.timeseries.generation * peak_generation).sum(
        axis=1))

    return residual_load.idxmax(), residual_load.idxmin()


import pandas as pd


def select_worstcase_snapshots(network):
    """
    Select two worst-case snapshots from time series

    Two time steps in a time series represent worst-case snapshots. These are

    1. Load case: refers to the point in the time series where the
        (load - generation) achieves its maximum and is greater than 0.
    2. Feed-in case: refers to the point in the time series where the
        (load - generation) achieves its minimum and is smaller than 0.

    These two points are identified based on the generation and load time
    series. In case load or feed-in case don't exist None is returned.

    Parameters
    ----------
    network : :class:`~.grid.network.Network`
        Network for which worst-case snapshots are identified.

    Returns
    -------
    :obj:`dict`
        Dictionary with keys 'load_case' and 'feedin_case'. Values are
        corresponding worst-case snapshots of type
        :pandas:`pandas.Timestamp<timestamp>` or None.

    """

    grids = [network.mv_grid] + list(network.mv_grid.lv_grids)

    gens = []
    loads = []
    for grid in grids:
        gens.extend(list(grid.graph.nodes_by_attribute('generator')))
        loads.extend(list(grid.graph.nodes_by_attribute('load')))

    generation_timeseries = pd.Series(0, index=network.timeseries.timeindex)
    for gen in gens:
        generation_timeseries += gen.timeseries.p

    load_timeseries = pd.Series(0, index=network.timeseries.timeindex)
    for load in loads:
        load_timeseries += load.timeseries.p

    residual_load = load_timeseries - generation_timeseries

    timestamp = {}
    timestamp['load_case'] = (
        residual_load.idxmax() if max(residual_load) > 0 else None)
    timestamp['feedin_case'] = (
        residual_load.idxmin() if min(residual_load) < 0 else None)

    return timestamp


Revise function to select worst-case snapshots