import cv2
import numpy as np
import matplotlib.pyplot as plt
import os, glob, shutil
import sys # process command line arguments
import math
import datetime

class MatchInfo:
  def __init__(self, match, score, error):
    self.match = match # True if the two shapes match, False if the two shapes do not match
    self.score = score # match score between the two shapes
    self.error = error # True if there was a computation error, False otherwise

# computes ratio of white pixels to total pixels, if less than a threshold value, then considers the image shapeless
# returns False if the ratio of white pixels:total pixels is lower than the threshold value, True otherwise
def shape_exists(img, white_threshold):
    img = img / 255.0

    num_white_pixels = 0 # number of white pixels in the image
    height, width = img.shape
    total_pixels = height * width # total number of pixels in the image

    for y in range(height):
        for x in range(width):
            if img[y][x] == 1.0:
                num_white_pixels += 1

    white_ratio = float(num_white_pixels) / float(total_pixels) # get ratio of white pixels to all pixels in the image
    if (white_ratio < white_threshold):
        return False
    else:
        return True

def img_to_sig(arr):
    """Convert a 2D array to a signature for cv2.EMD"""
    # format image data values to meet requirements of cv2.EMD
    # cv2.EMD requires single-precision, floating-point input
    sig = np.empty((arr.size, 3), dtype=np.float32)
    count = 0
    nonzero_value = False
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            sig[count] = np.array([arr[i,j], i, j])
            count += 1
            if arr[i,j] > 0:
                nonzero_value = True # a non-zero value in the signature
    if nonzero_value:
        return sig
    else:
        return None

# compute Earth mover's distance
def emd(reference_img, uncategorized_img, EMD_cutoff):
    # convert image matrices into numpy arrays
    arr1 = np.array(reference_img) # convert matrix into a numpy array
    arr2 = np.array(uncategorized_img) # convert matrix into a numpy array

    # convert numpy arrays of images into signatures for cv2.EMD
    sig1 = img_to_sig(arr1)
    sig2 = img_to_sig(arr2)
    print(sig1)
    print(sig2)

    if (sig1 is None) or (sig2 is None):
        # sys.exit("Error: signatures must contain at least one non-zero value")
        return MatchInfo(False, math.inf, True)

    # compute the amount of work required to match smaller weight distribution to larger weight distribution (have all white pixels in the image with fewer white pixels match all of the white pixels in the image with more white pixels)
    work, _, flow = cv2.EMD(sig1, sig2, cv2.DIST_L2)
    print("Work: " + str(work)) # output the amount of work required to transform one distribution into the other

    if work > EMD_cutoff:
        print("No match")
        return MatchInfo(False, work, False)
    else:
        print("Match")
        return MatchInfo(True, work, False)

# takes in two image histograms, a histogram cutoff score, and a histogram comparison method; returns whether the shapes in the two images match up (whether the two images have around the same histogram) and the match score as a MatchInfo object
def color_match(reference_hist, uncategorized_hist, histogram_cutoff, comparison_method):
    # compute whether the two images have similar histograms
    # if they do not, return False
    match_score = cv2.compareHist(reference_hist, uncategorized_hist, comparison_method) # compute the histogram match score between the reference image histogram and the uncategorized image histogram
    print("Histogram match result: " + str(match_score))
    if match_score < histogram_cutoff:
        return MatchInfo(False, match_score, False) # since the histograms of the two images differing significantly, they can't be the same shape
    else:
        return MatchInfo(True, match_score, False) # since the histograms of the two images are around the same, say that they're the same shape

def emd_shape_match(reference_img, uncategorized_img, EMD_cutoff):
    # compute EMD scores
    # if either EMD score is lower than a cutoff score, return True (doesn't require much work to transform one distribution into the other since the shapes are very similar); else return False
    # compute EMD between reference image and uncategorized image
    reference_uncategorized_match = emd(reference_img, uncategorized_img, EMD_cutoff)
    # print("EMD result between reference image and uncategorized image: " + str(reference_uncategorized_match))
    if reference_uncategorized_match.match:
        return reference_uncategorized_match # not a big difference between the reference image and uncategorized image, so consider them as having the same shape

    # compute EMD between reference image and the uncategorized image reflected across the y-axis
    reference_reflected_uncategorized_match = emd(reference_img, np.fliplr(uncategorized_img), EMD_cutoff)
    # print("EMD result between reference image and the uncategorized image reflected across the y-axis: " + str(reference_reflected_uncategorized_match))
    if reference_reflected_uncategorized_match.match:
        return reference_reflected_uncategorized_match # not a big difference between the reference image and the uncategorized image reflected across the y-axis, so consider them as having the same shape

    # determine lowest EMD score
    lowest_EMD = reference_reflected_uncategorized_match.score
    if reference_uncategorized_match.score < reference_reflected_uncategorized_match.score:
        lowest_EMD = reference_uncategorized_match.score
    return MatchInfo(False, lowest_EMD, False)  # EMD scores are too high, so classify the two shapes as not the same; return the lowest EMD score between the 2 images

