#!/usr/bin/python
import numpy as np
import scipy.sparse
import xgboost as xgb
### simple example
# load file from text file, also binary buffer generated by xgboostdtrain = xgb.DMatrix('../data/agaricus.txt.train')

def my_int(str1):
    if "NA" in str1:
        return 0
    else:
        return int(str1)

def my_float(str1):
    if "NA" in str1:
        return 0
    else:
        return float(str1)    
    
ftr_x = open('../train_.csv', 'r')
te_mat = []
for trxline in ftr_x:
    splited = trxline.split(",")
    te_mat.append([my_int(splited[0]), my_int(splited[1]), my_int(splited[2]), my_int(splited[3]), my_int(splited[4]), \
                                my_int(splited[5]), my_int(splited[6]), my_int(splited[7]), my_int(splited[8]), my_int(splited[9]), \
                                my_float(splited[10]), my_float(splited[11]), my_float(splited[12]), my_float(splited[13]), my_float(splited[14]), \
                                my_float(splited[15]), my_float(splited[16])])

te_np_arr = np.array(te_mat)
dtrain = xgb.DMatrix(te_np_arr[:,0:16], label=te_np_arr[:,16])
#dtrain = xgb.DMatrix(te_np_arr[:,8], label=te_np_arr[:,0])
#dtrain = xgb.DMatrix(te_np_arr[:,1:10], label=te_np_arr[:,0])
# specify parameters via map, definition are same as c++ version
param = {'max_depth':2, 'eta':1, 'silent':1, 'objective':'binary:logistic' }

# specify validations set to watch performance
watchlist  = [(dtrain,'train')]
num_round = 30
bst = xgb.train(param, dtrain, num_round, watchlist)

fts = open('../test_.csv', 'r')
enroll_ids = []
ts_mat = []
for tsline in fts:
    splited = tsline.split(",")
    enroll_ids.insert(len(enroll_ids), splited[0])
    ts_mat.append([my_int(splited[3]), my_int(splited[4]), my_int(splited[5]), my_int(splited[6]), my_int(splited[7]), \
                                my_int(splited[8]), my_int(splited[9]), my_int(splited[10]), my_int(splited[13]), my_float(splited[14]), \
                                my_float(splited[15]), my_float(splited[16]), my_float(splited[17]), my_float(splited[18]), my_float(splited[19]), \
                                my_float(splited[20])])

ts_np_arr = np.array(ts_mat)
dtest = xgb.DMatrix(ts_np_arr)

# this is prediction
preds = bst.predict(dtest)

frslt = open('../test/xgb_ninjasoul_1_30.csv', 'w')
for idx in xrange(len(enroll_ids)):
    frslt.write(enroll_ids[idx] + "," + str(preds[idx]) + "\n")
