import tkinter as tk
from tkinter.filedialog import askdirectory
import numpy as np
from PIL import ImageTk, Image
import os
import glob
from multiprocessing.dummy import Pool
import shutil
import cv2
import dlib as db
import skimage
import matplotlib.pyplot as plt
from scipy import ndimage
import copy
import scipy.spatial.distance as dist

def sorting_menu(location):
    # region sorting_menu
    # Window contains 5 tools
    # 1. View all images 2. Find duplicate images with K-means
    # 3. Find similar photos to a photo input of our choice
    # 4. Apply filter to an image of our choice 5. Face detection program
    print(location)

    # Tasks
    def view_photos():
        # Initializing window to display the images
        gtk = tk.Toplevel()

        def getFileName(image):
            # Gets the address of the image that we have selected
            global image_pressed
            print(str(image))
            image_pressed = image
            gtk.destroy()

        # Placement of all images on different buttons on the window
        row = 0
        column = 0
        for images in os.listdir(str(location)):
            if images.endswith("jpg"):
                im = Image.open(location+'/'+images)
                im.thumbnail((64, 64))
                tkimage = ImageTk.PhotoImage(im)
                handler = lambda img=images: getFileName(img)  # here modify
                imageButton = tk.Button(gtk, image=tkimage, command=handler)  # here
                imageButton.image = tkimage
                imageButton.grid(row=row, column=column)
                if column == 8:
                    row += 1
                    column = 0
                else:
                    column += 1

        # Wait for an input
        gtk.mainloop()

        return

    def find_duplicates():
        # Function classify images into k categories based on similarity among the images
        # Histogram of each images with 10 bins is produced and HAC is applied on these histogram data
        #   to classify images and so find duplicates

        # HAC Clustering
        def HAClustering(X):
            X = np.float32(X)
            m, n = X.shape
            plt.figure(1)

            num_clusters = m
            idx = np.arange(m)

            centroids = copy.deepcopy(X)
            cluster_sizes = np.ones(m)

            dists = dist.squareform(dist.pdist(centroids))
            np.fill_diagonal(dists, float('inf'))

            min_dist = np.min(dists)
            iteration = 0

            while min_dist < 40:
                iteration += 1
                min_dist = np.min(dists)
                print(min_dist)

                i = np.where(dists == min_dist)[0][0]
                j = np.where(dists == min_dist)[0][1]

                # Make sure that i < j
                if i > j:
                    t = i
                    i = j
                    j = t
                else:
                    pass

                # temp = np.array([dists[i], dists[j]])
                # dists[i] = np.min(temp, axis=0)
                # dists = np.delete(dists, i, axis=0)
                # dists = np.delete(dists, i, axis=1)
                temp = np.array([dists[i], dists[j]])
                dists[i, :] = np.mean(temp, axis=0)
                dists[:, i] = np.mean(temp, axis=0)
                dists[j, :] = float('inf')
                dists[:, j] = float('inf')

                centroids[i] = (centroids[i] + centroids[j]) / 2
                cluster_sizes[i] += cluster_sizes[j]
                cluster_sizes[j] = 0
                centroids[j] = float('inf')
                idx[idx == idx[j]] = idx[i]

                num_clusters -= 1

            # Reindexing clusters
            u = np.unique(idx)
            for i in range(len(u)):
                idx[idx == u[i]] = i

            return idx
        # K Means clustering (Not Used)

        def KMeansCLustering(X, k):
            X = np.float32(X)
            m, n = X.shape

            step = int(m/k)
            centers = []
            for a in range(k):
                centers.append(X[int(a*step)])

            centers = np.array(centers)
            # centers = np.array(random.choices(X, k=k))

            idx = np.zeros([m])

            iter = 0
            MAX_ITER = 100

            while True:
                old_idx = idx

                # Allotting clusters
                for i in range(m):
                    min_dist = float('inf')
                    for c in range(k):
                        distance = np.linalg.norm(X[i] - centers[c])
                        if distance < min_dist:
                            min_dist = distance
                            idx[i] = c

                # Updating cluster centers
                for cen in range(k):
                    new_centre_members = []
                    for i in range(m):
                        if idx[i] == cen:
                            new_centre_members.append(X[i])
                        else:
                            pass
                    new_centre_members = np.array(new_centre_members)
                    centers[cen] = np.mean(new_centre_members, axis=0)

                # Display
                # End condition
                if np.array_equal(idx, old_idx):
                    break

                # Stop early
                iter += 1
                if iter > MAX_ITER:
                    break
            return idx

        def produce_histogram(filename):
            # Converts the full size image to 64X64 image and returns the histogram of the pixel values in 10 bins
            image = Image.open(filename)
            image.thumbnail((64,64))
            img_shape = image.size
            hist = [float(distribution)/float(img_shape[0]*img_shape[1])*100 for distribution in np.histogram(np.asarray(image), bins=10)[0]]
            return hist

        def get_data(filenames):
            # Returns an array of histograms of all photos of a folder in which photos are located
            # All 12 processors are pooled together to make computation faster
            pool = Pool(12)
            mapping = pool.map(produce_histogram, filenames)
            pool.close()
            pool.join()
            return mapping

        # def find_dupes(hists):

        # find the filenames of all photos in the selected folder
        image_paths = glob.glob(location + '/*jpg')
        num_images = len(image_paths)

        # get histogram data
        hist_data = get_data(image_paths)
        print(hist_data[0])

        # get the index of each image after HA clustering
        # k_idx = np.int32(KMeansCLustering(hist_data, k_numbers))
        k_idx = np.int32(HAClustering(hist_data))

        # Make directories for each cluster
        for i in range(len(np.unique(k_idx))):
            os.makedirs(location+'/'+str(i+1))

        # Copy photos with same cluster index into same folder
        create = shutil.copy
        for i in range(num_images):
            create(image_paths[i], location + '/' + str(k_idx[i] + 1) + '/')

        print("Done")

        return

    def similar_images():

        def view_photos1():
            gtk1 = tk.Toplevel()

            def getFileName(image_filename):
                print(str(image_filename))
                gtk1.destroy()

                # Get all images in the folder
                image_paths = glob.glob(location + '/*jpg')
                images = image_all(image_paths)

                img_file_1 = location +'/'+ image_filename

                # create an instance
                image_1 = cv2.imread(img_file_1, 0)

                for i, thumb in enumerate(images):
                    thumb = np.asarray(thumb)
                    orb_index = cv2.ORB_create()
                    k1, d1 = orb_index.detectAndCompute(image_1, None)
                    k2, d2 = orb_index.detectAndCompute(thumb, None)

                    temp = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)

                    # matching objects
                    matching_index = temp.match(d1, d2)

                    matching_regions = []

                    for r in matching_index:
                        if (r.distance < 50):
                            matching_regions.append(r)

                    if len(matching_index) == 0:
                        print("similar_index", 0)

                    else:
                        similar_index = len(matching_regions) / len(matching_index)
                        print("similar_index:", similar_index)

                        if (similar_index > 0.1):
                            path = "similar_pics/"
                            create = shutil.copy
                            create(image_paths[i], path)

            # Display all images from which one image can be selected to find similar images of
            row = 0
            column = 0
            for images in os.listdir(str(location)):
                if images.endswith("jpg"):
                    im = Image.open(location + '/' + images)
                    im.thumbnail((64, 64))
                    tkimage = ImageTk.PhotoImage(im)
                    handler = lambda img=images: getFileName(img)  # here modify
                    imageButton = tk.Button(gtk1, image=tkimage, command=handler)  # here
                    imageButton.image = tkimage
                    imageButton.grid(row=row, column=column)
                    if column == 8:
                        row += 1
                        column = 0
                    else:
                        column += 1
            gtk1.mainloop()
            return

        def read_images(filename):
            # Read Image
            img = Image.open(filename)
            return img

        def image_all(filenames):
            # returns a list of all images
            # All 12 processors pooled to read images to increase performance
            pool = Pool(12)
            mapping = pool.map(read_images, filenames)
            pool.close()
            pool.join()

            return mapping

        # Run the program
        view_photos1()
        return

    def stylize_image():
        # Provides window to apply different types of filters on an image of the users choice
        def view_photos2():
            # Display all photos so that the user can select an image
            gtk2 = tk.Toplevel()

            def getFileName(image_filename):
                # Gets the address of the file user has selected
                print(str(image_filename))
                gtk2.destroy()

                # Read image
                img_file_1 = location + '/' + image_filename
                # rgb image
                img = cv2.imread(img_file_1)

                def vignette_filter():
                    # Displays the image with a black outline mask
                    rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

                    # vintage filter

                    img_shape = img.shape

                    row = img_shape[0]
                    col = img_shape[1]

                    # filter

                    kernel_1 = cv2.getGaussianKernel(col, 500)
                    kernel_2 = cv2.getGaussianKernel(row, 500)

                    main_kernel = kernel_1 * kernel_2.T

                    fil = 8000 * np.sin(np.linalg.norm(main_kernel)) * (255 * main_kernel) / (
                        np.linalg.norm(main_kernel))
                    fil = fil.T

                    print(fil.shape)
                    #
                    # vintage_image = np.zeros(img.shape)
                    vintage_image = np.copy(rgb_img)
                    vintage_image = np.int32(vintage_image)
                    print(vintage_image.shape)
                    for i in range(3):
                        vintage_image[:, :, i] = vintage_image[:, :, i] * fil

                    # Normalize
                    vintage_image[vintage_image > 255] = 255
                    vintage_image[vintage_image < 0] = 0
                    vintage_image = np.uint8(vintage_image)
                    plt.figure()
                    plt.title("Vignette Filter")
                    plt.imshow(vintage_image)
                    plt.show()

                def edges():
                    # Returns edge intensity image
                    # Edge is combination of horizontal and vertical sobel filter convolved through image
                    img = cv2.imread(img_file_1)
                    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

                    def gradient_x(img1):
                        kernel = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
                        img1 = cv2.GaussianBlur(img1, (5, 5), 0)
                        img1 = np.float32(img1)
                        grad_img = ndimage.convolve(img1, kernel)
                        return grad_img

                    def gradient_y(img1):
                        kernel = np.transpose(np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]))
                        img1 = cv2.GaussianBlur(img1, (5, 5), 0)
                        img1 = np.float32(img1)
                        grad_img = ndimage.convolve(img1, kernel)
                        return grad_img

                    def gradient_mag(gx, gy):
                        grad_img = np.hypot(gx, gy)
                        return grad_img

                    # Calculating horizontal and vertical gradient values
                    grad_x = gradient_x(img)
                    grad_y = gradient_y(img)

                    # total gradient
                    grad = gradient_mag(grad_x, grad_y)
                    plt.figure()
                    plt.title("Edges Filter")
                    plt.imshow(grad, cmap='gray')
                    plt.show()

                def gotham_filter():
                    image_name = img_file_1

                    # Reading the image
                    image = cv2.imread(image_name, 1)

                    # convert BGR TO RGB
                    img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

                    # converting the int type to float type
                    img = skimage.img_as_float(img)

                    # spitting the channels
                    red_pixels = img[:, :, 0]
                    green_pixels = img[:, :, 1]
                    blue_pixels = img[:, :, 2]

                    # function for sharpening the image
                    def sharpe_image(img, a, b):
                        # blur the image first
                        blur_image = cv2.GaussianBlur(img, (5, 5), 0)

                        # sharpening the image
                        sharp_image = np.clip(filtered_image * a - blur_image * b, 0, 1.0)

                        return sharp_image

                    # interpolation of red pixels

                    data_points = []
                    for i in range(11):
                        data_points.append(i * 0.1)

                    temp_red = np.interp(red_pixels.flatten(), np.linspace(0, 1, len(data_points)), data_points)
                    red_pixels = temp_red.reshape(red_pixels.shape)

                    # increase the color of blue
                    blue_pixels = np.clip(blue_pixels + 0.2, 0, 1.0)

                    # increse the color of green
                    # green_pixels = np.clip(green_pixels+0.1,0,1.0)

                    # filtered image
                    filtered_image = np.stack([red_pixels, green_pixels, blue_pixels], axis=2)

                    # sharpening the image
                    sharp_image = sharpe_image(filtered_image, 1.5, 0.5)

                    # spitting the channels
                    red_pixels = sharp_image[:, :, 0]
                    green_pixels = sharp_image[:, :, 1]
                    blue_pixels = sharp_image[:, :, 2]

                    # interpolation of blue pixels
                    data_points = [0, 0.047, 0.118, 0.251, 0.318, 0.392, 0.42, 0.439, 0.475, 0.561, 0.58, 0.627, 0.671,
                                   0.733, 0.847, 0.925, 1]

                    temp_blue = np.interp(blue_pixels.flatten(), np.linspace(0, 1, len(data_points)), data_points)

                    blue_pixels = temp_blue.reshape(blue_pixels.shape)

                    # final filter
                    final_filter_image = np.stack([red_pixels, green_pixels, blue_pixels], axis=2)

                    # show the final filtered image
                    plt.subplot(1, 2, 1)
                    plt.imshow(img)
                    plt.title("Original image")
                    plt.subplot(1, 2, 2)
                    plt.imshow(final_filter_image)
                    plt.title("Gotham Filter image")
                    plt.show()

                # Pop up window of different types of filter selection
                filters_window = tk.Toplevel()
                filters_window.title("Filters")

                filter_commands = [vignette_filter, edges, gotham_filter]
                filter_buttons = ["Vignette Filter", "Edges Filter", 'Gotham Filter']

                for button3 in range(len(filter_buttons)):
                    create1 = tk.Button(master=filters_window, relief=tk.RAISED, borderwidth=2, command=filter_commands[button3], height=5, bg='#291e1d', fg='white', text=filter_buttons[button3], font=('Helvetica 18 bold', 15))
                    create1.grid(row=0, column=button3)

                filters_window.mainloop()

            # Window to show all photos. User can select a photo to apply filter
            row = 0
            column = 0
            for images in os.listdir(str(location)):
                if images.endswith("jpg"):
                    im = Image.open(location + '/' + images)
                    im.thumbnail((64, 64))
                    tkimage = ImageTk.PhotoImage(im)
                    handler1 = lambda img=images: getFileName(img)  # here modify
                    imageButton = tk.Button(gtk2, image=tkimage, command=handler1)  # here
                    imageButton.image = tkimage
                    imageButton.grid(row=row, column=column)
                    if column == 8:
                        row += 1
                        column = 0
                    else:
                        column += 1
            gtk2.mainloop()

        view_photos2()
        return

    def people_photos():
        def read_images(filename):
            img = cv2.imread(filename, 0)
            # img.thumbnail((64, 64))
            return img

        def image_all(filenames):
            pool = Pool(12)
            mapping = pool.map(read_images, filenames)
            pool.close()
            pool.join()

            return mapping

        image_paths = glob.glob(location+'/*jpg')
        all_images = image_all(image_paths)

        for i, image_one in enumerate(all_images):
            face_detect = db.get_frontal_face_detector()

            all_faces = face_detect(image_one)

            print(all_faces)

            if len(all_faces) > 0:
                path = "Faces/"
                create = shutil.copy
                create(image_paths[i], path)

        return

    # Back and close button functions
    def go_back():
        sort_window.destroy()
        main_menu()

    def close():
        sort_window.destroy()

    # Displays window to select the type of function needed by the user to apply for their photos
    sort_window = tk.Tk()
    sort_window.title("Tools")
    logo1 = ImageTk.PhotoImage(file='Logo.png')

    logo_frame = tk.Frame(master=sort_window)
    logo_frame.pack()
    logo_label1 = tk.Label(master=logo_frame, image=logo1)
    logo_label1.pack(side=tk.TOP)

    m2_back_button = tk.Button(master=logo_frame, text="Go Back", command= go_back)
    m2_back_button.pack()
    m2_close_button = tk.Button(master=logo_frame, text="Close", command=close)
    m2_close_button.pack()

    m2_command_frame = tk.Frame(master=sort_window)
    m2_command_frame.pack()

    commands = [view_photos, find_duplicates, similar_images, stylize_image, people_photos]
    buttons = ["View All Photos", "Find Duplicate Photos", "Find similar images", "Stylize Images", "See all photos with face"]
    for button2 in range(len(buttons)):
        create = tk.Button(relief=tk.RAISED, borderwidth=2, command=commands[button2], master=m2_command_frame, height=5, bg='#291e1d', fg='white', text=buttons[button2], font=('Helvetica 18 bold', 15))
        create.grid(row=0, column=button2)
    sort_window.mainloop()

