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

# %% auto 0
__all__ = ['Label', 'RawData', 'load']

# %% ../../nbs/preprocessing/load.ipynb 1
from typing import Dict, Literal, Union, List, Optional
from typing_extensions import TypeAlias

import joblib
import nilmtk

import enilm.etypes
import enilm.etypes.ser
import enilm.datasets.loaders
import enilm.metrics
import enilm.appliances

Label: TypeAlias = Union[Literal["mains"], enilm.etypes.AppName]
RawData: TypeAlias = Dict[Label, enilm.etypes.ser.PDTimeSeries]


def load(
    ds: enilm.etypes.Datasets,
    house: int,
    selected_physical_quantity: str,
    selected_ac_type: Dict[Label, str],
    selected_apps_names: List[str | nilmtk.Appliance],
    memory: Optional[joblib.Memory] = None,
) -> 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 `MeterGroup`) 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
    """

    def _load(
        ds: enilm.etypes.Datasets,
        house: int,
        selected_physical_quantity: str,
        selected_ac_type: Dict[Label, str],
        selected_apps_names: List[str | nilmtk.Appliance],
    ):
        data = {}

        load_res = enilm.datasets.loaders.load(
            dataset=ds,
            building_nr=house,
        )
        house_elec = load_res.elec

        mains_elec: nilmtk.ElecMeter = house_elec.mains()  # type: ignore
        data["mains"] = mains_elec.power_series_all_data(
            physical_quantity=selected_physical_quantity,
            ac_type=selected_ac_type["mains"],
        )

        # load selected appliances as nilmtk.Appliance
        selected_apps_nilmtk: List[nilmtk.Appliance] = []
        for selected_app_name in selected_apps_names:
            if isinstance(selected_app_name, str):
                selected_apps_nilmtk.append(
                    enilm.appliances.get_appliance_by_label_name_from_nilmtk_metergroup(
                        label_name=selected_app_name,
                        elec=house_elec,
                    )
                )
            elif isinstance(selected_app_name, nilmtk.Appliance):
                selected_apps_nilmtk.append(selected_app_name)
            else:
                raise ValueError(f'Invalid type for selected_app_name {selected_app_name}: {type(selected_app_name)}')
        
        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] = app_elec.power_series_all_data(
                physical_quantity=selected_physical_quantity,
                ac_type=selected_ac_type["apps"],
            )

        # 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():  # type: ignore
                meter.store.close()  # type: ignore
        elif isinstance(mains_elec, nilmtk.ElecMeter):
            mains_elec.store.close()  # type: ignore

        for app_elec in selected_apps_elecs:
            app_elec.store.close()  # type: ignore

        return data

    if memory is not None:
        _load = memory.cache(_load)

    # IMPORTANT: pass params, otherwise joblib will always use the
    # same params even if the input data or params change
    return _load(
        ds=ds,
        house=house,
        selected_physical_quantity=selected_physical_quantity,
        selected_ac_type=selected_ac_type,
        selected_apps_names=selected_apps_names,
    )
