###
# Run as following:
# `python3 det.py -d ./images/prova -c newResult.csv`
# __  -d : the directory where you wanna to scan files
# __  -c : the output file name which you are going to make that contain the result of scanning
###


import os
import tempfile
import subprocess
import csv
from imutils.object_detection import non_max_suppression
import numpy as np
import argparse
import time
import cv2
from PIL import Image
import shutil
import pytesseract
from multiprocessing import Pool

ap = argparse.ArgumentParser()
ap.add_argument('-d', '--directory', type=str,
                help="the name of the directory to scan")
ap.add_argument('-c', '--csv', type=str,
                help="the name of the output file which have the result of scanning")
ap.add_argument('-tn', '--textnegative', action='store_true',
                help="skip files that begin with 'tn_'")
ap.add_argument('-r', '--recursive', type=int, 
                help="recursively called with defined concurrent operations")
args = vars(ap.parse_args())

directory = args['directory']
outputFile = args['csv']
textnegative = args['textnegative']
recursive = args['recursive']

detecion_net = cv2.dnn.readNet('models/detection.weights', 'models/detection.cfg')
recognition_net = cv2.dnn.readNet('models/recognition.weights', 'models/recognition.cfg')


image_width = 0
image_height = 0
number_classes = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']


def get_output_layers(net):
    layer_names = net.getLayerNames()

    output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]

    return output_layers


def check_min_value(a):
    if a < 0:
        return 0
    else:
        return a


def check_X_max_value(x):
    if x > image_width:
        return image_width
    else:
        return x


def check_Y_max_value(y):
    if y > image_height:
        return image_height
    else:
        return y


def crop_image(img, x, y, w, h):
    global crop_img_counter

    x1 = check_min_value(x - int(w * 0.1))
    x2 = check_X_max_value(x + w + int(w * 0.1))
    y1 = check_min_value(y - int(h * 0.1))
    y2 = check_Y_max_value(y + h + int(h * 0.1))
    crop_img = img[y1:y2, x1:x2]

    return crop_img


def draw_bounding_box(img, class_id, label, x, y, x_plus_w, y_plus_h):
    # label = str(plate_classes[class_id])

    color = (0, 255, 0)

    cv2.rectangle(img, (x, y), (x_plus_w, y_plus_h), color, 2)

    cv2.putText(img, label, (x - 10, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 1.5, color, 3)


def recog_number(image):
    scale = 0.00392

    # read pre-trained model and config file

    # create input blob
    blob = cv2.dnn.blobFromImage(image, scale, (140, 120), (0, 0, 0), True, crop=False)

    # set input blob for the network
    recognition_net.setInput(blob)

    # run inference through the network
    # and gather predictions from output layers
    outs = recognition_net.forward(get_output_layers(recognition_net))

    # initialization
    class_ids = []
    confidences = []
    boxes = []
    center_X = []
    conf_threshold = 0.5
    nms_threshold = 0.4

    # for each detetion from each output layer
    # get the confidence, class id, bounding box params
    # and ignore weak detections (confidence < 0.5)
    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > 0.5:
                center_x = int(detection[0] * image.shape[1])
                center_y = int(detection[1] * image.shape[0])
                w = int(detection[2] * image.shape[1])
                h = int(detection[3] * image.shape[0])
                x = center_x - w / 2
                y = center_y - h / 2
                class_ids.append(class_id)
                confidences.append(float(confidence))
                boxes.append([x, y, w, h])
                center_X.append(center_x)

    # apply non-max suppression
    indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)

    # go through the detections remaining
    # after nms and draw bounding box

    result = ''

    valid_boxes = []
    valid_classids = []
    valid_centerX = []


    for i in indices:
        i = i[0]
        box = boxes[i]
        x = box[0]
        valid_boxes.append(box)
        valid_classids.append(class_ids[i])
        valid_centerX.append(x)

    for i in range(0, len(valid_centerX)):
        for j in range(i + 1, len(valid_centerX)):
            if valid_centerX[i] > valid_centerX[j]:
                temp = valid_centerX[i]
                valid_centerX[i] = valid_centerX[j]
                valid_centerX[j] = temp

                tem = valid_classids[i]
                valid_classids[i] = valid_classids[j]
                valid_classids[j] = tem

    for i in range(0, len(valid_classids)):

        result += number_classes[valid_classids[i]]


    return result