# region main_menu of the GUI
def main_menu():
    # Command for open button
    def open_folder():
        mm_entry.delete(0, tk.END)
        folder_path = askdirectory()
        mm_entry.insert(0, folder_path)

    def copy_location():
        # Extract the folder location in the entry box
        photos_location = mm_entry.get()
        main_menu.destroy()
        sorting_menu(photos_location)

    def mm_destroy():
        # Command for cancel button
        main_menu.destroy()

    # Initializing the main menu window with logo
    main_menu = tk.Tk()
    logo = ImageTk.PhotoImage(file='Logo.png')

    main_menu.title("Select Directory")

    # Placement of labels buttons and entry boxes in the window
    open_frame = tk.Frame(relief=tk.SUNKEN, borderwidth=3)
    open_frame.pack(fill=tk.X)

    logo_label = tk.Label(master=open_frame, image=logo)
    logo_label.grid(row=0, column=0)
    logo_line = tk.Label(master=open_frame, text="Image Sorter", font='Helvetica 18 bold')
    logo_line.grid(row=0, column=1, sticky='w')

    mm_label = tk.Label(master=open_frame, text="Select the directory of Photos")
    mm_label.grid(row=1, column=1, sticky='w')
    load_button = tk.Button(master=open_frame, text="Open", command=open_folder)
    load_button.grid(row=2, column=0, sticky='ew')

    mm_entry = tk.Entry(master=open_frame, width=80)
    mm_entry.grid(row=2, column=1)

    command_frame = tk.Frame( borderwidth=3)
    command_frame.pack(fill=tk.X)

    mm_ok_button = tk.Button(master=command_frame, text='OK', command=copy_location)
    mm_ok_button.pack(side=tk.RIGHT)
    mm_cancel_button = tk.Button(master=command_frame, text='Cancel', command=mm_destroy)
    mm_cancel_button.pack(side=tk.RIGHT)

    # Wait for an input
    main_menu.mainloop()

    # endregion

# Run the program
main_menu()


import numpy as np
from scipy.spatial.distance import minkowski
from .base import ClusteringMethod

class KMeans(ClusteringMethod):

  def __init__(self, epochs, n_clusters, distance_method=minkowski):
    self.epochs     = epochs
    self.n_clusters = n_clusters
    self.distance_method = distance_method
    # Extra
    self.reuse = False

  def fit(self, X):
    self.init_centroids(X)
    grupos = None
    for e in range(self.epochs):
        if (e%20 == 0):
          print(f">> epoch: {e}")
        grupos = self.get_groups(X)
        new_c = self.new_centroids(X, grupos)
        #diff = np.array(self.centroids) - np.array(new_c)
        self.centroids = new_c
    return grupos

  def get_groups(self, X):
    groups = []
    for x in X:
        dist = np.array( [self.distance_method(c, x) for c in self.centroids])
        g    = np.argmin( dist )
        groups.append(g)
    groups = np.array(groups)
    return groups

  def new_centroids(self, X, grupos):
    new_c = []
    for c in range(self.n_clusters):
        elements = X[grupos == c]
        if len(elements) > 0:
          new_c.append(np.mean(elements, axis = 0))
        else:
          if not self.reuse:
            print("REUSE")
            self.reuse=True
          new_c.append(self.centroids[c])
    new_c = np.array(new_c)

    return new_c

  def init_centroids(self, X):
    # Select centoids
    self.centroids = X[np.random.choice(len(X), self.n_clusters, replace=False)]


import logging
import uuid

import configmanager as C
import errors as errs
import helpers
import lexconfig
import morphparser as mp
import parseparadigms as pp
import pextract as pex

def add_paradigm(lexicon, pid, paradigm, paradigms, identifier, pos, classes):
    """
    Add a new paradigm (update language model, save to karp).
    Args:
        lexicon (str): the lexicon name
        pid (string): the paradigms name (human readable id)
        paradigm (obj): the paradigm (Paradigm.py)
        paradigms (list): a list of all internal paradigms
        identifier (str): the tables identifier
        pos (str): the tables word class
        classes (dict): inflectional classes to add '{"bklass": [3,4]}'
    """
    presource = lexconfig.get_paradigmlexicon(lexicon)
    lresource = lexconfig.get_lexiconname(lexicon)
    logging.debug('id %s, para %s.\n classes %s, identifier %s',
                  pid, paradigm, classes, identifier)
    # Set name, id etc for the paradigm object
    paradigm.set_id(pid)
    puuid = str(uuid.uuid1())  # generate a uid and set it
    paradigm.set_uuid(puuid)
    paradigm.set_lexicon(presource)
    paradigm.set_pos(pos)
    paradigm._entries = 1
    for key, val in classes.items():
        paradigm.add_class(key, [val])
    logging.debug('uuid %s', puuid)

    # Add to our in memory paradigms and update language model
    if lresource not in paradigms:
        paradigms[lresource] = {}
    if pos not in paradigms[lresource]:
        paradigms[lresource][pos] = {}
    all_paras, numex, lms, alpha = paradigms[lresource].get(pos, ({}, 0, None))
    alpha = mp.extend_alphabet(paradigm, alpha)
    mp.lms_paradigm(paradigm, lms, alpha, lexconfig.get_ngramorder(lexicon),
                    lexconfig.get_ngramprior(lexicon))
    all_paras[puuid] = paradigm
    paradigms[lresource][pos] = (all_paras, numex+1, lms, alpha)

    # Save to karp
    helpers.karp_add(paradigm.jsonify(), resource=presource, _id=puuid)

