import time
import xgboost as xgb
import lightgbm as lgb
import catboost as cat
import pandas as pd
import numpy as np
from hyperopt import fmin, tpe, Trials, STATUS_OK, STATUS_FAIL, hp

import os
from os.path import join
from os import listdir
from sklearn.model_selection import StratifiedKFold, TimeSeriesSplit, KFold
from sklearn.utils import check_array, check_X_y
from copy import deepcopy


class ScikitDataset:
    def __init__(self, x, y=None):
        if isinstance(x, pd.DataFrame):
            self.x = check_array(x.values)
            self.features = x.columns
        elif isinstance(x, np.ndarray):
            self.x = check_array(x)
            self.features = ['f' + str(i+1) for i in range(0, x.shape[1])]
        else:
            raise TypeError('X should be an instance of pandas.DataFrame or numpy.ndarray')

        if isinstance(y, pd.Series):
            self.y = check_array(y.values, ensure_2d=False)
        elif isinstance(y, pd.DataFrame):
            self.y = check_array(y.iloc[:, 0].values, ensure_2d=False)
        elif isinstance(y, np.ndarray):
            self.y = check_array(y, ensure_2d=False)
        elif y is None:
            self.y = y
        else:
            raise TypeError('Y should be an instance of pandas.Series, numpy.ndarray or None for test sample')

        if y is not None:
            self.x, self.y = check_X_y(self.x, self.y)

    def get_label(self):
        return self.y


