# -*- coding: utf-8 -*-
"""
@author: kevinhikali
@email: hmingwei@gmail.com
"""

from __future__ import division
import os
import numpy as np
import sys
dataname = sys._getframe().f_code.co_filename
dataname = dataname.split('/')[-1]
dataname = dataname.split('.')[0]

LABEL_SIZE = 1
IMAGE_SIZE = 32
NUM_CHANNELS = 3
PIXEL_DEPTH = 255
NUM_CLASSES = 10

TRAIN_NUM = 10000
TEST_NUM = 10000

PATH = '/home/kevin/catkin_ws/src/res_kevin'


def read_file(filename):
    import cPickle
    with open(filename, 'rb') as fo:
        _dict = cPickle.load(fo)
    return _dict


def import_cifar10_train():
    print('Importing cifar10 data:')

    filename_queue = []
    for _filename in os.listdir(PATH + '/' + dataname + '/'):
        if _filename.startswith("data"):
            filename_queue.append(_filename)

    nBatch = len(filename_queue)

    _data = []
    _label = []
    for i in range(nBatch):
        pack = read_file(PATH + '/' + dataname + '/'
                                    + filename_queue[i])
        _data.append(pack['data'])
        _label.append(pack['labels'])

    data = np.concatenate(_data)
    label = np.concatenate(_label)
    label = np.array(label)

    return data, label


def import_cifar10_check():
    print('Importing cifar10 data:')

    filename_queue = 'test_batch'

    pack = read_file(PATH + '/' + dataname + '/'
                     + filename_queue)
    data = pack['data']
    label = pack['labels']
    label = np.array(label)

    return data, label
