# -*- coding: utf-8 -*-

import os, glob, shutil
import numpy as np

np.random.seed(2016)

# 元データ root_path
r_path = '../../data/101_ObjectCategories'

# データ配置用 root_path
o_path = '../../data/Caltech-101'

# データ読み込み
path = '%s/*/*.jpg'%r_path
files = sorted(glob.glob(path))
files = np.array(files)

#############################
## Keras, VGG16, ResNet 用 ##
#############################

# 使用するラベル ---①
use_labels = ['airplanes', 'Motorbikes', 'Faces_easy', 'watch', 'Leopards', 'bonsai']
labels_count = [0,0,0,0,0,0]
# 学習、評価、テストに使用する件数を指定　---③
train_nums = [80,80,44,24,20,13]
valid_nums = [80,80,43,24,20,13]
test_nums = [640,638,348,191,160,102]


# まず train+valid と testデータに分割する。

# ディレクトリが存在しなかったら作成する。
for i in range(0, len(use_labels)):
    if not os.path.exists('%s/train_org/%i'%(o_path, i)):
        os.makedirs('%s/train_org/%i'%(o_path, i))
    if not os.path.exists('%s/test/%i'%(o_path, i)):
        os.makedirs('%s/test/%i'%(o_path, i))
# ファイル数分だけ処理を繰り返す
for fl in files:
    # ファイル名取得
    filename = os.path.basename(fl)
    # 親ディレクトリ = ラベル 取得
    parent_dir =  os.path.split(os.path.split(fl)[0])[1]

    if parent_dir in use_labels:

        ind = use_labels.index(parent_dir)

        num = labels_count[ind]
        valid_num = valid_nums[ind]
        test_num = test_nums[ind]

        if num < train_nums[ind] + valid_nums[ind]:
            cp_path = '%s/train_org/%i/'%(o_path, ind)
            shutil.copy(fl, cp_path)
        else:
            cp_path = '%s/test/%i/'%(o_path, ind)
            shutil.copy(fl, cp_path)

        labels_count[ind] += 1

    else:
        # 今回使用するラベル以外は無視する。
        continue


# trainデータを train, valid に分割する。
# ho数分だけ繰り返す
for ho in range(0,2):

    for ii in range(0, len(use_labels)):

        # ディレクトリが存在しなかったら作成する。
        if not os.path.exists('%s/train/%i/%i'%(o_path, ho, ii)):
            os.makedirs('%s/train/%i/%i'%(o_path, ho, ii))
        if not os.path.exists('%s/valid/%i/%i'%(o_path, ho, ii)):
            os.makedirs('%s/valid/%i/%i'%(o_path, ho, ii))

        # データ読み込み
        path = '%s/train_org/%i/*.jpg'%(o_path,ii)
        files = sorted(glob.glob(path))
        files = np.array(files)

        perm = np.random.permutation(len(files))
        random_train = files[perm]

        train_files = random_train[:train_nums[ii]]
        valid_files = random_train[train_nums[ii]:]

        # trainデータを配置
        for file in train_files:
            # ファイル名取得
            filename = os.path.basename(file)
            # 親ディレクトリ = ラベル 取得
            p_dir =  os.path.split(os.path.split(file)[0])[1]

            shutil.copy(file, '%s/train/%i/%i/'%(o_path,ho,int(p_dir)))

        # validデータを配置
        for file in valid_files:
            # ファイル名取得
            filename = os.path.basename(file)
            # 親ディレクトリ = ラベル 取得
            p_dir =  os.path.split(os.path.split(file)[0])[1]

            shutil.copy(file, '%s/valid/%i/%i/'%(o_path,ho,int(p_dir)))

# ディレクトリとラベルの紐づけ書き出し

# ラベルとの紐づけ --- ②
str = '0:airplanes,1:Motorbikes,2:Faces_easy,3:watch,4:Leopards,5:bonsai'
f = open('%s/label.csv'%o_path, 'w')
f.write(str)
f.close()


#############################
## Yolo　用                ##
#############################

# データ配置用 root_path
yolo_path = '../../data/Yolo'
if not os.path.exists('%s/train/'%yolo_path):
    os.makedirs('%s/train/'%yolo_path)
if not os.path.exists('%s/test/'%yolo_path):
    os.makedirs('%s/test/'%yolo_path)