def main():
    # Read in command line arguments
    num_args = len(sys.argv)

    if num_args < 4:
        # Error
        print("ERROR: Program usage: python3 combination.py (-all, -one, or -specific) (image folder name) [specific flag parameters] (x-resolution) (y-resolution)")
        return -1

    # Determine program mode (-all, -one, or -specific)
    # -all categorizes all images in a folder; default -all
    # -specific determines whether two images have the same fractal shape
    # -one finds all images that have the same shape as the reference image
    program_mode = sys.argv[1] # the program mode: either -all or -specific
    folder_name = sys.argv[2] # the name of the folder within the program's directory in which all images to categorize are stored
    histogram_cutoff = 0.9997 # histogram difference cutoff
    comparison_method = cv2.HISTCMP_CORREL # histogram comparison method
    EMD_cutoff = 1.80 # EMD work cutoff
    white_threshold = 0.005 # white pixel:total pixel ratio cutoff for shapeless images

    hist_output_path = r'hist_initial_output' # output folder path
    if not os.path.exists(hist_output_path):
        os.makedirs(hist_output_path)
    EMD_output_path = r'EMD_final_output' # final output folder path
    if not os.path.exists(EMD_output_path):
        os.makedirs(EMD_output_path)
    info_file = open("category_info.txt", 'w+') # create a text file to store categorization information

    # save the command
    info_file.write("Running: ")
    for i in range(num_args):
        info_file.write(sys.argv[i] + " ")

    info_file.write("\nwhite_threshold: " + str(white_threshold) + "\n") # remember the white_threshold value
    info_file.write("histogram_cutoff: " + str(histogram_cutoff) + "\n") # remember the histogram_cutoff threshold value
    info_file.write("comparison_method: cv2.HISTCMP_CORREL\n") # remember the comparison method for cv2.compareHist()
    info_file.write("EMD_cutoff: " + str(EMD_cutoff) + "\n") # remember the EMD_cutoff threshold value

    start_time = datetime.datetime.now() # get program start time to later compute program total runtime

    if program_mode == "-specific": # compare two specific images using just histogram comparison
        # Program usage: python3 combination.py -specific (image folder name) (image 1 number) (image 2 number) (x-resolution) (y-resolution)
        if num_args != 7:
            # Error
            print("ERROR: Program usage: python3 combination.py (program mode) (image folder name) (image 1 number) (image 2 number) (x-resolution) (y-resolution)")
            return -1

        reference_image = sys.argv[3] # the file number of the image in the folder to use as the reference image
        uncategorized_image = sys.argv[4] # the file number of the image in the folder to use as the uncategorized image
        x_res = int(sys.argv[5]) # width of image in pixels for resizing
        y_res = int(sys.argv[6]) # height of image in pixels for resizing

        # read in image files to compare
        reference_img_fp = folder_name + '/frame.' + '{:0>6}'.format(reference_image) + '.ppm' # reference image filepath
        uncategorized_img_fp = folder_name + '/frame.' + '{:0>6}'.format(uncategorized_image) + '.ppm' # uncategorized image filepath

        assert os.path.isfile(reference_img_fp), 'file \'{0}\' does not exist'.format(reference_img_fp) # make sure reference image exists at the filepath specified
        assert os.path.isfile(uncategorized_img_fp), 'file \'{0}\' does not exist'.format(uncategorized_img_fp) # make sure the uncategorized image exists at the filepath specified
        reference_img = cv2.resize(cv2.imread(reference_img_fp, cv2.IMREAD_GRAYSCALE), (x_res, y_res)) # read reference image as an array of grayscale values so only getting one value per pixel instead of values for three color channels per pixel
        uncategorized_img = cv2.resize(cv2.imread(uncategorized_img_fp, cv2.IMREAD_GRAYSCALE), (x_res, y_res)) # read uncategorized image as an array of grayscale values so only getting one value per pixel instead of values for three color channels per pixel

        # print image sizes if they are successfully read
        if reference_img is not None:
            print('reference_img.size: ', reference_img.shape)
        else:
            print('imread({0}) -> None'.format(reference_img_fp))

        if uncategorized_img is not None:
            print('uncategorized_img.size: ', uncategorized_img.shape)
        else:
            print('imread({0}) -> None'.format(uncategorized_img_fp))

        reference_hist = cv2.calcHist([reference_img], [0], None, [256], [0, 256]) # compute histogram of the reference image
        uncategorized_hist = cv2.calcHist([uncategorized_img], [0], None, [256], [0, 256]) # compute histogram of the uncategorized image

        # write out resized images for reference purposes
        reference_img_output_fp = 'hist_initial_output/frame.' + '{:0>6}'.format(reference_image) + '.png' # output reference image filepath
        uncategorized_img_output_fp = 'hist_initial_output/frame.' + '{:0>6}'.format(uncategorized_image) + '.png' # output uncategorized image filepath
        cv2.imwrite(reference_img_output_fp, reference_img)
        cv2.imwrite(uncategorized_img_output_fp, uncategorized_img)

        initial_match = color_match(reference_hist, uncategorized_hist, histogram_cutoff, comparison_method).match # compute whether the two images have similar histograms
        print("Same shape results: " + str(initial_match)) # print out whether the shapes match

    elif program_mode == "-one": # create a category with all images that have similar histograms as a specific reference image
        # Program usage: python3 combination.py -one (image folder name) (image to use as reference) (x-resolution) (y-resolution)
        if num_args != 6:
            # Error
            print("ERROR: Program usage: python3 combination.py -one (image folder name) (image to use as reference) (x-resolution) (y-resolution)")
            return -1
        x_res = int(sys.argv[4]) # width of image in pixels for resizing
        y_res = int(sys.argv[5]) # height of image in pixels for resizing
        print("x_res: " + str(x_res) + ", y_res: " + str(y_res))

        # read in reference image file
        reference_image = sys.argv[3] # the file number of the image in the folder to use as the reference image
        reference_img_fp = folder_name + '/frame.' + '{:0>6}'.format(reference_image) + '.ppm' # reference image filepath
        assert os.path.isfile(reference_img_fp), 'file \'{0}\' does not exist'.format(reference_img_fp) # make sure reference image exists at the filepath specified
        reference_img = cv2.resize(cv2.imread(reference_img_fp, cv2.IMREAD_GRAYSCALE), (x_res, y_res)) # read reference image as an array of grayscale values so only getting one value per pixel instead of values for three color channels per pixel
        reference_hist = cv2.calcHist([reference_img], [0], None, [256], [0, 256]) # compute histogram of the reference image

        # create a new category with the reference image
        new_category_path = r'hist_initial_output/category.sample' # output category folder path
        if not os.path.exists(new_category_path):
            os.makedirs(new_category_path)
        reference_img_output_fp = new_category_path + '/' + 'frame.{:0>6}'.format(reference_image) + '.png'
        cv2.imwrite(reference_img_output_fp, reference_img)

        directory = r'{}'.format(folder_name) # read this folder for the images to look through to find the same shapes
        for filename in os.listdir(directory):
            if filename.endswith(".ppm") and filename != ('frame.{:0>6}'.format(reference_image) + '.ppm'):
                print(filename)
                # read in uncategorized image file
                uncategorized_img_fp = folder_name + '/' + filename # uncategorized image filepath
                assert os.path.isfile(uncategorized_img_fp), 'file \'{0}\' does not exist'.format(uncategorized_img_fp) # make sure the uncategorized image exists at the filepath specified
                uncategorized_img = cv2.resize(cv2.imread(uncategorized_img_fp, cv2.IMREAD_GRAYSCALE), (x_res, y_res)) # read uncategorized image as an array of grayscale values so only getting one value per pixel instead of values for three color channels per pixel
                uncategorized_hist = cv2.calcHist([uncategorized_img], [0], None, [256], [0, 256]) # compute histogram of the uncategorized image

                # compare the shapes in the uncategorized image and the reference image
                if (color_match(reference_hist, uncategorized_hist, histogram_cutoff, comparison_method).match):
                    print(str(reference_image) + " & " + str(filename) + " match")
                    # the shapes are the same, so categorize the uncategorized image in the same group as the reference image
                    # make a copy of the uncategorized image file to put in the shape category folder
                    categorized_img_output_fp = new_category_path + '/' + filename[:-3] + 'png'
                    cv2.imwrite(categorized_img_output_fp, uncategorized_img)

    else: # categorize all images using histogram comparison for initial groupings and EMD within these initial groupings to create final shape categories
        # Program usage: python3 combination.py -all (image folder name) (x-resolution) (y-resolution)
        if num_args != 5:
            # Error
            print("ERROR: Program usage: python3 combination.py -all (image folder name) (x-resolution) (y-resolution)")
            return -1
        x_res = int(sys.argv[3]) # width of image in pixels for resizing
        y_res = int(sys.argv[4]) # height of image in pixels for resizing
        print("x_res: " + str(x_res) + ", y_res: " + str(y_res))

        # holds the same shapeless shapes
        EMD_shapeless_category_path = r'EMD_final_output/category.' + '{:0>3}'.format(0) # output category folder path for shapeless images
        hist_shapeless_category_path = r'hist_initial_output/category.' + '{:0>3}'.format(0) # output category folder path for shapeless images

        if not os.path.exists(EMD_shapeless_category_path):
            os.makedirs(EMD_shapeless_category_path)
        if not os.path.exists(hist_shapeless_category_path):
            os.makedirs(hist_shapeless_category_path)

        EMD_uncategorized_category_path = r'EMD_final_output/uncategorized' # output category folder path for a copy of all shapes that are categorized alone
        if not os.path.exists(EMD_uncategorized_category_path):
            os.makedirs(EMD_uncategorized_category_path)

        images = {} # create a hashtable storing the categorization information of every image in the folder
        # key is the frame number and value is a dictionary {"filename": filename, "histogram": the image histogram, "hist_categorized": boolean True/False for whether the image has been categorized according to histogram matching, "EMD_categorized": boolean True/False for whether the image has been categorized according to EMD, "hist_category_num": shape category number in initial histogram groupings, "error": boolean True/False for whether there are errors in categorization}
        # iterate through every .ppm frame in the folder
        directory = r'{}'.format(folder_name) # read this folder for the images to categorize
        num_frames = 0 # keep track of the number of frames in the folder to categorize
        num_hist_categories = 0 # indices for initial categories, starting at 0
        histogram_categories = {} # hold categories created using histogram comparison; key is the category number and value is an array of all image filenames in that shape category
        for filename in os.listdir(directory):
            if filename.endswith(".ppm"):
                print(filename)

                # read in image file
                img_fp = folder_name + '/' + filename # reference image filepath
                assert os.path.isfile(img_fp), 'file \'{0}\' does not exist'.format(img_fp) # make sure image exists at the filepath specified
                img = cv2.resize(cv2.imread(img_fp, cv2.IMREAD_GRAYSCALE), (x_res, y_res)) # read image as an array of grayscale values so only getting one value per pixel instead of values for three color channels per pixel

                # print image sizes if they are successfully read
                if img is not None:
                    print('reference_img.size: ', img.shape)
                else:
                    print('imread({0}) -> None'.format(img_fp))

                file_info = {} # dictionary for every image to store the image's categorization info
                file_info["filename"] = filename
                file_info["hist_categorized"] = False # initialize "hist_categorized" to false
                file_info["EMD_categorized"] = False # initialize "EMD_categorized" to false
                file_info["hist_category_num"] = -1 # initialize to -1 when uncategorized using histogram
                file_info["EMD_category_num"] = -1 # initialize to -1 when uncategorized using EMD
                file_info["histogram"] =  cv2.calcHist([img], [0], None, [256], [0, 256]) # compute the image histogram
                file_info["error"] = False # initialize to no errors in categorization

                # remove images with super few white pixels from categorization consideration and place them all in category.000
                # shape_exists
                if (shape_exists(img, white_threshold) == False):
                    # copy shapeless image into shapeless category
                    # copy into initial histogram groupings
                    hist_categorized_output_fp = hist_shapeless_category_path + '/' + filename[-16:-4] + '.png'
                    cv2.imwrite(hist_categorized_output_fp, img)

                    # copy into EMD groupings
                    EMD_categorized_output_fp = EMD_shapeless_category_path + '/' + filename[-16:-4] + '.png'
                    cv2.imwrite(EMD_categorized_output_fp, img)

                    file_info["hist_categorized"] = True # categorized in shapeless category
                    file_info["EMD_categorized"] = True # categorized in shapeless category
                    file_info["hist_category_num"] = 0 # in shapeless category
                    file_info["EMD_category_num"] = 0 # in shapeless category

                images[num_frames] = file_info # save in overall images dictionary
                num_frames += 1 # increment total number of frames in the folder to categorize
            else:
                continue # not a .ppm file

        # Shape categorization
        print("num_frames: " + str(num_frames))

        # create initial groupings using histogram comparison
        for i in range(num_frames):
            if images[i]["hist_categorized"] == False: # go to the next uncategorized image and set that as the reference image for a new shape category
                # read in reference image file
                reference_img_fp = folder_name + '/' + images[i]["filename"] # reference image filepath
                assert os.path.isfile(reference_img_fp), 'file \'{0}\' does not exist'.format(reference_img_fp) # make sure reference image exists at the filepath specified
                reference_img = cv2.resize(cv2.imread(reference_img_fp, cv2.IMREAD_GRAYSCALE), (x_res, y_res)) # read reference image as an array of grayscale values so only getting one value per pixel instead of values for three color channels per pixel

                num_hist_categories += 1 # increment the number of shape categories
                # create new shape category with this image as reference, category name = num_hist_categories
                new_category_path = r'hist_initial_output/category.' + '{:0>3}'.format(num_hist_categories) # output category folder path
                if not os.path.exists(new_category_path):
                    os.makedirs(new_category_path)

                histogram_categories[str(num_hist_categories)] = [] # array of all image numbers in this category
                histogram_categories[str(num_hist_categories)].append(reference_img_fp) # save reference image filepath

                info_file.write("\ncategory." + '{:0>3}'.format(num_hist_categories) + "\n") # save category number in text file
                info_file.write("Reference image: " + images[i]["filename"] + "\n") # save the reference image for this shape category in text file

                # make a copy of the image file
                categorized_ref_output_fp = new_category_path + '/' + 'frame.{:0>6}'.format(i) + '.png' # copy of reference image categorized filepath
                cv2.imwrite(categorized_ref_output_fp, reference_img)

                images[i]["hist_categorized"] = True # remember that this image has been categorized
                images[i]["hist_category_num"] = num_hist_categories # remember the category this image is a part of

                if (i > num_frames):
                    continue # this is the last frame
                for j in range(i + 1, num_frames):
                    if images[j]["hist_categorized"] == False:
                        # read in uncategorized image file
                        uncategorized_img_fp = folder_name + '/' + images[j]["filename"] # uncategorized image filepath
                        assert os.path.isfile(uncategorized_img_fp), 'file \'{0}\' does not exist'.format(uncategorized_img_fp) # make sure the uncategorized image exists at the filepath specified
                        uncategorized_img = cv2.resize(cv2.imread(uncategorized_img_fp, cv2.IMREAD_GRAYSCALE), (x_res, y_res)) # read uncategorized image as an array of grayscale values so only getting one value per pixel instead of values for three color channels per pixel

                        initial_match = color_match(images[i]["histogram"], images[j]["histogram"], histogram_cutoff, comparison_method) # see if the two images have similar histograms
                        # compare the shapes in the two images
                        if (initial_match.match):
                            info_file.write(images[j]["filename"] + ": " + str(initial_match.score) + "\n") # save file number corresponding to the histogram match score
                            print(str(i) + " & " + str(j) + " match")
                            # the shapes have similar histograms, so categorize the two images together
                            # make a copy of the image file
                            categorized_img_output_fp = new_category_path + '/' + 'frame.{:0>6}'.format(j) + '.png'
                            cv2.imwrite(categorized_img_output_fp, uncategorized_img)

                            histogram_categories[str(num_hist_categories)].append(uncategorized_img_fp) # save image filepath for this new image added to the category

                            images[j]["hist_categorized"] = True # remember that this image has been categorized
                            images[j]["hist_category_num"] = num_hist_categories # save the category number the image was categorized in

        # create final shape categories using EMD
        num_final_categories = 0
        num_single_image_categories = 0 # hold the number of categories with only one image in them

        for category, image_list in histogram_categories.items():
            for i in range(len(image_list)):
                image = image_list[i] # get the filepath of an image in this shape category
                ref_digits = image[-10:-4] # get the 6 digits in the filename
                img_num = int(ref_digits[:-1].lstrip('0') + ref_digits[-1]) # get the image number, remove leading 0s
                print("IMAGE: " + str(image))
                print("img_num: " + str(img_num))
                if (images[img_num]["EMD_categorized"] == False): # go to the next uncategorized image in this shape category and set that as the reference image for a new shape category
                    print("reading in")
                    # read in reference image file
                    reference_img_fp = image # reference image filepath
                    assert os.path.isfile(reference_img_fp), 'file \'{0}\' does not exist'.format(reference_img_fp) # make sure reference image exists at the filepath specified
                    reference_img = cv2.resize(cv2.imread(reference_img_fp, cv2.IMREAD_GRAYSCALE), (x_res, y_res)) # read reference image as an array of grayscale values so only getting one value per pixel instead of values for three color channels per pixel

                    num_final_categories += 1 # increment the number of shape categories
                    # create new shape category with this image as reference, category name = num_final_categories
                    new_category_path = r'EMD_final_output/category.' + '{:0>3}'.format(num_final_categories) # output category folder path
                    if not os.path.exists(new_category_path):
                        os.makedirs(new_category_path)

                    info_file.write("\ncategory." + '{:0>3}'.format(num_final_categories) + "\n") # save category number in text file
                    info_file.write("Reference image: " + image[-16:] + "\n") # save the reference image for this shape category in text file

                    # make a copy of the image file
                    categorized_ref_output_fp = new_category_path + '/' + image[-16:-4] + '.png' # copy of reference image categorized filepath
                    cv2.imwrite(categorized_ref_output_fp, reference_img)

                    images[img_num]["EMD_categorized"] = True # remember that this image has been categorized
                    images[img_num]["EMD_category_num"] = num_final_categories # remember the category this image is a part of

                    num_category_images = 1 # keep track of the number of images in this category, one image in this category so far

                    for j in range(i, len(image_list)):
                        # read in uncategorized image file
                        uncategorized_img_fp = image_list[j] # uncategorized image filepath
                        uncategorized_digits = uncategorized_img_fp[-10:-4] # get the 6 digits in the filename
                        uncategorized_img_num = int(uncategorized_digits[:-1].lstrip('0') + uncategorized_digits[-1]) # get the image number, remove leading 0s
                        print("here: " + str(images[uncategorized_img_num]["EMD_categorized"]))
                        if (images[uncategorized_img_num]["EMD_categorized"] == False): # go to the next uncategorized image in this shape category
                            print("uncategorized")
                            assert os.path.isfile(uncategorized_img_fp), 'file \'{0}\' does not exist'.format(uncategorized_img_fp) # make sure the uncategorized image exists at the filepath specified
                            uncategorized_img = cv2.resize(cv2.imread(uncategorized_img_fp, cv2.IMREAD_GRAYSCALE), (x_res, y_res)) # read uncategorized image as an array of grayscale values so only getting one value per pixel instead of values for three color channels per pixel

                            shape_match = emd_shape_match(reference_img, uncategorized_img, EMD_cutoff) # see if the two images have low EMD score and should actually be in a group together
                            if (shape_match.error):
                                images[uncategorized_img_num]["error"] = True
                                images[uncategorized_img_num]["EMD_categorized"] = True # remember that this image has been categorized
                                continue
                            # compare the shapes in the two images
                            if (shape_match.match):
                                info_file.write(image_list[j]+ ": " + str(shape_match.score) + "\n") # save file number corresponding to the EMD score
                                print(str(i) + " & " + str(j) + " match")
                                # the shapes are the same, so categorize the two images together
                                # make a copy of the image file
                                categorized_img_output_fp = new_category_path + '/' + uncategorized_img_fp[-16:-4] + '.png'
                                cv2.imwrite(categorized_img_output_fp, uncategorized_img)

                                images[uncategorized_img_num]["EMD_categorized"] = True # remember that this image has been categorized
                                images[uncategorized_img_num]["EMD_category_num"] = num_final_categories # save the category number the image was categorized in

                                num_category_images += 1 # increment the number of images in this shape category

                    if (num_category_images == 1):
                        # if this is the only image in the category, copy it into a folder holding all of the images that are categorized alone
                        copy_uncategorized_output_fp = EMD_uncategorized_category_path + '/' + image[-16:-4] + '.png' # copy of reference image categorized filepath
                        cv2.imwrite(copy_uncategorized_output_fp, reference_img)
                        info_file.write("Single image category \n") # save the reference image for this shape category in text file
                        num_single_image_categories += 1 # increment the number of categories with only one image in them
                    else:
                        info_file.write("Number of images in this category: " + str(num_category_images) + "\n")

        info_file.write("\nnum_single_image_categories: " + str(num_single_image_categories) + "\n") # spacer and header for the number of categories with only one image in them

    end_time = datetime.datetime.now() # get program end time to compute program total runtime
    run_time = end_time - start_time # compute program total runtime

    info_file.write("\nRuntime: \n") # spacer and header for runtime info
    info_file.write("start_time: " + str(start_time) + "\n")
    info_file.write("end_time: " + str(end_time) + "\n")
    info_file.write("run_time: " + str(run_time) + "\n")

    info_file.close() # close the text file

