# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/datasets/loaders.ipynb.

# %% auto 0
__all__ = ['Cat', 'Sep', 'seps', 'DatasetInfo', 'load_dsinfo', 'LoadResult', 'load', 'LoadMultipleResult', 'load_multiple_seps']

# %% ../../nbs/datasets/loaders.ipynb 1
"""
Quick loaders for common datasets with frequently used setting
"""
import datetime
from pathlib import Path
from typing import Literal, NamedTuple, Tuple, Union, Iterable, Dict, Optional

import nilmtk
from nilmtk import DataSet, MeterGroup
from pydantic import PositiveInt
from pydantic.main import BaseModel
import pandas as pd
from tqdm import tqdm

import enilm
from . import get_nilmtk_dataset
from ..dt.nilmtktf import get_tzinfo_from_ds
import enilm.etypes


class DatasetInfo(BaseModel):
    dataset: enilm.etypes.DatasetID
    building_nr: PositiveInt


def load_dsinfo(ds_info: DatasetInfo) -> Tuple[nilmtk.DataSet, nilmtk.MeterGroup]:
    ds = get_nilmtk_dataset(ds_info.dataset)
    elec = ds.buildings[ds_info.building_nr].elec  # MeterGroup
    return ds, elec


class LoadResult(BaseModel):
    ds_info: DatasetInfo
    ds: DataSet
    elec: MeterGroup
    tz: datetime.tzinfo

    class Config:
        arbitrary_types_allowed = True


def load(dataset: enilm.etypes.DatasetID, building_nr: PositiveInt) -> LoadResult:
    ds_info = DatasetInfo(dataset=dataset, building_nr=building_nr)
    ds, elec = load_dsinfo(ds_info)
    return LoadResult(
        ds_info=ds_info,
        ds=ds,
        elec=elec,
        tz=get_tzinfo_from_ds(ds),
    )


Cat = Union[Literal["mains"], enilm.etypes.AppName]
Sep = Union[
    Literal["full"],
    Literal["ondays"],
    Literal["offdays"],
    Literal["busdays"],
    Literal["weekends"],
]
seps = ["full", "busdays", "weekends", "ondays", "offdays"]


class LoadMultipleResult(NamedTuple):
    data: Dict[enilm.etypes.DatasetID, Dict[enilm.etypes.HouseNr, Dict[Cat, Dict[Sep, pd.Series]]]]
    elec: Dict[enilm.etypes.DatasetID, Dict[enilm.etypes.HouseNr, Dict[Cat, nilmtk.ElecMeter]]]
    loaded: Dict[enilm.etypes.DatasetID, Dict[enilm.etypes.HouseNr, LoadResult]]