class BaseExperiment:
    def __init__(self, data_path, learning_task, eval_metric, validation_type, greater_is_better=False, n_folds=None,
                 datetime_feature=None, output_path='', validation_seed=0):
        self.learning_task = learning_task
        self.data_path = data_path
        self.eval_metric = eval_metric
        self.validation_type = validation_type
        self.n_folds = n_folds
        self.datetime_feature = datetime_feature
        self.validation_seed = validation_seed
        self.output_path = output_path
        self.greater_is_better = greater_is_better
        self.best_params = None

        if self.learning_task not in ['classification', 'regression']:
            raise ValueError('Task type must be "classification" or "regression"')
        if not callable(eval_metric):
            raise ValueError('Wrong metric provided: eval_metric must be callable')
        if validation_type == 'cv':
            if self.n_folds is None:
                raise ValueError('For "cv" validation type argument "n_folds" must be specified')
        elif validation_type == 'datetime':
            if self.datetime_feature is None:
                raise ValueError('For "datetime" validation type argument "holdout_size" must be specified')
            if self.n_folds is None:
                raise ValueError('For "datetime" validation type argument "n_folds" must be specified')
        else:
            raise ValueError('Validation type must be "cv" or "datetime"')

    def read_data(self):
        files_in_data_dir = listdir(self.data_path)
        if 'x_train.csv' not in files_in_data_dir:
            raise ValueError('File "x_train.csv" is no found in directory "{0}"'.format(self.data_path))
        else:
            self.x = pd.read_csv(join(self.data_path, 'x_train.csv'))

        if 'y_train.csv' not in files_in_data_dir:
            raise ValueError('File "y_train.csv" is no found in directory "{0}"'.format(self.data_path))
        else:
            self.y = pd.read_csv(join(self.data_path, 'y_train.csv'))
            if 'target' not in self.y.columns:
                raise ValueError('File "y_train.csv" does not have "target" columns')
            self.y = self.y.target

    def make_splits(self):
        if self.validation_type == 'cv':
            if self.learning_task == 'classification':
                self.cv = StratifiedKFold(n_splits=self.n_folds, shuffle=True, random_state=self.validation_seed)
            else:
                self.cv = KFold(n_splits=self.n_folds, shuffle=True, random_state=self.validation_seed)
        elif self.validation_type == 'datetime':
            self.x = self.x.sort_values(self.datetime_feature, ascending=True)
            self.cv = TimeSeriesSplit(n_splits=self.n_folds)
        else:
            raise ValueError('Validation type must be "cv" or "datetime"')

        cv_pairs = []
        for train_idx, test_idx in self.cv.split(self.x, self.y):
            x_train, x_test = self.x.iloc[train_idx], self.x.iloc[test_idx]
            y_train, y_test = self.y.iloc[train_idx], self.y.iloc[test_idx]
            dtrain = self._convert_to_dataset(x_train, y_train)
            dtest = self._convert_to_dataset(x_test, y_test)
            cv_pairs.append((dtrain, dtest))
        return cv_pairs

    def eval_metric(self, *args, **kwargs):
        raise NotImplementedError('Method "eval_metric" is not implemented.')

    def _fit_predict(self, dtrain, dtest, params):
        raise NotImplementedError('Method "fit" is not implemented.')

    def fit_predict(self, x_train, y_train, x_test):
        raise NotImplementedError('Method "fit" is not implemented.')

    def _convert_to_dataset(self, x, y):
        raise NotImplementedError('Method "to_dataset" is not implemented.')

    def _get_params(self):
        raise NotImplementedError('Method "_get_params" is not implemented.')

    def _get_seed_attribute_name(self):
        raise NotImplementedError('Method "_get_seed_attribute_name is not implemented"')

    def run(self, use_best_params=False, verbose=True):

        if use_best_params and self.best_params is None:
            raise AttributeError('There is no attribute "best_params". Use method "run_optimize" first')

        self.read_data()
        cv_pairs = self.make_splits()
        params = self.best_params if use_best_params else None
        result = self._run_cv(cv_pairs, None)
        if verbose:
            print('[SING]\t\teval_time={0:.2} sec\t{1}={2:.6f}'.format(
                result['eval_time'], self.eval_metric.__name__, result['loss']))
        return result

    def run_bagging(self, seeds, use_best_params=False, verbose=True):

        if use_best_params and self.best_params is None:
            raise AttributeError('There is no attribute "best_params". Use method "run_optimize" first')

        self.read_data()
        cv_pairs = self.make_splits()
        seed_name = self._get_seed_attribute_name()
        self.bagging_results = []
        self.bagging_preds = []

        for i, seed in enumerate(seeds):
            seed_params = self.best_params if use_best_params else {}
            seed_params.update({seed_name: seed})
            bag = self._run_cv(cv_pairs, seed_params)
            self.bagging_results.append(bag)
            self.bagging_preds.append(bag['preds'])
            if verbose:
                print('[{0}/{1}]\tseed={2}\teval_time={3:.2f} sec\t{4}={5:.6f}'.format(
                      i+1, len(seeds), seed, bag['eval_time'], self.eval_metric.__name__, bag['loss']))

        self.bagging_preds = np.array(self.bagging_preds).mean(axis=0).T
        self.bagging_loss = np.mean([self.eval_metric(dtest.get_label(), bag_pred)
                                     for (_, dtest), bag_pred in zip(cv_pairs, self.bagging_preds.T)])
        self.bagging_std = np.std([self.eval_metric(dtest.get_label(), bag_pred)
                                     for (_, dtest), bag_pred in zip(cv_pairs, self.bagging_preds.T)])
        self.bagging_eval_time = sum([bag['eval_time'] for bag in self.bagging_results])

        if verbose:
            print('-' * 80)
            print('[TOTAL]\t\teval_time={0:.2f} sec\t{1}={2:.6f}\tstd={3:.3f}'.format(
                  self.bagging_eval_time, self.eval_metric.__name__, self.bagging_loss, self.bagging_std))
        return self.bagging_results

    def run_optimize(self, space, max_evals=5, seed=0, verbose=True):
        self.read_data()
        cv_pairs = self.make_splits()
        self.space = space
        self.trials = Trials()
        self._hyperopt_eval_num, self.best_loss = 0, np.inf
        self._hyperopt_max_evals = max_evals
        start_time = time.time()
        _ = fmin(fn=lambda params: self._run_cv(cv_pairs, params, opti=True, opti_verbose=verbose),
                 space=space, algo=tpe.suggest, max_evals=max_evals, trials=self.trials,
                 rstate=np.random.RandomState(seed=seed))
        opti_time = time.time() - start_time
        if verbose:
            print('-' * 80)
            print('[TOTAL]\t\teval_time={0:.2f} sec\tbest={1:.6f}\n'.format(opti_time, self.best_loss))
            print(self.best_params)
        # self.best_params = self.trials.best_trial['result']['params']
        # self.best_n_estimators = self.trials.best_trial['result']['best_n_estimators']
        return [trial['result'] for trial in self.trials.trials]

    def _run_cv(self, cv_pairs, additional_params, opti=False, opti_verbose=False):
        params = self._get_params()
        if additional_params:
            params.update(additional_params)

        self.evals, start_time = [], time.time()
        preds = []
        for dtrain, dtest in cv_pairs:
            model, y_pred = self._fit_predict(dtrain, dtest, params)
            self.evals.append(self.eval_metric(dtest.get_label(), y_pred))
            preds.append(y_pred)
        eval_time = time.time() - start_time

        greater_is_better = -1 if self.greater_is_better else 1
        mean_loss = np.mean(self.evals) * greater_is_better
        result_cv = {
            'loss': mean_loss,
            'status': STATUS_FAIL if np.isnan(mean_loss) else STATUS_OK,
            'params': params,
            'eval_time': eval_time,
            'preds': preds
        }

        if opti:
            self._hyperopt_eval_num += 1
            if mean_loss < self.best_loss:
                self.best_loss = mean_loss
                self.best_params = params
            if opti_verbose:
                print('[{0}/{1}]\t\teval_time={2:.2f} sec\t{3}={4:.6f}\tbest={5:.6f}'.format(
                        self._hyperopt_eval_num, self._hyperopt_max_evals, eval_time,
                        self.eval_metric.__name__, mean_loss, self.best_loss))
        return result_cv