def add_word_to_paradigm(lexicon, paradigm, paradigms, identifier, pos,
                         classes, inst):
    """
    Add a word to extisting paradigm (update language model, save to karp)
    Args:
        lexicon (str): the lexicon name
        paradigm (obj): the paradigm (Paradigm.py)
        paradigms (dict): a dictionary with all paradigms
            '{"lexname": {"nn": [], "vb": []}'
        identifier (str): the tables identifier
        pos (str): the tables word class
        classes (dict): all inflectional classes to add '{"bklass": [3,4]}'
        inst (list): this table's variable instantiations for the paradigm
    """
    presource = lexconfig.get_paradigmlexicon(lexicon)
    lresource = lexconfig.get_lexiconname(lexicon)
    logging.debug('old count %s', paradigm.count)
    # Set variabel instannces, count, resource etc for the paradigm object
    var_inst = [('first-attest', identifier)]+list(enumerate(inst, 1))
    paradigm.add_var_insts(var_inst)
    paradigm.members.append(identifier)
    for key, val in classes.items():
        paradigm.add_class(key, [val])
    paradigm.count += 1
    paradigm.set_lexicon(presource)

    # Update our in memory paradigms and update language model
    all_paras, numex, lms, alpha = paradigms[lresource].get(pos, ({}, 0, None))
    alpha = mp.extend_alphabet(paradigm, alpha)
    paradigms[lresource][pos] = (all_paras, numex, lms, alpha)
    logging.debug('new count %s', paradigm.count)

    # Save to karp
    helpers.karp_update(paradigm.uuid, paradigm.jsonify(), resource=presource)

def remove_word_from_paradigm(lexicon, paradigm, paradigms, identifier, pos):
    """
    Remove a word from a paradigm  (update language model, save to karp)
    Args:
        lexicon (str): the lexicon name
        paradigm (obj): the paradigm (Paradigm.py)
        paradigms (dict): a dictionary with all paradigms
            '{"lexname": {"nn": [], "vb": []}'
        identifier (str): the tables identifier
        pos (str): the tables word class
    """
    logging.debug('old count %s', paradigm.count)
    # remove variable_instances
    for ix, var_inst in enumerate(paradigm.var_insts):
        if dict(var_inst).get('first-attest', '') == identifier:
            paradigm.var_insts.pop(ix)
            break

    # remove from members
    try:
        paradigm.members.remove(identifier)
    except Exception as e:
        logging.exception(e)
        logging.warning('identifier %s cannot be removed, since it was not present in %s',
                        identifier, paradigm.p_id)

    # remove the words classes from the paradigm if they are not used by other
    # words, by first removing all and then readding the all classes used by
    # other words
    paradigm.empty_classes()
    # get all classes from all other words, add to paradigm:
    q = 'extended||and|paradigm|equals|%s||not|%s|equals|%s'\
        % (paradigm.p_id, lexconfig.get_identifierfield(lexicon), identifier)
    for iclass in lexconfig.get_inflectionalclassnames(lexicon):
        res = helpers.karp_query('statlist',
                                 {'q': q, 'size': 1000,
                                  'buckets': '%s.bucket' % iclass},
                                 resource=lexconfig.get_lexiconname(lexicon),
                                 mode=lexconfig.get_lexiconmode(lexicon))
        paradigm.add_class(iclass, [cl[0] for cl in res.get('stattable', [])])

    # decrease paradigm member count
    paradigm.count -= 1
    logging.debug('new count %s', paradigm.count)
    if paradigm.count > 0:
        # TODO remove from alphabet
        helpers.karp_update(paradigm.uuid, paradigm.jsonify(),
                            resource=lexconfig.get_paradigmlexicon(lexicon))
    else:
        remove_paradigm(lexicon, paradigm, paradigms, pos)

def remove_paradigm(lexicon, paradigm, paradigms, pos):
    """
    Remove a  paradigm  (update language model, delete from karp)
    Args:
        lexicon (str): the lexicon name
        paradigm (obj): the paradigm (Paradigm.py)
        paradigms (dict): a dictionary with all paradigms
            '{"lexname": {"nn": [], "vb": []}'
        pos (str): the tables word class
    """
    helpers.karp_delete(paradigm.uuid, lexconfig.get_paradigmlexicon(lexicon))
    lresource = lexconfig.get_lexiconname(lexicon)
    all_paras, numex, lms, alpha = paradigms[lresource].get(pos, ({}, 0, None))
    del all_paras[paradigm.uuid]
    del lms[paradigm.uuid]
    alpha = mp.paradigms_to_alphabet(all_paras.values())
    # update the internal model, save the recomputed alphabet
    paradigms[lresource][pos] = all_paras, numex-1, lms, alpha

def inflect_table(lexicon, table, paradigms, identifier, pos, ppriorv=None,
                  kbest=10, match_all=False):
    """
    Find matching paradigms for an inflectiontable, possibly by adding
    word forms to the original table.
    Args:
        lexicon (str): the lexicon name
        table (list): the comma separated word forms, possibly with msds.
                      "katt,katter|pl indef nom,katts"
        paradigms (dict): a dictionary with all paradigms
            '{"lexname": {"nn": [], "vb": []}'
        identifier (str, optional): the word's identifier
        pos (str): the tables word class
        ppriorv (int, optional): setting for pextract.
            Defaults to the lexicon config's value.
        kbest (int, optional): how many suggestion to return. Defaults to 10.
        match_all (bool, optional): must the output strictly match the input?
             Defaults to false. Set to true if the input table is complete and
             no extra forms may be added.
    Returns:
       {"Results": [results]}
         where a result consits of:
           {score: float, paradigm: str, new: bool,
            identifier: str,
            baseform: str,
            variables: dict of variable instansiations,
            WordForms: [{writtenForm: str, msd: str}]
            partOfSpeech: str,
            paradigm_entries: int,
            }
        If the paradigm is new, the key 'pattern' (str) will show a compact
        represention of the paradigm.
        Possibly more, lexicon specific fields are added.
    """
    lexconf = lexconfig.get_lexiconconf(lexicon)
    restrict_baseform = helpers.read_restriction(lexconf)
    paras, numex, lms = helpers.relevant_paradigms(paradigms, lexicon, pos)
    fill_tags = '|' in table
    pex_table = helpers.tableize(table, add_tags=False, fill_tags=fill_tags)
    logging.debug('inflect forms %s msd %s. Restricted %s',
                  pex_table[0], pex_table[1], restrict_baseform)
    res = []
    if paras:
        # if there are known paradigms for the current lexicon and part of
        # speech, try these
        if ppriorv is None:
            ppriorv = lexconfig.get_pprior(lexicon)
        res = mp.test_paradigms(pex_table, paras, numex, lms,
                                C.config["print_tables"], C.config["debug"],
                                ppriorv, returnempty=False,
                                baseform=restrict_baseform,
                                match_all=match_all)
    if res:
        ans = {"Results": helpers.format_kbest(lexicon, res, kbest=kbest,
                                               pos=pos,
                                               lemgram=identifier)}
    else:
        # if no results are found (no matching paradigms), make a new paradigm
        # giving exactly the input forms
        logging.debug('invent! %s', len(res))
        pex_table = helpers.tableize(table, add_tags=True,
                                     identifier=identifier)
        paradigm = pex.learnparadigms([pex_table])[0]
        logging.debug('learned %s', paradigm)
        ans = {'Results': [helpers.karp_tableize(lexicon, table,
                                                 paradigm=paradigm,
                                                 pos=pos,
                                                 identifier=identifier)]}
    return ans

def make_new_table(lexicon, table, paradigm, paradigms, identifier, baseform,
                   pos, classes, ppriorv=None, newword=False, newpara=False):
    """
    Check that the given table and identifiers are ok, that the table matches
    given paradigm and then add the table to the paradigm,
    in karp and internally.
    Args:
        lexicon (str): the lexicon name
        table (str): the comma separated word forms, possibly with msds.
            "katt,katter|pl indef nom,katts"
        paradigm (str): a paradigm name
        paradigms (dict): a dictionary with all paradigms
            '{"lexname": {"nn": [], "vb": []}'
        identifier (str): the word's identifier
        baseform (str): the word's baseform
        pos (str): the tables word class
        classes (str): input classes 'bklass:2,cklass:3'
        ppriorv (int, optional): setting for pextract
        newword (bool, optional): is the word new? Defaults to false.
        newpara (bool, optional): is the paradigm new? Defaults to false.

    Returns:
       a tuple (identifier, wf_table, para, v, classes)
         where
         identifier (str): identifier of  the table
         wf_table (obj): the formatted inflection table
         para (obj): the paradigm object
         v (list): the variable instances
         classes (dict): the formatted classes of the word
    """
    # check that the table's identier is unique
    # if it is new: make sure that it doesn't already exists.
    #    otherwise: construct a new one
    # if it is not new: make sure that it does exists.
    #    otherwise: fail
    lresource = lexconfig.get_lexiconname(lexicon)
    ok = helpers.check_identifier(identifier,
                                  lexconfig.get_identifierfield(lexicon),
                                  lresource,
                                  lexconfig.get_lexiconmode(lexicon),
                                  unique=newword, fail=not newword)
    if not ok:
        logging.debug('identifier %s not ok', identifier)
        lexconf = lexconfig.get_lexiconconf(lexicon)  # TODO remove after refactoring helpers
        word = baseform or helpers.get_baseform(lexconf, identifier)
        identifier = helpers.make_identifier(lexicon, word, pos)
        logging.debug('\t...use %s', identifier)

    # make sure that all required fields are present
    required = [('identifier', identifier), ('paradigm', paradigm),
                ('partOfSpeech', pos)]
    for name, field in required:
        if not field:
            raise errs.MflException("Both identifier, partOfSpeech and paradigm must be given!",
                                    code="unknown_%s" % name)
    # parse input classes
    try:
        logging.debug('make classes: %s', classes)
        if classes:
            classes = dict([c.split(':') for c in classes.split(',')])
        else:
            classes = {}
    except Exception as e1:
        logging.warning('Could not parse classes')
        logging.error(e1)
        raise errs.MflException("Could not parse classes. Format should be 'classname:apa,classname2:bepa'",
                                classes,
                                code="unparsable_class")

    paras, numex, lms = helpers.relevant_paradigms(paradigms, lexicon, pos)
    if newpara and paradigm in [p.name for p in paras]:
        raise errs.MflException("Paradigm name %s is already used" % (paradigm),
                                code="unique_paradigm")

    # parse the input table into pextract and karp format
    pex_table = helpers.tableize(table, add_tags=False)
    wf_table = helpers.karp_wftableize(lexicon, paradigm, table, classes,
                                       baseform=baseform, identifier=identifier,
                                       pos=pos,
                                       resource=lresource)
    if newpara:
        # prepare to try all paradigm
        fittingparadigms = paras
        # check that this is a new name
        helpers.check_identifier(paradigm, pp.id_field,
                                 lexconfig.get_paradigmlexicon(lexicon),
                                 lexconfig.get_paradigmmode(lexicon))

    else:
        logging.debug('not new, look for %s', paradigm)
        # prepare to try only the given paradigm
        fittingparadigms = [p for p in paras if p.p_id == paradigm]
        if not fittingparadigms:
            raise errs.MflException("Could not find paradigm %s" % paradigm,
                                    code="unknown_paradigm")

    logging.debug('fitting %s', fittingparadigms)
    logging.debug('pex_table %s %s', pex_table[0], pex_table[1])
    if ppriorv is None:
        ppriorv = lexconfig.get_pprior(lexicon)

    ans = mp.test_paradigms(pex_table, fittingparadigms, numex, lms,
                            C.config["print_tables"], C.config["debug"],
                            ppriorv, returnempty=False,
                            match_all=True)

    if newpara and ans:
        logging.warning("Could inflect %s as %s", table, ans[0][1].p_id)
        raise errs.MflException("Table should belong to paradigm %s" % (ans[0][1].p_id),
                                code="inflect_problem")

    if not newpara and not ans:
        logging.warning("Could not inflect %s as %s", table, paradigm)
        raise errs.MflException("Table can not belong to paradigm %s" % (paradigm),
                                code="inflect_problem")

    if ans:
        # found a match, add the word to this paradig
        score, para, v = ans[0]
        add_word_to_paradigm(lexicon, para, paradigms, identifier, pos,
                             classes, v)

    else:
        # if no results are found (no matching paradigms), make a new paradigm
        # giving exactly the input forms
        pex_table = helpers.tableize(table, add_tags=False,
                                     identifier=identifier)
        para = pex.learnparadigms([pex_table])[0]
        v = [var for ix, var in para.var_insts[0][1:]]
        add_paradigm(lexicon, paradigm, para, paradigms, identifier, pos,
                     classes)

    return identifier, wf_table, para, v, classes


