#-----------------------------------------------------------------------------------------
# modules
#-----------------------------------------------------------------------------------------
from   sklearn.metrics import accuracy_score # convenient accuracy measure
import xgboost         as     xgb            # XGBoost software package
import h5py                                  # deal with .mat data files
import sys                                   # deal with command arguments
import gc; gc.enable();                      # garbage memory collection 


#-------------------------------------------------------------------------------
# functions
#-------------------------------------------------------------------------------
def serialise_model(__pne__, __model__):
    __model__.save_model(__pne__); gc.collect()

def deserialise_model(__pne__):
    __model__ = xgb.Booster(); xgb.Booster.load_model(__model__, __pne__)
    return __model__

def load_mat(__pne__):  # load .mat data and return xgboost.DMatrix
    print('load_mat(%s): Commence.' % __pne__, flush = 1)
    Mat = h5py.File(__pne__)
    X__ = (Mat.get('X')[()]).transpose();
    y__ = (Mat.get('y')[()]).transpose().flatten(); del Mat;
    Xy_ = xgb.DMatrix(data = X__, label = y__)
    __shape__ = X__.shape
    print('load_mat(): DMatrix ready. X__.shape = ', end = '', flush = 1)
    print(__shape__, flush = 1)
    del X__; gc.collect();
    return Xy_, y__, __shape__;

def set_model_params(__max_depth__, __num_class__): 
    xgb_params = {
        'process_type':               'default',        # \in {default, update} for the continuation of a tree.
        'tree_method':                'gpu_hist',
        'booster':                    'gbtree',
        'grow_policy':                'lossguide',      # lossguide: split at nodes with highest loss change.
        'num_parallel_tree':          1,                # number of trees in boosted random forest.
        'min_split_loss':             0,                # 0: split whenver it improves.
        'learning_rate':              1.0,              # 1: no decay. weight on previous trees.
        'max_depth':                  __max_depth__,    # \in [0, Inf], 0 accepted if (hist | lossguided). cf. memory (!)
        'max_leaves':                 0,
        'reg_lambda':                 0.0,              # L2-regularisation
        'reg_alpha':                  0.0,              # L1-regularisation
        'num_class':                  __num_class__,
        'objective':                  'multi:softmax',
        'eval_metric':                'merror',
        'predictor':                  'gpu_predictor',
        'verbosity':                  1,
        'validate_parameters':        0,
        'single_precision_histogram': 0,                # gpu_hist can fail with single precision
        'deterministic_histogram':    0                 # default 1: rounds losing accuracy
    }

    return xgb_params


#-----------------------------------------------------------------------------------------
# data
#-----------------------------------------------------------------------------------------
data_pne         = '/workspace/temp/data.mat'
my_model_path    = '/workspace/temp/'
num_classes      = 3
my_max_depth     = 21 # 21
print('from: %s' % (data_pne))
print('to:   %s' % (my_model_path))


#-----------------------------------------------------------------------------------------
# user parameters
#-----------------------------------------------------------------------------------------
if len(sys.argv) != 2: 
    print('Give number of rounds');
    exit()

dump_model    = 1
my_print_freq = 1
xgb_params    = set_model_params(my_max_depth, num_classes)


#-----------------------------------------------------------------------------------------
# ----------------------------------------------------- using xgboost-core
#-----------------------------------------------------------------------------------------
Xy, y, shape  = load_mat(data_pne)
progress      = dict()
my_num_rounds = int(sys.argv[1])
my_model_pne  = '%s/%03d.json' % (my_model_path, my_num_rounds)

try:
    M0 = deserialise_model(my_model_pne)
    print('Existing model detected.\n')

except:
    M0 = None; print('Model will be newly trained.\n')
    M0 = xgb.train(xgb_params, Xy, my_num_rounds, [(Xy, 'convergence')], 
        evals_result = progress,
        verbose_eval = my_print_freq,
        xgb_model    = M0)

## report result
y0  = np.ubyte(M0.predict(Xy))
ACC = accuracy_score(y0, np.ubyte(y))
print('\nAccuracy = %6.2f%%' % (ACC * 100.0) )

## secure memory as much as possible
del Xy; del y; gc.collect()

## saving the trained model
if dump_model == 1:
    print('Dumping model M0: ', end = '', flush = 1);
    try:
        serialise_model(my_model_pne, M0)
    except:
        print('serialisation failure.')