class ScikitExperiment(BaseExperiment):
    def __init__(self, skl_model, data_path, learning_task, eval_metric, validation_type, greater_is_better=False,
                 n_folds=None, datetime_feature=None, output_path='', validation_seed=0):
        super(ScikitExperiment, self).__init__(data_path, learning_task, eval_metric, validation_type, greater_is_better,
                                               n_folds, datetime_feature, output_path, validation_seed)
        self.model = skl_model

    def _convert_to_dataset(self, x, y=None):
        return ScikitDataset(x, y)

    def _fit_predict(self, dtrain, dtest, params=None):
        model = deepcopy(self.model)
        model = model.set_params(**params)
        model.fit(dtrain.x, dtrain.y)
        if self.learning_task == 'classification':
            y_pred = model.predict_proba(dtest.x)
            if len(set(dtrain.get_label())) == 2:
                y_pred = y_pred[:, 1]
        else:
            y_pred = model.predict(dtest.x)
        return model, y_pred

    def fit_predict(self, x_train, y_train, x_test, use_best_params=False, return_fitted_model=False):
        if use_best_params and self.best_params is None:
            raise AttributeError('There is no attribute "best_params". Use method "run_optimize" first')

        params = self.best_params if use_best_params else self._get_params()
        dtrain = self._convert_to_dataset(x_train, y_train)
        dtest = self._convert_to_dataset(x_test)
        model, y_pred = self._fit_predict(dtrain, dtest, params=params)
        if return_fitted_model:
            return y_pred, model
        else:
            return y_pred

    def _get_params(self):
        return self.model.get_params()

    def _get_seed_attribute_name(self):
        if hasattr(self.model, 'random_state'):
            return 'random_state'
        else:
            raise AttributeError('Provided SKL model doesn\'t have "random_state" attribute, can\'t perform bagging')


class BoostingExperiment(BaseExperiment):
    def __init__(self, bst_name, bst_params, bst_rounds, data_path, learning_task, eval_metric, validation_type,
                 greater_is_better=False, n_folds=None, datetime_feature=None,  output_path='', validation_seed=0):
        super(BoostingExperiment, self).__init__(data_path, learning_task, eval_metric, validation_type, greater_is_better,
                                                 n_folds, datetime_feature, output_path, validation_seed)
        self.params = bst_params
        self.name = bst_name
        self.early = int(bst_rounds / 10)
        self.rounds = bst_rounds
        self.model = None

        if bst_rounds <= 0:
            raise ValueError('bst_rounds must be >= 1')

        if bst_name not in ['xgboost', 'lightgbm', 'catboost']:
            raise ValueError('Unknown booster name {0}. Should be one of the folowing - xgboost, lightgbm, catboost')

    def _get_params(self):
        return self.params

    def _get_seed_attribute_name(self):
        return 'seed'

    def _convert_to_dataset(self, x, y=None):
        if self.name == 'xgboost':
            return xgb.DMatrix(x, label=y)
        elif self.name == 'lightgbm':
            return lgb.Dataset(x, label=y, free_raw_data=False)
        elif self.name == 'catboost':
            return cat.Pool(x, label=y)

    @property
    def _eval_metric(self):
        def __eval_metric(preds, train_dataset):
            labels = train_dataset.get_label()
            return self.eval_metric.__name__, self.eval_metric(labels, preds), self.greater_is_better
        return __eval_metric

    def _fit_predict(self, dtrain, dtest, params):
        model, y_pred = None, None
        if self.name == 'xgboost':
            model = xgb.train(params, dtrain, self.rounds, evals=(dtrain, dtest),
                              feval=self._eval_metric, early_stopping_rounds=self.early)
            y_pred = model.predict(dtest)

        elif self.name == 'lightgbm':
            evals_result = {}
            model = lgb.train(params, dtrain, self.rounds, valid_sets=dtest, valid_names='test',
                              feval=self._eval_metric, verbose_eval=0, evals_result=evals_result)
            y_pred = model.predict(dtest.data)

        elif self.name == 'catboost':
            pass

        return model, y_pred

        # if self.learning_task == 'classification':
        #     y_pred = model.predict_proba(dtest.x)
        #     if len(set(dtrain.get_label())) == 2:
        #         y_pred = y_pred[:, 1]
        # else:
        #     y_pred = model.predict(dtest.x)