from django.db import models

from product.models import ProductVariation

from account.models import Address
from django.conf import settings  # Here

# Create your models here.
class Order(models.Model):

    variation = models.ForeignKey(ProductVariation,
                                  on_delete=models.DO_NOTHING)
    created_at = models.DateField(auto_now_add=True)
    uploaded_at = models.DateField(auto_now=True)
    quantity = models.IntegerField()
    total = models.FloatField()
    user = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE)

    def unity_price(self):
        return self.total / self.quantity

class Orders(models.Model):
    STATUS_CHOICES = (('C', 'CREATED'), ('P', 'PAID_TO_DELIEVER'),('D', 'DELIVERED'))

    order = models.ManyToManyField(Order)
    created_at = models.DateField(auto_now_add=True)
    uploaded_at = models.DateField(auto_now=True)
    user = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE)

    address = models.ForeignKey(Address,
                                on_delete=models.DO_NOTHING,
                                blank=False,
                                null=False)
    status = models.CharField(max_length=1,
                              choices=STATUS_CHOICES,
                              blank=False,
                              null=False)
    total_paid = models.FloatField()


import numpy as np
import SimpleITK as sitk
from skimage import transform
import os
import tre3d
import random
from scipy import ndimage
import math

mode = 'E1'

BASE_64_DIR = '/mnt/hd3/new3dDataset_64/'
TRAIN_64 = BASE_64_DIR + 'train/'
TEST_64 = BASE_64_DIR + 'test/'
PATH = '/home/puwenbo/puwenbo/Dataset/new3dDataset/test/'

if mode == 'E1':
    TEST_64 = TEST_64 + 'E1/'
    OUTPUT = '/home/puwenbo/puwenbo/Result/E1/'

if mode == 'E2':
    TEST_64 = TEST_64 + 'E2/'
    OUTPUT = '/home/puwenbo/puwenbo/Result/E2/'

p = sitk.ParameterMap()
p['FixedInternalImagePixelType'] = ['float']
p['MovingInternalImagePixelType'] = ['float']
p['FixedImageDimension'] = ['3']
p['MovingImageDimension'] = ['3']
p['UseDirectionCosines'] = ['true']
p['Registration'] = ['MultiResolutionRegistration']
p['Interpolator'] = [ "BSplineInterpolator"]
p['ResampleInterpolator'] = ["FinalBSplineInterpolator"]
p['Resampler'] = ["DefaultResampler"]
p['FixedImagePyramid'] = ["FixedRecursiveImagePyramid"]
p['MovingImagePyramid'] = ["MovingRecursiveImagePyramid"]
p['Optimizer'] = ["AdaptiveStochasticGradientDescent"]
p['Transform'] =  ["EulerTransform"]
p['Metric'] = ["AdvancedMattesMutualInformation"]
p['AutomaticScalesEstimation'] = [ "true"]
p['AutomaticTransformInitialization'] = ["true"]
p['HowToCombineTransforms'] = ["Compose"]
p['NumberOfHistogramBins']=['32']
p['ErodeMask'] = ['false']
p['NumberOfResolutions'] = ['4']
p['MaximumNumberOfIterations'] = ['250']
p['NumberOfSpatialSamples'] = ['2048']
p['NewSamplesEveryIteration'] = ['true']
p['ImageSampler'] = ['Random']
p['BSplineInterpolationOrder'] = ['1']
p['FinalBSplineInterpolationOrder'] = ['3']
p['DefaultPixelValue'] = ['0']

def get64():
    if not os.path.exists(PATH):
        os.mkdir(PATH)
    if not os.path.exists(PATH):
        os.mkdir(PATH)
    for dir in os.listdir('/mnt/hd3/new3dDataset/test'):
        if not os.path.exists(PATH+dir):
            os.mkdir(PATH+dir)
        for CT in os.listdir('/mnt/hd3/new3dDataset/test/'+dir+'/CT'):
            imageCT = sitk.ReadImage('/mnt/hd3/new3dDataset/test/'+dir+'/CT/'+CT)
            imageData = sitk.GetArrayFromImage(imageCT)
            res = transform.resize(imageData, (64,64,64))
            if not os.path.exists(PATH+dir+'/CT/'):
                os.mkdir(PATH+dir+'/CT/')
            sitk.WriteImage(sitk.GetImageFromArray(res), PATH+dir+'/CT/'+CT)
        for MR in os.listdir('/mnt/hd3/new3dDataset/test/'+dir+'/MR'):
            imageMR = sitk.ReadImage('/mnt/hd3/new3dDataset/test/' + dir + '/MR/' + MR)
            imageData = sitk.GetArrayFromImage(imageMR)
            res = transform.resize(imageData, (64, 64, 64))
            if not os.path.exists(PATH + dir + '/MR/'):
                os.mkdir(PATH + dir + '/MR/')
            sitk.WriteImage(sitk.GetImageFromArray(res), PATH + dir + '/MR/' + MR)

def getMovingIMG():
    for patient in os.listdir(TEST_64 + 'ground'):
        os.mkdir(OUTPUT + patient) if not os.path.exists(OUTPUT + patient) else print(patient)
        image = sitk.ReadImage(TEST_64+'ground/'+patient)
        imageData = sitk.GetArrayFromImage(image)
        T = getRandomT()
        # print(T)
        res = ndimage.affine_transform(imageData, tre3d.get_transform_matrix(T))
        sitk.WriteImage(sitk.GetImageFromArray(normalization(res)), OUTPUT + patient + '/moving.nii')
        np.savetxt(OUTPUT + patient + '/T', T)

def getRandomT():
    shift_z = random.randint(-10, 10)
    shift_y = random.randint(-10, 10)
    shift_x = random.randint(-10, 10)
    rotation_z = random.randint(-20, 20)
    # rotation_y = random.randint(-45, 45)
    # rotation_z = random.randint(-45, 45)
    # rotation_x = 0
    rotation_y = 0
    rotation_x = 0
    T = np.array([shift_z, shift_y, shift_x, rotation_z, rotation_y, rotation_x, 1])
    return T

def normalization(new_img):
    # ary = sitk.GetArrayFromImage(img)
    # new_img = np.copy(img)
    new_img[new_img < 0] = 0
    new_img = (new_img - np.min(new_img)) / (np.max(new_img) - np.min(new_img))
    return new_img

def elastix():
    for patient in os.listdir(TEST_64 + '/fixed'):
        try:
            elastixImageFilter = sitk.ElastixImageFilter()
            fix = sitk.ReadImage(TEST_64 + '/fixed/' + patient)
            fix_data = sitk.GetArrayFromImage(fix)
            fix_norm = sitk.GetImageFromArray(normalization(fix_data))
            elastixImageFilter.SetFixedImage(fix_norm)
            elastixImageFilter.SetMovingImage(sitk.ReadImage(OUTPUT + patient + '/moving.nii'))
            elastixImageFilter.SetParameterMap(p)
            if not os.path.isdir(OUTPUT+patient):
                os.mkdir(OUTPUT+patient)
            elastixImageFilter.SetOutputDirectory(OUTPUT+patient)
            elastixImageFilter.Execute()
            res = normalization(sitk.GetArrayFromImage(elastixImageFilter.GetResultImage()))
            sitk.WriteImage(sitk.GetImageFromArray(res), OUTPUT+patient+'/output.nii')
        except Exception:
            print(patient)

def getTrans(patient):
    file = open(OUTPUT + patient + '/TransformParameters.0.txt')
    TransformParameters = []
    for line in file:
        if line.startswith('(TransformParameters'):
            matrix = line.split(')')[0]
            matrix = matrix.split(' ')
            matrix = matrix[-6:]
            for i in range(6):
                if i > 2:
                    TransformParameters.append(math.degrees(float(matrix[5 - i])))
                else:
                    TransformParameters.append(float(matrix[5 - i]))
            TransformParameters.append(1)
            # print(TransformParameters)
            break
    file.close()
    return  TransformParameters

def get_t(patient):
    T = []
    patient = patient.split('.')[0]
    path = TEST_64
    with open(path+'ground2movingparameters.txt', 'r') as file:
        for line in file.readlines():
            parts = line.split('\t')
            name = parts[0]
            if patient == name:
                t = parts[1].strip('\n').strip('[').strip(']')
                t = t.split(' ')
                for i in t:
                    if i != '':
                        T.append(float(i))
                break
    return T

def getTRE(patient):
    TransformParameters = getTrans(patient)
    # T = get_t(patient)
    T = np.loadtxt(OUTPUT + patient + '/T')
    print(T)
    trans = np.array(TransformParameters, dtype=np.float32)
    # print(math.degrees(TransformParameters[3]))
    print(TransformParameters)
    res = tre3d.cal_tre_3d([64, 64, 64], T, trans)
    return res
#
# def getDice(patient):
#


import functools
import logging
from collections import defaultdict
from os.path import commonpath
from pathlib import Path
from typing import Any, Dict, List, Optional, Set, Tuple

import click

from fuzzing_cli.fuzz.config import FuzzingOptions
from fuzzing_cli.fuzz.ide import IDEArtifacts
from fuzzing_cli.fuzz.lessons import FuzzingLessons
from fuzzing_cli.fuzz.rpc.rpc import CONTRACT_ADDRESS, RPCClient
from fuzzing_cli.fuzz.types import Contract, SeedSequenceTransaction

LOGGER = logging.getLogger("fuzzing-cli")

class NoTransactionFound(Exception):
    pass

def _uniq(lst: List[str]) -> List[str]:
    """
    Remove duplicates from a list while preserving order
    """
    seen = set()
    return [x for x in lst if not (x in seen or seen.add(x))]

