# Path: DSTK/GAM/gam.py
# def _recurse(tree, feature_vec):
# def _get_sum_of_gamma_correction(tree, data, labels, class_weights, feature_name):
# def _get_shape_for_attribute(attribute_data, labels, class_weights, feature_name, criterion, splitter,
#                              max_depth, min_samples_split, min_samples_leaf, min_weight_fraction_leaf,
#                              max_features, random_state, max_leaf_nodes, presort):
#     def __init__(self, **kwargs):
#     def _get_metadata_dict(self):
#     def _train_cost(self, data, labels):
#     def _get_pseudo_responses(self, data, labels):
#     def _init_shapes_and_data(self, data, labels):
#     def _update_learning_rate(self, dct, epoch):
#     def _initialize_class_weights(self, labels):
#     def _get_class_weights(self, labels):
#     def train(self, data, targets, **kwargs):
#     def _get_trimmed_record_indices(self, responses):
#     def _calculate_gradient_shape(self, data, labels, bag_indices=None, max_workers=1):
#     def __init__(self, gam):
#     def __getattr__(self, item):
#     def smoothen(self, data, penalty=None):
#     def _create_smooth_shape(shape, values, name, penalties):
#     def _fit_spline(values, target_vals, penalties):
# class GAM(BaseGAM):
# class SmoothGAM(BaseGAM):
# 
# Path: DSTK/tests/tests_gam/test_shape_function.py
# def _create_partition(lst_of_splits):
#     return np.append(lst_of_splits, np.PINF)
# 
# Path: DSTK/GAM/base_gam.py
# def load_from_tar(file_name):
# 
#     model_name = file_name.split('/')[-1].replace('.tar.gz', '')
# 
#     gam = DeserializedGAM()
# 
#     with tf.open(file_name, "r:gz") as tar:
#         f = tar.extractfile('{}/{}.json'.format(model_name, model_name))
#         content = js.loads(f.read())
#         gam._n_features = content['num_features']
#         gam.feature_names = content['feature_names']
# 
#         for member in tar.getmembers():
#             if member.isfile() and (member.name != '{}/{}.json'.format(model_name, model_name)):
#                 f = tar.extractfile(member.path)
#                 content = js.loads(f.read())
#                 gam.shapes.update({content['feature_name']: ShapeFunction(content['splits'],
#                                                                           content['values'],
#                                                                           content['feature_name'])})
# 
#     assert set(gam.shapes.keys()) == set(gam.feature_names), 'feature names and shape names do not match up'
#     return gam


from sklearn import datasets as ds
from DSTK.GAM.gam import GAM, ShapeFunction
from DSTK.tests.tests_gam.test_shape_function import _create_partition
from DSTK.GAM.base_gam import load_from_tar
import numpy as np
import os
import shutil
data = cancer_ds['data'][:, :20]
labels = 2 * cancer_ds['target'] - 1

assert_scores = [
    [0.5538394842641805, 0.44616051573581944],
    [0.49861290044203543, 0.5013870995579646],
    [0.5470227126670573, 0.4529772873329428],
    [0.513940794277825, 0.48605920572217504],
    [0.529758125364891, 0.470241874635109]
]


test_root_folder = '/tmp/test_gam_serialization'

def teardown():
    if os.path.exists(test_root_folder):
        shutil.rmtree(test_root_folder)

def test_gam_training():
    gam = GAM(max_depth=3, max_leaf_nodes=5, random_state=42, balancer_seed=42)
    gam.train(data, labels, n_iter=5, learning_rate=0.0025, num_bags=1, num_workers=3)

    for idx, vec in enumerate(data[:5, :]):
        gam_scores = gam.score(vec)
        np.testing.assert_almost_equal(np.sum(gam_scores), 1.0, 10)
        np.testing.assert_almost_equal(gam_scores, assert_scores[idx], 10)


def test_correct_scoring():
    func1 = ShapeFunction(_create_partition(np.linspace(1, 10, 10)), np.linspace(1, 10, 11), 'attr_1')

