# -*- coding: utf-8 -*-
#!/usr/bin/python

import os
import sys
##########################################################
# usage(run this file before any feature extract function file)
# pypy get_id.py  ../../data/train ./data/test ./data/trainLabels.csv

# xid_train.p is a list like ['loIP1tiwELF9YNZQjSUO',''....] to specify
# the order of samples in traing data
# ../../data/train is the path of original train data
# ./ins_train is where the local folder of ins.p, {address:ins}
# ./jump_train is where the local folder of jmp.p, {address:jump ins}
##########################################################

trainDatesetRoot = sys.argv[1]
testDatesetRoot = sys.argv[2]
trainLabelRoot = sys.argv[3]
testLabelRoot = sys.argv[4]

xid=[i.split('.')[0] for i in os.listdir(trainDatesetRoot) if '.asm' in i]
Xt_id=[i.split('.')[0] for i in os.listdir(testDatesetRoot) if '.asm' in i]

#train label
f=open(trainLabelRoot)
f.readline()
label={}
for line in f:
    xx=line.split(',')
    label[xx[0]]=int(xx[-1])
f.close()
y=[label[i] for i in xid]

#test lable
f=open(testLabelRoot)
f.readline()
tlabel={}
for line in f:
    xx=line.split(',')
    tlabel[xx[0]]=int(xx[-1])
f.close()
yt=[tlabel[i] for i in Xt_id]

import pickle
pickle.dump(xid,open('xid.p','wb'))
pickle.dump(Xt_id,open('Xt_id.p','wb'))
pickle.dump(y,open('y.p','wb'))  #实际用于测试的数据
pickle.dump(yt,open('yt.p','wb'))  #实际用于测试的数据

