#!/usr/bin/env python
# -- coding: utf-8 --
"""
Copyright (c) 2020. All rights reserved.
Created by C. L. Wang on 10.11.20
"""

import argparse
import os
from multiprocessing import Pool

import cv2


def traverse_dir_files(root_dir, ext=None):
    """
    列出文件夾中的文件, 深度遍歷
    :param root_dir: 根目錄
    :param ext: 后綴名
    :return: [文件路徑列表, 文件名稱列表]
    """
    names_list = []
    paths_list = []
    for parent, _, fileNames in os.walk(root_dir):
        for name in fileNames:
            if name.startswith('.'):  # 去除隱藏文件
                continue
            if ext:  # 根據(jù)后綴名搜索
                if name.endswith(tuple(ext)):
                    names_list.append(name)
                    paths_list.append(os.path.join(parent, name))
            else:
                names_list.append(name)
                paths_list.append(os.path.join(parent, name))
    return paths_list, names_list


def check_img(path, size):
    """
    檢查圖像
    """
    is_good = True
    try:
        img_bgr = cv2.imread(path)
        h, w, _ = img_bgr.shape
        if h < size or w < size:
            is_good = False
        _ = cv2.resize(img_bgr, (size, size))
    except Exception as e:
        is_good = False

    if path.endswith("jpg"):
        with open(path, 'rb') as f:
            check_chars = f.read()[-2:]
        if check_chars != b'\xff\xd9':
            print('[Info] Not complete jpg image')
            is_good = False

    if not is_good:
        print('[Info] error path: {}'.format(path))
        os.remove(path)


def check_error(img_dir, n_prc, size):
    """
    檢查錯(cuò)誤圖像的數(shù)量
    """
    print('[Info] 處理文件夾路徑: {}'.format(img_dir))
    paths_list, names_list = traverse_dir_files(img_dir)
    print('[Info] 數(shù)據(jù)總量: {}'.format(len(paths_list)))

    pool = Pool(processes=n_prc)  # 多線程下載
    for idx, path in enumerate(paths_list):
        pool.apply_async(check_img, (path, size))
        if (idx+1) % 1000 == 0:
            print('[Info] idx: {}'.format(idx+1))

    pool.close()
    pool.join()

    print('[Info] 數(shù)據(jù)處理完成: {}'.format(img_dir))


def parse_args():
    """
    處理腳本參數(shù)，支持相對(duì)路徑
    :return: in_folder 輸入文件夾, size 尺寸, n_prc 進(jìn)程數(shù)
    """
    parser = argparse.ArgumentParser(description='檢查圖片腳本')
    parser.add_argument('-i', dest='in_folder', required=True, help='輸入文件夾', type=str)
    parser.add_argument('-p', dest='n_prc', required=False, default=100, help='進(jìn)程數(shù)', type=str)
    parser.add_argument('-s', dest='size', required=False, default=50, help='最小邊長(zhǎng)', type=str)
    args = parser.parse_args()

    in_folder = args.in_folder
    size = int(args.size)
    n_prc = int(args.n_prc)
    print("[Info] 文件路徑：{}".format(in_folder))
    print("[Info] 進(jìn)程數(shù): {}".format(n_prc))
    print("[Info] 邊長(zhǎng): {}".format(size))

    return in_folder, n_prc, size


def main():
    arg_in, n_prc, size = parse_args()
    check_error(arg_in, n_prc, size)


if __name__ == '__main__':
    main()