# -*- coding:utf-8 -*-
import os
import numpy as np
import shutil
import configparser
from PIL import Image
import time

# 读取配置文件
config = configparser.ConfigParser()
config.read('config.ini', encoding='utf-8')
data_dir = config.get('section_1', 'DATA_DIR')
data_list = os.listdir(data_dir)
img_size = int(config.get('section_2', 'IMG_SIZE'))


def remove_old(DIR):
    """
    清空DIR目录下所有文件
    :param DIR:目标目录
    :return:
    """
    list = os.listdir(DIR)
    print("清空旧文件，目标目录：{}".format(DIR))
    if not list:
        print("目标目录已为空\n")
        return 0
    for f in list:
        path = os.path.join(DIR, f)
        os.remove(path)
    print("目标目录已清空\n")

def update_image(src_dir, update_list, label):
    """
    具体实现图片素材更新到项目data中
    :param src_dir:图片源目录，项目外源素材,../orange_data/xxx
    :param update_list:待更新图片目录，项目内data
    :param label:标签，指明调用类型，bad.test,bad.train,good.test,good.train
    :return:
    """

    # 确定目标目录
    list_label = label.split(".")
    print("update : " + label)
    if list_label[1] == "test":
        des_dir = "data/test"
    else:
        des_dir = "".join("data/train".format())

    # 目标目录已有文件列表
    list = os.listdir(des_dir)
    list_len = len(list)
    # 遍历更新列表
    for image in update_list:
        # 求image在更新列表中的索引值，以用于命名新文件
        index = update_list.index(image)
        if list_label[1] == "test":
            image_new = "".join("test.{}.jpg".format(list_len + index))         # test在已有文件index上追加test.index.jpg
        else:
            image_new = "".join("{}.{}.jpg".format(list_label[0], index))       # bad.index.jpg/goog.index.jpg

        # 图片完整路径
        image_src = os.path.join(src_dir, image)
        image_des = os.path.join(des_dir, image_new)
        shutil.copy(image_src, image_des)           # 图片复制

        # 图片尺寸修改
        img = Image.open(image_des)
        img_resize = img.resize((img_size, img_size), Image.ANTIALIAS)          # Image.ANTIALIAS抗锯齿
        img_resize.save(image_des)
    print(label + "更新成功\n")


def cut_image(image_dir):
    """
    将源路径下的图片分成test和train两部分
    :param image_dir:../orange_data/xxx
    :return:
    :param image_dir:
    :return:
    """

    print("分割" + image_dir + "下图片为test和train部分")
    image_list = os.listdir(image_dir)          # 得到image_dir下所有文件列表
    image_num = len(image_list)
    print("image num : {}".format(image_num))

    # 打乱文件顺序
    # 用numpy的方法，根据打乱的索引random_index来打乱image_list文件顺序
    random_index = np.arange(image_num)
    np.random.shuffle(random_index)
    image_list = np.array(image_list)
    image_list = image_list[random_index]
    image_list = image_list.tolist()

    # 取前10%做test，后90%做train
    test_num = int(image_num / 10)
    print("{} to test,{} to train\n".format(test_num,image_num - test_num))
    # 切片得test和train图片列表
    test_list = image_list[:test_num]
    train_list = image_list[test_num:]
    return test_list,train_list


def update():
    """
    更新项目data
    :return:
    """
    update_list = ["bad","good"]
    for i in range(2):
        label = update_list[0]
        src_dir = "{}/{}".format(data_dir, data_list[i])
        # 获得test列表，train列表
        test_list, train_list = cut_image(src_dir)
        update_image(src_dir, test_list, label + ".test")
        update_image(src_dir, train_list, label + ".train")

def main():
    # 移除项目下data旧图片
    remove_old("data/test")
    remove_old("data/train")
    # 更新图片至项目内
    update()


if __name__ == '__main__':
    main()