class CorpusRepository:
    def __init__(
        self,
        rpc: RPCClient,
        artifacts: IDEArtifacts,
        options: FuzzingOptions,
        corpus_target: Optional[str],
    ):
        self._corpus_target = corpus_target
        self._rpc = rpc
        self._artifacts = artifacts
        self._options = options
        (
            self._address_to_contract_mapping,
            self._contract_to_address_mapping,
            self._source_file_to_address_mapping,
        ) = self._construct_address_contract_mapping()
        self._initialize()

    def _initialize(
        self,
        addresses_under_test: Optional[List[str]] = None,
        targets: Optional[List[str]] = None,
    ) -> None:
        (
            self.contract_targets,
            self.source_targets,
        ) = self._construct_targets(addresses_under_test, targets)

        self.validation_errors = []
        self._validate()

    def apply_auto_fix(self, suggested_fixes: List[Dict[str, Any]]) -> None:
        """
        Apply the suggested fixes to the config file. Suggested fixes are of the following types:
            - Add a new address target
            - Add a new source target
            - Remove an address target
            - Remove a source target
        """
        addresses_under_test = self.contract_targets[:]  # make shallow copy of the list
        targets = self.source_targets[:]  # make shallow copy of the list
        for fix in suggested_fixes:
            if fix["type"] == "add_addresses":
                addresses_under_test.extend(fix["data"])
                continue
            if fix["type"] == "add_targets":
                targets.extend(fix["data"])
                continue
            if fix["type"] == "remove_addresses":
                addresses_under_test = [
                    address
                    for address in addresses_under_test
                    if address not in fix["data"]
                ]
                continue
            if fix["type"] == "remove_targets":
                targets = [target for target in targets if target not in fix["data"]]
                continue
        # Re-initialize the repository with the new targets and addresses, and validate
        self._initialize(addresses_under_test, targets)

    @property
    def _fuzzing_lessons(self) -> Tuple[Set[str], List[List[SeedSequenceTransaction]]]:
        """Get the fuzzing lessons to be added to the campaign's seed state"""
        blocks_to_skip: Set[str] = set([])
        suggested_seed_seqs: List[List[SeedSequenceTransaction]] = []
        for lesson in FuzzingLessons.get_lessons():
            click.secho(
                f"Lesson \"{lesson['description']}\" will be added to the campaign's seed state"
            )
            LOGGER.debug(
                f"Adding lesson \"{lesson['description']}\" to the campaign's seed state"
            )
            blocks_to_skip.update(
                {b["blockNumber"] for s in lesson["transactions"] for b in s}
            )
            suggested_seed_seqs.extend(lesson["transactions"])
        return blocks_to_skip, suggested_seed_seqs

    def _get_contract_by_address(
        self,
        contract_address: CONTRACT_ADDRESS,
    ) -> Optional[Contract]:
        """Get the artifacts of the contracts at the given addresses"""
        deployed_bytecode = self._rpc.get_code(contract_address)
        if deployed_bytecode is None:  # it's unknown contract
            LOGGER.warning(
                f'No deployed bytecode is found in an RPC node for contract: "{contract_address}"'
            )
            return None
        contract = self._artifacts.get_contract(deployed_bytecode)
        if not contract or contract.get("mainSourceFile", None) is None:
            LOGGER.warning(
                f'Contract "{contract_address}" could not be found in sources.'
                f" You can try to manually set the sources using the targets option. "
                f"More at: https://fuzzing-docs.diligence.tools/getting-started/configuring-the-cli#configuration"
            )
            return None
        return contract

    @property
    @functools.lru_cache(maxsize=1)
    def all_deployed_contracts_addresses(self) -> List[str]:
        """Get all the deployed contracts addresses from an RPC node excluding transactions from the fuzzing lessons"""
        blocks_to_skip, _ = self._fuzzing_lessons
        contracts_addresses = self._rpc.get_all_deployed_contracts_addresses(
            blocks_to_skip
        )
        return contracts_addresses

    @staticmethod
    def _contract_key(contract: Contract):
        return contract["mainSourceFile"], contract["contractName"]

    def _construct_address_contract_mapping(
        self,
    ) -> Tuple[Dict[str, Contract], Dict[Tuple[str, str], str], Dict[str, List[str]]]:
        address_to_contract_mapping = {}
        contract_to_address_mapping = {}
        # source file to contract addresses mapping. One source file can have multiple contracts, so we need to
        # keep track of all the contracts addresses for a given source file
        source_file_to_address_mapping = defaultdict(list)
        for contract_address in self.all_deployed_contracts_addresses:
            contract = self._get_contract_by_address(contract_address)
            if contract is None:
                address_to_contract_mapping[contract_address] = None
                continue
            address_to_contract_mapping[contract_address] = contract
            contract_to_address_mapping[self._contract_key(contract)] = contract_address
            source_file_to_address_mapping[contract["mainSourceFile"]].append(
                contract_address
            )

        return (
            address_to_contract_mapping,
            contract_to_address_mapping,
            source_file_to_address_mapping,
        )

    @staticmethod
    def _path_inclusion_checker(paths: List[str]):
        """Construct a function that checks if a given path is in the list of paths.
        The paths can be files or folders"""
        directory_paths: List[str] = []
        file_paths: List[str] = []
        for _path in paths:
            if Path(_path).is_dir():
                directory_paths.append(_path)
            else:
                file_paths.append(_path)

        def inner_checker(path: str):
            if path in file_paths:
                # we have found exact file match
                return True
            # try to find folder match
            for dir_path in directory_paths:
                if commonpath([dir_path, path]) == dir_path:
                    # file is in the directory
                    return True
            return False

        return inner_checker

    def _construct_targets(
        self,
        addresses_under_test: Optional[List[str]] = None,
        targets: Optional[List[str]] = None,
    ) -> Tuple[List[CONTRACT_ADDRESS], List[str]]:
        """
        Construct the targets from the addresses under test and the targets options
        or from provided addresses under test and targets (after prompting for automatic fixes).
        """
        if addresses_under_test is None:
            addresses_under_test = self._options.addresses_under_test
        if targets is None:
            targets = self._options.targets[
                :
            ]  # make a copy of the targets to not modify the original list

        return _uniq(addresses_under_test), _uniq(targets)

    @property
    def seed_state(self) -> Dict[str, any]:
        try:
            blocks_to_skip, suggested_seed_seqs = self._fuzzing_lessons
            processed_transactions = self._rpc.get_transactions(
                block_numbers_to_skip=list(blocks_to_skip)
            )

            if len(processed_transactions) == 0:
                raise NoTransactionFound(
                    "No transactions were found in an Ethereum RPC node"
                )

            setup = {
                "address-under-test": list(self.contract_targets)[0],
                "steps": processed_transactions,
                # make None if there are no other addresses under test to avoid sending empty list
                "other-addresses-under-test": list(self.contract_targets)[1:] or None,
            }

            """Get a seed state for the target contract to be used by Harvey"""
            if self._corpus_target:
                setup["target"] = self._corpus_target
            if len(suggested_seed_seqs) > 0:
                setup["suggested-seed-seqs"] = suggested_seed_seqs
            return {
                "discovery-probability-threshold": 0.0,
                "assertion-checking-mode": 1,
                "num-cores": self._options.number_of_cores,
                "analysis-setup": setup,
            }
        except NoTransactionFound:
            raise
        except Exception as e:
            raise Exception(
                "Unable to generate the seed state. Please check configuration and try again."
            ) from e

    def _validate(self):
        """
        Here we perform the following checks:
        1. All contracts deployed at the addresses specified in the addresses_under_test are found on the RPC.
        2. All contract deployed on the RPC are found in artifacts of the project.
        3. All contracts deployed on the RPC are provided as targets and as addresses under test.
        4. All contracts provided as addresses under test have a corresponding target (source file name)
           provided as well.
        5. All contracts provided as targets (source file name) have a corresponding deployed contract's
           provided as an address under test.
        6. All contracts provided as targets (source file name) has been deployed on the RPC.
        """
        unknown_contracts = []
        contracts_with_no_artifact = []
        # contract address was set as address under test, but its source file was not provided as a target
        source_target_not_set = []
        contract_target_not_set = []
        not_deployed_contracts = []

        processed_source_targets = [
            self._artifacts.normalize_path(s) for s in self.source_targets
        ]
        path_check = self._path_inclusion_checker(processed_source_targets)

        # These are the contracts that are deployed on the RPC, but we can't find them using the
        # `rpc.get_all_deployed_contracts_addresses`. This can happen if the contract was deployed
        # from another contract (factory contract).
        discovered_contracts = {}
        discovered_contract_to_address_mapping = {}

        for contract_address in self.contract_targets:
            # 1. All contracts deployed at the addresses specified in the addresses_under_test are found on the RPC.
            if contract_address not in self.all_deployed_contracts_addresses:
                contract = self._get_contract_by_address(contract_address)
                if contract:
                    discovered_contracts[contract_address] = contract
                    discovered_contract_to_address_mapping[
                        self._contract_key(contract)
                    ] = contract_address
                    continue
                unknown_contracts.append(contract_address)
                continue

            contract = self._address_to_contract_mapping.get(
                contract_address, None
            ) or discovered_contracts.get(contract_address, None)
            # 2. All contract deployed on the RPC are found in artifacts of the project.
            if contract is None:
                LOGGER.debug(
                    f'No artifact found for contract with address: "{contract_address}"'
                )
                contracts_with_no_artifact.append(contract_address)
                continue
            # 4. All contracts provided as addresses under test have a corresponding target provided as well.
            if not path_check(
                self._artifacts.normalize_path(contract["mainSourceFile"])
            ):
                # This is a contract that's been provided as address under test,
                # but its source file was not provided as a target
                source_target_not_set.append(
                    (
                        contract_address,
                        self._artifacts.normalize_path(contract["mainSourceFile"]),
                    )
                )

        for source_file_name in self.source_targets:
            # here we use path_check because target could be a directory,
            # so we need to check if the source file is in the directory
            # if the target is a file, then path_check will just check if the source file is the same as the target
            source_file_name = self._artifacts.normalize_path(source_file_name)
            _path_check = self._path_inclusion_checker([source_file_name])
            for contract in self._artifacts.contracts:
                if not _path_check(
                    self._artifacts.normalize_path(contract["mainSourceFile"])
                ):
                    continue
                contract_address = self._contract_to_address_mapping.get(
                    self._contract_key(contract), None
                ) or discovered_contract_to_address_mapping.get(
                    self._contract_key(contract), None
                )
                # 6. All contracts provided as targets (source file name) has been deployed on the RPC.
                if contract_address is None:
                    not_deployed_contracts.append(self._contract_key(contract))
                    continue
                # 5. All contracts provided as targets (source file name) have a corresponding
                # deployed contract's address provided as an address under test.
                if contract_address not in self.contract_targets:
                    contract_target_not_set.append(
                        (contract_address, source_file_name, contract["contractName"])
                    )
                    continue

        # 3. All contracts deployed on the RPC are provided as targets and as addresses under test.
        not_targeted_contracts = []
        targets_in_use = {
            *self.contract_targets,
            *[addr for addr, _, _ in contract_target_not_set],
            *[addr for addr, _ in source_target_not_set],
            *contracts_with_no_artifact,
        }
        for contract_address in self.all_deployed_contracts_addresses + list(
            discovered_contracts.keys()
        ):
            if contract_address in targets_in_use:
                continue
            contract = self._get_contract_by_address(contract_address)
            if contract is None:
                # This is the contract without the artifact which should be caught by the previous checks,
                # so we skip it here
                continue
            not_targeted_contracts.append(
                (
                    contract_address,
                    contract.get("mainSourceFile"),
                    contract.get("contractName"),
                )
            )

        if unknown_contracts:
            self.validation_errors.append(
                {
                    "type": "unknown_contracts",
                    "data": unknown_contracts,
                }
            )

        if contracts_with_no_artifact:
            self.validation_errors.append(
                {
                    "type": "contracts_with_no_artifact",
                    "data": contracts_with_no_artifact,
                }
            )

        if contract_target_not_set:
            self.validation_errors.append(
                {
                    "type": "contract_target_not_set",
                    "data": contract_target_not_set,
                }
            )

        if source_target_not_set:
            self.validation_errors.append(
                {
                    "type": "source_target_not_set",
                    "data": source_target_not_set,
                }
            )

        if not_deployed_contracts:
            self.validation_errors.append(
                {
                    "type": "not_deployed_contracts",
                    "data": not_deployed_contracts,
                }
            )

        if not_targeted_contracts:
            self.validation_errors.append(
                {
                    "type": "not_targeted_contracts",
                    "data": not_targeted_contracts,
                }
            )


