import importlib
import inspect
import copy
import random
import string

import torch
import numpy as np
import pandas as pd
from torch.utils.data import Dataset

from lightwood.mixers.helpers.transformer import Transformer
from lightwood.encoders.time_series.helpers.common import generate_target_group_normalizers
from lightwood.config.config import CONFIG
from lightwood.constants.lightwood import ColumnDataTypes
from lightwood.encoders import (
    NumericEncoder,
    TsNumericEncoder,
    CategoricalAutoEncoder,
    MultiHotEncoder,
    DistilBertEncoder,
    DatetimeEncoder,
    Img2VecEncoder,
    TsRnnEncoder,
    TimeSeriesPlainEncoder,
    ShortTextEncoder,
    VocabularyEncoder
)


class SubSet(Dataset):
    def __init__(self, data_source, indexes):
        self.data_source = data_source
        self.index_mapping = {}
        for i in range(len(indexes)):
            self.index_mapping[i] = indexes[i]

    def __len__(self):
        return int(len(self.index_mapping.keys()))

    def __getitem__(self, idx):
        return self.data_source[self.index_mapping[idx]]

    def get_feature_names(self, where='input_features'):
        return self.data_source.get_feature_names(where)

    def __getattribute__(self, name):
        if name in ['config', 'encoders', 'transformer', 'training',
                    'output_weights', 'dropout_dict', 'disable_cache', 'out_types', 'out_indexes']:
            return self.data_source.__getattribute__(name)
        else:
            return object.__getattribute__(self, name)

    def __setattr__(self, name, value):
        if name in ['config', 'encoders', 'transformer', 'training',
                    'output_weights', 'dropout_dict', 'disable_cache']:
            return dict.__setattr__(self.data_source, name, value)
        else:
            super().__setattr__(name, value)


