"""deals with loading daily data using config found in [yaml files](enilm.yaml.config)"""

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

# %% auto 0
__all__ = ['DailyData', 'DailyDataNP', 'load', 'chop', 'resample', 'clean', 'scale', 'NDaysOnOff', 'get_ndays_onoff']

# %% ../../../nbs/yaml/daily.ipynb 2
from typing import Dict, Iterable
from typing_extensions import TypeAlias
from dataclasses import dataclass
from pathlib import Path
import datetime

import numpy as np
import pandas as pd
from tqdm.notebook import tqdm
from loguru import logger
import tslearn.preprocessing

import enilm.etypes.ser
import enilm.dicthelpers.save_load

import enilm.yaml.config
from ..data import Label, DayDate, overlapping

# %% ../../../nbs/yaml/daily.ipynb 4
DailyData: TypeAlias = Dict[Label, Dict[DayDate, enilm.etypes.ser.PDTimeSeries]]
DailyDataNP: TypeAlias = Dict[Label, Dict[DayDate, np.ndarray]]

# %% ../../../nbs/yaml/daily.ipynb 6
def load(config: enilm.yaml.config.Config) -> DailyData:
    def get_days_dates_of_ser(ser: pd.Series):
        # how to get the date from the index?
        get_date_from_index_entry = None
        if isinstance(ser.index[0].date, datetime.date):
            get_date_from_index_entry = lambda dt: dt.date  # noqa: E731
        elif isinstance(ser.index[0].date(), datetime.date):
            get_date_from_index_entry = lambda dt: dt.date()  # noqa: E731
        else:
            raise ValueError("I don't know how to get the date from the index!")

        return pd.Series(ser.index.map(get_date_from_index_entry)).drop_duplicates()

    def get_data_split_by_day(cat_data: pd.Series):
        each_day = {}
        assert isinstance(cat_data.index, pd.DatetimeIndex)
        for day in tqdm(get_days_dates_of_ser(cat_data)):
            each_day[day.strftime("%Y-%m-%d")] = cat_data[cat_data.index.date == day]
        return each_day

    data_overlapping = overlapping(config)
    data_path = config.data_path / "data" / "each_day"
    each_day = {}
    for cat, cat_data in tqdm(data_overlapping.items()):
        each_day[cat] = enilm.dicthelpers.save_load.compute_or_load(
            fn=get_data_split_by_day,
            fn_params=dict(cat_data=cat_data),
            dict_path=data_path / f"{cat}.pkl",
        )
    return each_day

# %% ../../../nbs/yaml/daily.ipynb 20
def chop(config: enilm.yaml.config.Config, show_stats=True):
    if config.selected_n_samples_per_day is None:
        raise ValueError("config.selected_n_samples_per_day must be set")
    
    each_day = load(config)
    selected_n_samples_per_day = config.selected_n_samples_per_day
    data_path = config.data_path / "data" / "each_day_chopped"

    def fn(cat_data):
        res = {}
        for day_date, day_data in cat_data.items():
            if day_data.size == selected_n_samples_per_day:
                res[day_date] = day_data.copy()
        return res

    each_day_chopped = {}
    for cat, cat_data in tqdm(each_day.items()):
        each_day_chopped[cat] = enilm.dicthelpers.save_load.compute_or_load(
            fn=fn,
            fn_params=dict(cat_data=cat_data),
            dict_path=data_path / f"{cat}.pkl",
        )
    
    # stats (compared to the original data)
    before_ndays = len(each_day['mains'].keys())
    after_ndays = len(each_day_chopped['mains'])
    removed_ndays = before_ndays - after_ndays
    removed_pct = removed_ndays / before_ndays
    left_pct = 1 - removed_pct
    if show_stats:
        print(f"Number of days before: {before_ndays}")
        print(f"Number of days after chopping: {after_ndays}")
        print(f"Removed {removed_ndays} days ({removed_pct*100:.2f}%)")
    
    # warrning if many days were dropped
    if left_pct <= 0.5:
        logger.warning("More than 50% of the data has been removed! Do you really want to continue?")

    if after_ndays < 10:
        logger.warning("Less than 10 days left! Do you really want to continue?")

    return each_day_chopped