import collections
import itertools
import random

from model.clusters import AbstractClusterRules
from utils import stats_utils, bdsa_utils

def aggregate_clusters(edges: list, eid2source2element: dict, score:int, cluster_rules: AbstractClusterRules,
                                  compare_previous_limit=None):
    """
    Aggregate former cluster (eid2source2element) with new edges
    :param edges:
    :param eid2source2element:
    :param cluster_rules:
    :param compare_previous_limit:
    :return:
    """
    for source2elements in eid2source2element.values():
        for source_pair in itertools.combinations(source2elements.keys(), 2):
            for element_pair in itertools.product(source2elements[source_pair[0]], source2elements[source_pair[1]]):
                edges.append(WeightedEdge(element_pair[0], element_pair[1], score))
    partition_using_agglomerative(edges, eid2source2element, cluster_rules, compare_previous_limit)

def partition_using_agglomerative(edges: list, eid2source2element: dict, cluster_rules: AbstractClusterRules,
                                  compare_previous_limit=None):
    """
    Cluster graph in this way:
    - Initially each node is a cluster
    - Foreach edge (sorted by decreasing weight):
     - If nodes for this edge are in different clusters, if a specific precondition for merging is met (tipically there
        should be elements from same source)
     - Stop at MIN_EDGE_WEIGHT
    Inspired (but different) by paper 'Clustering large probabilistic graph'
    :param edges: list of edges with nodes and weight. Each node must have an attribute 'source'.
    :param eid2source2element: the dictionary that represents the output
    :param cluster_rules: class that says if elements can be merged, and build cluster objects.
    :param: compare_previous_limit: Compute detailed similarity (costly) if only if gross similarity is bigger \
    than this. If None, never compute detailed similarity. Cf stats_utils.compare_clusters
    :return:
    """

    edges_sorted = sorted(edges, key=lambda e: (e.weight, e.node1, e.node2), reverse=True)
    node2cid = {}
    cid2cluster = {}
    cluster_counter = 0
    for edge in edges_sorted:
        node1 = edge.node1
        node2 = edge.node2
        c1 = node2cid[node1] if node1 in node2cid else None
        c2 = node2cid[node2] if node2 in node2cid else None
        if c1 is None and c2 is None:
            if cluster_rules.is_pair_mergeable(node1, node2):
                node2cid[node1] = cluster_counter
                node2cid[node2] = cluster_counter
                cid2cluster[cluster_counter] = cluster_rules.build_cluster(node1, node2)
                cluster_counter += 1
        elif c2 is None:
            added = cid2cluster[c1].add_element_if_allowed(node2)
            if added:
                node2cid[node2] = c1
        elif c1 is None:
            added = cid2cluster[c2].add_element_if_allowed(node1)
            if added:
                node2cid[node1] = c2
        # If 2 nodes already belong to the same cluster, there is nothing to do.
        elif c1 != c2:
            cmin = min(c1, c2)
            cmax = max(c1, c2)
            added = cid2cluster[cmin].merge_cluster(cid2cluster[cmax])
            if added:
                for node in cid2cluster[cmax].nodes():
                    node2cid[node] = cmin
                del cid2cluster[cmax]
    if compare_previous_limit:
        old_clusters_copy = {eid: frozenset(element for elements in source2element.values() for element in elements) for eid, source2element in eid2source2element.items()}
        eid2source2element.clear()
        new_cluster_copy = {cid: frozenset(cluster.nodes()) for cid, cluster in cid2cluster.items()}
        res = compare_clusters(old_clusters_copy, new_cluster_copy, compare_previous_limit)
        del old_clusters_copy
        del new_cluster_copy
    else:
        eid2source2element.clear()
        res = 0
    for eid, cluster in cid2cluster.items():
        for node in cluster.nodes():
            eid2source2element[eid][node.source].add(node)
    return res

class WeightedEdge:
    """
    Weighted edge, 1st and 2nd
    """
    def __init__(self, node_a, node_b, weight:float):
        self.weight = weight
        self.node1 = min(node_a, node_b)
        self.node2 = max(node_a, node_b)

    def __repr__(self):
        return "%s --- %s [%f]" % (self.node1, self.node2, self.weight)

    def __eq__(self, other):
        return self.node1 == other.node1 and self.node2 == other.node2

    def __hash__(self):
        return hash(self.node1) + hash(self.node2)

def nodes_flattener(source2elements):
    """
    Flatten nodes grouped in source2elements
    >>> nodes_flattener({10: {7,3,8}, 20: {6,4,2}})
    [2, 3, 4, 6, 7, 8]
    """

    return sorted(element for elements in source2elements.values() for element in elements)

def build_all_pairs_generic(clusters, nodes_retriever=nodes_flattener):
    """
    Build all pairs of nodes in same cluster.
    :param: clusters: list of clusters
    :param: nodes_retriever: method that retrieve nodes from a cluster
    >>> list(build_all_pairs_generic([[4,1,3,2], set([7,9,8])], lambda cluster:cluster))
    [(1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4), (7, 8), (7, 9), (8, 9)]
    >>> list(build_all_pairs_generic([{30:{100,200}},{10: {7,3,8}, 20: {6,4,2}}]))
    [(100, 200), (2, 3), (2, 4), (2, 6), (2, 7), (2, 8), (3, 4), (3, 6), (3, 7), (3, 8), (4, 6), (4, 7), (4, 8), (6, 7), (6, 8), (7, 8)]

    """
    for cluster in clusters:
        all_nodes_in_cluster = sorted(nodes_retriever(cluster))
        for pair in itertools.combinations(all_nodes_in_cluster, 2):
            yield pair

def add_edges_bipartite(edges, source2sas1, source2sas2, weight:int):
    """
    Add edges between all pairs from sas1 and sas2
    :param edges:
    :param sas1:
    :param sas2:
    :param weight:
    :return:
    """
    for source1, source2 in itertools.product(source2sas1.keys(), source2sas2.keys()):
        if source1 != source2:
            for sa1, sa2 in itertools.product(source2sas1[source1], source2sas2[source2]):
                edges.append(WeightedEdge(sa1, sa2, weight))

def add_group_edges(edges, sas_groups, weight: int):
    """
    Add edges between all elements in sas_groups (ie full graph)
    :param edges:
    :param sas_groups:
    :param weight:
    :param exclude_rule:
    :param group_all if True, group everything in a single clique, otherwise make separate cliques per group
    :return:
    """
    for sas in sas_groups:
        for sa_pair in itertools.combinations(sas, 2):
            edges.append(WeightedEdge(sa_pair[0], sa_pair[1], weight))

ClusterDifference = collections.namedtuple('ClusterDifference',
                                           'original_pairs added_pairs deleted_pairs sample_added sample_deleted')

def _build_flat_clusters(eid2source2elements:dict):
    """
    Build a list of clusters, represented as sorted tuple of elements
    >>> eid2source2elements = {1: {'s1':  {10,20},'s2':{30}}, 2: {'s1': {10}}, 3: {'s4': {50,60,70,80}}}
    >>> _build_flat_clusters(eid2source2elements)
    {(50, 60, 70, 80), (10,), (10, 20, 30)}
    """
    res = set()
    for source2elements in eid2source2elements.values():
        res.add(tuple(sorted(element for elements in source2elements.values() for element in elements)))
    return res

def cluster_differences(eid2source2element_old: dict, eid2source2element_new: dict,
                        sample_function=lambda x: x, overlapping_cluster=False) -> ClusterDifference:
    """
    Compute difference between 2 clustering. In particular, detect new and deleleted pairs
    :param overlapping_cluster: True if cluster are overlapping. Notice that alignment cluster are not overlapping per-se:
    there may be different virtual atts from an original att that pertain to different cluster, but technically they are
    different.

    >>> eid2source2element1 = {1: {'s1': {10,20,30}, 's2':  {40,50}}, \
                               2: {'s1': {60, 70}}, \
                               3: {'s3': {80}}, \
                               4: {'s5': {100,110}}}
    >>> eid2source2element2 = {9: {'s1': {10,20}, 's2': {40,50}}, \
                               5: {'s1': {60, 30, 70}}, \
                               6: {'s3': {80, 90}}, \
                               4: {'s5': {100,110}}}
    >>> cluster_differences(eid2source2element1, eid2source2element2)
    ClusterDifference(new_pairs=3, deleted_pairs=4, new_pairs_percent=0.25, deleted_pairs_percent=0.3333333333333333)
    """
    clusters_old = _build_flat_clusters(eid2source2element_old)
    clusters_new = _build_flat_clusters(eid2source2element_new)
    total_pairs_old = sum(bdsa_utils.nb_pairs_in_cluster(len(cluster)) for cluster in clusters_old)

    # Now we need to compute pairs of elements in same cluster, and check if some pairs are new and/or some is old.
    # If clusters are non-overlapping, we can throw away identical clusters, as they share exactly the same pairs.
    clusters_only_old = clusters_old if overlapping_cluster else clusters_old - clusters_new
    clusters_only_new = clusters_new if overlapping_cluster else clusters_new - clusters_old
    old_pairs = set().union(*(set(itertools.combinations(cluster, 2)) for cluster in clusters_only_old))
    new_pairs = set().union(*(set(itertools.combinations(cluster, 2)) for cluster in clusters_only_new))
    added_pairs = new_pairs - old_pairs
    deleted_pairs = old_pairs - new_pairs
    sample_added = stats_utils.safe_sample(added_pairs, 300)
    sample_deleted = stats_utils.safe_sample(deleted_pairs, 300)

    return ClusterDifference \
        (total_pairs_old, len(added_pairs), len(deleted_pairs), sample_added, sample_deleted)