def load_multiple_seps(
    datasets: Union[Iterable[enilm.etypes.DatasetID], enilm.etypes.DatasetID],
    physical_quantity_and_ac_type: Dict[enilm.etypes.DatasetID, Dict[str, Tuple[str, str]]],
    houses_per_ds: Dict[enilm.etypes.DatasetID, Iterable[enilm.etypes.HouseNr]],
    apps_per_ds: Dict[enilm.etypes.DatasetID, Iterable[nilmtk.Appliance]],
    datasets_path: Path,
) -> LoadMultipleResult:
    """
    example for physical_quantity_and_ac_type: {HIPE: {'mains': ('power', 'apparent'), 'app': ('power', 'active')}}
    """
    loaded = {}
    elec = {}
    data = {}

    if not isinstance(datasets, list):
        datasets = [datasets]

    for ds in datasets:
        print("Loading", ds)
        mains_physical_quantity, mains_ac_type = physical_quantity_and_ac_type[ds]["mains"]
        app_physical_quantity, app_ac_type = physical_quantity_and_ac_type[ds]["app"]
        apps_names = enilm.appliances.to_str(apps_per_ds[ds])

        base = datasets_path / ds.name
        loaded[ds] = {}
        elec[ds] = {}
        data[ds] = {}
        for house in houses_per_ds[ds]:
            loaded[ds][house] = enilm.datasets.loaders.load(ds, house)
            elec[ds][house] = {}
            data[ds][house] = {}

            elec[ds][house]["mains"]: nilmtk.Electric = loaded[ds][house].elec.mains()
            data[ds][house]["mains"] = {}
            mains_path_base = base / mains_physical_quantity / mains_ac_type / str(house)
            mains_path = mains_path_base / "mains.feather"
            if mains_path.exists():
                df = pd.read_feather(mains_path).set_index("index")
                data[ds][house]["mains"]["full"] = df["full"]
                # MUST dropna because all data is merged into one dataframe, otherwise busdays and weekends would have the same number of rows as full but with na's
                data[ds][house]["mains"]["busdays"] = df["busdays"].dropna()
                data[ds][house]["mains"]["weekends"] = df["weekends"].dropna()
                data[ds][house]["mains"]["ondays"] = df["ondays"].dropna()
                data[ds][house]["mains"]["offdays"] = df["offdays"].dropna()
            else:
                data[ds][house]["mains"]["full"] = elec[ds][house]["mains"].power_series_all_data(
                    physical_quantity=mains_physical_quantity,
                    ac_type=mains_ac_type,
                )
                data[ds][house]["mains"]["full"].index.name = "index"  # in REFIT the index name is unix!
                (
                    data[ds][house]["mains"]["busdays"],
                    data[ds][house]["mains"]["weekends"],
                ) = enilm.split.busday.split_pd_busday_weekend(data[ds][house]["mains"]["full"])
                (
                    data[ds][house]["mains"]["ondays"],
                    data[ds][house]["mains"]["offdays"],
                ) = enilm.split.busday.split_pd_on_offday(data[ds][house]["mains"]["full"], ds)

                # save
                path = mains_path_base
                path.mkdir(parents=True, exist_ok=True)
                pd.DataFrame(data[ds][house]["mains"]).reset_index().to_feather(mains_path)

            apps_iter = list(zip(apps_names, apps_per_ds[ds])) # for tqdm to work, since zip on its own is unbounded
            for app_name, nilmtk_app in tqdm(apps_iter):
                elec[ds][house][app_name]: nilmtk.Electric = enilm.appliances.get_elec(nilmtk_app, loaded[ds][house].elec)
                data[ds][house][app_name] = {}
                app_path_base = base / app_physical_quantity / app_ac_type / str(house)
                app_path = app_path_base / f"{app_name}.feather"
                if app_path.exists():
                    df = pd.read_feather(app_path).set_index("index")
                    data[ds][house][app_name]["full"] = df["full"]
                    data[ds][house][app_name]["busdays"] = df["busdays"].dropna()
                    data[ds][house][app_name]["weekends"] = df["weekends"].dropna()
                    data[ds][house][app_name]["ondays"] = df["ondays"].dropna()
                    data[ds][house][app_name]["offdays"] = df["offdays"].dropna()
                else:
                    data[ds][house][app_name]["full"] = elec[ds][house][app_name].power_series_all_data(
                        physical_quantity=app_physical_quantity,
                        ac_type=app_ac_type,
                    )
                    data[ds][house][app_name]["full"].index.name = "index"  # in REFIT the index name is unix!
                    (
                        data[ds][house][app_name]["busdays"],
                        data[ds][house][app_name]["weekends"],
                    ) = enilm.split.busday.split_pd_busday_weekend(data[ds][house][app_name]["full"])
                    (
                        data[ds][house][app_name]["ondays"],
                        data[ds][house][app_name]["offdays"],
                    ) = enilm.split.busday.split_pd_on_offday(data[ds][house][app_name]["full"], ds)

                    # save
                    path = app_path_base
                    path.mkdir(parents=True, exist_ok=True)
                    pd.DataFrame(data[ds][house][app_name]).reset_index().to_feather(app_path)

            # check: full = busdays + weekends = ondays + offdays
            for cat in ["mains", *apps_names]:
                if data[ds][house][cat]["full"].shape[0] != data[ds][house][cat]["busdays"].shape[0] + data[ds][house][cat]["weekends"].shape[0]:
                    print(f"ERR: full != busdays + weekends in {ds}, {house}, {cat}")
                if data[ds][house][cat]["full"].shape[0] != data[ds][house][cat]["ondays"].shape[0] + data[ds][house][cat]["offdays"].shape[0]:
                    print(f"ERR: full != ondays + offdays in {ds}, {house}, {cat}")

    return LoadMultipleResult(
        data=data,
        elec=elec,
        loaded=loaded,
    )