class DataSource(Dataset):
    def __init__(self,
                 data_frame,
                 config,
                 prepare_encoders=True,
                 initialize_transformer=True):
        """
        Create a lightwood datasource from the data frame
        :param data_frame:
        :param config
        """
        self.subsets = {}
        self.data_frame = data_frame
        self.config = config
        self.training = False  # Flip this flag if you are using the datasource while training
        self.output_weights = None
        self.output_weights_offset = None
        self.dropout_dict = {}
        self.enable_dropout = False
        self.enable_cache = self.config['data_source']['cache_transformed_data']
        self.out_indexes = None

        self.out_types = [feature['type'] for feature in self.config['output_features']]
        self.output_feature_names = [feature['name'] for feature in self.config['output_features']]
        self.input_feature_names = [feature['name'] for feature in self.config['input_features']]
        self.output_features = self.config['output_features']
        self.input_features = self.config['input_features']

        for col in self.input_features:
            if len(self.input_features) > 1:
                dropout = 0.1
            else:
                dropout = 0.0

            if 'dropout' in col:
                dropout = col['dropout']

            self.dropout_dict[col['name']] = dropout

        for col in self.config['output_features']:
            prev_col = [c for c in self.config['input_features'] if c['name'] == f'__mdb_ts_previous_{col["name"]}']
            # len(prev_col) <= 1 always holds
            if prev_col:
                col['group_info'] = {
                    prev_col['name']: self.get_column_original_data(prev_col['name'])
                    for prev_col in self.config['input_features'] if prev_col['grouped_by']
                }
        
        if initialize_transformer:
            self.transformer = Transformer(self.input_feature_names, self.output_feature_names)
        else:
            self.transformer = None

        if prepare_encoders:
            self.encoders = self._prepare_encoders()
        else:
            self.encoders = None
        
        self._clear_cache()

    def extend(self, df):
        """
        :df: pandas.DataFrame or DataSource
        """
        if isinstance(df, pd.DataFrame):
            self.data_frame.append(df)
        else:
            raise TypeError(':df: must be pandas.DataFrame')

    def create_subsets(self, nr_subsets):
        subsets_indexes = {}
        np.random.seed(len(self.data_frame))

        subset_nr = 1
        for i in range(len(self.data_frame)):
            if subset_nr not in subsets_indexes:
                subsets_indexes[subset_nr] = []
            subsets_indexes[subset_nr].append(i)

            subset_nr += 1
            if subset_nr > nr_subsets:
                subset_nr = 1

        for subset_nr in subsets_indexes:
            self.subsets[subset_nr] = SubSet(self, subsets_indexes[subset_nr])

    def _clear_cache(self):
        self.encoded_cache = {}
        self.transformed_cache = None

    def subset(self, percentage):
        """
        Removes :percentage: of data randomly from itself and returns a new
        datasource made from that data

        :param percentage: float

        :return: DataSource
        """
        np.random.seed(int(round(percentage * 100000)))

        msk = np.random.rand(len(self.data_frame)) < (1 - percentage)

        sub_df = self.data_frame[~msk]
        self.data_frame = self.data_frame[msk]

        self._clear_cache()

        ds = DataSource(
            sub_df,
            self.config,
            prepare_encoders=False,
            initialize_transformer=False
        )
        ds.encoders = self.encoders
        ds.transformer = self.transformer
        ds.output_weights = self.output_weights

        return ds

    def __len__(self):
        """
        return the length of the datasource (as in number of rows)
        :return: number of rows
        """
        return int(self.data_frame.shape[0])

    def __getitem__(self, idx):
        sample = {}

        dropout_features = None

        if self.training and random.randint(0, 3) == 1 and self.enable_dropout and CONFIG.ENABLE_DROPOUT:
            dropout_features = [feature['name'] for feature in self.config['input_features'] if random.random() > (1 - self.dropout_dict[feature['name']])]

            # Make sure we never drop all the features, since this would make the row meaningless
            if len(dropout_features) > len(self.config['input_features']):
                dropout_features = dropout_features[:-1]
            #logging.debug(f'\n-------------\nDroping out features: {dropout_features}\n-------------\n')

        if self.enable_cache:
            if self.transformed_cache is None:
                self.transformed_cache = [None] * len(self)

            if dropout_features is None or len(dropout_features) == 0:
                cached_sample = self.transformed_cache[idx]
                if cached_sample is not None:
                    return cached_sample

        for feature_set in ['input_features', 'output_features']:
            sample[feature_set] = {}
            for feature in self.config[feature_set]:
                col_name = feature['name']
                col_config = self.get_column_config(col_name)

                if col_name not in self.encoded_cache:  # if data is not encoded yet, encode values
                    if not ((dropout_features is not None and col_name in dropout_features) or not self.enable_cache):
                        self.get_encoded_column_data(col_name)

                # if we are dropping this feature, get the encoded value of None
                if (dropout_features is not None and col_name in dropout_features):
                    custom_data = {col_name: [None]}
                    # if the dropout feature depends on another column, also pass a None array as the dependant column
                    if 'depends_on_column' in col_config:
                        for col in col_config['depends_on_column']:
                            custom_data[col] = [None]
                    sample[feature_set][col_name] = self.get_encoded_column_data(col_name, custom_data=custom_data)[0]
                elif not self.enable_cache:
                    if col_name in self.data_frame:
                        custom_data = {col_name: [self.data_frame[col_name].iloc[idx]]}
                        if 'depends_on_column' in col_config:
                            for col in col_config['depends_on_column']:
                                custom_data[col] = [self.data_frame[col].iloc[idx]]
                    else:
                        custom_data = {col_name: [None]}

                    sample[feature_set][col_name] = self.get_encoded_column_data(col_name, custom_data=custom_data)[0]
                else:
                    sample[feature_set][col_name] = self.encoded_cache[col_name][idx]

        # Create weights if not already created
        if self.output_weights is None:
            for col_config in self.config['output_features']:
                if 'weights' in col_config:

                    weights = col_config['weights']
                    new_weights = None

                    for val in weights:
                        encoded_val = self.get_encoded_column_data(
                            col_config['name'],
                            custom_data={
                                col_config['name']: [val]
                            }
                        )

                        if new_weights is None:
                            mean = np.mean(list(weights.values()))
                            new_weights = [mean] * len(encoded_val[0])

                        for value_index, hot in enumerate(encoded_val[0]):
                            if hot:
                                new_weights[value_index] = weights[val] / sum(encoded_val[0]).item()

                    if self.output_weights is None or self.output_weights == False:
                        self.output_weights = new_weights
                    else:
                        self.output_weights.extend(new_weights)
                else:
                    self.output_weights = False

        sample = self.transformer.transform(sample)
        if self.out_indexes is None:
            self.out_indexes = self.transformer.out_indexes

        if self.enable_cache:
            self.transformed_cache[idx] = sample

        if self.output_weights_offset is None:
            self.output_weights_offset = {}
            for idx, (otype, oidxs) in enumerate(zip(self.out_types, self.out_indexes)):
                self.output_weights_offset[idx] = self.output_weights_offset.get(idx-1, 0)
                if otype not in (ColumnDataTypes.CATEGORICAL, ColumnDataTypes.MULTIPLE_CATEGORICAL):
                     self.output_weights_offset[idx] += (oidxs[1] - oidxs[0])

        return sample

    def get_column_original_data(self, column_name):
        if column_name not in self.data_frame:
            nr_rows = self.data_frame.shape[0]
            return [None] * nr_rows

        return self.data_frame[column_name].tolist()

    def _lookup_encoder_class(self, column_config, is_target):
        default_encoder_classes = {
            ColumnDataTypes.NUMERIC: NumericEncoder,
            ColumnDataTypes.CATEGORICAL: CategoricalAutoEncoder,
            ColumnDataTypes.MULTIPLE_CATEGORICAL: MultiHotEncoder,
            ColumnDataTypes.DATETIME: DatetimeEncoder,
            ColumnDataTypes.IMAGE: Img2VecEncoder,
            ColumnDataTypes.TEXT: DistilBertEncoder,
            ColumnDataTypes.SHORT_TEXT: ShortTextEncoder,
            ColumnDataTypes.TIME_SERIES: TsRnnEncoder,
            # ColumnDataTypes.AUDIO: AmplitudeTsEncoder
        }

        target_encoder_classes = {
            ColumnDataTypes.TEXT: VocabularyEncoder
        }

        column_type = column_config['type']
        if is_target and column_type in target_encoder_classes:
            encoder_class = target_encoder_classes[column_type]
        else:
            encoder_class = default_encoder_classes[column_type]

        # specialized dispatches
        if is_target and column_config.get('additional_info', {}).get('time_series_target', False) and \
                column_type == ColumnDataTypes.NUMERIC:
            encoder_class = TsNumericEncoder
        if column_type == ColumnDataTypes.TIME_SERIES and '__mdb_ts_previous' in column_config['name']:
            encoder_class = TimeSeriesPlainEncoder

        return encoder_class

    def _make_column_encoder(self, encoder_class, encoder_attrs=None, is_target=False):
        encoder_instance = encoder_class(is_target=is_target)
        encoder_attrs = encoder_attrs or {}
        for attr in encoder_attrs:
            if hasattr(encoder_instance, attr):
                setattr(encoder_instance, attr, encoder_attrs[attr])
        return encoder_instance

    def _prepare_column_encoder(self, config, is_target=False, training_data=None):
        column_data = self.get_column_original_data(config['name'])
        encoder_class = config.get(
            'encoder_class',
            self._lookup_encoder_class(config, is_target)
        )
        encoder_attrs = config.get('encoder_attrs', {})
        encoder_attrs['original_type'] = config.get('original_type', None)
        encoder_attrs['secondary_type'] = config.get('secondary_type', None)

        encoder_instance = self._make_column_encoder(
            encoder_class,
            encoder_attrs,
            is_target=is_target
        )

        if training_data and 'training_data' in inspect.getfullargspec(encoder_instance.prepare).args:
            encoder_instance.prepare(
                column_data,
                training_data=training_data
            )
        else:
            # joint column data augmentation for time series
            if config['type'] == ColumnDataTypes.TIME_SERIES and not is_target and '__mdb_ts_previous' not in config['name']:
                encoder_instance.prepare(column_data, previous_target_data=training_data['previous'])
            else:
                encoder_instance.prepare(column_data)

        return encoder_instance

    def _prepare_encoders(self):
        """
        Get the encoder for all the output and input column and prepare them
        with all available data for that column.
        """
        encoders = {}
        input_encoder_training_data = {'targets': [], 'previous': []}

        # handle time series previous target columns
        previous_cols = []
        for config in self.config['input_features']:
            column_name = config['name']
            if column_name.startswith('__mdb_ts_previous_'):
                column_data = self.get_column_original_data(column_name)
                previous_cols.append(column_name)
                col_info = {'data': column_data,
                            'name': column_name,
                            'original_type': config['original_type'],
                            'group_info': {conf['name']: self.get_column_original_data(conf['name'])
                                           for conf in self.config['input_features'] if conf['grouped_by']},
                            'output_type': config['type']
                            }
                col_info = generate_target_group_normalizers(col_info)
                input_encoder_training_data['previous'].append(col_info)
                config['additional_info']['normalizers'] = col_info['normalizers']
                config['additional_info']['group_combinations'] = col_info['group_combinations']

        # train output encoders (except for time-series related)
        for config in self.config['output_features']:
            if config['type'] not in ColumnDataTypes.TIME_SERIES:
                column_name = config['name']
                column_data = self.get_column_original_data(column_name)
                encoder_instance = self._prepare_column_encoder(config, is_target=True)

                input_encoder_training_data['targets'].append({
                    'encoded_output': encoder_instance.encode(column_data),
                    'unencoded_output': column_data,
                    'output_encoder': encoder_instance,
                    'output_type': config['type']
                })

                encoders[column_name] = encoder_instance

        # train input encoders
        for config in self.config['input_features']:
            column_name = config['name']
            encoder_instance = self._prepare_column_encoder(config,
                                                            is_target=False,
                                                            training_data=input_encoder_training_data)
            encoders[column_name] = encoder_instance

            if column_name not in previous_cols:
                if config['type'] == ColumnDataTypes.TIME_SERIES and len(input_encoder_training_data['previous']) > 0:
                    for d in input_encoder_training_data['previous']:
                        try:
                            if not d['name'] in config['depends_on_column']:
                                config['depends_on_column'].append(d['name'])
                        except KeyError:
                            config['depends_on_column'] = [d['name']]

        # train time series output encoder
        for config in self.config['output_features']:
            column_name = config['name']

            prev_col = [c for c in self.config['input_features'] if
                        c['name'] == f'__mdb_ts_previous_{column_name}']
            if prev_col:
                # 0-indexed as we know it can only be one
                config['depends_on_column'] = [prev_col[0]['name']]
                config['encoder_attrs']['normalizers'] = prev_col[0]['additional_info']['normalizers']
                config['encoder_attrs']['group_combinations'] = prev_col[0]['additional_info'][
                    'group_combinations']

            encoder_instance = self._prepare_column_encoder(config, is_target=True)
            encoders[column_name] = encoder_instance

        return encoders

    def make_child(self, df):
        """
        :param df: DataFrame

        :return: DataSource
        """
        child = DataSource(
            df,
            self.config,
            prepare_encoders=False,
            initialize_transformer=False
        )
        child.transformer = self.transformer
        child.encoders = self.encoders
        child.output_weights = self.output_weights
        return child

    def get_encoded_column_data(self, column_name, custom_data=None):
        if column_name in self.encoded_cache and custom_data is None:
            return self.encoded_cache[column_name]

        # The first argument of encoder is the data, if no custom data is specified,
        # use all the datasource's data for this column
        if custom_data is not None:
            args = [custom_data[column_name]]
        else:
            args = [self.get_column_original_data(column_name)]

        config = self.get_column_config(column_name)

        # Pass dependency info when applicable
        encoder_arity = len(inspect.signature(self.encoders[column_name].encode).parameters)
        if encoder_arity >= 2 and 'depends_on_column' in config:
            arg2 = []
            for col in config['depends_on_column']:
                sublist = {'group_info': {conf['name']: self.get_column_original_data(conf['name'])
                                          for conf in self.config['input_features'] if conf['grouped_by']},
                           'name': col}
                if custom_data is None:
                    sublist['data'] = self.get_column_original_data(col)
                else:
                    sublist['data'] = custom_data.get(col, None)
                arg2.append(sublist)
            args.append(arg2)

        encoded_vals = self.encoders[column_name].encode(*args)
        if isinstance(encoded_vals, list) and isinstance(encoded_vals[0], torch.Tensor):
            encoded_vals = torch.stack(encoded_vals)

        # Cache the encoded data so we don't have to run the encoding,
        # Don't cache custom_data
        # (custom_data is usually used when running without cache or dropping out a feature for a certain pass)
        if column_name not in self.encoded_cache and custom_data is None:
            self.encoded_cache[column_name] = encoded_vals
        return encoded_vals

    def get_decoded_column_data(self, column_name, encoded_data, decoder_instance=None):
        """
        :param column_name: column names to be decoded
        :param encoded_data: encoded data of tensor type
        :return decoded_data : Dict :Decoded data of input column
        """
        if decoder_instance is None:
            if column_name not in self.encoders:
                raise ValueError("""
                    Data must have been encoded before at some point.
                    You should not decode before having encoding at least once
                    """)
            decoder_instance = self.encoders[column_name]

        decoded_data = {}
        if getattr(decoder_instance, 'predict_proba', False):
            # return complete belief distribution
            preds, pred_probs, labels = decoder_instance.decode(encoded_data)
            decoded_data['predictions'] = preds
            decoded_data['class_distribution'] = pred_probs
            decoded_data['class_labels'] = labels
        elif isinstance(decoder_instance, TsNumericEncoder):
            group_info = {conf['name']: self.get_column_original_data(conf['name'])
                          for conf in self.config['input_features'] if conf.get('grouped_by', False)}
            decoded_data['predictions'] = decoder_instance.decode(encoded_data, group_info=group_info)
        else:
            decoded_data['predictions'] = decoder_instance.decode(encoded_data)

        return decoded_data

    def get_feature_names(self, where='input_features'):
        return [feature['name'] for feature in self.config[where]]

    def get_column_config(self, column_name):
        """
        Get the config info for the feature given a config as defined in data_schemas definition.py
        :param column_name:
        :return:
        """
        for feature_set in ['input_features', 'output_features']:
            for feature in self.config[feature_set]:
                if feature['name'] == column_name:
                    return feature
        return None

    def train(self):
        self.training = True
    
    def eval(self):
        self.training = False