def recog_text(image):
    scale = 0.00392

    # read pre-trained model and config file

    # create input blob
    blob = cv2.dnn.blobFromImage(image, scale, (416, 416), (0, 0, 0), True, crop=False)

    # set input blob for the network
    detecion_net.setInput(blob)

    # run inference through the network
    # and gather predictions from output layers
    outs = detecion_net.forward(get_output_layers(detecion_net))

    # initialization
    class_ids = []
    confidences = []
    boxes = []
    center_Y_list = []
    conf_threshold = 0.5
    nms_threshold = 0.4

    # for each detetion from each output layer
    # get the confidence, class id, bounding box params
    # and ignore weak detections (confidence < 0.5)
    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > 0.5:
                center_x = int(detection[0] * image.shape[1])
                center_y = int(detection[1] * image.shape[0])
                w = int(detection[2] * image.shape[1])
                h = int(detection[3] * image.shape[0])
                x = center_x - w / 2
                y = center_y - h / 2
                class_ids.append(class_id)
                confidences.append(float(confidence))
                boxes.append([x, y, w, h])
                center_Y_list.append(center_y)

    # apply non-max suppression
    indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)

    # go through the detections remaining
    # after nms and draw bounding box

    text = ""

    for i in indices:
        i = i[0]
        box = boxes[i]
        x = box[0]
        y = box[1]
        w = box[2]
        h = box[3]

        # if center_Y_list[i] in range(int(image.shape[0] * 0.2), int(image.shape[0] * 0.8)):

        plate_img = crop_image(image, round(x), round(y), round(w), round(h))

        license_str = recog_number(plate_img)

        draw_bounding_box(image, class_ids[i], license_str, round(x), round(y), round(x + w), round(y + h))

        print(license_str)

        text += license_str + ","

    return text[:-1]


class detector:
    def __init__(self, subdir=None, outputName=None):
        super().__init__()
        self.superdir = ""
        self.subdir = subdir if subdir else directory
        self.outputName = outputName if outputName else outputFile
        self.ocr()

    def finddir(self):
        currentPath = os.getcwd()
        self.superdir = currentPath + "/images"
        self.subdir = [self.superdir + "/" +
                       sd for sd in os.listdir(self.superdir) if not ("." in sd)]

    def ocr(self):
        if os.path.isfile(self.outputName):
            os.remove(self.outputName)
        content = os.listdir(self.subdir)
        files = [fil for fil in content if (".jpg" in fil or ".JPG" in fil) and os.path.isfile(os.path.join(self.subdir,fil))]
        if recursive:
            with Pool(processes=recursive) as pool:
                pool.map(self.process_file, files)
        else:
            for fil in files:
                self.process_file(fil)
                
        # Check for the -r or --recursive argument
        if recursive:
            directories = [item for item in content if os.path.isdir(os.path.join(self.subdir, item))]
            for directory in directories:
                # Create a new instance of detector for each subdirectory
                detector(subdir=os.path.join(self.subdir, directory), outputName=self.outputName)
    
    def process_file(self, fil):
        if textnegative and fil.startswith('tn_'):
            return
        filepath = self.subdir + "/" + fil
        text = self.text_detect(filepath, fil)
        csvPath = self.outputName

        file_exists = os.path.isfile(csvPath)
        if text == "":
            return
        with open(csvPath, mode='a') as csv_file:
            fieldnames = ['filename', 'text']
            writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
            if not file_exists:
                writer.writeheader()
            writer.writerow({'filename': fil, 'text': text})

    def text_detect(self, filepath, file):
        global image_width, image_height

        image = cv2.imread(filepath)

        image_width = image.shape[1]
        image_height = image.shape[0]

        text = recog_text(image)
        print(text)

        #image = cv2.resize(image, (1000, 800))
        #cv2.imshow("Result", image)
        #cv2.waitKey(0)

        return text


if __name__ == "__main__":
    if directory == None or outputFile == None:
        print('\tInvalid argument inputed. You should input correct arguments.\n\t--e.g.\t `python3 det.py -d ./images/prova -c prova.csv')
    else:
        if not os.path.isdir(directory) or not '.csv' in outputFile:
            if not os.path.isdir(directory):
                print('\tSuch directory doesn\'t exist.')
            else:
                print('\tInvalid argument inputed. You should input correct arguments.\n\t--e.g.\t `python3 det.py -d ./images/prova -c prova.csv')
        else:
            detector()