# -*- coding: utf-8 -*-
'''
@authot:李华清
@brief：根据lfw网站所给的图片（lfw，lfwcrop_grey,lfwcrop_color)和pairs.txt,写出需要比较的两幅图像的路径(left.txt,right.txt)以及人脸匹配信息(label.txt)
@brief:选择性进行图像尺度变换和灰度化
'''
import os
import numpy as np
import cv2


# 得到一个文件夹下的子文件夹
def dir_subdir(dirname):
    dirname = str(dirname)
    if dirname == "":
        return []
    if dirname[-1] != "/":
        dirname = dirname + "/"
    dirlist = os.listdir(dirname)
    subdir_list = [];
    for x in dirlist:
        if os.path.isdir(dirname + x):
            subdir_list.append(x)
    return subdir_list


# 得到一个文件夹下，指定后缀名的文件
def dir_file(dirname, ext):
    dirname = str(dirname)
    if dirname == "":
        return []
    if dirname[-1] != "/":
        dirname = dirname + "/"
    dirlist = os.listdir(dirname)
    filelist = []
    for x in dirlist:
        if not ('.' in x):
            continue
        if x.split('.')[-1] == ext:
            filelist.append(x)
    return filelist


# 读取LFW的pairs.txt保存到result中
# result是列表，paris是字典，字典元素flag，img1，img2，num1，num2
def read_paris(filelist="pairs.txt"):
    filelist = str(filelist)
    fp = open(filelist, 'r')
    result = []
    for lines in fp.readlines():
        lines = lines.replace("\n", "").split("\t")
        if len(lines) == 2:
            print "lenth=2:" + str(lines)
            continue
        elif len(lines) == 3:
            pairs = {
                'flag': 1,
                'img1': lines[0],
                'img2': lines[0],
                'num1': lines[1],
                'num2': lines[2],
            }
            result.append(pairs)
            continue
        elif len(lines) == 4:
            pairs = {
                'flag': 2,
                'img1': lines[0],
                'num1': lines[1],
                'img2': lines[2],
                'num2': lines[3],
            }
            result.append(pairs)
        else:
            print "read file Error!"
            exit()
    fp.close
    print "Read paris.txt DONE!"
    return result


# 找到文件夹下后缀名为ext的文件,可包含子文件夹
def dirdir_list(lfwdir, ext='jpg'):
    subdirlist = dir_subdir(lfwdir)
    filelist = []
    # 含有子文件夹
    if len(subdirlist):
        for i in subdirlist:
            list = dir_file(lfwdir + i, ext)
            for j in list:
                j = lfwdir + i + '/' + j
                filelist.append(j)
    # 不包含子文件夹
    else:
        list = dir_file(lfwdir, ext)
        for line in list:
            line = lfwdir + line
            filelist.append(line)
    return filelist


def grey_resize(lfwdir, filelist, desdir='lfw_change/', grey=False, resize=False, height=64, width=64):
    if (grey == False and resize == False):
        return []
    # 创建目的文件夹
    if not os.path.exists(desdir):
        os.makedirs(desdir)

    if desdir[-1] != '/':
        desdir = desdir + '/'

    # 判断文件夹下是否含有子文件夹
    flag = 0
    subdir = dir_subdir(lfwdir)
    if len(subdir):
        flag = 1

    # 创建文件夹
    if flag == 1:
        for path in filelist:
            path = path.split('/')
            if not os.path.exists(desdir + path[-2]):
                os.makedirs(desdir + path[-2])

    # 处理后图像所放路径
    filelistnew = []
    for path in filelist:
        path = path.split('/')
        if flag == 1:
            filelistnew.append(desdir + path[-2] + '/' + path[-1])
        else:
            filelistnew.append(desdir + path[-1])
    # 进行灰度化和尺度变换

    num = 0
    nums = len(filelistnew)
    for line in filelist:
        img = cv2.imread(line)

        if grey == True:
            if num % 100 == 0:
                print "grey:" + str(num) + '/' + str(nums)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        if resize == True:
            if num % 100 == 0:
                print "resize:" + str(num) + '/' + str(nums)
            img = cv2.resize(img, (height, width))
        cv2.imwrite(filelistnew[num], img)
        num = num + 1
    return filelistnew


