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

# %% auto 0
__all__ = ['valid_datasets', 'DayDate', 'ResampleParams', 'ACTypes', 'Config', 'load']

# %% ../../nbs/yaml/config.ipynb 1
from pathlib import Path
from typing import Optional, Dict, List
from typing_extensions import TypeAlias

from pydantic import BaseModel, field_validator
import yaml
import pandas as pd
from loguru import logger

import enilm
import enilm.etypes

# %% ../../nbs/yaml/config.ipynb 5
class ResampleParams(BaseModel):
    rule: str  # e.g. '6s' = every 6 seconds
    how: str  # e.g. 'mean'
    fill: str  # e.g. 'ffill' for forward fill

    @field_validator("rule")
    @classmethod
    def rule_must_be_valid_duration(cls, v) -> str:
        try:
            pd.to_timedelta(v)
        except Exception as e:
            raise ValueError(f"Cannot convert {v} to pd.Timedelta: {e}")
        return v

    def n_samples_per_day(self) -> int:
        return int(24 * 60 * 60 / pd.to_timedelta(self.rule).seconds)

# %% ../../nbs/yaml/config.ipynb 6
valid_datasets = list(map(lambda ds: ds.name, enilm.etypes.Datasets))


class ACTypes(BaseModel):
    mains: str
    apps: str

DayDate: TypeAlias = str  # e.g. "2013-04-01"


class Config(BaseModel):
    dataset: str
    house: int

    # can be added later
    data_path: Optional[Path] = None
    selected_physical_quantity: Optional[str] = None
    selected_ac_type: Optional[ACTypes] = None
    selected_apps: Optional[List[str]] = None
    resample_params: Optional[ResampleParams] = None
    selected_n_samples_per_day: Optional[int] = None
    manually_deleted_days: Optional[List[DayDate]] = None
    n_days: Optional[int] = None
    selected_train_percent: Optional[float] = 0.7
    selected_daily_scaling_method: Optional[str] = "mean_var"
    selected_gpu: Optional[int] = None

    def __init__(self, **data):
        if "data_path" in data and data["data_path"] is not None:
            # data_path to Path
            data["data_path"] = Path(data["data_path"])
            # create data_path if it does not exist
            data["data_path"].mkdir(parents=True, exist_ok=True)

        if (
            "selected_n_samples_per_day" not in data
            or data["selected_n_samples_per_day"] is None
        ):
            # compute n_samples_per_day if not given

            # attempt 1: from resample_params.rule
            if "resample_params" in data:
                assert isinstance(data["resample_params"], ResampleParams)
                resample_sec = pd.to_timedelta(data["resample_params"].rule).seconds
                data["selected_n_samples_per_day"] = int(24 * 60 * 60 / resample_sec)
                logger.info(
                    f"selected_n_samples_per_day not given, computing it resampling rule:"
                    f"{data['resample_params'].rule} -> {data['selected_n_samples_per_day']}"
                )
            else:
                logger.warning(
                    "selected_n_samples_per_day not given and cannot be computed"
                )

        # if selected_n_samples_per_day is given, resample_params.rule can be computed
        if (
            "selected_n_samples_per_day" in data
            and data["selected_n_samples_per_day"] is not None
            and "resample_params" not in data
        ):
            data["resample_params"] = ResampleParams(
                rule=f"{24*60*60/data['selected_n_samples_per_day']}s",
                how="mean",
                fill="ffill",
            )
            logger.info(
                f"resample_params.rule not given, computing it from selected_n_samples_per_day:"
                f"{data['selected_n_samples_per_day']} -> {data['resample_params'].rule}"
            )

        # warning if a resample-rule with `m` (or `M`) is used, since most likely it was meant to be in minutes but `m` refers to months!
        if "resample_params" in data:
            assert isinstance(data["resample_params"], ResampleParams)
            if "m" in data["resample_params"] or "M" in data["resample_params"]:
                logger.info(
                    "resample_params.rule contains 'm' or 'M' which refers to months, not minutes!"
                    "use 'min' instead if you meant minutes"
                )

        # hanlde dataset of type enilm.etypes.Datasets
        if isinstance(data["dataset"], enilm.etypes.Datasets):
            data["dataset"] = data["dataset"].name

        super().__init__(**data)

    def hash(self) -> str:
        import base64
        import hashlib

        _hash = hashlib.md5()
        _hash.update(str(self).encode())
        return base64.b64encode(_hash.digest()).decode()

    @field_validator("dataset")
    def valid_dataset(cls, v):
        if v not in valid_datasets:
            raise ValueError(f"dataset {v} not in {valid_datasets}")
        return v

# %% ../../nbs/yaml/config.ipynb 8
from pathlib import Path
from typing import Union
from io import StringIO

# %% ../../nbs/yaml/config.ipynb 9
def load(config: Union[StringIO, Path]) -> 'Config':
    if isinstance(config, Path):
        config = config.open()
    data_dict = yaml.safe_load(config)
    return Config(**data_dict)
