from __future__ import print_function

import sys

sys.path.append("../../tests/python")
import xgboost as xgb
import numpy as np
import unittest
from nose.plugins.attrib import attr
from sklearn.datasets import load_digits, load_boston, load_breast_cancer, make_regression

rng = np.random.RandomState(1994)


def non_increasing(L, tolerance):
    return all((y - x) < tolerance for x, y in zip(L, L[1:]))

#Check result is always decreasing and final accuracy is within tolerance
def assert_accuracy(res, tree_method, comparison_tree_method, tolerance):
    assert non_increasing(res[tree_method], tolerance)
    assert np.allclose(res[tree_method][-1], res[comparison_tree_method][-1], 1e-3, 1e-2)


def train_boston(param_in, comparison_tree_method):
    data = load_boston()
    dtrain = xgb.DMatrix(data.data, label=data.target)
    param = {}
    param.update(param_in)
    res_tmp = {}
    res = {}
    num_rounds = 10
    xgb.train(param, dtrain, num_rounds, [(dtrain, 'train')], evals_result=res_tmp)
    res[param['tree_method']] = res_tmp['train']['rmse']
    param["tree_method"] = comparison_tree_method
    xgb.train(param, dtrain, num_rounds, [(dtrain, 'train')], evals_result=res_tmp)
    res[comparison_tree_method] = res_tmp['train']['rmse']

    return res


def train_digits(param_in, comparison_tree_method):
    data = load_digits()
    dtrain = xgb.DMatrix(data.data, label=data.target)
    param = {}
    param['objective'] = 'multi:softmax'
    param['num_class'] = 10
    param.update(param_in)
    res_tmp = {}
    res = {}
    num_rounds = 10
    xgb.train(param, dtrain, num_rounds, [(dtrain, 'train')], evals_result=res_tmp)
    res[param['tree_method']] = res_tmp['train']['merror']
    param["tree_method"] = comparison_tree_method
    xgb.train(param, dtrain, num_rounds, [(dtrain, 'train')], evals_result=res_tmp)
    res[comparison_tree_method] = res_tmp['train']['merror']
    return res


def train_cancer(param_in, comparison_tree_method):
    data = load_breast_cancer()
    dtrain = xgb.DMatrix(data.data, label=data.target)
    param = {}
    param['objective'] = 'binary:logistic'
    param.update(param_in)
    res_tmp = {}
    res = {}
    num_rounds = 10
    xgb.train(param, dtrain, num_rounds, [(dtrain, 'train')], evals_result=res_tmp)
    res[param['tree_method']] = res_tmp['train']['error']
    param["tree_method"] = comparison_tree_method
    xgb.train(param, dtrain, num_rounds, [(dtrain, 'train')], evals_result=res_tmp)
    res[comparison_tree_method] = res_tmp['train']['error']
    return res


def train_sparse(param_in, comparison_tree_method):
    n = 5000
    sparsity = 0.75
    X, y = make_regression(n, random_state=rng)
    X = np.array([[np.nan if rng.uniform(0, 1) < sparsity else x for x in x_row] for x_row in X])
    dtrain = xgb.DMatrix(X, label=y)
    param = {}
    param.update(param_in)
    res_tmp = {}
    res = {}
    num_rounds = 10
    bst = xgb.train(param, dtrain, num_rounds, [(dtrain, 'train')], evals_result=res_tmp)
    res[param['tree_method']] = res_tmp['train']['rmse']
    param["tree_method"] = comparison_tree_method
    bst = xgb.train(param, dtrain, num_rounds, [(dtrain, 'train')], evals_result=res_tmp)
    res[comparison_tree_method] = res_tmp['train']['rmse']
    return res


def assert_updater_accuracy(tree_method, comparison_tree_method, variable_param, tolerance):
    param = {'tree_method': tree_method}
    for k, set in variable_param.items():
        for val in set:
            param_tmp = param.copy()
            param_tmp[k] = val
            print(param_tmp, file=sys.stderr)
            assert_accuracy(train_boston(param_tmp, comparison_tree_method), tree_method, comparison_tree_method, tolerance)
            assert_accuracy(train_digits(param_tmp, comparison_tree_method), tree_method, comparison_tree_method, tolerance)
            assert_accuracy(train_cancer(param_tmp, comparison_tree_method), tree_method, comparison_tree_method, tolerance)
            assert_accuracy(train_sparse(param_tmp, comparison_tree_method), tree_method, comparison_tree_method, tolerance)


@attr('gpu')
class TestGPU(unittest.TestCase):
    def test_gpu_hist(self):
        variable_param = {'max_depth': [2, 6, 11], 'max_bin': [2, 16, 1024], 'n_gpus': [1, -1]}
        assert_updater_accuracy('gpu_hist', 'hist', variable_param, 0.02)

    def test_gpu_exact(self):
        variable_param = {'max_depth': [2, 6, 15]}
        assert_updater_accuracy('gpu_exact', 'exact', variable_param, 0.02)

    def test_gpu_hist_experimental(self):
        variable_param = {'max_depth': [2, 6], 'max_leaves': [255, 4], 'max_bin': [2, 16, 1024]}
        assert_updater_accuracy('gpu_hist_experimental', 'hist', variable_param, 0.01)