# 根据pairs.list,灰度化或者统一图像尺度后，得到left.txt,right.txt,label.txt
def split_pairs(pairslist, lfwdir, ext='jpg'):
    num = 0
    sum = len(pairslist)

    # lfw图像组织形式，只有图像：0，文件夹+图像：1
    flag = 0
    # dir_subdir得到一个文件夹下的子文件夹
    subdir = dir_subdir(lfwdir)
    if len(subdir):
        flag = 1

    path_left = ""
    path_right = ""
    label = ""

    # left.txt and right.txt
    # 文件夹+图像形式
    print "split pairs.txt"
    if flag == 1:
        for lines in pairslist:
            num = num + 1
            if num % 100 == 0:
                print str(num) + "/" + str(sum)

            dir_left = lfwdir + lines['img1'] + '/'
            dir_right = lfwdir + lines['img2'] + '/'

            file_left = lines['img1'] + '_' + str("%04d" % int(lines["num1"])) + '.' + ext
            file_right = lines['img2'] + '_' + str("%04d" % int(lines["num2"])) + '.' + ext

            path_left = path_left + dir_left + file_left + "\n"
            path_right = path_right + dir_right + file_right + "\n"

    # 图像
    else:
        for lines in pairslist:
            num = num + 1
            print str(num) + "/" + str(sum)
            path_left = path_left + lfwdir + lines['img1'] + '_' + str("%04d" % int(lines["num1"])) + '.' + ext + "\n"
            path_right = path_right + lfwdir + lines['img2'] + '_' + str("%04d" % int(lines["num2"])) + '.' + ext + "\n"
    # label.txt
    for lines in pairslist:
        if int(lines["flag"]) == 1:
            label = label + "0\n"
        else:
            label = label + "1\n"

    result = {
        'path_left': path_left,
        'path_right': path_right,
        'label': label}
    return result


def testdeal(pairs='pairs.txt', lfwdir='lfw/', ext='jpg', changdir='lfw_chang/', grey=False, resize=False, height=64,
             width=64):
    # pairslist is a list which contains 2 kinds dicts(flag1 flag2).
    # 读取LFW的pairs.txt保存到result中
    # result是列表，paris是字典，字典元素flag，img1，img2，num1，num2
    pairslist = read_paris(pairs)
    if grey == True or resize == True:
        filelist = dirdir_list(lfwdir, ext)
        filelist = grey_resize(lfwdir, filelist, changdir, grey, resize, height, width)
        lfwdir = changdir
    pairs_result = split_pairs(pairslist, lfwdir, ext)
    return pairs_result


def write_pairs(pairs_result, savepath):
    if not os.path.exists(savepath):
        os.makedirs(savepath)
    left = savepath + 'left.txt'
    right = savepath + 'right.txt'
    label = savepath + 'label.txt'

    fp_left = open(left, 'w')
    fp_right = open(right, 'w')
    fp_label = open(label, 'w')

    fp_left.write(pairs_result['path_left'])
    fp_right.write(pairs_result['path_right'])
    fp_label.write(pairs_result['label'])

    fp_left.close()
    fp_right.close()
    fp_label.close()


def demo_lfw():
    caffe_dir = '/home/herbert/caffe/'
    pairs = caffe_dir + 'data/deepID/pairs.txt'
    lfwdir = caffe_dir + 'data/deepID/lfwcrop_color/faces/'
    ext = 'ppm'
    # pairs_result : a list contains pairs
    pairs_result = testdeal(pairs, lfwdir, ext)
    savepath = caffe_dir + 'examples/deepID/lfw_'
    write_pairs(pairs_result, savepath)


def grey_pairs():
    DEEPID = 'data/deepID_grey/'
    pairs = DEEPID + 'pairs.txt'
    lfwdir = DEEPID + 'lfwcrop_grey/faces/'
    ext = 'pgm'
    pairs_result = testdeal(pairs, lfwdir, ext)
    return pairs_result


def demo_webface_resize():
    DEEPID = '/home/herbert/caffe/data/deepID/'
    pairs = DEEPID + 'pairs.txt'
    #    lfwdir=DEEPID+'lfwcrop_color/faces/'
    lfwdir = '/media/ikode/Document/big_materials/document/deep_learning/caffe/face_datasets/webface/croped/'
    ext = 'jpg'
    lfw_chang = '/media/ikode/Document/big_materials/document/deep_learning/caffe/face_datasets/webface/change/'

    pairs_result = testdeal(pairs, lfwdir, ext, lfw_chang, False, True)
    savepath = '/home/herbert/caffe/' + 'examples/deepID/webface_'
    write_pairs(pairs_result, savepath)


def demo_color_resize():
    DEEPID = '/home/herbert/caffe/data/deepID/'
    pairs = DEEPID + 'pairs.txt'
    #    lfwdir=DEEPID+'lfwcrop_color/faces/'
    lfwdir = '/media/ikode/Document/big_materials/document/deep_learning/caffe/face_datasets/webface/croped/'
    ext = 'jpg'
    lfw_chang = '/media/ikode/Document/big_materials/document/deep_learning/caffe/face_datasets/webface/change/'

    pairs_result = testdeal(pairs, lfwdir, ext, lfw_chang, False, True)
    return pairs_result


if __name__ == '__main__':
    demo_lfw()
