from __future__ import annotations

from datetime import datetime
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Set, Tuple, Union

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

from learninghouse.models.sensor import Sensors

if TYPE_CHECKING:
    from learninghouse.models.brain import Brain


class DatasetPreprocessing:
    CATEGORICAL_KEY = "categorical"
    NUMERICAL_KEY = "numerical"

    @classmethod
    def sensorsconfig(cls) -> Tuple[List[str], List[str]]:
        sensors = Sensors.load_config()
        categoricals = sensors.categoricals
        numericals = sensors.numericals

        categoricals.append("month_of_year")
        numericals.append("day_of_month")
        categoricals.append("day_of_week")
        numericals.append("hour_of_day")
        numericals.append("minute_of_hour")

        return categoricals, numericals

    @staticmethod
    def add_time_information(data: Dict[str, Any]) -> Dict[str, Any]:
        if "timestamp" not in data:
            data["timestamp"] = datetime.now().timestamp()

        date = datetime.fromtimestamp(data["timestamp"])
        data["datetime"] = date.strftime("%Y-%m-%d %H:%M:%s")
        data["month_of_year"] = date.month
        data["day_of_month"] = date.day
        data["day_of_week"] = date.strftime("%A")
        data["hour_of_day"] = date.hour
        data["minute_of_hour"] = date.minute

        return data

    @classmethod
    def get_x_selected_and_numerical_columns(
        cls, brain: Brain, data: pd.DataFrame, only_features: bool
    ) -> Tuple[pd.DataFrame, List[str]]:
        categoricals, numericals = cls.sensorsconfig()

        categoricals = cls.columns_intersection(categoricals, data)

        used_columns = categoricals + cls.columns_intersection(numericals, data)

        if len(categoricals) > 0:
            x_temp = pd.get_dummies(data[used_columns], columns=categoricals)

            if only_features:
                features_in_dataframe = cls.columns_intersection(
                    x_temp, brain.dataset.features
                )

                x_selected = x_temp[features_in_dataframe]
            else:
                x_selected = x_temp
        else:
            if only_features:
                features_in_dataframe = cls.columns_intersection(
                    data, brain.dataset.features
                )

                x_selected = data[features_in_dataframe]
            else:
                x_selected = data[used_columns]

        x_selected = cls.sort_columns(x_selected)

        return x_selected, numericals

    @classmethod
    def prepare_training(
        cls, brain: Brain, data: pd.DataFrame, only_features: bool
    ) -> Tuple[Brain, pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
        x_vector, numericals = cls.get_x_selected_and_numerical_columns(
            brain, data, only_features
        )

        y_vector = data[brain.configuration.name]

        if brain.configuration.dependent_encode:
            y_vector = brain.dataset.dependent_encoder.fit_transform(y_vector)

        x_train, x_test, y_train, y_test = train_test_split(
            x_vector, y_vector, test_size=brain.configuration.test_size, random_state=0
        )

        if only_features:
            numericals = cls.columns_intersection(numericals, brain.dataset.features)
        else:
            numericals = cls.columns_intersection(numericals, data)

        x_train = cls.transform_columns(
            brain.dataset.imputer.fit_transform, x_train, numericals
        )
        x_test = cls.transform_columns(
            brain.dataset.imputer.transform, x_test, numericals
        )

        x_train = cls.sort_columns(x_train)
        x_test = cls.sort_columns(x_test)

        return brain, x_train, x_test, y_train, y_test

    @classmethod
    def prepare_prediction(cls, brain: Brain, data: pd.DataFrame) -> pd.DataFrame:
        x_vector, numericals = cls.get_x_selected_and_numerical_columns(
            brain, data, True
        )

        numericals = cls.columns_intersection(brain.dataset.columns, numericals)

        missing_columns = set.difference(
            cls.set_of_columns(numericals), cls.set_of_columns(x_vector)
        )

        for missing_column in missing_columns:
            x_vector.insert(0, missing_column, [np.nan])

        x_vector = x_vector.reindex(columns=brain.dataset.columns, fill_value=0)
        x_vector = cls.sort_columns(x_vector)

        x_vector = cls.transform_columns(
            brain.dataset.imputer.transform, x_vector, numericals
        )

        return cls.sort_columns(x_vector)

    @staticmethod
    def transform_columns(
        func: Callable, data: pd.DataFrame, columns: List[str]
    ) -> pd.DataFrame():
        data_temp = data.copy()
        data_temp[columns] = func(data[columns])
        return data_temp

    @staticmethod
    def sort_columns(data: pd.DataFrame) -> pd.DataFrame:
        data_temp = data.copy()
        return data_temp.reindex(sorted(data.columns), axis=1)

    @classmethod
    def columns_intersection(
        cls,
        list_or_dataframe1: Union[pd.DataFrame, List[str]],
        list_or_dataframe2: Union[pd.DataFrame, List[str]],
    ) -> List[str]:
        set1 = cls.set_of_columns(list_or_dataframe1)
        set2 = cls.set_of_columns(list_or_dataframe2)
        return sorted(list(set.intersection(set1, set2)))

    @staticmethod
    def set_of_columns(list_or_dataframe: Union[pd.DataFrame, List[str]]) -> Set[str]:
        set_of_columns = None
        if isinstance(list_or_dataframe, pd.DataFrame):
            set_of_columns = set(list_or_dataframe.columns.values.tolist())
        elif isinstance(list_or_dataframe, list):
            set_of_columns = set(list_or_dataframe)

        return set_of_columns
