# main opencv dependency
import cv2

# numpy is a library for large, multi-dimensional arrays and matrices
# along with a large collection of high-level math functions
import numpy as np

# this library provides a way of using operating system dependent functionality
# in this case it is reading/writing files
import os


# `Plate` class is created for convenience purposes
class Plate:
    def __init__(self, image, filename):
        self.image = image
        self.filename = filename

# it's already pre-trained haar cascade model
plateCascade = cv2.CascadeClassifier("haarcascade_russian_plate_number.xml")
data_path = "data/"
result_thresh_path = "result/thresh/"
result_canny_path = "result/canny/"

# image is converted into resulting form using two image processing techniques:
# thresholding and canny edge detection
def save(plate):
    if not plate: return
    # convert to gray colorspace
    gray_image = cv2.cvtColor(plate.image, cv2.COLOR_BGR2GRAY)

    # apply thresholding tehcnique
    thresh_image = cv2.threshold(
        gray_image, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU
    )[1]

    # apply canny edge detection
    canny_image = cv2.Canny(gray_image, 100, 200, apertureSize=3)

    # save results
    filename = result_thresh_path + plate.filename + ".png"
    cv2.imwrite(filename, thresh_image)

    filename = result_canny_path + plate.filename + ".png"
    cv2.imwrite(filename, canny_image)


# determines the contours of the plate and its angle
# calculates rotation matrix based on given angle and coordinates
# and applies affine transform (rotation) to the given image
def rotate(plate):
    if not plate: return None
    # convert to gray colorspace
    gray_image = cv2.cvtColor(plate.image, cv2.COLOR_BGR2GRAY)
    # invert the image
    inverse_gray = cv2.bitwise_not(gray_image)

    # apply image thresholding to inverted image
    ret, thresh = cv2.threshold(
        inverse_gray, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU
    )

    # determine coordinates of the corresponding region
    coords = np.column_stack(np.where(thresh > 0))
    # determine angle of the given rect
    angle = cv2.minAreaRect(coords)[-1]

    if angle < -45:
        angle = -(90 + angle)
    else:
        angle = -angle

    # calculate width, height, center
    (height, width) = plate.image.shape[:2]
    center = (width // 2, height // 2)
    # calculate rotation matrix & apply affine transform (rotation) with suggested properties
    matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated = cv2.warpAffine(
        plate.image,
        matrix,
        (width, height),
        flags=cv2.INTER_CUBIC,
        borderMode=cv2.BORDER_REPLICATE,
    )
    plate.image = rotated
    return plate


# before cropping image is grayed
# recognizes plate on image using cascade classifier and correspondingly
# crops image with given width and determined height
def crop(plate):
    image = plate.image
    # we change colorspace of the image to gray
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # try to detect plates on the image using cascade classifier
    detected_plates = plateCascade.detectMultiScale(gray_image, 1.3, 5)
    # pre-defined resulting plate image width
    image_width = 400
    enumerated_plates = list(enumerate(detected_plates))
    # if no plates were detected on the image return `None`
    if not enumerated_plates:
        print("error with:" + plate.filename)
        return None
    i, (x, y, w, h) = enumerated_plates[0]
    # determine the ROI (region of interest) in the image
    roi_color = image[y : y + h, x : x + w]
    r = image_width / roi_color.shape[1]
    dim = (image_width, int(roi_color.shape[0] * r))
    # we just crop the image appropriately using suggested interpolation method = cv2.INTER_AREA
    resized_image = cv2.resize(roi_color, dim, interpolation=cv2.INTER_AREA)
    resized_width = resized_image.shape[0]
    resized_height = resized_image.shape[1]
    image[100 : 100 + resized_width, 100 : 100 + resized_height] = resized_image
    cv2.destroyAllWindows()
    plate.image = resized_image
    return plate

# read all images from `data/` source directory and ignore any system files
# construct `Plate` object from each entry and append it to resulting array
def get_plates():
    plates = []
    for entry in sorted(os.listdir(data_path)):
        filename = os.path.splitext(os.path.basename(entry))[0]
        if not filename.startswith('.'):
            image = cv2.imread(data_path + filename  + ".jpg")
            plates.append(Plate(image, filename))
    return plates

for plate in get_plates():
    save(rotate(crop(plate)))