# -*- coding: utf-8 -*-
import os
import cv2
import numpy as np
import onMouse
import txtProc

workDir = "./images/"
winName = "label window"
totalImg = np.ndarray([480, 640 * 2, 3], dtype='uint8')
readMeImg = cv2.imread("pre.jpg")


def init():
    """
    mkdir etc.
    :return:
    """
    try:
        if not os.path.exists("./label.txt"):
            file = open("./label.txt", "w")
            file.close()
        if not os.path.exists("./checkImages"):
            os.mkdir("./checkImages")
        return True
    except:
        print("err: make dirs failure.")
        return False


def labelPic(filePath, subFolderNum, subFolderIdx, fileNumInSubFolder, idx):
    '''
    label current pic
    :param filePath:
    :return:
    '''
    img = cv2.imread(filePath)
    img2 = cv2.resize(img, (720, 480))
    totalImg[:, 0:560, :] = readMeImg
    cv2.putText(totalImg, str(subFolderNum), (250, 304), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
    cv2.putText(totalImg, str(subFolderIdx), (250, 337), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
    cv2.putText(totalImg, str(fileNumInSubFolder), (250, 370), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
    cv2.putText(totalImg, str(idx), (250, 403), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
    totalImg[:, 560:560 + 720, :] = img2

    stus = onMouse.run(totalImg, filePath, img.shape)
    return stus


def run():
    """
    main process
    :return:
    """
    status = init()
    if not status:
        print("init failed. exit.")
        return

    startLable = False
    lastimgfile = txtProc.lastLabeledImg()

    subFolderNum = 0
    subFolderIdx = 0
    ifQuit = False  # quit program by user
    for parent, dirNames, fileNames in os.walk(workDir):
        if len(dirNames) != 0:
            subFolderNum = len(dirNames)

        fileNumInSubFolder = len(fileNames)
        i = 0
        while i < fileNumInSubFolder:
            filePath = os.path.join(parent, fileNames[i])

            if startLable or lastimgfile == None:
                status = labelPic(filePath, subFolderNum, subFolderIdx, fileNumInSubFolder, i)
                print("labelPic status:", status, "i: ", i)
                if status == 1:
                    i += 1
                elif status == -1:
                    ifQuit = True
                    break
            else:
                i += 1

            print(filePath)
            print(lastimgfile)
            if (lastimgfile != None) and (lastimgfile in filePath):
                startLable = True

        if ifQuit:
            print("successful quit by user.")
            break

        subFolderIdx += 1

    return


if __name__ == "__main__":
    """
    label tool for Retinaface
    lhf 2022.3.10
    """
    run()
    # str0 = "./images/folder0/WIN_20220303_12_09_16_Pro.jpg"
    # str1 = "folder0/WIN_20220303_12_09_16_Pro.jpg"
    # print(str1 in str0)