# Path: DSTK/GAM/gam.py
# def _recurse(tree, feature_vec):
# def _get_sum_of_gamma_correction(tree, data, labels, class_weights, feature_name):
# def _get_shape_for_attribute(attribute_data, labels, class_weights, feature_name, criterion, splitter,
#                              max_depth, min_samples_split, min_samples_leaf, min_weight_fraction_leaf,
#                              max_features, random_state, max_leaf_nodes, presort):
#     def __init__(self, **kwargs):
#     def _get_metadata_dict(self):
#     def _train_cost(self, data, labels):
#     def _get_pseudo_responses(self, data, labels):
#     def _init_shapes_and_data(self, data, labels):
#     def _update_learning_rate(self, dct, epoch):
#     def _initialize_class_weights(self, labels):
#     def _get_class_weights(self, labels):
#     def train(self, data, targets, **kwargs):
#     def _get_trimmed_record_indices(self, responses):
#     def _calculate_gradient_shape(self, data, labels, bag_indices=None, max_workers=1):
#     def __init__(self, gam):
#     def __getattr__(self, item):
#     def smoothen(self, data, penalty=None):
#     def _create_smooth_shape(shape, values, name, penalties):
#     def _fit_spline(values, target_vals, penalties):
# class GAM(BaseGAM):
# class SmoothGAM(BaseGAM):
# 
# Path: DSTK/tests/tests_gam/test_shape_function.py
# def _create_partition(lst_of_splits):
#     return np.append(lst_of_splits, np.PINF)
# 
# Path: DSTK/GAM/base_gam.py
# def load_from_tar(file_name):
# 
#     model_name = file_name.split('/')[-1].replace('.tar.gz', '')
# 
#     gam = DeserializedGAM()
# 
#     with tf.open(file_name, "r:gz") as tar:
#         f = tar.extractfile('{}/{}.json'.format(model_name, model_name))
#         content = js.loads(f.read())
#         gam._n_features = content['num_features']
#         gam.feature_names = content['feature_names']
# 
#         for member in tar.getmembers():
#             if member.isfile() and (member.name != '{}/{}.json'.format(model_name, model_name)):
#                 f = tar.extractfile(member.path)
#                 content = js.loads(f.read())
#                 gam.shapes.update({content['feature_name']: ShapeFunction(content['splits'],
#                                                                           content['values'],
#                                                                           content['feature_name'])})
# 
#     assert set(gam.shapes.keys()) == set(gam.feature_names), 'feature names and shape names do not match up'
#     return gam


# Path: DSTK/tests/tests_gam/test_gam.py
from sklearn import datasets as ds
from DSTK.GAM.gam import GAM, ShapeFunction
from DSTK.tests.tests_gam.test_shape_function import _create_partition
from DSTK.GAM.base_gam import load_from_tar
import numpy as np
import os
import shutil

data = cancer_ds['data'][:, :20]
labels = 2 * cancer_ds['target'] - 1

assert_scores = [
    [0.5538394842641805, 0.44616051573581944],
    [0.49861290044203543, 0.5013870995579646],
    [0.5470227126670573, 0.4529772873329428],
    [0.513940794277825, 0.48605920572217504],
    [0.529758125364891, 0.470241874635109]
]


test_root_folder = '/tmp/test_gam_serialization'

def teardown():
    if os.path.exists(test_root_folder):
        shutil.rmtree(test_root_folder)

def test_gam_training():
    gam = GAM(max_depth=3, max_leaf_nodes=5, random_state=42, balancer_seed=42)
    gam.train(data, labels, n_iter=5, learning_rate=0.0025, num_bags=1, num_workers=3)

    for idx, vec in enumerate(data[:5, :]):
        gam_scores = gam.score(vec)
        np.testing.assert_almost_equal(np.sum(gam_scores), 1.0, 10)
        np.testing.assert_almost_equal(gam_scores, assert_scores[idx], 10)


def test_correct_scoring():
    func1 = ShapeFunction(_create_partition(np.linspace(1, 10, 10)), np.linspace(1, 10, 11), 'attr_1')

    func2 = ShapeFunction(_create_partition(np.linspace(1, 10, 10)), -1 * np.linspace(1, 10, 11), 'attr_2')