"""Split daily data"""

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

# %% auto 0
__all__ = ['DailySplit', 'train_test', 'XYNP', 'traintest_xy', 'XYPD', 'traintest_xy_pd']

# %% ../../../nbs/yaml/daily_split.ipynb 2
from dataclasses import dataclass
from pathlib import Path
from typing import Dict
import math

import numpy as np
import pandas as pd
from tqdm.notebook import tqdm

import enilm.dicthelpers.save_load
from ...etypes import AppName
from . import DailyData
from ..config import Config

# %% ../../../nbs/yaml/daily_split.ipynb 4
@dataclass
class DailySplit:
    train: DailyData
    test: DailyData

# %% ../../../nbs/yaml/daily_split.ipynb 6
def train_test(daily_data: DailyData, config: Config) -> DailySplit:
    if config.selected_train_percent is None:
        raise ValueError("selected_train_percent must be set in config")
    
    train_percent = config.selected_train_percent

    def split_each_day_train(each_day, train_percent):
        n_each_day_train = math.floor(len(each_day.keys()) * train_percent)
        return dict(list(each_day.items())[:n_each_day_train])

    def split_each_day_test(each_day, train_percent):
        n_each_day_train = math.floor(len(each_day.keys()) * train_percent)
        return dict(list(each_day.items())[n_each_day_train:])

    each_day_train = {}
    each_day_test = {}
    data_path = config.data_path / "data" / "each_day_cleaned_split"
    for cat, cat_data in daily_data.items():
        each_day_train[cat] = enilm.dicthelpers.save_load.compute_or_load(
            fn=split_each_day_train,
            fn_params=dict(
                each_day=cat_data,
                train_percent=train_percent,
            ),
            dict_path=data_path / "train" / f"{cat}.pkl",
        )
        each_day_test[cat] = enilm.dicthelpers.save_load.compute_or_load(
            fn=split_each_day_test,
            fn_params=dict(
                each_day=cat_data,
                train_percent=train_percent,
            ),
            dict_path=data_path / "test" / f"{cat}.pkl",
        )
    return DailySplit(
        train=each_day_train,
        test=each_day_test,
    )

# %% ../../../nbs/yaml/daily_split.ipynb 21
@dataclass
class XYNP:
    train_x: np.ndarray
    test_x: np.ndarray
    train_y: Dict[AppName, np.ndarray]
    test_y: Dict[AppName, np.ndarray]

def traintest_xy(each_day_data: DailySplit) -> XYNP:
    def each_day_train_test_x_y_split(all_each_day_train, all_each_day_test):
        apps_names = list(all_each_day_train.keys())
        apps_names.remove("mains")
        print(f"Apps: {apps_names}")

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

        # input: mains
        train_x = each_day_dict_to_np_arr(all_each_day_train["mains"])
        test_x = each_day_dict_to_np_arr(all_each_day_test["mains"])

        # gt: apps
        train_y = {}
        test_y = {}
        for app_name in tqdm(apps_names):
            train_y[app_name] = each_day_dict_to_np_arr(all_each_day_train[app_name])
            test_y[app_name] = each_day_dict_to_np_arr(all_each_day_test[app_name])

        return XYNP(train_x=train_x, train_y=train_y, test_x=test_x, test_y=test_y)

    return each_day_train_test_x_y_split(
        all_each_day_train=each_day_data.train,
        all_each_day_test=each_day_data.test,
    )

# %% ../../../nbs/yaml/daily_split.ipynb 33
@dataclass
class XYPD:
    train_x: np.ndarray
    test_x: np.ndarray
    train_y: Dict[AppName, np.ndarray]
    test_y: Dict[AppName, np.ndarray]

def traintest_xy_pd(each_day_data: DailySplit) -> XYPD:
    def each_day_train_test_x_y_split(all_each_day_train, all_each_day_test):
        apps_names = list(all_each_day_train.keys())
        apps_names.remove("mains")
        print(f"Apps: {apps_names}")

        def concatenate(each_day_dict):
            return pd.concat(list(each_day_dict.values()))

        # input: mains
        train_x = concatenate(all_each_day_train["mains"])
        test_x = concatenate(all_each_day_test["mains"])

        # gt: apps
        train_y = {}
        test_y = {}
        for app_name in tqdm(apps_names):
            train_y[app_name] = concatenate(all_each_day_train[app_name])
            test_y[app_name] = concatenate(all_each_day_test[app_name])

        return XYPD(train_x=train_x, train_y=train_y, test_x=test_x, test_y=test_y)

    return each_day_train_test_x_y_split(
        all_each_day_train=each_day_data.train,
        all_each_day_test=each_day_data.test,
    )
