# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/preprocessing/drop_non_overlapping.ipynb.

# %% auto 0
__all__ = ['overlapping']

# %% ../../nbs/preprocessing/drop_non_overlapping.ipynb 1
from typing import Iterable, Optional
import pandas as pd
import joblib
from . import RawData


def overlapping(
    raw_data: RawData,
    memory: Optional[joblib.Memory] = None,
) -> RawData:
    def get_datetime_index_intersection(
        datetime_indices: Iterable[pd.DatetimeIndex],
    ):
        curr_datetime_index = None
        for datetime_index in datetime_indices:
            if curr_datetime_index is None:
                curr_datetime_index = datetime_index
            else:
                curr_datetime_index = curr_datetime_index.intersection(datetime_index)
        return curr_datetime_index

    def _overlapping(raw_data: RawData):
        datetime_indices: Iterable[pd.DatetimeIndex] = [
            raw_data[cat].index for cat in raw_data
        ]
        datetime_indices_intersection = get_datetime_index_intersection(
            datetime_indices
        )
        data_overlapping = {}
        for cat, cat_data in raw_data.items():
            data_overlapping[cat] = cat_data[datetime_indices_intersection]
        return data_overlapping

    if memory:
        _overlapping = memory.cache(_overlapping)

    # IMPORTANT: pass params, otherwise joblib will always use the
    # same params even if the input data or params change
    return _overlapping(raw_data)