def compare_clusters(cl1:dict, cl2:dict, min_similarity_to_continue=0.95):
    """
    Compare 2 clustering of same data. Typically used to decide when to stop iteration
    If ratio of element in identical clusters are less than min_similarity_to_continue, then stop and return this ratio.
        - For instance, [12, 345, 67, 8] and [1, 2345, 67, 8] have a ratio of 3/8 (678 over 1,2,3,4,5,6,7,8)
    This method is used to check whether 2 clusters are enough similar, so if with this quick analysis they are quite different,
    then we should stop in order to avoid losing too much time.
    If ratio is > min_similarity_to_continue, then continues:
    - build pair of "correspondent clusters", from cl1 and cl2, i.e.:
        - compute jaccard containment similarity between all cluster pairs that share at least an element
        - take the pair with biggest JS --> label as 'correspondent' clusters
        - continue iteratively last point (unless a cluster is already in a pair)
    - compute score as #elements in cl1 that are in their correspondent cluster in cl2 / total elements

    :param cl1: dict of FROZEN set
    :param cl2: same
    :param min_similarity_to_continue: see description
    :return:
    >>> cl1 = {100:frozenset([1,2,3]), 200:frozenset([4,5,6]), 300:frozenset([7,8,9]), 400:frozenset([15,16])}
    >>> cl2 = {500:frozenset([1,2]), 600:frozenset([3,4,5,6]), 700:frozenset([7,8,9,10]), 800:frozenset([15,16])}
    >>> res = compare_clusters(cl1, cl2, 0.05)
    >>> round(res, 3)
    0.833
    >>> cl1
    {100: frozenset({1, 2, 3}), 200: frozenset({4, 5, 6}), 300: frozenset({8, 9, 7}), 400: frozenset({16, 15})}
    >>> cl2
    {500: frozenset({1, 2}), 600: frozenset({3, 4, 5, 6}), 700: frozenset({8, 9, 10, 7}), 800: frozenset({16, 15})}
    """
    len1 = sum(len(group) for group in cl1.values())
    len2 = sum(len(group) for group in cl2.values())
    # Join clusters
    intersect = set(cl1.values()) & set(cl2.values())
    len_intersec = sum(len(group) for group in intersect)
    similarity_rough = len_intersec / max(len1, len2) if len1 != 0 and len2 != 0 else 1
    if similarity_rough < min_similarity_to_continue:
        return similarity_rough

    # Now compare the remaining
    cl1_filtered = {key: values for key, values in cl1.items() if values not in intersect}
    cl2_filtered = {key: values for key, values in cl2.items() if values not in intersect}
    cl2_inverted = collections.defaultdict(set)
    for cid, elements in cl2_filtered.items():
        for element in elements:
            cl2_inverted[element].add(cid)
    element_pairs = []
    for cid1, elements in cl1_filtered.items():
        counter_cid2 = collections.Counter()
        for element in elements:
            counter_cid2.update(cl2_inverted[element])
        for cid2, occs in counter_cid2.items():
            # Jaccard similarity
            element_pairs.append((cid1, cid2, occs / max(len(elements), len(cl2_filtered[cid2])), occs))
    counted_c1 = set()
    counted_c2 = set()
    total_correct = len_intersec
    for pair in sorted(element_pairs, key=lambda pair: pair[2], reverse=True):
        if pair[0] not in counted_c1 and pair[1] not in counted_c2:
            total_correct += pair[3] # TODO fix
            counted_c1.add(pair[0])
            counted_c2.add(pair[1])
    return total_correct / max(len1, len2) if len1 != 0 and len2 != 0 else 1


from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.action_chains import ActionChains
from time import sleep
from selenium.webdriver.chrome.options import Options
from Pages.Login import LoginPage
from Pages.Desktop import Desktop

import unittest

o = Options()
o.add_argument('--no-sandbox')

s = Service('d:\chromedriver_win32\chromedriver.exe')
# driver = webdriver.Chrome(ChromeDriverManager().install())

class Test_Pinvest(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.driver = webdriver.Chrome(service=s, options=o)
        cls.driver.maximize_window()
        cls.driver.implicitly_wait(5)

    def test01_login_register(self):
        self.driver.get("https://admin-ha1.pinvest.ir/signin")
        login = LoginPage(driver=self.driver)
        login.enter_username("superadmin")
        login.enter_password("1234kaKA!")
        login.enter_show_password()
        login.enter_btn_submit()
        sleep(3)
        login.enter_site_header()
        # login.enter_logout()
        # login.enter_logout_no()
        # login.enter_logout()
        # login.enter_logout_yes()

# # # desktop

    def test02_desktop(self):
        desktop = Desktop(driver=self.driver)
        desktop.enter_click_desktop()
        # assert desktop
        desktop.enter_desktop_show()
        desktop.enter_desktop_sub_percent()
        desktop.enter_desktop_sub_value()
        desktop.enter_desktop_sub_active_units()
        desktop.enter_desktop_sub_price_unit()
        desktop.enter_desktop_all_users()
        desktop.enter_desktop_direct_customers()
        desktop.enter_desktop_indirect_customers()

    @classmethod
    def tearDownClass(cls) -> None:
        sleep(5)
        cls.driver.close()
        cls.driver.quit()


from datetime import datetime, timedelta
import os
from airflow import DAG
from airflow.operators.dummy_operator import DummyOperator
from airflow.operators import FeatureExtractorOperator
from airflow.operators import FeatureLabelOperator
from airflow.operators import FeatureVerificationOperator
from airflow.operators import CsvVerificationOperator
from airflow.operators import s3UploadOperator

default_args = {
    'owner': 'ahmad-chaiban',
    'start_date': datetime(2019, 1, 12),
    'depends_on_past': False,
    'email': ['ahmadchaiban@gmail.com'],
    'email_on_retry': True,
    'retries': 3,
    'retry_delay': timedelta(minutes=5),
    'catchup': False
}

dag = DAG(
    'image_pipeline',
    default_args=default_args,
    description='Extract features from images, create, load, transform and upload to DB',
    schedule_interval='0 * * * *'
)

start_operator = DummyOperator(
    task_id='Begin_execution',
    dag=dag
)

extract_img_features_plant = FeatureExtractorOperator(
    task_id = 'feature_extraction_plant',
    dag = dag,
    path_features = '../../../../capstone data/imgFeatures/*/*.*',
    path_images = '../../../../capstone data',
    path_output = '../../../../capstone data/imgFeatures/',
    category = 'plant'
)

extract_img_features_animal = FeatureExtractorOperator(
    task_id = 'feature_extraction_animal',
    dag = dag,
    path_features = '../../../../capstone data/imgFeatures/*/*.*',
    path_images = '../../../../capstone data',
    path_output = '../../../../capstone data/imgFeatures/',
    category = 'animal'
)

extract_img_features_human = FeatureExtractorOperator(
    task_id = 'feature_extraction_human',
    dag = dag,
    path_features = '../../../../capstone data/imgFeatures/*/*.*',
    path_images = '../../../../capstone data',
    path_output = '../../../../capstone data/imgFeatures/',
    category = 'human'
)

preVerificationStandbyOperator = DummyOperator(
    task_id='Verification_Standby',
    dag=dag
)

Feature_verify_plant = FeatureVerificationOperator(
    task_id = 'feature_verify_plant',
    dag = dag,
    path = '../../../../capstone data/imgFeatures',
    category = 'plant'
)

Feature_verify_animal = FeatureVerificationOperator(
    task_id = 'feature_verify_animal',
    dag = dag,
    path = '../../../../capstone data/imgFeatures',
    category = 'animal'
)

Feature_verify_human = FeatureVerificationOperator(
    task_id = 'feature_verify_human',
    dag = dag,
    path = '../../../../capstone data/imgFeatures',
    category = 'human'
)

preDataStorageStandbyOperator = DummyOperator(
    task_id='CSV_Storage_Standby',
    dag=dag
)

create_csv_plant = FeatureLabelOperator(
    task_id = 'feature_labeler_plant',
    dag = dag,
    main_path = '../../../../capstone data/imgFeatures',
    category = 'plant'
)

create_csv_animal = FeatureLabelOperator(
    task_id = 'feature_labeler_animal',
    dag = dag,
    main_path = '../../../../capstone data/imgFeatures',
    category = 'animal'
)

create_csv_human = FeatureLabelOperator(
    task_id = 'feature_labeler_human',
    dag = dag,
    main_path = '../../../../capstone data/imgFeatures',
    category = 'human'
)

preCSVVerificationStandbyOperator = DummyOperator(
    task_id='CSV_Verification_Standby',
    dag=dag
)

csv_verification_plant = CsvVerificationOperator(
    task_id = 'plant_csv_verification',
    dag = dag,
    path = '../../../dataFunctions/final_data',
    category = 'plant'

)

csv_verification_animal = CsvVerificationOperator(
    task_id = 'animal_csv_verification',
    dag = dag,
    path = '../../../dataFunctions/final_data',
    category = 'animal'

)

csv_verification_human = CsvVerificationOperator(
    task_id = 'human_csv_verification',
    dag = dag,
    path = '../../../dataFunctions/final_data',
    category = 'human'

)

preUploadStandbyOperator = DummyOperator(
    task_id = 'Pre_upload_Standby',
    dag = dag
)

s3_upload_plant = s3UploadOperator(
    task_id = 'plant_upload',
    dag = dag,
    path = '../../../dataFunctions/final_data',
    category = 'plant',
    bucket_name = 'capstone-project-2187'
)

s3_upload_animal = s3UploadOperator(
    task_id = 'animal_upload',
    dag = dag,
    path = '../../../dataFunctions/final_data',
    category = 'animal',
    bucket_name = 'capstone-project-2187'
)

s3_upload_human = s3UploadOperator(
    task_id = 'human_upload',
    dag = dag,
    path = '../../../dataFunctions/final_data',
    category = 'human',
    bucket_name = 'capstone-project-2187'
)

end_operator = DummyOperator(
    task_id='end_execution',
    dag=dag
)

start_operator >> [extract_img_features_plant, extract_img_features_animal, extract_img_features_human] >> preVerificationStandbyOperator

preVerificationStandbyOperator >> [Feature_verify_plant, Feature_verify_animal, Feature_verify_human] >> preDataStorageStandbyOperator

preDataStorageStandbyOperator >> [create_csv_plant, create_csv_animal, create_csv_human] >> preCSVVerificationStandbyOperator

preCSVVerificationStandbyOperator >> [csv_verification_plant, csv_verification_animal, csv_verification_human] >> preUploadStandbyOperator

preUploadStandbyOperator >> [s3_upload_plant, s3_upload_animal, s3_upload_human] >> end_operator


import argparse
import tensorflow as tf

import models_tf as models
import utils

def optimistic_restore(session, save_file):
    reader = tf.train.NewCheckpointReader(save_file)
    saved_shapes = reader.get_variable_to_shape_map()
    var_names = sorted(
        [(var.name, var.name.split(':')[0]) for var in tf.global_variables() if var.name.split(':')[0] in saved_shapes])
    restore_vars = []

    with tf.variable_scope('', reuse=True):

        for var_name, saved_var_name in var_names:
            curr_var = tf.get_variable(saved_var_name)
            var_shape = curr_var.get_shape().as_list()

            if var_shape == saved_shapes[saved_var_name]:
                restore_vars.append(curr_var)

    saver = tf.train.Saver(restore_vars)
    saver.restore(session, save_file)

def inference(parameters, verbose=True):
    tf.set_random_seed(7)

    with tf.Graph().as_default():
        with tf.device('/' + parameters['device_type']):
            # initialize input holders
            if parameters["model_type"] == 'cnn':
                x_l_cc = tf.placeholder(tf.float32,
                                        shape=[None, parameters['input_size'][0], parameters['input_size'][1], 1])
                x_r_cc = tf.placeholder(tf.float32,
                                        shape=[None, parameters['input_size'][0], parameters['input_size'][1], 1])
                x_l_mlo = tf.placeholder(tf.float32,
                                         shape=[None, parameters['input_size'][0], parameters['input_size'][1], 1])
                x_r_mlo = tf.placeholder(tf.float32,
                                         shape=[None, parameters['input_size'][0], parameters['input_size'][1], 1])
                x = (x_l_cc, x_r_cc, x_l_mlo, x_r_mlo)
                model_class = models.BaselineBreastModel
            elif parameters["model_type"] == 'histogram':
                x = tf.placeholder(tf.float32, shape=[None, parameters['bins_histogram'] * 4])
                model_class = models.BaselineHistogramModel
            else:
                raise RuntimeError(parameters["model_type"])

            # holders for dropout and Gaussian noise
            nodropout_probability = tf.placeholder(tf.float32, shape=())
            gaussian_noise_std = tf.placeholder(tf.float32, shape=())

            # construct models
            model = model_class(parameters, x, nodropout_probability, gaussian_noise_std)
            y_prediction_density = model.y_prediction_density

        # allocate computation resources
        if parameters['device_type'] == 'gpu':
            session_config = tf.ConfigProto()
            session_config.gpu_options.visible_device_list = str(parameters['gpu_number'])
        elif parameters['device_type'] == 'cpu':
            session_config = tf.ConfigProto(device_count={'GPU': 0})
        else:
            raise RuntimeError(parameters['device_type'])

        with tf.Session(config=session_config) as session:
            session.run(tf.global_variables_initializer())

            # loads the pre-trained parameters if it's provided
            optimistic_restore(session, parameters['model_path'])

            # load input images
            datum_l_cc = utils.load_images(parameters['image_path'], 'L-CC')
            datum_r_cc = utils.load_images(parameters['image_path'], 'R-CC')
            datum_l_mlo = utils.load_images(parameters['image_path'], 'L-MLO')
            datum_r_mlo = utils.load_images(parameters['image_path'], 'R-MLO')

            # populate feed_dict for TF session
            # No dropout and no gaussian noise in inference
            feed_dict_by_model = {nodropout_probability: 1.0, gaussian_noise_std: 0.0}
            if parameters["model_type"] == 'cnn':
                feed_dict_by_model[x_l_cc] = datum_l_cc
                feed_dict_by_model[x_r_cc] = datum_r_cc
                feed_dict_by_model[x_l_mlo] = datum_l_mlo
                feed_dict_by_model[x_r_mlo] = datum_r_mlo
            elif parameters["model_type"] == 'histogram':
                feed_dict_by_model[x] = utils.histogram_features_generator(
                    [datum_l_cc, datum_r_cc, datum_l_mlo, datum_r_mlo],
                    parameters,
                )

            # run the session for a prediction
            prediction_density = session.run(y_prediction_density, feed_dict=feed_dict_by_model)

            if verbose:
                # nicely prints out the predictions
                print('Density prediction:\n' +
                      '\tAlmost entirely fatty (0):\t\t\t' + str(prediction_density[0, 0]) + '\n' +
                      '\tScattered areas of fibroglandular density (1):\t' + str(prediction_density[0, 1]) + '\n' +
                      '\tHeterogeneously dense (2):\t\t\t' + str(prediction_density[0, 2]) + '\n' +
                      '\tExtremely dense (3):\t\t\t\t' + str(prediction_density[0, 3]) + '\n')

            return prediction_density[0]

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='Run Inference')
    parser.add_argument('model_type')
    parser.add_argument('--bins-histogram', default=50)
    parser.add_argument('--model-path', default=None)
    parser.add_argument('--device-type', default="cpu")
    parser.add_argument('--gpu-number', default=0)
    parser.add_argument('--image-path', default="images/")
    args = parser.parse_args()

    parameters_ = {
        "model_type": args.model_type,
        "bins_histogram": args.bins_histogram,
        "model_path": args.model_path,
        "device_type": args.device_type,
        "image_path": args.image_path,
        "gpu_number": args.gpu_number,
        "input_size": (2600, 2000),
    }

    if parameters_["model_path"] is None:
        if args.model_type == "histogram":
            parameters_["model_path"] = "saved_models/BreastDensity_BaselineHistogramModel/model.ckpt"
        elif args.model_type == "cnn":
            parameters_["model_path"] = "saved_models/BreastDensity_BaselineBreastModel/model.ckpt"
        else:
            raise RuntimeError(parameters_['model_class'])

    inference(parameters_)

