# -*- coding: utf-8 -*-
#!/usr/bin/python
'''
    此文件的功能是对数据库进行抽取，形成一个小的测试集
'''
from csv import DictReader
import os

DataRoot = r'E:\kiggle\kiggleData'

labelRoot = DataRoot + '\\'+ 'subTrainLabels.csv'

import shutil
#生成新的label文件
from csv import writer
count = [0,0,0]
with open(DataRoot + '\\sample\\subTrainLabels.csv', 'w',newline='') as f:
    fw = writer(f)
    colnames = ['Id','Class']
    fw.writerow(colnames)
    for row in DictReader(open(labelRoot)):
        label = int(row['Class'])
        if(count[label-1] < 4):
            count[label -1] += 1
            fw.writerow([row['Id'],label])
            shutil.move(DataRoot + '\\train\\'+row['Id'] + '.asm', DataRoot + '\\sample\\train\\'+row['Id'] + '.asm')
            shutil.move(DataRoot + '\\train\\' + row['Id'] + '.bytes',DataRoot + '\\sample\\train\\' + row['Id'] + '.bytes')









