# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/yaml/data.ipynb.

# %% auto 0
__all__ = ['Label', 'RawData', 'raw', 'resample', 'overlapping']

# %% ../../nbs/yaml/data.ipynb 1
from typing import Dict, List, Iterable
from pathlib import Path
from typing import Union, Literal
from typing_extensions import TypeAlias

import nilmtk
import pandas as pd

import enilm
import enilm.datasets
import enilm.appliances
import enilm.yaml.config
import enilm.cache
import enilm.cache.pdser
import enilm.etypes
import enilm.etypes.ser
from .config import DayDate

# %% ../../nbs/yaml/data.ipynb 3
Label: TypeAlias = Union[Literal["mains"], enilm.etypes.AppName]  # mains, fridge, etc.
RawData: TypeAlias = Dict[Label, enilm.etypes.ser.PDTimeSeries]

# %% ../../nbs/yaml/data.ipynb 5
def _raw(
    house_elec: nilmtk.MeterGroup,
    data_path: Path,
    selected_physical_quantity: str,
    selected_ac_type: enilm.yaml.config.ACTypes,
    selected_apps_names: List[enilm.etypes.AppName],
) -> RawData:
    """
    - Input: Based on selected parameters:
        1. `selected_apps`
        2. `selected_physical_quantity`
        3. `selected_ac_type`
        4. `ds` and `house` for `house_elec` to get the `ElecMeter` (more specifcally `MeterGroupt`) for mains and appliances

    - Uses:
        - `cache.data.raw.by_ac_type` to load the data from the cache

    - Output:
        - Raw Data: `data` as double nested dict:
            - $level_1$: category (either `mains` or the applinace name from label)
            - $level_2$: the raw data generated from calling `power_series_all_data` on the corresponding electric meter as `pd.Series` with `DatetimeIndex` index
    """
    data = {}

    mains_elec: nilmtk.ElecMeter = house_elec.mains()
    data["mains"] = enilm.cache.pdser.compute_or_load(
        mains_elec.power_series_all_data,
        dict(
            physical_quantity=selected_physical_quantity,
            ac_type=selected_ac_type.mains,
        ),
        data_path / selected_ac_type.mains / "mains.h5",
    )

    selected_apps_nilmtk = enilm.appliances.as_nilmtk(house_elec, selected_apps_names)
    selected_apps_elecs = enilm.appliances.as_elecs(house_elec, selected_apps_nilmtk)
    for app_name, app_elec in zip(selected_apps_names, selected_apps_elecs):
        data[app_name] = enilm.cache.pdser.compute_or_load(
            app_elec.power_series_all_data,
            dict(
                physical_quantity=selected_physical_quantity,
                ac_type=selected_ac_type.apps,
            ),
            data_path / selected_ac_type.apps / f"{app_elec.instance()}.h5",
        )

    # close (note the the mains_elec may be a MeterGroup if the house has multiple aggregation meters, e.g. multiple phases)
    if isinstance(mains_elec, nilmtk.MeterGroup):
        for meter in house_elec.mains().all_meters():
            if hasattr(meter, "store"):
                meter.store.close()
    elif isinstance(mains_elec, nilmtk.ElecMeter):
        if hasattr(mains_elec, "store"):
            mains_elec.store.close()

    for app_elec in selected_apps_elecs:
        if hasattr(app_elec, "store"): # otherwise not working for REDD/1
            app_elec.store.close()

    return data

def raw(config: enilm.yaml.config.Config) -> RawData:
    ds = enilm.datasets.get_dataset_by_name(config.dataset)
    house = config.house
    nilmtkds = enilm.datasets.get_nilmtk_dataset(ds)
    house_elec = nilmtkds.buildings[house].elec
    
    # check if the config is valid for loading (i.e. has all the required fields)
    if config.data_path is None:
        raise ValueError("data_path is required")
    if config.selected_physical_quantity is None:
        raise ValueError("selected_physical_quantity is required")
    if config.selected_ac_type is None:
        raise ValueError("selected_ac_type is required")
    if config.selected_apps is None:
        raise ValueError("selected_apps is required")   

    return _raw(
        selected_apps_names=config.selected_apps,
        selected_physical_quantity=config.selected_physical_quantity,
        selected_ac_type=config.selected_ac_type,
        house_elec=house_elec,
        data_path=Path(config.data_path / "data" / "raw" / "by_ac_type"),
    )

# %% ../../nbs/yaml/data.ipynb 16
from tqdm import tqdm

# %% ../../nbs/yaml/data.ipynb 17
def resample(config: enilm.yaml.config.Config) -> RawData:
    if not config.resample_params:
        raise ValueError("resample_params must be set in config")
    
    data_resampled = {}
    data = raw(config)
    resampled_data_path = config.data_path / "data" / "resampled"

    for cat, cat_data in tqdm(data.items()):
        data_resampled[cat] = enilm.cache.pdser.compute_or_load(
            fn=nilmtk.utils.safe_resample,
            fn_params=dict(
                data=cat_data,
                rule=config.resample_params.rule,
                fill_method=config.resample_params.fill,
                how=config.resample_params.how,
            ),
            ser_path=resampled_data_path / f"{cat}.h5",
            store_format=enilm.cache.pdser.PandasSeriesStoreFormat.HDF5,
        )
    return data_resampled

# %% ../../nbs/yaml/data.ipynb 24
def overlapping(config: enilm.yaml.config.Config) -> RawData:
    data_resampled = resample(config)

    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

    datetime_indices_intersection = get_datetime_index_intersection(
        [data_resampled[cat].index for cat in data_resampled]
    )

    overlapping_data_path = config.data_path / "data" / "overlapping"
    data_overlapping = {}
    for cat, cat_data in data_resampled.items():
        data_overlapping[cat] = enilm.cache.pdser.compute_or_load(
            fn=lambda cat_data: cat_data[datetime_indices_intersection],
            fn_params=dict(
                cat_data=cat_data,
            ),
            ser_path=overlapping_data_path / f"{cat}.h5",
            store_format=enilm.cache.pdser.PandasSeriesStoreFormat.HDF5,
        )
    return data_overlapping