"""
python density_model_tf.py histogram
python density_model_tf.py cnn
"""


import pickle
import gymnasium as gym
import torch

from torch import nn
from torch.optim import Adam
from torch.distributions import Categorical
from utils import *
from gymnasium.utils.save_video import save_video

# Class for training an RL agent within an environment
class PGTrainer:
    def __init__(self, params):
        self.params = params
        self.env = gym.make(self.params['env_name'])
        self.agent = Agent(env=self.env, params=self.params)
        self.actor_policy = PGPolicy(input_size=self.env.observation_space.shape[0], output_size=self.env.action_space.n, hidden_dim=self.params['hidden_dim']).to(get_device())
        self.optimizer = Adam(params = self.actor_policy.parameters(), lr=self.params['lr'])

    def run_training_loop(self):
        list_ro_reward = list()

        for ro_idx in range(self.params['n_rollout']):
            trajectory = self.agent.collect_trajectory(policy=self.actor_policy)
            #print(trajectory)
            loss = self.estimate_loss_function(trajectory)
            self.update_policy(loss)
            # DONE TODO: Calculate avg reward for this rollout
            # HINT: Add all the rewards from each trajectory. There should be "ntr" trajectories within a single rollout.
            avg_ro_reward = sum(sum(traj_rewards) for traj_rewards in trajectory['reward']) / self.params['n_trajectory_per_rollout']
            print(f'End of rollout {ro_idx}: Average trajectory reward is {avg_ro_reward: 0.2f}')
            # Append average rollout reward into a list
            list_ro_reward.append(avg_ro_reward)
        # Save avg-rewards as pickle files
        pkl_file_name = self.params['exp_name'] + '.pkl'
        with open(pkl_file_name, 'wb') as f:
            pickle.dump(list_ro_reward, f)
        # Save a video of the trained agent playing
        self.generate_video()
        # Close environment
        self.env.close()

    def estimate_loss_function(self, trajectory):
        loss = list()
            # DONE TODO: Compute loss function
            # HINT 1: You should implement equation of policy gradient, reward to go and reward discounting here. Which will be used based on the flags set from the main function
            # Get trajectory action log-prob
            # Calculate discounted and reward-to-go
            # Calculate loss
            # HINT 2: Get trajectory action log-prob
            # HINT 3: Calculate Loss function and append to the list
        gamma = 0.99
        for t_idx in range(self.params['n_trajectory_per_rollout']):
            log_probs = trajectory['log_prob'][t_idx]
            rewards = trajectory['reward'][t_idx]
            if self.params['reward_to_go']:
                reward_sum = 0
                for i in range(len(rewards) - 1, -1, -1):
                    reward_sum = rewards[i] + gamma * reward_sum
                    loss.append(-log_probs[i] * reward_sum)
            else:
                reward_sum = sum(rewards)
                if self.params['reward_discount']:
                    for i in range(len(rewards)):
                        loss.append(-log_probs[i] * (reward_sum * gamma ** i))
                else:
                    for i in range(len(rewards)):
                        loss.append(-log_probs[i] * reward_sum)
        print("Discounted reward to go:", apply_discount(rewards))
        print("Reward-to-go:", apply_reward_to_go(rewards))
        print("Return-Reward: ", apply_return(rewards))

        loss = torch.stack(loss).mean()
        return loss

    def update_policy(self, loss):
        loss.backward()
        self.optimizer.step()
        self.optimizer.zero_grad()

    def generate_video(self, max_frame=1000):
        self.env = gym.make(self.params['env_name'], render_mode='rgb_array_list')
        obs, _ = self.env.reset()
        for _ in range(max_frame):
            action_idx, log_prob = self.actor_policy(torch.tensor(obs, dtype=torch.float32, device=get_device()))
            obs, reward, terminated, truncated, info = self.env.step(self.agent.action_space[action_idx.item()])
            if terminated or truncated:
                break
        save_video(frames=self.env.render(), video_folder=self.params['env_name'][:-3], fps=self.env.metadata['render_fps'], step_starting_index=0, episode_index=0)

# CLass for policy-net
class PGPolicy(nn.Module):
    def __init__(self, input_size, output_size, hidden_dim):
        super(PGPolicy, self).__init__()
        # TODO: Define the policy net
        # HINT: You can use nn.Sequential to set up a 2 layer feedforward neural network.
        self.policy_net = nn.Sequential(
            nn.Linear(input_size,hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim,output_size),
            nn.Softmax(dim = -1)
        )

    def forward(self, obs):
        # TODO: Forward pass of policy net
        # HINT: (use Categorical from torch.distributions to draw samples and log-prob from model output)
        logits = self.policy_net(obs)
        dist = Categorical(logits=logits)
        action_index = dist.sample()
        log_prob = dist.log_prob(action_index)
        return action_index, log_prob

# Class for agent
class Agent:
    def __init__(self, env, params=None):
        self.env = env
        self.params = params
        self.action_space = [action for action in range(self.env.action_space.n)]

    def collect_trajectory(self, policy):
        obs, _ = self.env.reset(seed=self.params['rng_seed'])
        rollout_buffer = list()
        for _ in range(self.params['n_trajectory_per_rollout']):
            trajectory_buffer = {'log_prob': list(), 'reward': list()}
            while True:
                # TODO: Get action from the policy (forward pass of policy net)
                action_idx, log_prob = policy(torch.tensor(obs, dtype=torch.float32, device=get_device()))
                # TODO: Step environment (use self.env.step() function)
                obs, reward, terminated, truncated, info = self.env.step(self.action_space[action_idx.item()])
                # Save log-prob and reward into the buffer
                trajectory_buffer['log_prob'].append(log_prob)
                trajectory_buffer['reward'].append(reward)
                # Check for termination criteria
                if terminated or truncated:
                    obs, _ = self.env.reset()
                    rollout_buffer.append(trajectory_buffer)
                    break
        rollout_buffer = self.serialize_trajectory(rollout_buffer)
        return rollout_buffer

    # Converts a list-of-dictionary into dictionary-of-list
    @staticmethod
    def serialize_trajectory(rollout_buffer):
        serialized_buffer = {'log_prob': list(), 'reward': list()}
        for trajectory_buffer in rollout_buffer:
            serialized_buffer['log_prob'].append(torch.stack(trajectory_buffer['log_prob']))
            serialized_buffer['reward'].append(trajectory_buffer['reward'])
        return serialized_buffer