# %% ../../../nbs/yaml/daily.ipynb 26
def resample(config: enilm.yaml.config.Config) -> DailyData:
    selected_n_samples_per_day = config.selected_n_samples_per_day
    data_path = config.data_path / "data" / "each_day_resampled"
    each_day_chopped = chop(config)

    def each_day_dict_to_np_arr(each_day_dict):
        return np.concatenate(list(each_day_dict.values()))

    for cat_data in each_day_chopped.values():
        flattened_daily_data = each_day_dict_to_np_arr(cat_data)
        n_days = len(cat_data.keys())
        claim = (flattened_daily_data.size / selected_n_samples_per_day) == n_days
        if claim == False:
            print("Claim failed!!!")
            print("i.e. samples are note aligned accross days")
            print("flattened_daily_data.size", flattened_daily_data.size)
            print("selected_n_samples_per_day", selected_n_samples_per_day)
            print("n_days", n_days)
            print(
                "flattened_daily_data.size / selected_n_samples_per_day",
                flattened_daily_data.size / selected_n_samples_per_day,
            )
            print(
                "flattened_daily_data.size / selected_n_samples_per_day == n_days",
                flattened_daily_data.size / selected_n_samples_per_day == n_days,
            )
            break

    each_day_resampled = {}
    if claim == False:
        # each_day_resampled = (
        #     tslearn.preprocessing
        #     .TimeSeriesResampler(sz=selected_n_samples_per_day)
        #     .fit_transform(day_data)
        #     .squeeze()
        # )
        raise NotImplementedError()
    else:
        for cat, cat_data in each_day_chopped.items():
            each_day_resampled[cat] = enilm.dicthelpers.save_load.compute_or_load(
                fn=lambda cat: each_day_chopped[cat],
                fn_params=dict(cat=cat),
                dict_path=data_path / f"{cat}.pkl",
            )
    return each_day_resampled

# %% ../../../nbs/yaml/daily.ipynb 32
def clean(config: enilm.yaml.config.Config) -> DailyData:
    if config.manually_deleted_days is None:
        raise ValueError("config.manually_deleted_days must be specified for manual cleaning")
    
    days_to_delete = config.manually_deleted_days
    data_path = config.data_path / "data" / "each_day_cleaned"
    each_day_resampled = resample(config)

    each_day_cleaned = {}
    def fn(cat_data):
        res = {}
        for day in cat_data:
            if not day in days_to_delete:
                res[day] = cat_data[day].copy()
        return res

    for cat, cat_data in tqdm(each_day_resampled.items()):
        each_day_cleaned[cat] = enilm.dicthelpers.save_load.compute_or_load(
            fn=fn,
            fn_params=dict(cat_data=cat_data),
            dict_path=data_path / f"{cat}.pkl",
        )
    return each_day_cleaned

# %% ../../../nbs/yaml/daily.ipynb 40
def scale(config: enilm.yaml.config.Config) -> DailyDataNP:
    if config.selected_daily_scaling_method is None:
        raise ValueError("config.selected_daily_scaling_method must be specified for scaling")
    
    def scale_each_day(each_day, scaler):
        if scaler == "min_max":
            days_data = np.array(list(each_day.values()))
            scalled = tslearn.preprocessing.TimeSeriesScalerMinMax().fit_transform(days_data).squeeze()
            each_day = {day: scalled[i] for i, day in enumerate(each_day.keys())}
        elif scaler == "mean_var":
            days_data = np.array(list(each_day.values()))
            scalled = tslearn.preprocessing.TimeSeriesScalerMeanVariance().fit_transform(days_data).squeeze()
            each_day = {day: scalled[i] for i, day in enumerate(each_day.keys())}
        return each_day

    selected_scaler = config.selected_daily_scaling_method
    data_path = config.data_path / "data" / "each_day_scaled"
    each_day_cleaned = clean(config)

    each_day_scaled = {}
    for cat, cat_data in tqdm(each_day_cleaned.items()):
        each_day_scaled[cat] = enilm.dicthelpers.save_load.compute_or_load(
            fn=lambda cat_data: scale_each_day(cat_data, selected_scaler),
            fn_params=dict(cat_data=cat_data),
            dict_path=data_path / f"{cat}.pkl",
        )
    return each_day_scaled


# %% ../../../nbs/yaml/daily.ipynb 49
@dataclass
class NDaysOnOff:
    on_days: int
    off_days: int

def get_ndays_onoff(config: enilm.yaml.config.Config, day_dates: Iterable[str]):
    ds = enilm.datasets.get_dataset_by_name(config.dataset)
    holidays_cal = enilm.nilmdt.get_holidays_calendar_from_ds(ds)
    
    n_ondays = 0
    n_offdays = 0
    for day in day_dates:
        if isinstance(day, str):
            day = pd.Timestamp(day)
        is_weekend = day.weekday() >= 5
        is_holiday = day.date() in holidays_cal
        if is_weekend or is_holiday:
            n_offdays += 1
        else:
            n_ondays += 1
    return NDaysOnOff(on_days=n_ondays, off_days=n_offdays)
