"""Helpers to get datasets paths and avoid navigating to the config file `config.toml`"""

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

# %% auto 0
__all__ = ['platform', 'hostname', 'get_all_datasets_paths', 'id_to_str', 'get_dataset_by_name', 'get_dataset_path',
           'get_nilmtk_dataset', 'get_buildings', 'get_country_code_of_dataset', 'overview']

# %% ../../nbs/datasets/main.ipynb 2
import getpass
import os
import socket
import sys
import io
from pathlib import Path
from typing import Dict, OrderedDict, Optional

import nilmtk
import yaml
import pandas as pd
from loguru import logger

import enilm.etypes
import enilm.config_file

# %% ../../nbs/datasets/main.ipynb 3
# get device info
platform = sys.platform
hostname = socket.gethostname()
try:
    username = getpass.getuser()
except KeyError:
    username = None

# %% ../../nbs/datasets/main.ipynb 5
def get_all_datasets_paths(config_file_path: Optional[Path] = None) -> Dict[str, str]:
    """
    Read config.toml are return the path to the dataset.
    If paths[.PLATFORM[.HOSTNAME[.USERNAME]]] are specified, returns platform specific paths

    config_file_path: Path to the config file. If None, will try to read from ENILM_CONFIG_FILE_PATH env var, then default to config.yaml.

    Returns
    -------
    Dictionary which keys are the datasets names and values are the paths
    """
    if config_file_path is None:
        config_file_path = enilm.config_file.get_config_file_path()
    assert isinstance(config_file_path, Path)
    ds_paths = yaml.safe_load(config_file_path.read_text())["paths"]["datasets"]

    # check for submatch
    paths = {}
    if platform in ds_paths:
        paths = ds_paths[platform]

        # accept default if no hostname is specified
        found_hostname = False
        if not hostname in paths and "default" in paths:
            paths = paths["default"]
            found_hostname = True
        elif hostname in paths:
            paths = paths[hostname]
            found_hostname = True

        if found_hostname:
            if username in paths:
                paths = paths[username]

    return paths

# %% ../../nbs/datasets/main.ipynb 7
def id_to_str(dataset_id: enilm.etypes.DatasetID) -> str:
    if isinstance(dataset_id, enilm.etypes.Datasets):
        return dataset_id.name
    return dataset_id


def get_dataset_by_name(ds_name: str) -> enilm.etypes.Datasets:
    return enilm.etypes.Datasets[ds_name]


def get_dataset_path(
    dataset: enilm.etypes.DatasetID, config_file_path: Optional[Path] = None
) -> str:
    """
    Read config.toml are return the path to the dataset.
    If paths[.PLATFORM[.HOSTNAME[.USERNAME]]] are specified, returns platform specific paths

    Parameters
    ----------
    dataset: desired dataset

    Returns
    -------
    Dataset path from config file
    """
    paths = get_all_datasets_paths(config_file_path)
    dataset_name = id_to_str(dataset)

    if dataset_name in paths:
        path: str = paths[dataset_name]
        if not os.path.exists(path):
            raise ValueError(
                f"File in {path} for dataset {dataset_name} does not exist"
            )
        return path
    raise ValueError(f"Cannot find the dataset {dataset_name} in the config file")

# %% ../../nbs/datasets/main.ipynb 9
def get_nilmtk_dataset(
    dataset: enilm.etypes.DatasetID, config_file_path: Optional[Path] = None
) -> nilmtk.DataSet:
    """
    Parameters
    ----------
    dataset: desired dataset

    Returns
    -------
    NILMTK DataSet
    """
    return nilmtk.DataSet(get_dataset_path(dataset, config_file_path))

# %% ../../nbs/datasets/main.ipynb 11
def get_buildings(
    dataset: enilm.etypes.DatasetID, config_file_path: Optional[Path] = None
) -> OrderedDict[int, nilmtk.Building]:
    """
    Returns:
    --------
    buildings : OrderedDict
        Each key is an integer, starting from 1.
        Each value is a nilmtk.Building object.
    """
    nilmtk_dataset = get_nilmtk_dataset(dataset, config_file_path)
    return nilmtk_dataset.buildings

# %% ../../nbs/datasets/main.ipynb 13
def get_country_code_of_dataset(
    ds: enilm.etypes.DatasetID, 
    config_file_path: Optional[Path] = None,
) -> str:
    if ds == enilm.etypes.Datasets.DRED or ds == "DRED":
        return "NL"
    ds: nilmtk.DataSet = get_nilmtk_dataset(ds, config_file_path)
    # from metadata: geo_location->country
    if "geo_location" in ds.metadata and "country" in ds.metadata["geo_location"]:
        logger.info(f"Found country in metadata: {ds.metadata['geo_location']['country']}")
        return ds.metadata["geo_location"]["country"]
    # from metadata: timezone
    if "timezone" in ds.metadata:
        logger.warning("No country found in metadata, using timezone")
        logger.info(f"Found timezone in metadata: {ds.metadata['timezone']}")
        return ds.metadata["timezone"].split("/")[0]
    raise ValueError("Unknown country for dataset")

# %% ../../nbs/datasets/main.ipynb 15
def overview() -> pd.DataFrame:
    data = """
    Dataset ,Country Code,Year,Type                 ,#H   ,#A  ,Summed up Duration [d],Agg Sampling  , Appl Sampling
    UK-DALE ,GBR         ,2017,residential          ,5    ,109 ,2247                  ,"6 s, 16 kHz" ,6 s
    REDD    ,USA         ,2011,residential          ,6    ,92  ,119                   ,"1 Hz, 15 kHz",1 Hz
    AMPds(2),CAN         ,2016,residential          ,1    ,20  ,730                   ,1 min         ,1 min
    REFIT   ,GBR         ,2016,residential          ,20   ,177 ,14600                 ,8 s           ,8 s
    dataport,USA         ,2015,residential          ,1200+,8598,1376120               ,"1 Hz, 1 min" ,"1 Hz, 1 min"
    ECO     ,CHE         ,2016,residential          ,6    ,45  ,1227                  ,1 Hz          ,1 Hz
    DRED    ,NLD         ,2014,residential          ,1    ,12  ,183                   ,1 Hz          ,1 Hz
    Enertalk,KOR         ,2019,residential          ,22   ,75  ,1714                  ,15 Hz         ,15 Hz
    HES     ,GBR         ,2010,residential          ,251  ,5860,15976                 ,2-10 min      ,2-10 min
    IDEAL   ,GBR         ,-   ,residential          ,-    ,-   ,-                     ,1 Hz          ,1 or 5 Hz
    IMD     ,BRA         ,2020,industrial           ,1    ,8   ,111                   ,1 Hz          ,1 Hz
    PLAID   ,USA         ,2014,residential          ,65   ,1876,1-20 s                ,-             ,30 kHz
    SynD    ,AUT         ,2020,synthetic residential,1    ,21  ,180                   ,5 z           ,5 Hz
    HIPE    ,GER         ,   ?,industrial           ,?    ,?   ,?                     ,?             ,?
    """
    return pd.read_csv(io.StringIO(data), sep=",")