if __name__ == "__main__":
    main()


import evokit.chromosome as chr
import evokit.crossover as cr
import evokit.mutation as mu
import evokit.ga as ga

def fitness(chromosome):
    x = chromosome.getFloat(0, 32, -5000, 5000)
    y = chromosome.getFloat(32, 64, -5000, 5000)
    z = chromosome.getFloat(64, 96, -5000, 5000)

    value = ((x - 1) ** 2) + (y ** 2) + ((z + 1) ** 2)
    return value

def main():
    chromosomeFactory = chr.BinaryChromosomeFactory(96)
    crossover = cr.UniformCrossover()
    mutation = mu.BitFlipMutation(0.01)
    result = ga.geneticAlgorithm(fitness, chromosomeFactory, 100,
                crossover, mutation, 1000)

    print("Best fitness:", result["fitness"])
    print("Best solution:",
            result["solution"].getFloat(0, 32, -5000, 5000),
            result["solution"].getFloat(32, 64, -5000, 5000),
            result["solution"].getFloat(64, 96, -5000, 5000))
    result["observer"].plot()

if __name__ == "__main__":
    main()


#!/usr/bin/env python

from widget import Widget
from collections import deque
import logging

logger = logging.getLogger('neoled')

class GraphWidget(Widget):
    def callback(self, bus, argument):
        self.called_bus = bus
        self.data.append(argument)
        self.data.popleft()

    def __init__(self, matrix, bus, offscreenCanvas, x=0, y=10, width=50, height=10, lineColor="0xFFFFFF",
                 borderColor="0x0A0A0A", bgColor="0x00000A", listen="event.key",
                 ranges=[(0, 3, "0x00FF00"), (3, 5, "0xFF8800"), (5, 10, "0xFF0000")], barWidth=2, stacked=True,
                 type=__name__):
        super(GraphWidget, self).__init__(matrix, bus, offscreenCanvas, x, y, width, height)
        self.lineColor = super(GraphWidget, self).color_from_hex(int(lineColor, 0))
        self.borderColor = super(GraphWidget, self).color_from_hex(int(borderColor, 0))
        self.bgColor = super(GraphWidget, self).color_from_hex(int(bgColor, 0))

        self.bus.subscribe(listen, self.callback)

        self.ranges = ranges
        self.bar_width = barWidth
        self.stacked = stacked
        self.data = deque([0] * (self.width // self.bar_width))
        self.scale = 1.0
        logger.info("Started Graph Widget at x: " + str(self.x) + " y: " + str(self.y))

    def Display(self):

        super(GraphWidget, self).drawBackground(self.bgColor)

        m = max(list(self.data))
        scale = 1.0

        if (m > 0):
            scale = self.height * 1.0 / m

        if scale < self.scale:
            self.scale = scale

        count = 0
        for d in list(self.data):

            for i in range(0, self.height, 1):
                color = super(GraphWidget, self).color_from_hex(0x000000)

                for r in self.ranges:
                    if self.stacked:
                        if (i < d * scale) and (r[0] <= i < r[1]):
                            color = super(GraphWidget, self).color_from_hex(int(r[2], 0))
                    else:
                        if (i < d * scale) and (r[0] <= d < r[1]):
                            color = super(GraphWidget, self).color_from_hex(int(r[2], 0))

                for w in range(self.bar_width):
                    # print "d: "+str(d)+"w:" + str(w) + " count:" + str(count) + " = " + str(self.x + w + (count * self.bar_width))
                    self.offscreenCanvas.SetPixel(self.x + w + (count * self.bar_width), self.y - i, color.red,
                                                  color.green, color.blue)

            count += 1

        super(GraphWidget, self).drawBorder(self.borderColor)


import os

class XO:
    occupied = [" " for _ in range(9)]
    winner = None
    leng = 0
    def __init__(self, name):
        self.name = name
        self.vis = []

    def play(self):
        print("\n\t{}'s Turn:".format(self.name), end = "")
        n = int(input())
        if XO.occupied[n-1] != " ":
            print("Wrong Input, Enter correctly")
            self.play()
            return
        self.vis.append(n)
        XO.occupied[n-1] = self.name
        XO.leng += 1

    def check(self):
        if self.contains(self.vis, [1,2,3]) or self.contains(self.vis, [4,5,6]) or self.contains(self.vis, [7,8,9]):
            XO.winner = self
        elif self.contains(self.vis, [1,4,7]) or self.contains(self.vis, [2,5,8]) or self.contains(self.vis, [3,6,9]):
            XO.winner = self
        elif self.contains(self.vis, [1,5,9]) or self.contains(self.vis, [3,5,7]):
            XO.winner = self

        return XO.winner

    def contains(self, arr, sub):
        if((set(arr) & set(sub)) == set(sub)):
            return True
        return False

    @staticmethod
    def printBoard():
        print("{}{}|{}|{} \t{}|{}|{}".format("\t"*2,XO.occupied[0],XO.occupied[1],XO.occupied[2],1,2,3))
        print("{}-+-+- \t-+-+-".format("\t"*2))
        print("{}{}|{}|{} \t{}|{}|{}".format("\t"*2, XO.occupied[3],XO.occupied[4],XO.occupied[5],4,5,6))
        print("{}-+-+- \t-+-+-".format("\t"*2))
        print("{}{}|{}|{} \t{}|{}|{}".format("\t"*2, XO.occupied[6],XO.occupied[7],XO.occupied[8],7,8,9))

if __name__ == '__main__':
    cls = lambda:os.system("cls")
    p1 = XO("X")
    p2 = XO("O")
    XO.printBoard()
    while not XO.winner and XO.leng != 9:
        p1.play()
        p1.check()
        cls()
        XO.printBoard()
        if not XO.winner and XO.leng != 9:
            p2.play()
            p2.check()
            cls()
            XO.printBoard()
    if XO.leng != 9:
        print("\n\tThe Winner is {}".format(XO.winner.name))
    else:
        print("\n\tMatch Tied")


#http://jecvay.com/2015/02/python3-web-bug-series5.html
#零基础自学用Python 3开发网络爬虫(五): 使用第三方模块快速抓取与解析
import requests
from bs4 import BeautifulSoup
response = requests.get("http://jecvay.com")
# print(response.text)
soup = BeautifulSoup(response.text, "html.parser")

print(soup.title.text)
# print(soup.body.text)

for x in soup.findAll("a"):
    print(x['href'])


"""Plotting relational data using Flareplots."""
from __future__ import annotations

from abc import ABC, abstractmethod
from itertools import chain
from typing import NamedTuple
import sys
import warnings

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from .coordinates import bezier, Polar
from .utils import closest, Mark, Label

class MarkupStrategy(ABC):
    def __init__(self, gap_scale: float=0.3):
        self._mark = list()
        self._colors = list()
        self._labels = list()
        self._data = None
        self.gap_scale = gap_scale

    def marks(self, positions, marking, on='number'):
        positions = positions.reset_index().loc[:,[on, "tick"]]
        positions = positions.set_index(on).squeeze()
        for label, color, selection in marking:
            self.selector(positions, label, color, selection)
        return self.data

    def gap(self, width):
        gap = width * self.gap_scale
        side = width / 2 - gap

        return np.array([-side, side])

    @abstractmethod
    def selector(self, positions, label, color, selection) -> None:
        pass

    @abstractmethod
    def as_df(self) -> None:
        pass

    @property
    def data(self):
        df = self.as_df()
        df["color"] = self._colors
        df["label"] = self._labels
        self._data = df
        return self._data

class RangeMarkup(MarkupStrategy):
    def selector(self, positions, label, color, selection) -> None:
        try:
            values = positions.loc[
                    [closest(min(selection), positions.index, lower=True),
                        closest(max(selection), positions.index, lower=False)]].values
        except IndexError:
            warnings.warn(f"{label} at {selection} not found")
            values = np.nan
        self._mark.append(values + self.gap(positions.values[1]))
        self._colors.append(color)
        self._labels.append(label)

    def as_df(self) -> None:
        return pd.DataFrame(self._mark)

class SelectionMarkup(MarkupStrategy):
    def selector(self, positions, label, color, selection) -> None:
        values = positions.loc[selection].values
        self._mark.append(values.reshape(-1,1) + self.gap(positions.values[1]))
        for _ in values:
            self._colors.append(color)
            label = None if label in self._labels else label
            self._labels.append(label)

    def as_df(self) -> None:
        return pd.DataFrame(np.concatenate(self._mark))

class IndividualMarkup(MarkupStrategy):
    def selector(self, positions, label, color, selection) -> None:
        values = positions.loc[min(selection): max(selection)].values
        self._mark.append(values.reshape(-1,1) + self.gap(positions.values[1]))
        for _ in values:
            self._colors.append(color)
            label = None if label in self._labels else label
            self._labels.append(label)

    def as_df(self) -> None:
        return pd.DataFrame(np.concatenate(self._mark))

class Flareplotter:
    def __init__(self, data, fig=None, ax=None, sort_kw={}):
        self.data = data

        if not ax:
            fig = plt.figure()
            ax = fig.add_subplot(polar=True)
        else:
            fig = plt.gcf()

        self.fig = fig
        self.ax = ax
        self.highlighted = list()
        self.labels = None
        self.ticks = None
        self.sort_kw = sort_kw

    def plot(self, *, cmap: str = "Greys", cbar: bool=False,
            vmin: float=0, vmax: float=1, linewidth: float=1.5,
            scale_thickness: bool=True, **line_kwargs) -> Fingerprinter:
        """
:param cmap: str
        colormap to use for the connecting lines. Standard matplotlib colormaps
        are supported.
    :default: "Greys"
        which is a colormap from white to black.
:param cbar: bool
        wether a cbar should be added to the plot or not.
    :default: False
        which plots a colorbar
:param vmin: float
        low point for mapping the colors of :param cmap: to the
        cells values.
    :default: 0
:param vmax: float
        high point for mapping the colors of :param cmap: to the
        cells values.
    :default: 1
:param linewidth: int|float
        base thickness of the connecting lines.
    :default: 1.5
:param scale_thickness: bool
        wether the thickness of the connecting lines scale with the value.
    :default: True

:returns: Fingerprinter
        """
        self.ax.set_theta_zero_location("E", offset=0)
        self.ax.grid(False)

        info = self.label_info()
        self.ax.set_xticks(self.ticks)
        self.ax.set_xticklabels(self.labels)
        self.ax.set_rlim(top=1)
        self.ax.set_yticklabels([])

        sm_cmap = plt.cm.ScalarMappable(
                cmap=cmap,
                norm=plt.Normalize(vmin=vmin, vmax=vmax),
                )

        if cbar:
            cax = self.fig.colorbar(sm_cmap, pad=0.2, ax=self.ax)
            cax.set_label("Frequency")

        for (left, right), value in self.data.items():
            curve = bezier(
                    Polar(theta=info[left].tick, r=1),
                    Polar(theta=info[right].tick, r=1),
                    )
            curve_points = np.array([point.convert().values for point in curve])
            lw = (linewidth + value)**2 if scale_thickness else linewidth
            self.ax.plot(
                    *curve_points.T,
                    color=sm_cmap.get_cmap()(value),
                    lw=lw,
                    zorder=10-10*(1-value),
                    **line_kwargs
                    )
            self.ax.spines['polar'].set_visible(True)

        return self

    def mark(self, marking: dict[str, Mark], *, strategy_name: str = "range", **strategy_kwargs) -> None:
        self.reset_marking()
        strategies = {
                "range": RangeMarkup,
                "individual": IndividualMarkup,
                "selection": SelectionMarkup,
                }
        strat = strategies[strategy_name](**strategy_kwargs)
        strat.marks(self.positions, marking)

        for _, row in strat.data.iterrows():
            if not np.isnan(row[0]):
                thetas = np.arange(row[0], row[1]+0.01, 0.01)
                visibility = True
            else:
                thetas = [0]
                visibility = False
            fill = self.ax.fill_between(
                    thetas,
                    np.ones_like(thetas),
                    1.1,
                    color=row["color"],
                    label=row["label"],
                    zorder=11
                    )
            fill.set_visible(visibility)
            self.highlighted.append(fill)

        self.ax.set_rlim(top=1.1)
        self.ax.spines['polar'].set_visible(False)

        return self.ax

    def reset_marking(self) -> Flareplotter:
        for n in range(len(self.highlighted)):
            self.highlighted[n].remove()
        self.highlighted = []

        try:
            self.ax.get_legend().remove()
        except AttributeError:
            # None if no legend
            pass

        for _ in range(len(self.fig.legends)):
            self.fig.legends[0].remove()

        self.ax.set_rlim(top=1)
        self.ax.spines['polar'].set_visible(True)
        return self

    def replace_labels(self, new_labels: list[str]):
        """Replace rotated "labels" with new labels.

        :param new_labels: list[str]
            List with new labels as string

        """
        if len(new_labels) != len(self.ax.texts):
            msg = ("Make sure that the length of new labels matches the number of ticks. "
                   "Further the labels should have already been replaced by texts throught the "
                   "use of `rotate_labels()`, if not you might want to use ax.set_xticklabels()")
            raise ValueError(msg)

        for old, new in zip(self.ax.texts, new_labels):
            old.set_text(new)

    def rotate_labels(self, shift: float=-0.1) -> Flareplotter:
        self.reset_labels()
        self.fig.canvas.draw()

        for label in self.ax.get_xticklabels():
            text = label.get_text()
            if text.isnumeric():
                text = int(text)
            deg = np.rad2deg(self.label_info()[text].tick)

            if 90 <= deg <= 270:
                deg -= 180

            lab = self.ax.text(*label.get_position()+np.array([0, shift]),
                    text,
                    transform=label.get_transform(),
                    ha=label.get_ha(),
                    va=label.get_va(),
                    )
            lab.set_rotation(deg)
        self.ax.set_xticklabels([])
        return self

    def reset_labels(self):
        for _ in range(len(self.ax.texts)):
            self.ax.texts[0].remove()

        labels, ticks = zip(*[(text, label.tick) for text, label in self.label_info().items()])
        self.ax.set_xticks(ticks)
        self.ax.set_xticklabels(labels)

    def label_info(self):
        unique_labels = sorted(set(chain(*self.data.index)), **self.sort_kw)
        # int(''.join(
        #       (str(label) for label in unique_labelsgc
        #        if str.isdigit(label))
        # ))
        numbers = map(
                lambda lab: int(''.join(filter(str.isdigit, str(lab)))),
                unique_labels,
                )
        ticks = np.linspace(0, 2*np.pi, endpoint=False, num=len(unique_labels))

        try:
            info = dict()
            for (number, label), tick in zip(
                    sorted(zip(numbers, unique_labels)),
                    ticks,
                    ):
                info[label] = Label(number, tick)
        except ValueError:
            info = dict()
            for number, (label, tick) in enumerate(zip(unique_labels, ticks)):
                info[label] = Label(number, tick)

        self.labels, self.ticks = zip(*[(text, label.tick) for text, label in info.items()])
        return info

    @property
    def positions(self) -> pd.Series:
        df = pd.DataFrame(self.label_info()).T
        df[0] = df[0].astype(int)
        df = (df.set_index(0, append=True)
                .swaplevel())
        df.index.names = ["number", "label"]
        df.columns = ["tick"]

        return df

    def _repr_png_(self):
        return display(self.fig)

def flareplot(data, *, cmap="Greys", ax=None, sort_kw={}, **line_kwargs):
    """Wrapper around Flareplotter class.

:param data: pandas.DataFrame
        data of merged DataFrame to be plotted.
:param cmap: str
        colormap to use for the connecting lines. Standard matplotlib colormaps
        are supported.
    :default: "Greys"
        which is a colormap from white to black.
:param cbar: bool
        wether a cbar should be added to the plot or not.
    :default: False
        which plots a colorbar
:param vmin: float
        low point for mapping the colors of :param cmap: to the
        cells values.
    :default: 0
:param vmax: float
        high point for mapping the colors of :param cmap: to the
        cells values.
    :default: 1
:param linewidth: int|float
        base thickness of the connecting lines.
    :default: 1.5
:param scale_thickness: bool
        wether the thickness of the connecting lines scale with the value.
    :default: True

:returns: Fingerprinter
        """

    flare = Flareplotter(data, ax=ax, sort_kw=sort_kw)
    flare.plot(cmap=cmap, **line_kwargs)
    flare.rotate_labels()
    return flare

if __name__ == "__main__":
    sys.exit()


import argparse
import os
import numpy as np
import torch
import torch.nn as nn
from torch.nn import functional as F
import matplotlib.pyplot as plt
from sklearn import metrics
from statistics import mean
import pickle
import math
import sys
import torchvision.models as models

sys.path.insert(0, 'utils/')
sys.path.insert(0, 'data/')
sys.path.insert(0, 'models/')
sys.path.insert(0, 'layers/')

from vgg_utils import _nanargmin, set_parameter_requires_grad
from data_vgg import get_data_loaders

#%%

EPOCHS = 100

loss_fn = F.cross_entropy

LEARNING_RATE = 1e-04
REG = 1e-04

def fit(feature_extractor, task_classifier, adv_classifier, data, device, model_path, output):
    # train and validation loaders
    train_loader, valid_loader = data
    print("Train/Val batches: {}/{}".format(len(train_loader),
                                            len(valid_loader)))

    # Set the optimizer
    task_opt = torch.optim.Adam(list(feature_extractor.parameters()) +
                                list(task_classifier.parameters()),
                                lr=LEARNING_RATE,
                                weight_decay=REG)

    adv_opt = torch.optim.Adam(list(adv_classifier.parameters()),
                               lr=LEARNING_RATE,
                               weight_decay=REG)

    # Start training
    train_history = {'train_loss': [], 'train_acc': [], 'train_apcer': [], 'train_bpcer': [], 'train_eer': [], 'train_bpcer_apcer1': [], 'train_bpcer_apcer5': [], 'train_bpcer_apcer10': [], 'train_apcer1': [], 'train_apcer5': [], 'train_apcer10': [],
                     'val_loss': [], 'val_acc': [], 'val_apcer': [], 'val_bpcer': [], 'val_eer': [], 'val_bpcer_apcer1': [], 'val_bpcer_apcer5': [], 'val_bpcer_apcer10': [], 'val_apcer1': [], 'val_apcer5': [], 'val_apcer10': []}

    # Best validation params
    best_val = -float('inf')
    best_epoch = 0

    for epoch in range(EPOCHS):
        print('\nEPOCH {}/{}\n'.format(epoch + 1, EPOCHS))

        # TRAINING
        # set model to train
        feature_extractor.train()
        task_classifier.train()
        adv_classifier.train()
        for i, (x, y, f, f_norm) in enumerate(train_loader):  # iterations loop
            # send mini-batch to gpu
            x = x.to(device)

            y = y.type(torch.LongTensor)
            y = y.to(device)

            f = f.type(torch.LongTensor)
            f = f.to(device)

            f_norm = f_norm.type(torch.LongTensor)
            f_norm = f_norm.to(device)

            # split real and fake observations
            x_real = x[f == -1]
            x_fake = x[f != -1]
            x_concat = torch.cat([x_real, x_fake], dim=0)

            y_real = y[f == -1]
            y_fake = y[f != -1]
            y_concat = torch.cat([y_real, y_fake], dim=0)

            f_norm_fake = f_norm[f != -1]

            if x_real.shape[0] == 0 or x_fake.shape[0] == 0  :
                continue

            # forward pass
            h_conv = feature_extractor(x_concat)
            h_fake = feature_extractor(x_fake)

            h_conv_flat = h_conv.view(h_conv.size(0), -1)
            h_fake_flat = h_fake.view(h_fake.size(0), -1)

            y_task = task_classifier(h_conv_flat)
            y_adv = adv_classifier(h_fake_flat)

            # Compute task-specific loss
            task_loss = loss_fn(y_task, y_concat)

            # Compute adversial loss
            adv_loss = loss_fn(y_adv, f_norm_fake)

            task_opt.zero_grad()
            task_loss.backward(retain_graph=True)
            task_opt.step()

            adv_opt.zero_grad()
            adv_loss.backward()
            adv_opt.step()

            # display the mini-batch loss
            sys.stdout.write("\r" + '........{}-th mini-batch loss: {:.3f}'.format(i, task_loss.item()))
            sys.stdout.flush()

        # Validation
        tr_loss, tr_acc, tr_apcer, tr_bpcer, tr_eer, tr_bpcer_apcer1, tr_bpcer_apcer5, tr_bpcer_apcer10, tr_apcer1, tr_apcer5, tr_apcer10 = eval_model(feature_extractor, task_classifier, adv_classifier, train_loader, device)
        train_history['train_loss'].append(tr_loss.item())
        train_history['train_acc'].append(tr_acc)
        train_history['train_apcer'].append(tr_apcer)
        train_history['train_bpcer'].append(tr_bpcer)
        train_history['train_eer'].append(tr_eer)
        train_history['train_bpcer_apcer1'].append(tr_bpcer_apcer1)
        train_history['train_bpcer_apcer5'].append(tr_bpcer_apcer5)
        train_history['train_bpcer_apcer10'].append(tr_bpcer_apcer10)
        train_history['train_apcer1'].append(tr_apcer1)
        train_history['train_apcer5'].append(tr_apcer5)
        train_history['train_apcer10'].append(tr_apcer10)

        val_loss, val_acc, val_apcer, val_bpcer, val_eer, val_bpcer_apcer1, val_bpcer_apcer5, val_bpcer_apcer10, val_apcer1, val_apcer5, val_apcer10 = eval_model(feature_extractor, task_classifier, adv_classifier, valid_loader, device)
        train_history['val_loss'].append(val_loss.item())
        train_history['val_acc'].append(val_acc)
        train_history['val_apcer'].append(val_apcer)
        train_history['val_bpcer'].append(val_bpcer)
        train_history['val_eer'].append(val_eer)
        train_history['val_bpcer_apcer1'].append(val_bpcer_apcer1)
        train_history['val_bpcer_apcer5'].append(val_bpcer_apcer5)
        train_history['val_bpcer_apcer10'].append(val_bpcer_apcer10)
        train_history['val_apcer1'].append(val_apcer1)
        train_history['val_apcer5'].append(val_apcer5)
        train_history['val_apcer10'].append(val_apcer10)

        # save best validation model
        if best_val < val_acc:
            #torch.save(model.state_dict(), model_path + 'vgg_fpad.pth')
            best_val = val_acc
            best_epoch = epoch

        # display the training loss
        print()
        print('\n>> Train loss: {:.3f}  |'.format(tr_loss.item()) + ' Train Acc: {:.3f}'.format(tr_acc) + '\n   Train APCER: {:.3f} |'.format(tr_apcer) + ' Train BPCER: {:.3f}'.format(tr_bpcer) + '\n   Train EER: {:.3f}'.format(tr_eer))

        print('\n>> Valid loss: {:.3f}  |'.format(val_loss.item()) + ' Valid Acc: {:.3f}'.format(val_acc) + '\n   Valid APCER: {:.3f} |'.format(val_apcer) + ' Valid BPCER: {:.3f}'.format(val_bpcer) + '\n   Valid EER: {:.3f}'.format(val_eer))

        print('\n>> Best model: {} / Acc={:.3f}'.format(best_epoch+1, best_val))
        print()

    # save train/valid history
    plot_fn = output + 'vgg_fpad_history.png'
    plot_train_history(train_history, plot_fn=plot_fn)

    # return best validation model
    #model.load_state_dict(torch.load(model_path + 'vgg_fpad.pth'))

    #return model, train_history, valid_loader, best_epoch+1
    return train_history, valid_loader, best_epoch+1

def plot_train_history(train_history, plot_fn=None):
    plt.switch_backend('agg')

    best_val_epoch = np.argmin(train_history['val_loss'])
    best_val_acc = train_history['val_acc'][best_val_epoch]
    best_val_loss = train_history['val_loss'][best_val_epoch]
    plt.figure(figsize=(7, 5))
    epochs = len(train_history['train_loss'])
    x = range(epochs)
    plt.subplot(211)
    plt.plot(x, train_history['train_loss'], 'r-')
    plt.plot(x, train_history['val_loss'], 'g-')
    plt.plot(best_val_epoch, best_val_loss, 'bx')
    plt.xlabel('Epoch')
    plt.ylabel('Train/Val loss')
    plt.legend(['train_loss', 'val_loss'])
    plt.axis([0, epochs, 0, max(train_history['train_loss'])])
    plt.subplot(212)
    plt.plot(x, train_history['train_acc'], 'r-')
    plt.plot(x, train_history['val_acc'], 'g-')
    plt.plot(best_val_epoch, best_val_acc, 'bx')
    plt.xlabel('Epoch')
    plt.ylabel('Train/Val acc')
    plt.legend(['train_acc', 'val_acc'])
    plt.axis([0, epochs, 0, 1])
    if plot_fn:
        #plt.show()
        plt.savefig(plot_fn)
        plt.close()
    else:
        plt.show()

def eval_model(feature_extractor, task_classifier, adv_classifier, data_loader, device, debug=False):
    with torch.no_grad():

        feature_extractor.eval()
        task_classifier.eval()
        adv_classifier.eval()

        loss_eval = 0
        N = 0
        n_correct = 0

        TP = 0
        TN = 0
        FP = 0
        FN = 0

        PA = 0
        BF = 0

        eer_list = []

        BPCER_APCER1_list = []
        BPCER_APCER5_list = []
        BPCER_APCER10_list = []

        APCER1_list = []
        APCER5_list = []
        APCER10_list = []

        for i, (x, y, f, f_norm) in enumerate(data_loader):
            # send mini-batch to gpu
            x = x.to(device)

            y = y.type(torch.LongTensor)
            y = y.to(device)

            f = f.type(torch.LongTensor)
            f = f.to(device)

            f_norm = f_norm.type(torch.LongTensor)
            f_norm = f_norm.to(device)

            # split real and fake observations
            x_real = x[f == -1]
            x_fake = x[f != -1]
            x_concat = torch.cat([x_real, x_fake], dim=0)

            y_real = y[f == -1]
            y_fake = y[f != -1]
            y_concat = torch.cat([y_real, y_fake], dim=0)

            f_norm_fake = f_norm[f != -1]

            if x_real.shape[0] == 0 or x_fake.shape[0] == 0  :
                continue

            # forward pass
            h_conv = feature_extractor(x_concat)
            h_fake = feature_extractor(x_fake)

            h_conv_flat = h_conv.view(h_conv.size(0), -1)
            h_fake_flat = h_fake.view(h_fake.size(0), -1)

            y_task = task_classifier(h_conv_flat)
            y_adv = adv_classifier(h_fake_flat)

            # Compute task-specific loss
            task_loss = loss_fn(y_task, y_concat)

            # Compute adversial loss
            adv_loss = loss_fn(y_adv, f_norm_fake)

            loss_eval += task_loss * x.shape[0]

            y_pred = y_task

            # Compute Acc
            N += x.shape[0]
            ypred_ = torch.argmax(y_pred, dim=1)
            n_correct += torch.sum(1.*(ypred_ == y)).item()

            y = y.cpu().numpy()
            ypred_ = ypred_.cpu().numpy()

            # Biometric metrics

            TP += np.sum(np.logical_and(ypred_, y))
            TN += np.sum(np.logical_and(1-ypred_, 1-y))

            FP += np.sum(np.logical_and(ypred_, 1-y))
            FN += np.sum(np.logical_and(1-ypred_, y))

            PA += np.sum(y == 0)
            BF += np.sum(y == 1)

            probs = torch.nn.functional.softmax(y_pred, dim=0)

            probs = probs.cpu().numpy()

            probs = probs[:, 1]

            fpr, tpr, threshold = metrics.roc_curve(y, probs)
            fnr = 1 - tpr

            BPCER_APCER1_list.append(fpr[(np.abs(fnr - 0.01)).argmin()])
            BPCER_APCER5_list.append(fpr[(np.abs(fnr - 0.05)).argmin()])
            BPCER_APCER10_list.append(fpr[(np.abs(fnr - 0.1)).argmin()])

            APCER1_list.append(fnr[(np.abs(fnr - 0.01)).argmin()])
            APCER5_list.append(fnr[(np.abs(fnr - 0.05)).argmin()])
            APCER10_list.append(fnr[(np.abs(fnr - 0.1)).argmin()])

            index = _nanargmin(np.absolute((fnr - fpr)))
            if math.isnan(index) == False:
                eer_list.append(fpr[index])

        loss_eval = loss_eval / N
        acc = n_correct / N
        APCER = (FP * 1.) / (FP + TN)
        BPCER = (FN * 1.) / (FN + TP)

        BPCER_APCER1=mean(BPCER_APCER1_list)
        BPCER_APCER5=mean(BPCER_APCER5_list)
        BPCER_APCER10=mean(BPCER_APCER10_list)

        APCER1=mean(APCER1_list)
        APCER5=mean(APCER5_list)
        APCER10=mean(APCER10_list)

        if eer_list != []:
            EER = mean(eer_list)
        else:
            EER = -1000000000

        return loss_eval, acc, APCER, BPCER, EER, BPCER_APCER1, BPCER_APCER5, BPCER_APCER10, APCER1, APCER5, APCER10

def main():

    IMG_PATH = "/ctm-hdd-pool01/DB/LivDet2015/train/"

    print()
    if torch.cuda.is_available():
        DEVICE = torch.device("cuda:1")  # you can continue going on here, like cuda:1 cuda:2....etc.
        print("Running on the GPU")
    else:
        DEVICE = torch.device("cpu")
        print("Running on the CPU")

    mode = input("Enter the mode [train/test]: ")
    #DATASET = input("Dataset [CrossMatch/Digital_Persona/GreenBit/Hi_Scan/Time_Series]: ")
    unseen_ = input("Unseen attack? [y/n]: ")

    if unseen_ == "y":
        unseen = True
        NUM_ITERATIONS = 1
        attack_txt = "UA"
    elif unseen_ == "n":
        unseen = False
        NUM_ITERATIONS = 3
        attack_txt = "OA"
    else:
        sys.exit("Error ('Unseen attack?'): incorrect input!")

    sensors = ["CrossMatch", "Digital_Persona", "GreenBit", "Hi_Scan", "Time_Series"]

    for DATASET in sensors:

        print("\nRunning in " + DATASET + "...\n")

        if DATASET == "CrossMatch" or DATASET=="Time_Series":
            NUM_MATERIALS = 3
        else:
            NUM_MATERIALS = 4

        # For LOOP - Test splits
        train_results_ = []
        results = []
        best_epochs = []

        for iteration in range(NUM_ITERATIONS):

            print("\n-- ITERATION {}/{} --".format(iteration+1, NUM_ITERATIONS))

            for test_material in range(NUM_MATERIALS):

                output_fn = "results/" + DATASET + "/" + DATASET + "_" + str(test_material) + "_"
                model_path = "/ctm-hdd-pool01/afpstudents/jaf/VGG19bn_" + DATASET + "_" + str(test_material) + "_"

                model = models.vgg19_bn(pretrained=True)
                set_parameter_requires_grad(model, feature_extracting=False)
                num_ftrs = model.classifier[6].in_features
                model.classifier[6] = nn.Linear(num_ftrs,2)

                feature_extractor = torch.nn.Sequential(*list(model.children())[0:2])

                task_classifier = torch.nn.Sequential(*list(model.children())[2])

                adv_classifier = torch.nn.Sequential(*list(model.children())[2])
                num_ftrs = model.classifier[6].in_features
                adv_classifier[6] = nn.Linear(num_ftrs,3)

                feature_extractor = feature_extractor.to(DEVICE)
                task_classifier = task_classifier.to(DEVICE)
                adv_classifier = adv_classifier.to(DEVICE)

                # Train or test
                if mode == 'train':

                    (train_loader, valid_loader, test_loader) = get_data_loaders(IMG_PATH, DATASET, test_material, croped=True, unseen_attack=unseen)

                    # Fit model
                    train_history, _, best_epoch = fit(feature_extractor=feature_extractor,
                                                  task_classifier = task_classifier,
                                                  adv_classifier = adv_classifier,
                                                  data=(train_loader, valid_loader),
                                                  device=DEVICE,
                                                  model_path = model_path,
                                                  output=output_fn)

                    # save train history
                    train_res_fn = output_fn + "history.pckl"
                    pickle.dump(train_history, open(train_res_fn, "wb"))

                elif mode == 'test':
                    sys.exit("Error: in construction yet!")
                    '''
                    model.load_state_dict(torch.load(
                                          os.path.join(*(output_fn, 'mlp_fpad.pth'))))

                    # load train history
                    res_fn = os.path.join(*(output_fn, '_history.pckl'))
                    train_history = pickle.load(open(res_fn, "rb"))
                    plot_fn = os.path.join(*(output_fn, 'mlp_fpad_history.png'))
                    plot_train_history(train_history, plot_fn=plot_fn)
                    '''
                else:
                    sys.exit("Error: incorrect mode!")

                #Train results
                train_results = pickle.load(open(train_res_fn, "rb"))
                train_results_.append([train_results['train_acc'][EPOCHS-1], train_results['train_apcer'][EPOCHS-1], train_results['train_bpcer'][EPOCHS-1], train_results['train_eer'][EPOCHS-1], train_results['train_bpcer_apcer1'][EPOCHS-1], train_results['train_bpcer_apcer5'][EPOCHS-1], train_results['train_bpcer_apcer10'][EPOCHS-1], train_results['train_apcer1'][EPOCHS-1], train_results['train_apcer5'][EPOCHS-1], train_results['train_apcer10'][EPOCHS-1]])

                # Test results
                test_loss, test_acc, test_apcer, test_bpcer, test_eer, test_bpcer_apcer1, test_bpcer_apcer5, test_bpcer_apcer10, test_apcer1, test_apcer5, test_apcer10 = eval_model(feature_extractor, task_classifier, adv_classifier, test_loader, DEVICE)
                print('\nTest loss: {:.3f}            |'.format(test_loss.item()) + ' Test Acc: {:.3f}'.format(test_acc) + '\nTest APCER: {:.3f}           |'.format(test_apcer) + ' Test BPCER: {:.3f}'.format(test_bpcer))
                print('Test BPCER@APCER=1%: {:.3f}  | Test APCER1: {:.3f}'.format(test_bpcer_apcer1, test_apcer1))
                print('Test BPCER@APCER=5%: {:.3f}  | Test APCER5: {:.3f}'.format(test_bpcer_apcer5, test_apcer5))
                print('Test BPCER@APCER=10%: {:.3f} | Test APCER10: {:.3f}'.format(test_bpcer_apcer10, test_apcer10))
                print('Test EER: {:.3f}'.format(test_eer))
                results.append((test_loss.item(), test_acc, test_apcer, test_bpcer, test_eer, test_bpcer_apcer1, test_bpcer_apcer5, test_bpcer_apcer10, test_apcer1, test_apcer5, test_apcer10))

                best_epochs.append(best_epoch)

                # save results
                res_fn = output_fn + 'results.pckl'
                pickle.dump(results, open(res_fn, "wb"))
                results = pickle.load(open(res_fn, "rb"))

            # Compute average and std
            acc_array = np.array([i[1] for i in results])
            apcer_array = np.array([i[2] for i in results])
            bpcer_array = np.array([i[3] for i in results])
            eer_array = np.array([i[4] for i in results])
            bpcer_apcer1_array = np.array([i[5] for i in results])
            bpcer_apcer5_array = np.array([i[6] for i in results])
            bpcer_apcer10_array = np.array([i[7] for i in results])
            apcer1_array = np.array([i[8] for i in results])
            apcer5_array = np.array([i[9] for i in results])
            apcer10_array = np.array([i[10] for i in results])

            print('\n\n\n---------------------------------\n-------------------------------------------')
            print('Average Acc: {:.3f}             |'.format(np.mean(acc_array)) + ' Std: {:.3f}'.format(np.std(acc_array)))
            print('Average APCER: {:.3f}           |'.format(np.mean(apcer_array)) + ' Std: {:.3f}'.format(np.std(apcer_array)))
            print('Average BPCER: {:.3f}           |'.format(np.mean(bpcer_array)) + ' Std: {:.3f}'.format(np.std(bpcer_array)))
            print('Average EER: {:.3f}             |'.format(np.mean(eer_array)) + ' Std: {:.3f}'.format(np.std(eer_array)))
            print('Average BPCER@APCER=1%: {:.3f}  |'.format(np.mean(bpcer_apcer1_array)) + ' Std: {:.3f}'.format(np.std(bpcer_apcer1_array)))
            print('Average BPCER@APCER=5%: {:.3f}  |'.format(np.mean(bpcer_apcer5_array)) + ' Std: {:.3f}'.format(np.std(bpcer_apcer5_array)))
            print('Average BPCER@APCER=10%: {:.3f} |'.format(np.mean(bpcer_apcer10_array)) + ' Std: {:.3f}'.format(np.std(bpcer_apcer10_array)))
            print('Average APCER1: {:.3f}          |'.format(np.mean(apcer1_array)) + ' Std: {:.3f}'.format(np.std(apcer1_array)))
            print('Average APCER5: {:.3f}          |'.format(np.mean(apcer5_array)) + ' Std: {:.3f}'.format(np.std(apcer5_array)))
            print('Average APCER10: {:.3f}         |'.format(np.mean(apcer10_array)) + ' Std: {:.3f}'.format(np.std(apcer10_array)))

            #Best epochs
            print('\nBest epochs:', end=" ")
            for epoch in best_epochs:
                print(epoch, end="   ")

            #Results of all loops (train and test)
            np.set_printoptions(formatter={'float': '{: 0.3f}'.format})
            print()
            print()
            print("*** " + DATASET + " RESULTS ***")
            print()
            print(">>Train results [Acc, APCER, BPCER, EER, BPCER@APCER=1%, BPCER@APCER=5%, BPCER@APCER=10%, APCER1, APCER5, APCER10]")
            print()
            for k in range(NUM_MATERIALS):
                print(*train_results_[k], sep = ", ")

            print()
            print()
            print(">>Test results")

            results_test = []

            for j in range(len(list(acc_array))):
                res = []
                res.append(acc_array[j])
                res.append(apcer_array[j])
                res.append(bpcer_array[j])
                res.append(eer_array[j])
                res.append(bpcer_apcer1_array[j])
                res.append(bpcer_apcer5_array[j])
                res.append(bpcer_apcer10_array[j])
                res.append(apcer1_array[j])
                res.append(apcer1_array[j])
                res.append(apcer10_array[j])

                print(*res, sep = ", ")

                results_test.append(res)

            if iteration == NUM_ITERATIONS-1:
                np.savetxt(DATASET + '_' + attack_txt + '_test.txt', results_test, fmt='%.3f', delimiter=',')

            '''
            print()
            print()
            print(">>Test results")
            print()
            print("Acc:             {}".format(acc_array))
            print("APCER:           {}".format(apcer_array))
            print("BPCER:           {}".format(bpcer_array))
            print("EER:             {}".format(eer_array))
            print("BPCER@APCER=1%:  {}".format(bpcer_apcer1_array))
            print("BPCER@APCER=5%:  {}".format(bpcer_apcer5_array))
            print("BPCER@APCER=10%: {}".format(bpcer_apcer10_array))
            print("APCER1:          {}".format(apcer1_array))
            print("APCER5:          {}".format(apcer5_array))
            print("APCER10:         {}".format(apcer10_array))
            '''

    print("\n\nDONE!")

if __name__ == '__main__':
    main()


import typing

from aiogram.dispatcher.filters import Text
from aiogram.types import Message
from sqlalchemy import select, and_
from sqlalchemy.ext.asyncio import AsyncSession

from keyboards.default import kb_edit_master, kb_return_main_menu
from loader import dp
from misc import state_add_master
from model import database

@dp.message_handler(Text(equals="Удалить мастера"))
async def start_delete_master_handler(m: Message, session: AsyncSession):
    masters = await database.master_barber.MasterBarber.all(session)
    markup = kb_edit_master.get_master_btn(masters)
    markup.add(
        kb_return_main_menu.get_back()
    )
    await m.answer(
        text="Выберите мастера.\n<i>При нажатии на кнопку - мастер удалится</i>",
        reply_markup=markup
    )
    await state_add_master.StateDeleteMaster.start.set()

@dp.message_handler(state=state_add_master.StateDeleteMaster.start)
async def delete_master_handler(m: Message, session: AsyncSession):
    master_name = m.text

    stmt = select(
        database.master_barber.MasterBarber, database.schedule_barber.ScheduleBarber
    ).join(database.schedule_barber.ScheduleBarber,
           and_(database.schedule_barber.ScheduleBarber.master_id == database.master_barber.MasterBarber.id,
                database.schedule_barber.ScheduleBarber.status == database.schedule_barber.ScheduleBarber.typing_status.progress),
           isouter=True,
           ).where(
        database.master_barber.MasterBarber.master == master_name,
        database.master_barber.MasterBarber.delete_status == 0
    )

    response = await session.execute(stmt)
    master, schedule_master = response.all()[0]

    if schedule_master:
        return await m.answer("Вы не можете удалить мастера! У него есть действующие записи")

    master.delete_status = 1
    session.add(master)
    await session.commit()

    masters = await database.master_barber.MasterBarber.all(session)
    markup = kb_edit_master.get_master_btn(masters)
    markup.add(
        kb_return_main_menu.get_back()
    )
    await m.answer("Мастер успешно удален.", reply_markup=markup)


import re
import json
from .scraper import Scraper

class MetasploitParser(Scraper):
    def __init__(self, filename=None, name=None, exploit_type=None, title=None, platform=None, exploit=None, mongoclient=None, date=None):
        ext = ['.rb']
        super().__init__(filename, name, exploit_type, title, platform, exploit, mongoclient, date, ext)
        self.brackets = {
            '{': '}',
            '[': ']',
            '(': ')',
            '}': '{',
            ']': '[',
            ')': '('
        }

    def parse_infos(self):
        cves = self.db['cves']

        self.exploit = self.exploit.replace('\\', '\\\\')

        query = self.parsed_col.find_one({"filename": self.filename})
        if query is not None:
            parsed = query['parsed']
            if parsed:
                return

        error = False
        parsed_file = True
        try:
            description = re.findall(r"[\'\"]Description[\'\"]\s*=>\s*(.*)", self.exploit, re.DOTALL)
            if description:
                description = description[0]
                description = re.sub(r'\s+', ' ', description)
                start, end, is_q = self.find_desc(description)
                if start == -1:
                    desc = re.findall(r"[\'\"]Description[\'\"]\s*=>\s*[\'\"](.*?)[\'\"],", self.exploit, re.DOTALL)[0]
                    desc, = desc[0]
                else:
                    desc = description[start:end]
                    description = description.replace("'", '"')
                    if is_q:
                        desc = desc[start + 3:end - 1]
                    else:
                        desc = desc[start + 2:end - 1]
            else:
                desc = ""

            array = re.findall(r"\\?'(.*?)\\?\\?'[^\w]", self.exploit)  # Get all quoted strings
            for string in array:
                string_1 = string
                if '#' in string:
                    if len(string) == 1:
                        string_1 = '# '
                    else:
                        string_1 = string.replace('#',
                                                  'dash')  # Remove '#' inside them to not lose data when eliminate comments
                    self.exploit = re.sub(re.escape(string), "{}".format(string_1), self.exploit)
                string_1 = string_1.strip('"')  # Replace '"' with '`' to avoid conflict in JSON
                if '"' in string_1:
                    self.exploit = re.sub(re.escape(string_1), string_1.replace('"', '`'), self.exploit)
            self.exploit = re.sub(r'(?:[,\s](#[^{\n]?.*\n{1}\w.*|#[^{\n]?.*)|(#[\'"].*?[\'"]?\n))', '',
                                  self.exploit)  # Delete comments
            self.exploit = re.sub(r'=begin.*?=end', '', self.exploit, flags=re.DOTALL)  # Delete =begin and =end
            self.exploit = re.sub(r'@.*\s*=\s*.*', '', self.exploit)  # Delete annotations

            lista = re.findall(
                r'\bdef initialize(?:\(\s*i|\s*super|\(\s*.*?\s*\)?\s*super)(.*?)(?:\bregister_options\b|\bend\b)\s*(?:def|[A-Za-z_0-9]+\s*=)',
                self.exploit, re.DOTALL)[0]  # Get the main initialize function
            nospace = lista + 'end of the body'  # Add a flag
            nospace1 = re.findall(
                r",?\s*({?\s*[\'\"](?:Name|Description)[\'\"].*?),?\s*\)\s*\)?\s*(?:(?:opt.*)?(?:de)?register_(?:advanced_)?(?:evasion_)?options|end of the body)",
                nospace, re.M | re.S)  # Get rid of the begin and the end of the function, keep only the 'variables'
            if nospace1:
                nospace = nospace1[0]
            title = re.findall(r"[\'\"]Name[\'\"]\s*=>\s*['\"](.*?[\s]*.*?)['\"],", nospace)[0]  # Get title
            nospace = re.sub(r'=>', ':', nospace)  # Replace '=>' with ':'
            nospace = re.sub(r'\s+', ' ', nospace)  # Replace all two+ spaces with only one
            nospace = nospace.replace('\\', '\\\\')  # Encode '\'

            if '{' not in nospace[:4]:
                nospace = '{' + nospace + '}'  # Add brackets

            platform = re.findall(r'(%w{|%w\(|%w\[)(.*?)(}|\)|\]),?', nospace)  # Find platform

            nospace = re.sub(r"('BadChars' : )(?:'.*?'(?=\s?,\s|\s?})|\".*?\"?(?=\s?,\s'|\s?})|.*?(, '))",
                             r"\g<1>''\g<2>", nospace)  # Replace 'BadChars'

            nospace = re.sub(r'0x0*[0-9a-fA-F][0-9a-fA-F]*', '0', nospace)  # Replace hex numbers with '0'
            nospace = re.sub(r"'EncoderType' : (.*?)(,|}|\])", r"'EncoderType' : ''\g<2>",
                             nospace)  # Replace 'EncoderType'
            nospace = re.sub(r"\]\s*\[", "], [", nospace)  # Add ',' between '] ['
            authors = re.findall(r"'Authors?' : \[\s*(.*?)\],? '", nospace)  # Find the 'Author' (array)
            if not authors:
                authors = re.findall(r"'Authors?' : \s*(.*?[\s']),? '", nospace)  # Find the 'Author' (only one)
            if authors:
                authors = re.split(r", (?=(?:[^\"']*[\"'][^\"']*[\"'])*[^\"']*$)",
                                   authors[0])  # Split authors by the most outer ','
                for author in authors:
                    author = author.strip()
                    if not ((author.startswith('"') and author.endswith('"')) or (
                            author.startswith("'") and author.endswith("'"))):
                        nospace = re.sub(re.escape(author), "", nospace)  # Delete wrong name format
                    else:
                        author_rep = author[1:-1].replace('"', '')  # Delete '"' from 'O"Shea'
                        author_rep = author_rep.replace("'", '')  # Delete '"' from 'O"Shea'
                        nospace = re.sub(re.escape(author[1:-1]), author_rep, nospace)
            desc = '"' + desc.replace('"', "'") + '"'  # Replace '"' with "'"
            title = title.replace('"', "'")
            nospace = re.sub(r"(?!\B'[^']*)(\w*_LICENSE|ARCH_\w*)(\s*\]*?)(?![^']*'\B)", r"'\g<1>'",
                             nospace)  # Add quotation marks to 'License' and 'Arch' values
            nospace = re.sub(r'\'', '"', nospace)  # Transform single quotes to duble quotes
            nospace = re.sub(r'(.*?)" "(.*)', r'\g<1>", "\g<2>', nospace)  # Add ',' between 'a_string" "another_one'
            nospace = re.sub(r'(,)\s*,', r' \g<1>', nospace)  # Remove two consecutive ','
            nospace = re.sub(r'(,)\s*([}\]\)])', r' \g<2>', nospace)  # Remove ',' from ', }' (not allowed on JSON)
            nospace = re.sub(r'"\\\\\\\\x[a-zA-Z0-9\\x"\.\(\)+\s]+,', r'"",',
                             nospace)  # Remove some encoded hex values
            nospace = re.sub(r'(\[|{)\s*,', r' \g<1>', nospace)  # Remove ',' from '{, ' (not allowed on JSON)
            nospace = re.sub(r'("BrowserRequirements" : )({.*})\s*([},\]])(?=\s*["{\[]+|$)', r'\g<1>""\g<3>',
                             nospace)  # Delete 'BrowserRequirements'
            nospace = re.sub(r'"Stance" : (.*?)(,|})', r'"Stance" : "\g<1>"\g<2>',
                             nospace)  # Add quotation marks to 'Stance' value

            nospace = re.sub(r'"\s*:\s*(\(.*?\)\s*)([,}])', r'" : "\g<1>"\g<2>',
                             nospace)  # Add quotation marks to ': (some_unquoted_value)'
            nospace = re.sub(r'([^:,\[{\(\)+\.]) ([{"\[\'])', r'\g<1>, \g<2>', nospace)  # Add ',' to ') ('

            nospace = re.sub(r'(\s*?{\s*?|\s*?,\s*?)(["])?([a-zA-Z0-9_]+)(["])?\s?:\s+([^\/\[]+?)',
                             r'\g<1>"\g<3>" : \g<5>', nospace)  # Add quotation marks to 'lowercase_string: value'
            description = re.findall(r'"Description"\s*:\s*(.*)', nospace)
            if description:
                description = description[0]
                start, end, _ = self.find_desc(description)
                if start != -1:
                    nospace = nospace.replace(description[start:end], desc.replace('::', ':~:'))
            nospace = re.sub(
                r'("PrependEncoder" :\s+)(.*\.[\s\n\r]*[\w]+\(.*?\).*?(}?,?\s*")|[^"\(]\S*[^"},]|\(.*\)|".*?"\s*([,}]))',
                r'\g<1>""\g<3>\g<4>', nospace)  # Delete 'PrependEncoder'
            nospace = re.sub(r'("Targets"\s*:\s*)(\d+.*?\d+)', r'\g<1>"\g<2>"',
                             nospace)  # Add quotation marks to 'Targets: 1..2'
            nospace = re.sub(r'(\s*[{:,\[])\s+(:\w+)', r'\g<1> "\g<2>"', nospace)  # Add quotation marks to ':some_value'
            nospace = re.sub(r'"Name"\s*:\s*"(.*?)", "',
                             '"Name" : "' + title.replace('\n', ' ').replace('::', ':~:') + '", "',
                             nospace)  # Replace title
            if platform:
                _, platform, _ = platform[0]
                nospace = re.sub(r'(%w{|%w\(|%w\[)(.*?)(}|\)|\]),?', '"{}",'.format(platform),
                                 nospace)  # Replace platform
            nospace = re.sub(r'("Space" : )[^"][0-9_,+A-Za-z\*\.\)\(\s\-\$]*', r'\g<1> 1024, ',
                             nospace)  # Add constant value to 'Space' adress

            nospace = re.sub(
                r'"Prepend" : (?:.*\.[\s\n\r]*[\w]+\(.*?\).*?(}?,?\s*")|"?.*?"?\s*([,}]\s*[",])|"?.*?"\s*?)',
                r'\g<1>\g<2>', nospace)  # Remove 'Prepend' value (may contain method calls - first regex)
            nospace = re.sub(r'(,)\s*,', r' \g<1>', nospace)
            nospace = re.sub(r'(,)\s*([}\]\)])', r' \g<2>', nospace)
            nospace = re.sub(r'(\[|{)\s*,', r' \g<1>', nospace)
            nospace = re.sub(r',\s*\d+\s*[\-\+\*]\s*\d+', '', nospace)

            nospace = re.sub(r'(:\s+\[\s*)([A-Z_]+[^a-z0-9-"\]}{\[]+)(\s*\])', r'\g<1>"\g<2>"\g<3>',
                             nospace)  # Add quotation marks to ': CAPITAL_LETTERS'
            nospace = re.sub(r'"Signature" : \/.*?\/.*?,', '', nospace)  # Remove Signature
            arithmetics = re.findall(r'(:\s*)(\d+)(\s*)(\*|\+|-)(\s*)(\d+)(,?)',
                                     nospace)  # Find arithmetic operations to replace with actual value
            while (arithmetics):
                for operation in arithmetics:
                    beg, a, sp1, sign, sp2, b, ending = operation
                    result = 0
                    if sign == '+':
                        result = int(a) + int(b)
                    elif sign == '*':
                        result = int(a) * int(b)

                    nospace = re.sub(re.escape('{}{}{}{}{}{}{}'.format(beg, a, sp1, sign, sp2, b, ending)),
                                     beg + str(result) + ending, nospace)
                arithmetics = re.findall(r'(:\s*)(\d+)(\s*)(\*|\+|-)(\s*)(\d+)(,?)',
                                         nospace)  # Find arithmetic operations to replace with actual value

            hex_values = re.findall(r'(?:\\x[0-9a-fA-F]+)+', nospace)  # Find hex escaped values
            for group in hex_values:
                nospace = re.sub(re.escape(group), "Some string", nospace)  # Replace them with a constant
            nospace = re.sub(r': nil', ': null', nospace)  # Replace 'nil'
            targets = re.findall(r'("Targets" : \[\s*\[?.*?[,\s]+.*\s*\]\s*\])', nospace)  # Edit the 'Targets' array
            if targets:
                to_delete = ['References', 'Author', 'Disclosure']
                values = [x in targets[0] for x in to_delete]
                if len(set(values)) == 1 and values[0] == False:
                    nospace = re.sub(r'("Targets" : \[\s*\[?.*?[,\s]+.*\s*\]\s*\])', r'', nospace)
                    nospace = re.sub(r'(,)\s*,', r' \g<1>', nospace)
                    nospace = re.sub(r'(,)\s*([}\]\)])', r' \g<2>', nospace)
                    nospace = re.sub(r'(\[|{)\s*,', r' \g<1>', nospace)
                    nospace = re.sub(r',\s*\d+\s*[\-\+\*]\s*\d+', r'', nospace)
                else:
                    to_edit = re.findall(r'\[\s*\[?\s*(.*?\s*})\s*\]\s*\]?', targets[0])  # Get the arrays
                    if not to_edit:
                        targets[0] = re.sub(r'"References" : \[\s*(.*?)\s*\]\s*\]', '', targets[0])
                        to_edit = re.findall(r'\[\s*\[?\s*(.*?\s*\])\s*\]?\s*\]?', targets[0])  # Can be only one
                    for i, match in enumerate(to_edit):
                        if match.endswith(']'):  # Some problematic endings
                            if not re.findall(r'"Arch" : \[', match):
                                match = match[:-1]
                        check_brackets = re.findall(r'^".*?"\s*,\s*(.*)', match)[
                            0]  # Check if after "some string" exist a '{' eg. ["some_string", {}]
                        if not check_brackets.startswith(
                                '{'):  # If not, add them, otherwise JSON will throw an exception (can't have "key" : "value" inside arrays only if it is an object)
                            to_insert = '{ ' + check_brackets + ' }'
                            match1 = match.replace(check_brackets, to_insert)
                            nospace = re.sub(re.escape(match), match1, nospace)
            nospace = re.sub(r'\\"([^,])', r"'\g<1>",
                             nospace)  # Some problematic endings from Windows paths eg. "C:\\"(treated like escape for '"')
            references = re.findall(r'"References" : \[\s*(.*?)\s*\]\s*,\s*"',
                                    nospace)  # Find references eg. ([ ["CVE", "YYYY-NR"] ])
            if references:
                references = references[0]
                if references:
                    if references.startswith('[') and not references.endswith(']'):
                        nospace = re.sub(re.escape(references), references + ' ]', nospace)
                    if not references.startswith(
                            '['):  # Some anomalies like '"URL" : "SRC"' or '"URL", "SRC", "URL2", "SRC2"' so trying to create the correct format
                        new_array_1 = references.split(', ')
                        new_array = []

                        for string in new_array_1:
                            if re.findall(r'(".*?")\s*:\s*(".*?")', string):
                                new_array.append(string)
                            else:
                                values = string.strip('"').split(' ')
                                new_array.extend('"' + string + '"' for string in values)
                        alt_array = []
                        alt_array.extend(
                            '[ {} ]'.format(arr) for arr in new_array if re.findall(r'(".*?")\s*:\s*(".*?")', arr))  #
                        non_array = [arr for arr in new_array if not re.findall(r'(".*?")\s*:\s*(".*?")', arr)]
                        if len(non_array) % 2 == 0:
                            new_array = ['[ {}, {} ]'.format(non_array[i], non_array[i + 1]) for i in
                                         range(0, len(non_array), 2)]
                            alt_array.extend(new_array)
                        references = ', '.join(alt_array)
                    references_replaced = re.sub(r'("\s*(?:\\\\\\\\|\+),?\s*")', '',
                                                 references)  # Eliminate multiline strings
                    references_replaced = re.sub(r'(\[\s*".*?")\s*:\s*(".*?"\s*\])', r'\g<1>, \g<2>',
                                                 references_replaced)  # Replace ':' with ','
                    nospace = re.sub(re.escape(references), references_replaced,
                                     nospace)  # Replace the old format with the newly created on

            nospace = re.sub(r':?\s([^"]\w*::.*?[^"])([,}\]])', r': ""\g<2>',
                             nospace)  # Add quotation marks to ': MSF::CONSOLE' like values
            nospace = re.sub(r'("DECODERSTUB"\s*:\s*).*?\)', r'\g<1>""', nospace)
            myfile = json.dumps(nospace.replace('\\', '\\\\'))  # Escape the 'backslash' and dumps to JSON

            jsonf = json.loads(json.loads(myfile))
            title = ''

            if self.collection.find_one({"filename": self.name}) is not None:
                title = self.collection.find_one({"filename": self.name})['cve']
            else:
                if 'References' in jsonf:
                    for arr in jsonf.get('References'):
                        if arr and arr[0] == 'CVE':
                            title = title + ' ' + arr[0] + '-' + arr[1]
                    title = title.lstrip()

            if not title:
                title = re.sub(r'\s', '_', self.title)
                title = re.sub(r'\.', '@', title)
                title = self.title + '_' + title

            myDict = {
                "EDB-ID": self.name,
                "Vulnerability": title,
                "Date": self.date,
                "Versions": json.loads(self.get_version_from_name(jsonf.get('Description'))),
            }

            arr = ['Name', 'Description', 'Platform', 'References', 'Targets']

            for key in arr:
                if key in jsonf.keys():
                    if key == 'Name':
                        myDict[key] = self.title
                    else:
                        myDict[key] = jsonf.get(key)
                else:
                    if key == 'Name':
                        myDict[key] = self.title
                    elif key == 'Platform':
                        myDict[key] = self.platform

            myDict['Type'] = self.exploit_type

            URIs = self.parse_url()

            myDict['URI'] = list(set(URIs))
            cves.update({"EDB-ID": self.name}, myDict, upsert=True)
        except Exception as e:
            error, parsed_file = self.found_error(e)
        finally:
            self.update_parsed_obj_db(parsed_file, error)

    def parse_url(self):

        register_options = re.findall(r'register_options\((.*?)],?(?:\s*self\.class)?\)?\s*(?:end)', self.exploit,
                                      re.DOTALL)
        if register_options:
            register_options = register_options[0]

        self.register_options = register_options

        URIs = re.findall(r'(GET|POST|PATCH|PUT) (.*) HTTP', self.exploit)
        if URIs:
            URIs = ['/' + URIs[i][1].lstrip('/') for i in range(len(URIs)) if URIs[i][1] != '/']
        uris = re.findall(r'\s*[\'\"](\/[a-zA-Z0-9_.\-~#{}]*?)[\'\"]\s*', self.exploit)
        URIs.extend(self.construct_url(uris))
        normalize = re.findall(r'normalize_uri\((.*)\)', self.exploit)
        if normalize:
            urls = {}
            for string in normalize:

                to_search_1 = re.findall(r'[?!^\'"](.*?)[?!^\'"]', string)
                splitted = string.replace(' ', '').replace("'", '').split(',')
                to_search = [word for word in splitted if word not in to_search_1]
                if to_search:
                    for i in range(len(to_search)):
                        if 'uri' in to_search[i] or 'path' in to_search[i] or 'datastore' in to_search[i].lower():
                            word = re.findall(r"'TARGETURI',\s*\[(.*?)\]", self.exploit)
                            if word:
                                word = word[0].replace(' ', '').split(',')[2]
                            else:
                                word = 'to_del'
                        elif 'rand' in to_search[i]:
                            continue
                        else:
                            word = re.findall(re.escape(to_search[i]) + r'\s*=\s*[\'"]?(.*?)[\'"]\n', self.exploit)
                            if not word:
                                word = 'to_del'
                            else:
                                word = word[0]
                        urls[to_search[i]] = word
            normaliz = []
            for string in normalize:
                for key in urls.keys():
                    if key in string:
                        string = string.replace(key, urls[key])
                normaliz.append(string.replace('to_del,', ''))
            URIs.extend(self.construct_url(['/' + '/'.join(
                s.strip('/').replace("'", '').replace(',', '').replace('"', '').replace(' ', '') for s in
                normaliz[i].split(',') if 'uri' not in s and 'nil' not in s) for i in range(len(normaliz))]))
        uris = re.findall(r"'uri'\s*=>\s*[\'\"]\/?#{.*?}(\S*)\?\S*[\'\"]", self.exploit)
        URIs.extend(self.construct_url(uris))

        for i, uri in enumerate(URIs):
            regex = re.findall(r'(.*?)%[qQwW][\(\[{](.*?)[\]\)}]', uri)
            if regex:
                regex = regex[0]
                URIs[i] = regex[0] + regex[1]
        return URIs

    # There are some problems with regex so parsing the description was the best way to get all the necessary info
    def find_desc(self, description):
        start = -1
        end = -1
        count = 0
        found = False
        is_q = False
        delimiter = ''
        for i, word in enumerate(description):

            if word == '%' and start == -1:  # Find the %[qQ] and the delimiter '(', '[' or '{'
                start = i
                if description[i + 1] in 'qQ':
                    is_q = True
                    delimiter = description[i + 2]
                else:
                    delimiter = description[i + 1]
            if word in '{[(':
                found = True
                count += 1
            elif word in '}])':
                count -= 1

            if (count == 0 and found and word == self.revert(
                    delimiter)):  # If there's a matching paranthesis and a odd number of them
                return (start, i + 1, is_q)
            elif (description[i - 1] == self.revert(delimiter) and word == ',' and "".join(description[i + 1:].split())[
                0] in "\'\""):  # There can be some anomalies and try to find the coresponding paranthesis in construction like ")'Author"
                return (start, i, is_q)
        return (start, end, is_q)

    def revert(self, char):
        return self.brackets.get(char)

    def construct_url(self, uris):
        URIs = []
        register_options = self.register_options
        if uris:
            for i in range(len(uris)):

                if '/bin/' in uris[i] or 'cmd' in uris[i] or '/c' == uris[i] or '/>' in uris[i] or '\", ' in uris[i] or "/', " in uris[i] or '/' == uris[i]:  # If we found '/bin/sh' or '/cmd' continue
                    continue
                elif 'path' in uris[i] or 'target' in uris[i] or 'base' in uris[
                    i]:  # If there's a construction like targeturi.path continue
                    continue
                elif '\\\\r' in uris[i] or '\\\\n' in uris[i] or '/@' in uris[i] or 'rand' in uris[i]:
                    continue
                elif 'datastore' in uris[i].lower():  # If url = datastore[TARGETURI]
                    word = re.findall(r'datastore\[(.*)\]', uris[i])  # Get the datastore key
                    if word:
                        word = word[0]
                    else:
                        continue
                    if word not in register_options:  # If the key is not in register_options, then the value is manually introduced in metasploit framework
                        continue
                    to_replace_arr = re.findall(r"'.*?{}',\s*\[(.*)\]".format(re.escape(word)), register_options)[
                        0]  # Get the array from construction like 'TARGETURI',[boolean, COMMENT', '/PATH']
                    to_replace_word = to_replace_arr.replace(' ', '').split(',')
                    if len(to_replace_word) > 2:
                        to_replace_word = to_replace_word[2][1:-1]  # Get the value from the array and strip the "'"
                    else:
                        continue
                    if to_replace_word == '/':
                        continue
                    URIs.extend(
                        [re.sub(r'datastore\[(.*)\]', to_replace_word, uris[i])])  # Replace the datastore with the value
                elif 'dash{' not in str(uris[i]) and '#{' not in str(uris[i]):
                    if '/' != uris[i]:  # If the path is '/', don't care
                        URIs.extend(['/' + uris[i].lstrip('/')])  # add the '/' at the begining of the path
                elif 'dash{' in str(uris[i]) or '#{' in str(uris[i]):
                    to_search = re.findall('(?:dash|#){(.*)}', uris[i])  # The url is like '/dash{url}/'
                    if to_search:
                        urls = []
                        for i in range(len(to_search)):
                            urls.extend(re.findall(to_search[i] + r'\s*=\s*[\'"]?(.*?)[\'"]\n',
                                                   self.exploit))  # Try to find that variable
                        if urls:
                            URIs.extend(['/' + re.sub(r'(?:dash|#){(.*)}', url, uris[i]).lstrip('/') for url in urls if
                                         url != '/'])  # Add the new uri to the main array
        return URIs