import logging

import numpy as np

from autogluon.common.features.types import R_CATEGORY, R_FLOAT, R_INT, R_OBJECT, S_BOOL

from .abstract_model import AbstractModel

logger = logging.getLogger(__name__)


class AbstractNeuralNetworkModel(AbstractModel):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._types_of_features = None

    # TODO: v0.1 clean method
    def _get_types_of_features(
        self, df, skew_threshold=None, embed_min_categories=None, use_ngram_features=None, needs_extra_types=True
    ):
        """Returns dict with keys: : 'continuous', 'skewed', 'onehot', 'embed', 'language', values = ordered list of feature-names falling into each category.
        Each value is a list of feature-names corresponding to columns in original dataframe.
        TODO: ensure features with zero variance have already been removed before this function is called.
        """
        if self._types_of_features is not None:
            Warning("Attempting to _get_types_of_features for Model, but previously already did this.")

        continuous_featnames = self._feature_metadata.get_features(
            valid_raw_types=[R_INT, R_FLOAT], invalid_special_types=[S_BOOL]
        )
        categorical_featnames = self._feature_metadata.get_features(valid_raw_types=[R_CATEGORY, R_OBJECT])
        bool_featnames = self._feature_metadata.get_features(required_special_types=[S_BOOL])

        language_featnames = []  # TODO: not implemented. This should fetch text features present in the data
        valid_features = categorical_featnames + continuous_featnames + bool_featnames + language_featnames

        if len(valid_features) < df.shape[1]:
            unknown_features = [feature for feature in df.columns if feature not in valid_features]
            logger.log(15, f"Model will additionally ignore the following columns: {unknown_features}")
            df = df.drop(columns=unknown_features)
            self._features_internal = list(df.columns)

        self.features_to_drop = df.columns[
            df.isna().all()
        ].tolist()  # drop entirely NA columns which may arise after train/val split
        if self.features_to_drop:
            logger.log(15, f"Model will additionally ignore the following columns: {self.features_to_drop}")
            df = df.drop(columns=self.features_to_drop)

        if needs_extra_types is True:
            types_of_features = {"continuous": [], "skewed": [], "onehot": [], "embed": [], "language": [], "bool": []}
            # continuous = numeric features to rescale
            # skewed = features to which we will apply power (ie. log / box-cox) transform before normalization
            # onehot = features to one-hot encode (unknown categories for these features encountered at test-time are encoded as all zeros). We one-hot encode any features encountered that only have two unique values.
            features_to_consider = [feat for feat in self._features_internal if feat not in self.features_to_drop]
            for feature in features_to_consider:
                feature_data = df[feature]  # pd.Series
                num_unique_vals = len(feature_data.unique())
                if feature in bool_featnames:
                    types_of_features["bool"].append(feature)
                elif num_unique_vals == 2:  # will be onehot encoded regardless of proc.embed_min_categories value
                    types_of_features["onehot"].append(feature)
                elif feature in continuous_featnames:
                    if np.abs(feature_data.skew()) > skew_threshold:
                        types_of_features["skewed"].append(feature)
                    else:
                        types_of_features["continuous"].append(feature)
                elif feature in categorical_featnames:
                    if (
                        num_unique_vals >= embed_min_categories
                    ):  # sufficiently many categories to warrant learned embedding dedicated to this feature
                        types_of_features["embed"].append(feature)
                    else:
                        types_of_features["onehot"].append(feature)
                elif feature in language_featnames:
                    types_of_features["language"].append(feature)
        else:
            types_of_features = []
            for feature in valid_features:
                if feature in categorical_featnames:
                    feature_type = "CATEGORICAL"
                elif feature in continuous_featnames or feature in bool_featnames:
                    feature_type = "SCALAR"
                elif feature in language_featnames:
                    feature_type = "TEXT"
                else:
                    raise ValueError(f"Invalid feature: {feature}")

                types_of_features.append({"name": feature, "type": feature_type})

        return types_of_features, df
