import os
import json
import pickle
import datetime
import tkinter as tk
from tkinter import (
    ttk,
    OptionMenu,
    StringVar,
    messagebox
)

class Component(object):
    def __init__(self, serial_number=None, component_type='Winglet Strut', size_fit='A320 Series', status='Manufactured-Unfinished'):
        """Constructor"""
        self._serial_number = serial_number         # 1 to 99
        # Winglet Dtrut, Door handle, Rudder Pin
        self._component_type = component_type
        # A320 Series, A380 Series, 10x75, 12x100, 16x150
        self._size_fit = size_fit
        self._status = status                       # Polish, Paint

    # get_serial_number() function is accessible to class Batch
    def get_serial_number(self):
        return self._serial_number

    # set_serial_number is accessible outside the class
    def set_serial_number(self, serial_number):
        self._serial_number = serial_number

class Batch(Component):
    def __init__(self, batch_number, serial_number, location="Not Allocated"):
        """Constructor"""
        super(Batch, self).__init__(serial_number)
        self._date_manufacture = datetime.datetime.today().strftime('%Y%m%d')
        self._batch_create = datetime.datetime.today().strftime('%H%M')
        self._batch_number = batch_number
        self._location = location

    def create_sn_component(self, component_type, size_fit):
        sn = []
        x = 1

        sn_component = int(self.get_serial_number())
        while x <= sn_component:
            sn.append({
                "sn_component": self.create_batch_number() + '-' + self.add_zero(x),
                'type_components': component_type,
                'size_fitment_comps': size_fit,
                'status': self._status
            })
            x += 1
        return sn

    def create_batch_number(self):
        batch_num = int(self._batch_number)

        return self._date_manufacture + self.add_zero(batch_num)

    # I have used a static method because I don't need  the self to be passed as the first argument
    @staticmethod
    def add_zero(arg):
        x = 1
        batch = None
        while x <= arg:
            if x < 10:
                batch = '000' + str(x)
            elif 10 <= x < 100:
                batch = '00' + str(x)
            elif 100 <= x < 1000:
                batch = '0' + str(x)
            else:
                batch = str(x)

            x += 1

        return batch

class App(tk.Frame):
    # Initialize Method
    def __init__(self, master):
        """Constructor"""
        tk.Frame.__init__(self, master)
        # Set the Font and color of Nee Batch Frame
        s = ttk.Style()
        s.configure('Blue.TLabelframe.Label', font=('courier', 13, 'bold'))
        s.configure('Blue.TLabelframe.Label', foreground='#374e9b')
        s.configure('Red.TLabelframe.Label', font=('courier', 13, 'bold'), foreground='#990000')
        # get screen width and height
        self.master.geometry('800x820')
        # Add a tittle on window
        self.master.title("PPEC - Inventory System")

        # Create instance for menu frame, this is the Menu Frame in the Main Window
        self._menu_frame = ttk.LabelFrame(
            self.master, text=' Menu ', style="Blue.TLabelframe")
        # Create instance frames from menu choice, this is the New Frames in the Main Window
        self._create_batch_frame = ttk.LabelFrame(
            self.master, text=' Create New Batch ', style="Red.TLabelframe")
        self._list_batch_frame = ttk.LabelFrame(
            self.master, text=' List All Batches ', style="Red.TLabelframe")
        self._detail_batch_frame = ttk.LabelFrame(
            self.master, text=' View Details Batch ', style="Red.TLabelframe")
        self._detail_comp_frame = ttk.LabelFrame(
            self.master, text=' View Details Component ', style="Red.TLabelframe")
        self._show_details_frame = ttk.LabelFrame(
            self.master, text=' Batch Details', style="Red.TLabelframe")
        self._comp_allocation_frame = ttk.LabelFrame(
            self.master, text=' Allocation Stock ', style="Red.TLabelframe")
        self._search_product_type_frame = ttk.LabelFrame(
            self.master, text=' Search By product Type ', style="Red.TLabelframe")
        self._finish_comp_frame = ttk.LabelFrame(
            self.master, text=' Finish a Component ', style="Red.TLabelframe")
        # the below tkinter LabelFrame is used to display the scrollbar
        self.scroll_comp_finished_frame = ttk.LabelFrame(
            self.master, text='Finished Products')
        self.scroll_comp_unfinished_frame = ttk.LabelFrame(
            self.master, text='Unfinished Products')

        # Initialize path and filename of json file
        self.path = '../Data'
        self.filename = 'BatchIndex'
        self.file_path_name_w = './' + self.path + '/' + self.filename + '.json'

        # Dictionary for dropdown components type
        self.choices_comp = {'Door handle': ['A320 Series', 'A380 Series'],
                             'Rudder Pin': ['A320 Series', 'A380 Series','10x75', '12x100', '16x150'],
                             'Winglet Strut': ['A320 Series', 'A380 Series']
                            }

        # Place Dropdown into the container element create a new status
        self.choices_status = {'Paint': ['AJ74','AB76', 'BT98', 'DF56', 'BE43'],
                                'Polished': ['Matte', 'Gloss']
                            }

        # Create a Tkinter variable for dropdown
        self.tk_var = StringVar(self)
        self.entry_type_comp = StringVar(self)
        self.entry_size_fitment = StringVar(self)
        self.entry_choose_location = StringVar(self)
        self.entry_choose_status = StringVar(self)
        self.entry_choose_paint = StringVar(self)
        self.entry_code_paint = StringVar(self)
        # Create a Tkinter variable for input
        self.entry_batch_no = StringVar(self)
        self.entry_components = StringVar(self)

        # Call methods
        self.create_menu()
        self._inventory = self.create_json_file()
        self.create_batch_pkl_file()
        self.create_comp_pkl_file()

    # Create json file and save all batch created
    def create_json_file(self):
        # Input the try except in a while loop, because try to see if file exist if not exist give an error that is catch by except,
        # the exception take the data and save into file and create, and start the loop again, this time check that file exist so open the file
        while True:
            try:
                fh = open(self.file_path_name_w, 'rb+')
                data_file = json.load(fh)
                return data_file
                break
            except FileNotFoundError:
                data = {'Batch': []}
                self.save_data(data)
                continue

    def save_data(self, datafile):
        # Write Data
        with open(self.file_path_name_w, 'w+') as fp:
            fp.write(json.dumps(datafile, indent=2, sort_keys=False))
        fp.close()

    # Create batch pickle file from json file
    def create_batch_pkl_file(self):
        # list to collect data from json file  and save into pickle file
        with open('../Data/BatchIndex.json') as json_file:
            # load the json file to check data saved inside
            data = json.load(json_file)

            # Check json file inside the Batch List all Ojects and save each object in a single pickle file with value of key:number
            for batch_obj in data['Batch']:
                number = batch_obj['number']
                qty_comp = batch_obj['quantity_comp']
                manufacture_date = batch_obj['manufacture_date']
                location = batch_obj['location']
                component = batch_obj['component']

                i = 0
                comps = []

                while i < len(component):
                    sn_component = component[i]['sn_component']
                    type_components = component[i]['type_components']
                    size_fitment_comps = component[i]['size_fitment_comps']
                    status = component[i]['status']
                    comps.append([sn_component,type_components,size_fitment_comps,status])

                    i += 1

                self.file_path_name_pkl_w = './' + self.path + '/' + batch_obj['number'] + '.pkl'

                batch = [number, qty_comp, manufacture_date, location, comps]
                batch_dict = { batch[0]: batch }
                self.save_batch_pkl_data(batch_dict)

    def save_batch_pkl_data(self, pkl_batch_file):
        # Write Data
        with open(self.file_path_name_pkl_w, 'wb') as fhp:
            pickle.dump(pkl_batch_file, fhp)
        fhp.close()

    # Create components pickle file from json file
    def create_comp_pkl_file(self):
         # list to collect data from json file  and save into pickle file
        with open('../Data/BatchIndex.json') as json_file:
            # load the json file to check data saved inside
            data = json.load(json_file)

            # Check json file inside the Batch List all Ojects and save each object in a single pickle file with value of key:number
            for batch_obj in data['Batch']:

                for comp_obj in batch_obj['component']:
                    sn_component = comp_obj['sn_component']
                    type_components = comp_obj['type_components']
                    size_fitment_comps = comp_obj['size_fitment_comps']
                    manufacture_date = batch_obj['manufacture_date']
                    status = comp_obj['status']
                    number = batch_obj['number']
                    location = batch_obj['location']

                    self.file_path_name_comp_pkl_w = './' + self.path + '/' + comp_obj['sn_component'] + '.pkl'

                    # batch = [sn_component, type_components, size_fitment_comps, manufacture_date, status, number]
                    # data[batch_obj['sn_component'][comp_obj]] = batch
                    comp = [sn_component, type_components, size_fitment_comps, manufacture_date, status, location, number]
                    comp_dict = { comp[0]: comp }
                    self.save_comp_pkl_data(comp_dict)

    def save_comp_pkl_data(self, pkl_comp_file):
        # Write Data
        with open(self.file_path_name_comp_pkl_w, 'wb') as fhcp:
            pickle.dump(pkl_comp_file, fhcp)
        fhcp.close()

    """ this method manage the choice you make on the dropdown options menu,
        display in the frame below the menu the choice you made"""
    def change_dropdown(self, *args):
        option = self.tk_var.get()
        # set a dictionary with coice for menu
        choice_dict = {
            '1': self.create_new_batch,
            '2': self.list_all_batches,
            '3': self.view_details_batch,
            '4': self.view_details_comp,
            '5': self.show_frame_allocation,
            '6': self.search_product_type,
            '7': self.finish_component,
            '8': self.close_app
        }
        # call the methods that user choose from the dropdown menu
        choice_dict.get(option, self.create_menu)()

    """ this method manage the choice you do for the component, for one type component you choose,
        display the size/fit for the this specific component """
    def search_product_options(self, *args):
        comp_types = self.choices_comp[self.entry_type_comp.get()]
        self.entry_size_fitment.set(comp_types[0])

        menu = self.optionmenu_size_fitment['menu']
        menu.delete(0, 'end')

        for comp_type in comp_types:
            menu.add_command(label=comp_type, command=lambda c_type=comp_type: self.entry_size_fitment.set(c_type))

    """ this method manage the choice you do for the status, for the painted finished you choose the code in the select,
        and polished no have code """
    def update_status_options(self, *args):
        comp_statuses = self.choices_status[self.entry_choose_status.get()]
        self.entry_code_paint.set(comp_statuses[0])

        menu = self.optionmenu_code_paint['menu']
        menu.delete(0, 'end')

        for comp_status in comp_statuses:
            menu.add_command(label=comp_status, command=lambda c_status=comp_status: self.entry_code_paint.set(c_status))

    # Check if number of batch exist
    def check_batch_number_exist(self):
        entry_batch_no = self.entry_batch_no.get()
        # Check if the bacth insert is a number
        if self.isNumber(entry_batch_no):
            new_batch = Batch(entry_batch_no, None)
            batch_num = new_batch.create_batch_number()

            try:
                data = self.create_json_file()
                temp_list = []
                # Access data
                for dict_obj in data['Batch']:
                    temp_list.append(dict_obj['number'])

                if batch_num in temp_list:
                    messagebox.showinfo('Oooops!!', 'The batch number ' + batch_num + ' already exist')
                else:
                    self.show_msg_before_save()
            except (ValueError, KeyError, TypeError):
                messagebox.showerror('Comp Json Error', 'Is not a number, enter a valid number.')
        else:
            messagebox.showerror('Warning', 'Is not a number, enter a valid number.')

    # Check the input is a number
    def isNumber(self, value):
        try:
            val = int(value) # check if the parameter passed is numeric
            return True # if so return True
        except (TypeError, ValueError): # if the params is not numeric the exception is a TypeError or ValueError
            return False # return False

    # This method create a menu with a dropdown to choice options
    def create_menu(self):
        # Create a container to hold all label for menu
        self._menu_frame.grid(column=0, row=9, padx=10, pady=10, ipadx=40, sticky=tk.W)

        # Dictionary with options
        choices = {'1', '2', '3', '4', '5', '6', '7', '8'}
        # set the default label in the dropdown option
        self.tk_var.set('Options')

        # Create the Dropdown Menu Option
        OptionMenu(self._menu_frame, self.tk_var, *sorted(choices)).grid(column=1, row=8, ipadx=5)

        # Place labels into the container element menu_frame
        ttk.Label(self._menu_frame, text="1. Create a New Batch").grid(
            column=0, row=0, sticky=tk.W)
        ttk.Label(self._menu_frame, text="2. List All Batches").grid(
            column=0, row=1, sticky=tk.W)
        ttk.Label(self._menu_frame, text="3. View details of a Batch").grid(
            column=0, row=2, sticky=tk.W)
        ttk.Label(self._menu_frame, text="4. View Details of Components").grid(
            column=0, row=3, sticky=tk.W)
        ttk.Label(self._menu_frame, text="5. Allocate Manufactured Stock").grid(
            column=0, row=4, sticky=tk.W)
        ttk.Label(self._menu_frame, text="6. Search By Product Type").grid(
            column=0, row=5, sticky=tk.W)
        ttk.Label(self._menu_frame, text="7. Finish a Component").grid(
            column=0, row=6, sticky=tk.W)
        ttk.Label(self._menu_frame, text="8. Quit").grid(
            column=0, row=7, sticky=tk.W)
        ttk.Label(self._menu_frame, text="Choose the option:").grid(
            column=0, row=8, ipady=10, sticky=tk.W)

        # link function to change dropdown and pass the choice as an argument
        self.tk_var.trace('w', self.change_dropdown)

        for child in self._menu_frame.winfo_children():
            child.grid_configure(padx=10, pady=4)

    # This method create a frame with all input to create a new batch
    def create_new_batch(self):
        self._show_details_frame.grid_forget()
        self._list_batch_frame.grid_forget()
        self._detail_batch_frame.grid_forget()
        self._detail_comp_frame.grid_forget()
        self._comp_allocation_frame.grid_forget()
        self._finish_comp_frame.grid_forget()
        self._search_product_type_frame.grid_forget()
        self.scroll_comp_finished_frame.grid_forget()
        self.scroll_comp_unfinished_frame.grid_forget()

        # Create a new Object from Batch class
        new_batch = Batch(batch_number='0', serial_number='0')
        # remove the button for exit to application
        self._show_details_frame.grid_remove()
        # Create a container to hold all label for menu
        self._create_batch_frame.grid(
            column=0, row=10, padx=10, pady=10, ipadx=40, sticky=(tk.N, tk.W, tk.E, tk.S))
        # Place labels into the container element create new batch
        ttk.Label(self._create_batch_frame, text="Batch number:").grid(
            column=0, row=0, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._create_batch_frame, text="{}".format(new_batch._date_manufacture)).grid(
            column=1, row=0, padx=0, pady=10, sticky=tk.W)
        ttk.Label(self._create_batch_frame, text="How many components\n in this Batch (1 to 9999):").grid(
            column=0, row=1, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._create_batch_frame, text="Select component type:").grid(
            column=0, row=2, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._create_batch_frame, text="Select Size/Fitment type:").grid(
            column=0, row=3, padx=10, pady=10, sticky=tk.W)
        # Place Entry into the container element create new batch
        ttk.Entry(self._create_batch_frame, textvariable=self.entry_batch_no).grid(
            column=1, row=0, padx=52, pady=10)
        ttk.Entry(self._create_batch_frame, textvariable=self.entry_components).grid(
            column=1, row=1)
        # Place Dropdown into the container element create a new batch
        self.entry_type_comp.trace('w', self.search_product_options)

        self.optionmenu_type_comp = OptionMenu(self._create_batch_frame, self.entry_type_comp, *sorted(self.choices_comp.keys()))
        self.optionmenu_size_fitment = OptionMenu(self._create_batch_frame, self.entry_size_fitment, '')

        self.entry_type_comp.set('Door handle')

        self.optionmenu_type_comp.grid(column=1, row=2, padx=20, pady=10, ipadx=5)
        self.optionmenu_size_fitment.grid(column=1, row=3, padx=20, pady=10, ipadx=5)

        # Place Buttons into the container element create new batch and call the methods check_batch_number_exist
        ttk.Button(self._create_batch_frame, text="Create Batch", command=self.check_batch_number_exist).grid(
            column=0, row=4, padx=10, pady=10, sticky=tk.W)

    # Show the message on separated windows of the batch created.
    def show_msg_before_save(self):
        entry_batch_no = self.entry_batch_no.get()
        entry_no_comps = self.entry_components.get()
        entry_type_comps = self.entry_type_comp.get()
        entry_size_fit = self.entry_size_fitment.get()
        # Call Date Class
        new_batch = Batch(entry_batch_no, entry_no_comps)
        sn_comps = new_batch.create_sn_component(entry_type_comps, entry_size_fit)
        batch_num = new_batch.create_batch_number()
        # Save in a variable the value return from Batch Object
        date_manufacture = new_batch._date_manufacture
        # Setting up the message box
        show_msg_saving = messagebox.askyesno(
            'Warning!',
            'This batch No {}\ncreated at {}\ncontains {} {} {}\n\n\tis this correct!'.format(
                batch_num, date_manufacture, len(sn_comps), entry_size_fit, entry_type_comps, _icon='warning')
        )
        # If click "yes" the details insert is saved on BatchIndex.json file
        if show_msg_saving:
            # Setting up the messagebox for show batch detail
            show_msg_details = messagebox.askyesno(
                'Batch Created',
                'Batch and Component records created at {} on {}\n\nShow Batch details?'.format(
                    new_batch._batch_create, date_manufacture, _icon='info')
            )
            # If click "Yes" The batch details are displayed
            if show_msg_details:
                self.show_details_batch()
            # Call the method that Save the detail on json file
            self.save_batch_to_json_file(self._inventory)
            self.create_batch_pkl_file()
            self.create_comp_pkl_file()

    """ The method save_batch_to_json_file() is called when you click "yes" on the message box that appear
        when click "Create batch" button and save all data in BatchIndex.json file.
    """
    def save_batch_to_json_file(self, inventory):
        entry_batch_no = self.entry_batch_no.get()
        entry_no_comps = self.entry_components.get()
        entry_type_comps = self.entry_type_comp.get()
        entry_size_fit = self.entry_size_fitment.get()
        # Call Date Class
        new_batch = Batch(entry_batch_no, entry_no_comps)
        sn_comps = new_batch.create_sn_component(entry_type_comps,entry_size_fit)
        batch_num = new_batch.create_batch_number()
        manufacture_date = new_batch._date_manufacture
        # Dictionary to collect data and save into json file
        a_dict = {}
        #data = self.create_json_file()
        temp_list = []
        for dic_obj in inventory["Batch"]:
            temp_list.append(dic_obj)
        qty_comp = entry_no_comps
        temp_list.append(
            {'number': batch_num,
            'quantity_comp': qty_comp,
            'location': new_batch._location,
            'manufacture_date': manufacture_date,
            'component': sn_comps
            })
        inventory["Batch"] = temp_list
        a_dict["Batch"] = inventory["Batch"]
        self.save_data(a_dict)
        # Clear the entry one the file is created to set with empty string
        self.entry_batch_no.set('')
        self.entry_components.set('')

    # Show the details of the batch created
    def show_details_batch(self):
        print('New Batch Details Method Called')
        entry_batch_no = self.entry_batch_no.get()
        entry_no_comps = self.entry_components.get()
        entry_type_comps = self.entry_type_comp.get()
        entry_size_fit = self.entry_size_fitment.get()
        # Call Date Class
        new_batch = Batch(entry_batch_no, entry_no_comps)
        sn_comps = new_batch.create_sn_component(entry_type_comps,entry_size_fit)
        batch_num = new_batch.create_batch_number()
        comp_location = new_batch._location

        sn = []
        comps = []
        for item in sn_comps:
            sn.append(item['sn_component'] )
            comps.append(item['sn_component'] + ' ' + item['status'])

        comps = "\n\t".join(comps)

        # Save in a variable the value return from Batch Object
        date_manufacture = new_batch._date_manufacture
        # remove the button for exit to application
        self._create_batch_frame.grid_remove()
        # Create a container to hold all label for menu
        self._show_details_frame.grid(column=0, row=10, padx=10, pady=10, ipadx=40, sticky=(tk.N, tk.W, tk.E, tk.S))
        # Place labels into the container element create new batch
        ttk.Label(self._show_details_frame, text="Batch number:\t{}".format(batch_num)).grid(
            column=0, row=0, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._show_details_frame, text="Manufacture Date:\t{}".format(date_manufacture)).grid(
            column=0, row=1, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._show_details_frame, text="Component type:\t{}".format(entry_type_comps)).grid(
            column=0, row=2, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._show_details_frame, text="Component Size/Fitment type: {}".format(entry_size_fit)).grid(
            column=0, row=3, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._show_details_frame, text="SerialNumbers: {}".format(sn)).grid(
            column=0, row=4, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._show_details_frame, text="Component Status: [\n\t{}\n\t]".format(comps)).grid(
            column=0, row=5, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._show_details_frame, text="Component Location: {}".format(comp_location)).grid(
            column=0, row=6, padx=10, pady=10, sticky=tk.W)

    # This method show all batches
    def list_all_batches(self):
        self._show_details_frame.grid_forget()
        self._create_batch_frame.grid_forget()
        self._detail_batch_frame.grid_forget()
        self._detail_comp_frame.grid_forget()
        self._comp_allocation_frame.grid_forget()
        self._search_product_type_frame.grid_forget()
        self._finish_comp_frame.grid_forget()
        self.scroll_comp_finished_frame.grid_forget()
        self.scroll_comp_unfinished_frame.grid_forget()

        # Open file json for retrieving all batch saved
        fh = open(self.file_path_name_w, 'r')
        data_file = json.load(fh)
        fh.close()

        # Loop for display all list Batches
        r = 1
        for list_batch in data_file['Batch']:
            batch_number = list_batch['number']
            quantity_comp = list_batch['quantity_comp']
            allocated_location = list_batch['location']
            for comp in list_batch['component']:
                type_components = comp['type_components']
                size_fitment_comps = comp['size_fitment_comps']

            # Place labels into the container element create row for each batch
            ttk.Label(self._list_batch_frame, text=batch_number).grid(column=0, row=r, padx=10, pady=10, sticky=tk.W)
            ttk.Label(self._list_batch_frame, text=type_components).grid(column=1, row=r, padx=10, pady=10, sticky=tk.W)
            ttk.Label(self._list_batch_frame, text=size_fitment_comps).grid(column=2, row=r, padx=10, pady=10, sticky=tk.W)
            ttk.Label(self._list_batch_frame, text=quantity_comp).grid(column=3, row=r, padx=10, pady=10, sticky=tk.W)
            ttk.Label(self._list_batch_frame, text=allocated_location).grid(column=4, row=r, padx=10, pady=10, sticky=tk.W)

            r += 1

        # Create a container to hold all label for menu
        self._list_batch_frame.grid(column=0, row=10, padx=10, pady=10, ipadx=40, sticky=(tk.N, tk.W, tk.E, tk.S))
        # Place labels into the container element create new batch
        ttk.Label(self._list_batch_frame, text="Batch number").grid(column=0, row=0, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._list_batch_frame, text="Type").grid(column=1, row=0, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._list_batch_frame, text="Size/Fitment").grid(column=2, row=0, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._list_batch_frame, text="Quantity Made").grid(column=3, row=0, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._list_batch_frame, text="Location").grid(column=4, row=0, padx=10, pady=10, sticky=tk.W)

    # Show detail of the batch
    def view_details_batch(self):
        self._show_details_frame.grid_forget()
        self._create_batch_frame.grid_forget()
        self._detail_comp_frame.grid_forget()
        self._list_batch_frame.grid_remove()
        self._comp_allocation_frame.grid_forget()
        self._finish_comp_frame.grid_forget()
        self._search_product_type_frame.grid_forget()
        self.scroll_comp_finished_frame.grid_forget()
        self.scroll_comp_unfinished_frame.grid_forget()

        # Create a container to hold all label for menu
        self._detail_batch_frame.grid(
            column=0, row=10, padx=10, pady=10, ipadx=40, sticky=(tk.N, tk.W, tk.E, tk.S))
        # Place labels into the container element create new batch
        ttk.Label(self._detail_batch_frame, text="Batch number:").grid(
            column=0, row=0, padx=10, pady=10, sticky=tk.W)
        # Place Entry into the container element create new batch
        ttk.Entry(self._detail_batch_frame, textvariable=self.entry_batch_no).grid(
            column=1, row=0, padx=52, pady=10)
        # Place Buttons into the container element create new batch
        ttk.Button(self._detail_batch_frame, text="Show Details", command=lambda : self.check_batch_exist(self.entry_batch_no.get())).grid(
            column=0, row=9, padx=10, pady=10, sticky=tk.W)

    # Check if batch exist, receive parameter from view_details_batch()
    def check_batch_exist(self, code):
        # Check if the bacth insert is a number
        if self.isNumber(code):

            if os.path.exists('../Data/' + code + '.pkl'):
                fhb = open('../Data/' + code + '.pkl', 'rb+')
                data = pickle.load(fhb)
                comps = []
                compStatus = []

                if code in data.keys():
                    # if the code is correct and exist execute the statement
                    batchDetail = data[code]  # pass to variable the element value to collect the right elements from the same list
                    ttk.Label(self._detail_batch_frame, text="Batch Number:\t{}".format(batchDetail[0])).grid(column=0, row=1, padx=10, pady=10, sticky=tk.W)
                    ttk.Label(self._detail_batch_frame, text="Location:\t{}".format(batchDetail[3])).grid(column=0, row=2, padx=10, pady=10, sticky=tk.W)
                    ttk.Label(self._detail_batch_frame, text="Manufacture date:\t{}".format(batchDetail[2])).grid(column=0, row=3, padx=10, pady=10, sticky=tk.W)
                    ttk.Label(self._detail_batch_frame, text="Type:\t{}".format(batchDetail[4][0][1])).grid(column=0, row=4, padx=10, pady=10, sticky=tk.W)
                    ttk.Label(self._detail_batch_frame, text="Size/Fit:\t{}".format(batchDetail[4][0][2])).grid(column=0, row=5, padx=10, pady=10, sticky=tk.W)
                    ttk.Label(self._detail_batch_frame, text="Number Components in Batch:\t{}".format(batchDetail[1])).grid(column=0, row=6, padx=10, pady=10, sticky=tk.W)
                    n_comps = int(batchDetail[1])
                    i = 0
                    for i in range(n_comps):
                        comps.append(batchDetail[4][i][0])
                        compStatus.append(batchDetail[4][i][0] + ' ' + batchDetail[4][i][3])

                    ttk.Label(self._detail_batch_frame, text="SN Component:\t{}".format(comps)).grid(column=0, row=7, padx=10, pady=10, sticky=tk.W)
                    ttk.Label(self._detail_batch_frame, text="Component Status: [\n\t{}\n\t]".format("\n\t".join(compStatus))).grid(column=0, row=8, padx=10, pady=10, sticky=tk.W)

                return True

            else:
                messagebox.showwarning('Warning','Batch not found')
                return False
        else:
            messagebox.showerror('Warning', 'Is not a number, enter a valid number.')

    # Show detail of components
    def view_details_comp(self):
        self._show_details_frame.grid_forget()
        self._create_batch_frame.grid_forget()
        self._detail_batch_frame.grid_forget()
        self._list_batch_frame.grid_remove()
        self._comp_allocation_frame.grid_forget()
        self._finish_comp_frame.grid_forget()
        self._search_product_type_frame.grid_forget()
        self.scroll_comp_finished_frame.grid_forget()
        self.scroll_comp_unfinished_frame.grid_forget()

        # Create a container to hold all label for menu
        self._detail_comp_frame.grid(
            column=0, row=10, padx=10, pady=10, ipadx=40, sticky=(tk.N, tk.W, tk.E, tk.S))
        # Place labels into the container element create new batch
        ttk.Label(self._detail_comp_frame, text="Component SN:").grid(
            column=0, row=0, padx=10, pady=10, sticky=tk.W)
        # Place Entry into the container element create new batch
        ttk.Entry(self._detail_comp_frame, textvariable=self.entry_components).grid(
            column=1, row=0, padx=52, pady=10)
        # Place Buttons into the container element create new batch
        ttk.Button(self._detail_comp_frame, text="Show Details", command=lambda : self.check_comp_exist(self.entry_components.get())).grid(
            column=0, row=7, padx=10, pady=10, sticky=tk.W)

    # Check if component exist by receiving parameter from view_details_comp()
    def check_comp_exist(self, code):
        # Check if the bacth insert is a number
        #if self.isNumber(code):

        if os.path.exists('../Data/' + code + '.pkl'):
            fhb = open('../Data/' + code + '.pkl', 'rb+')
            data = pickle.load(fhb)

            if code in data.keys():
                # if the code is correct and exist execute the statement
                compDetail = data[code]  # pass to variable the element value to collect the right elements from the same list
                ttk.Label(self._detail_comp_frame, text="Component Detail for:\t{}".format(compDetail[0])).grid(column=0, row=1, padx=10, pady=10, sticky=tk.W)
                ttk.Label(self._detail_comp_frame, text="Type:\t{}".format(compDetail[1])).grid(column=0, row=2, padx=10, pady=10, sticky=tk.W)
                ttk.Label(self._detail_comp_frame, text="Size/Fit:\t{}".format(compDetail[2])).grid(column=0, row=3, padx=10, pady=10, sticky=tk.W)
                ttk.Label(self._detail_comp_frame, text="Date of Manufacture:\t{}".format(compDetail[3])).grid(column=0, row=4, padx=10, pady=10, sticky=tk.W)
                ttk.Label(self._detail_comp_frame, text="Current Status:\t{}".format(compDetail[4])).grid(column=0, row=5, padx=10, pady=10, sticky=tk.W)
                ttk.Label(self._detail_comp_frame, text="Part of Batch:\t{}".format(compDetail[6])).grid(column=0, row=6, padx=10, pady=10, sticky=tk.W)

            return True

        else:
            messagebox.showwarning('Warning','Component not found')
            return False
        #else:
            #messagebox.showerror('Warning', 'Is not a number, enter a valid number.')

    # This method allocate a batch
    def show_frame_allocation (self):
        self._show_details_frame.grid_forget()
        self._create_batch_frame.grid_forget()
        self._list_batch_frame.grid_forget()
        self._detail_batch_frame.grid_forget()
        self._detail_comp_frame.grid_forget()
        self._finish_comp_frame.grid_forget()
        self._search_product_type_frame.grid_forget()
        self.scroll_comp_finished_frame.grid_forget()
        self.scroll_comp_unfinished_frame.grid_forget()

        # Dictionary for dropdown location
        choices_location = {'Dubai','Paisley'}
        # set the default label in the dropdown option
        self.entry_choose_location.set('Paisley')

        # Create a container to hold all label for menu
        self._comp_allocation_frame.grid(
            column=0, row=10, padx=10, pady=10, ipadx=40, sticky=(tk.N, tk.W, tk.E, tk.S))
        # Place labels into the container element create new batch
        ttk.Label(self._comp_allocation_frame, text="Enter Batch Number:").grid(
            column=0, row=0, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._comp_allocation_frame, text="Choose Location:").grid(
            column=0, row=1, padx=10, pady=10, sticky=tk.W)
        # Place Entry into the container element create new batch
        ttk.Entry(self._comp_allocation_frame, textvariable=self.entry_batch_no).grid(
            column=1, row=0, padx=52, pady=10)

        # Place Dropdown into the container element create a new batch
        OptionMenu(self._comp_allocation_frame, self.entry_choose_location, *
                   sorted(choices_location)).grid(column=1, row=1, padx=20, pady=10, ipadx=5)

        # Place Buttons into the container element create new batch
        ttk.Button(self._comp_allocation_frame, text="Allocate", command=lambda : self.check_allocation(self.entry_choose_location.get(), self.entry_batch_no.get())).grid(
            column=0, row=7, padx=10, pady=10, sticky=tk.W)

    def check_allocation(self, loc, code):
        # Check if the bacth insert is a number
        if self.isNumber(code):

            if os.path.exists('../Data/' + code + '.pkl'):
                fhb = open('../Data/' + code + '.pkl', 'rb+')
                data = pickle.load(fhb)
                fhb.close()

                if code in data.keys():
                    # if the code is correct and exist execute the statement
                    batchDetail = data[code]  # pass to variable the element value to collect the right elements from the same list
                    if (batchDetail[3] != 'Not Allocated'):
                        messagebox.showwarning('warning', 'Batch Already Allocated')
                    else:
                        yes_allocate_comp = messagebox.askyesno('Allocation', 'Batch Not Allocated, do you like to allocate batch num: {}?'.format(batchDetail[0]))
                        if yes_allocate_comp:
                            self.update_allocation_batch(loc, code)
            else:
                messagebox.showwarning('Warning','Batch not found')
        else:
            messagebox.showerror('Warning', 'Is not a number, enter a valid number.')

    def update_allocation_batch(self, loc, batchN):
        with open(self.file_path_name_w, 'r+') as f:
            data = json.load(f)
            for list_batch in data['Batch']:
                if (list_batch['number'] == batchN):
                    list_batch['location'] = loc

            f.seek(0)        # should reset file position to the beginning.
            json.dump(data, f, indent=2)
            f.truncate()     # remove remaining part
            f.close()
            # update also pickle file
            self.create_batch_pkl_file()
            self.create_comp_pkl_file()

            messagebox.showinfo('Allocate', 'This Batch is now allocated and will be shipped to the {} location'.format(loc))

    # method for search by product
    def search_product_type(self):
        self._show_details_frame.grid_forget()
        self._create_batch_frame.grid_forget()
        self._detail_batch_frame.grid_forget()
        self._list_batch_frame.grid_forget()
        self._detail_comp_frame.grid_forget()
        self._comp_allocation_frame.grid_forget()
        self._finish_comp_frame.grid_forget()
        self.scroll_comp_finished_frame.grid_forget()
        self.scroll_comp_unfinished_frame.grid_forget()
        # This statement destroy the old scrollbar when you call a new searching
        for child in self.scroll_comp_finished_frame.winfo_children(): child.destroy()
        for child in self.scroll_comp_unfinished_frame.winfo_children(): child.destroy()

        self._search_product_type_frame.grid(column=0, row=10, padx=10, pady=10, ipadx=40, sticky=tk.W)

        # Place labels into the container element search_product_type_frame
        ttk.Label(self._search_product_type_frame, text="Select component type:").grid(
            column=0, row=0, padx=10, pady=10, sticky=tk.W)
        ttk.Label(self._search_product_type_frame, text="Select Size/Fitment type:").grid(
            column=0, row=1, padx=10, pady=10, sticky=tk.W)

        # Dictionary for dropdown components type
        self.entry_type_comp.trace('w', self.search_product_options)

        self.optionmenu_type_comp = OptionMenu(self._search_product_type_frame, self.entry_type_comp, *sorted(self.choices_comp.keys()))
        self.optionmenu_size_fitment = OptionMenu(self._search_product_type_frame, self.entry_size_fitment, '')

        self.entry_type_comp.set('Door handle')

        self.optionmenu_type_comp.grid(column=1, row=0, padx=20, pady=10, ipadx=5)
        self.optionmenu_size_fitment.grid(column=1, row=1, padx=20, pady=10, ipadx=5)

        # Place Buttons into the container element create new batch
        ttk.Button(self._search_product_type_frame, text="Search", command=lambda: self.search_product(self.entry_type_comp.get(), self.entry_size_fitment.get())).grid(
            column=0, row=10, padx=10, pady=10, sticky=tk.W)

    def search_product(self, ctype, fit):
        # Open file json for retrieving all batch saved
        fh = open(self.file_path_name_w, 'r')
        data = json.load(fh)
        fh.close()

        unfinish_product = []
        finish_product = []

        # Loop and append in a new list all Batches founded
        msg_search = messagebox.askyesno('Info Component', 'You Selected {} {} is this Correct?'.format(fit, ctype))

        if msg_search:
            # this loop the json file and collect some data that I need for display in my scrollbar
            for list_product in data['Batch']:
                location = list_product['location']
                manufacture_date = list_product['manufacture_date']

                # this loop the json file and collect some data that I need for display in my scrollbar and check in a condition if finished or not
                for list_comp in list_product['component']:
                    serial_number = list_comp['sn_component']
                    status = list_comp['status']
                    # check if the component found are finished if so append in finish product list
                    if list_comp['type_components'] == ctype and list_comp['size_fitment_comps'] == fit and list_comp['status'] != 'Manufactured-Unfinished':
                        all_finished = [ctype,fit,serial_number,location,status,manufacture_date]
                        finish_product.append(all_finished)
                    # if not finished append into unfinish product list
                    elif list_comp['type_components'] == ctype and list_comp['size_fitment_comps'] == fit and list_comp['status'] == 'Manufactured-Unfinished':
                        all_unfinished = [ctype,fit,serial_number,location,status,manufacture_date]
                        unfinish_product.append(all_unfinished)

            # Check if the both list are empty, if so show a message and return False
            if len(unfinish_product) == 0 and len(finish_product) == 0:
                    messagebox.showwarning('Warning', 'No Stock Available!')
                    show_product = False
            # if not empty return True
            else:
                show_product = True

        # if show product is True call a function to display in a scrollbar all component founded
        if show_product:
            self.create_finished_scroll(finish_product, unfinish_product)

    """
        Start Methods to create 2 Scrolbars
        This methods create 2 scrollbar for finish and unfinished product
    """
    def create_finished_scroll(self, finished, unfinished):
        self.scroll_comp_finished_frame.grid(row=11, column=0, padx=10, pady=10, sticky=tk.W)
        self.create_frame_finished_scroll(self.scroll_comp_finished_frame, finished)

        self.scroll_comp_unfinished_frame.grid(row=12, column=0, padx=10, pady=10, sticky=tk.W)
        self.create_frame_unfinished_scroll(self.scroll_comp_unfinished_frame, unfinished)

    def create_frame_finished_scroll(self, box, finished):
        self._show_details_frame.grid_forget()
        self._create_batch_frame.grid_forget()
        self._detail_batch_frame.grid_forget()
        self._list_batch_frame.grid_forget()
        self._detail_comp_frame.grid_forget()
        self._comp_allocation_frame.grid_forget()
        self._finish_comp_frame.grid_forget()
        self._search_product_type_frame.grid_forget()

        canvas = tk.Canvas(box, borderwidth=0)
        self.finished_comp_frame = tk.Frame(canvas)
        vsb = tk.Scrollbar(box, orient="vertical", command=canvas.yview)
        canvas.configure(yscrollcommand=vsb.set, width=750, heigh=80)

        vsb.pack(side="right", fill="y")
        canvas.pack(side="left", fill="both", expand=True)
        canvas.create_window((4,4), window=self.finished_comp_frame, anchor="w", tags="frame")

        # be sure that we call OnFrameConfigure on the right canvas
        self.finished_comp_frame.bind("<Configure>", lambda event, canvas=canvas: self.on_scroll_frame_configure(canvas))

        self.fill_finished_window(self.finished_comp_frame, finished)

    def fill_finished_window(self, frame, finished):
        sn_components = ttk.Label(self.finished_comp_frame, text="Component Serial Number", font=('courier', 11, 'bold'))
        sn_components.grid(row=0, column=0, padx=10, sticky=tk.W)
        location_comp = ttk.Label(self.finished_comp_frame, text="Location", font=('courier', 11, 'bold'))
        location_comp.grid(row=0, column=1, padx=10, sticky=tk.W)
        finish = ttk.Label(self.finished_comp_frame, text="Finish", font=('courier', 11, 'bold'))
        finish.grid(row=0, column=2, padx=10, sticky=tk.W)
        manufacture_date = ttk.Label(self.finished_comp_frame, text="Manufacture Date", font=('courier', 11, 'bold'))
        manufacture_date.grid(row=0, column=3, padx=10, sticky=tk.W)

        for i in range(len(finished)):
            sn_components_w1 = ttk.Label(self.finished_comp_frame, text=str(finished[i][2:3]))
            sn_components_w1.grid(row=i+1, column=0, padx=10, sticky=tk.W)
            location_comp_w1 = ttk.Label(self.finished_comp_frame, text=str(finished[i][3:4]))
            location_comp_w1.grid(row=i+1, column=1, padx=10, sticky=tk.W)
            finish_w1 = ttk.Label(self.finished_comp_frame, text=str(finished[i][4:5]))
            finish_w1.grid(row=i+1, column=2, padx=10, sticky=tk.W)
            manufacture_date_w1 = ttk.Label(self.finished_comp_frame, text=str(finished[i][5:]))
            manufacture_date_w1.grid(row=i+1, column=3, padx=10, sticky=tk.W)

    def create_frame_unfinished_scroll(self, box, unfinished):
        self._show_details_frame.grid_forget()
        self._create_batch_frame.grid_forget()
        self._detail_batch_frame.grid_forget()
        self._list_batch_frame.grid_forget()
        self._detail_comp_frame.grid_forget()
        self._comp_allocation_frame.grid_forget()
        self._finish_comp_frame.grid_forget()
        self._search_product_type_frame.grid_forget()

        canvas = tk.Canvas(box, borderwidth=0)
        self.unfinished_comp_frame = tk.Frame(canvas)
        vsb = tk.Scrollbar(box, orient="vertical", command=canvas.yview)
        canvas.configure(yscrollcommand=vsb.set, width=750, heigh=80)

        vsb.pack(side="right", fill="y")
        canvas.pack(side="left", fill="both", expand=True)
        canvas.create_window((4,4), window=self.unfinished_comp_frame, anchor="w", tags="frame")

        # be sure that we call OnFrameConfigure on the right canvas
        self.unfinished_comp_frame.bind("<Configure>", lambda event, canvas=canvas: self.on_scroll_frame_configure(canvas))

        self.fill_unfinished_window(self.unfinished_comp_frame, unfinished)

    def fill_unfinished_window(self, frame, unfinished):
        sn_components = ttk.Label(self.unfinished_comp_frame, text="Component Serial Number", font=('courier', 11, 'bold'))
        sn_components.grid(row=0, column=0, padx=10, sticky=tk.W)
        location_comp = ttk.Label(self.unfinished_comp_frame, text="Location", font=('courier', 11, 'bold'))
        location_comp.grid(row=0, column=1, padx=10, sticky=tk.W)
        unfinish = ttk.Label(self.unfinished_comp_frame, text="Finish", font=('courier', 11, 'bold'))
        unfinish.grid(row=0, column=2, padx=10, sticky=tk.W)
        manufacture_date = ttk.Label(self.unfinished_comp_frame, text="Manufacture Date", font=('courier', 11, 'bold'))
        manufacture_date.grid(row=0, column=3, padx=10, sticky=tk.W)

        for i in range(len(unfinished)):
            # split the fifth element from unfinished list and take the first element to show on scrollbar
            split_unfinish = unfinished[i][4:5]
            split_result = [i.split('-', 1)[1] for i in split_unfinish]

            sn_components_w1 = ttk.Label(self.unfinished_comp_frame, text=str(unfinished[i][2:3]))
            sn_components_w1.grid(row=i+1, column=0, padx=10, sticky=tk.W)
            location_comp_w1 = ttk.Label(self.unfinished_comp_frame, text=str(unfinished[i][3:4]))
            location_comp_w1.grid(row=i+1, column=1, padx=10, sticky=tk.W)
            unfinish_w1 = ttk.Label(self.unfinished_comp_frame, text=str(split_result))
            unfinish_w1.grid(row=i+1, column=2, padx=10, sticky=tk.W)
            manufacture_date_w1 = ttk.Label(self.unfinished_comp_frame, text=str(unfinished[i][5:]))
            manufacture_date_w1.grid(row=i+1, column=3, padx=10, sticky=tk.W)

    def on_scroll_frame_configure(self, canvas):
        canvas.configure(scrollregion=canvas.bbox("all"))

    """
        =============== END Scrollbar Methods =========================

    """

    # Method that finish component
    def finish_component(self, fcomp='No'):
        self._show_details_frame.grid_forget()
        self._create_batch_frame.grid_forget()
        self._detail_batch_frame.grid_forget()
        self._list_batch_frame.grid_forget()
        self._detail_comp_frame.grid_forget()
        self._comp_allocation_frame.grid_forget()
        self._search_product_type_frame.grid_forget()
        self.scroll_comp_finished_frame.grid_forget()
        self.scroll_comp_unfinished_frame.grid_forget()

        # Create a container to hold all label for finish component
        self._finish_comp_frame.grid(column=0, row=10, padx=10,
            pady=10, ipadx=40, sticky=tk.W)

        # Place labels into the container element finish_comp_frame
        ttk.Label(self._finish_comp_frame, text="Enter SN of the Component to Finish").grid(
            column=0, row=0, padx=10, pady=10, sticky=tk.W)

        # Place Entry into the container element create new batch
        ttk.Entry(self._finish_comp_frame, textvariable=self.entry_components).grid(
            column=1, row=0, padx=52, pady=10)

        # Place Buttons into the container element create new batch
        if (fcomp == 'Yes'):
            ttk.Label(self._finish_comp_frame, text="Select Type Finish").grid(
                column=0, row=1, padx=10, pady=10, sticky=tk.W)
            ttk.Label(self._finish_comp_frame, text="Select Paint Code").grid(
                column=0, row=2, padx=10, pady=10, sticky=tk.W)

            self.entry_choose_status.trace('w', self.update_status_options)

            self.optionmenu_type_status = OptionMenu(self._finish_comp_frame, self.entry_choose_status, *sorted(self.choices_status.keys()))
            self.optionmenu_code_paint = OptionMenu(self._finish_comp_frame, self.entry_code_paint, '')

            self.entry_choose_status.set('Polished')

            self.optionmenu_type_status.grid(column=1, row=1, padx=20, pady=10, ipadx=5)
            self.optionmenu_code_paint.grid(column=1, row=2, padx=20, pady=10, ipadx=5)

            ttk.Button(self._finish_comp_frame, text="Save Status", command=lambda: self.update_finished_comp(
                self.entry_components.get(),self.entry_choose_status.get(),self.entry_code_paint.get())).grid(column=0, row=4, padx=11, pady=10, sticky=tk.W)
        else:
            ttk.Button(self._finish_comp_frame, text="Check Status", command=lambda: self.check_sn_comp_finished(self.entry_components.get())).grid(
                column=0, row=4, padx=10, pady=10, sticky=tk.W)

    def update_finished_comp(self, sn, stat, code):
        # Open json file and read and load data in a variable
        with open(self.file_path_name_w, 'r+') as f:
            data = json.load(f)
            i = 0
            for list_batch in data['Batch']:
                for i in list_batch['component']:
                    if (i['sn_component'] == sn):
                        i['status'] = stat + ': ' + code

            f.seek(0)        # <--- should reset file position to the beginning.
            json.dump(data, f, indent=2)
            f.truncate()     # remove remaining part
            f.close()
            # update also pickle file
            self.create_batch_pkl_file()
            self.create_comp_pkl_file()

            messagebox.showinfo('Finished', 'Component Number {} will be finished using {}: {}'.format(sn,stat,code))

    def check_sn_comp_finished(self, sncomp):
        # Check if file exist, if so open and read
        if os.path.exists('../Data/' + sncomp + '.pkl'):
            fhc = open('../Data/' + sncomp + '.pkl', 'rb+')
            data = pickle.load(fhc)

            if sncomp in data.keys():
                compDetails = data[sncomp]  # pass to variable the element value to collect the right elements from the same list
                if (compDetails[4] == 'Manufactured-Unfinished' and compDetails[5] != 'Not Allocated'):
                    yes_correct_comp = messagebox.askyesno('Info', 'You selected {} {} {} is this correct?'.format(compDetails[2], compDetails[1], compDetails[5]))
                    if yes_correct_comp:
                        self.finish_component('Yes')
                elif (compDetails[4] == 'Manufactured-Unfinished' and compDetails[5] == 'Not Allocated'):
                    yes_allocate_comp = messagebox.askyesno('Allocation', 'Batch Not Allocated, do you like to allocate batch num: {}?'.format(compDetails[6]))
                    if yes_allocate_comp:
                        self.show_frame_allocation()
                else:
                    messagebox.showwarning('Warning', 'This comp is already finished!')
        else:
            messagebox.showerror('Error', 'Component not found')

    # Method that close Application
    def close_app(self):
        # Below line is just for testing if everything work
        msgbox_exit = messagebox.askyesno(
            'Exit Program', 'Are you sure want to EXIT!', icon='warning')
        if msgbox_exit:
            # Below line is just for testing if everything work
            print('Close App method Called')
            self.quit()
        else:
            # Below line is just for testing if everything work
            print('Return to menu')

# ====================
# Start GUI
# ====================
if __name__ == "__main__":
    _frame = tk.Tk()
    app = App(_frame)
    app.mainloop()


import os
from argparse import ArgumentParser
import operator

class File:
    """File object that can store file attributes(file size, file path, etc).
    This has class dictionary that will store summary of all file types. """

    file_dic = {}

    def __init__(self, file):
        self._path, self._filename = os.path.split(file)
        self._file = file
        self._size = os.path.getsize(file)
        self._extension = os.path.splitext(self._filename)[1][1:]
        self.count_files()

    def get_file_extension(self):
        '''Since git have multiple dots(.) this function checks it.
        and returns the file extension.'''
        if self.file().count('.') > 1:
            file_extension = 'git'
            return file_extension
        else:
            return self._extension

    def path(self):
        return self._path

    def file(self):
        return self._file

    def size(self):
        return self._size

    def filename(self):
        return self._filename

    def count_files(self):
        '''Count files with the same extension, sum its size then store both
        count and size in a list and store the list in a dictionary
        with the key of file extension.'''
        if self._extension not in File.file_dic:
            File.file_dic[self._extension] = [1, self.size()]
        else:
            File.file_dic[self._extension][0] += 1
            File.file_dic[self._extension][1] += self.size()

def format_size(size):
    '''Format the different file sizes into B, KB, MB and GB.
    Returns the formatted size like '8.5 GB' and '5 MB'.'''
    size = remove_zfill(size)
    if len(size) <= 3:
        size = size + ' B'
        return size.rjust(9)
    elif len(size) <= 6:
        size = round(int(size)/1_000, 2)
        size = str(size) + ' KB'
        return size.rjust(9)
    elif len(size) <= 9:
        size = round(int(size)/1_000_000, 2)
        size = str(size) + ' MB'
        return size.rjust(9)
    else:
        size = round(int(size)/1_000_000_000, 2)
        size = str(size) + ' GB'
        return size.rjust(9)

def remove_zfill(size):
    '''Remove all the zeros added that were help for sorting.'''
    if not size.startswith('0'):
        return size
    for char in size:
        if len(size) == 1:
            return size
        if char == '0':
            return remove_zfill(size[1:])

def file_gen(folder_path):
    '''Generate all files complete with each path in the given folder and its subfolders.'''
    for path, _, files, in os.walk(folder_path):
        for file in files:
            file_ = os.path.join(path, file)
            if not os.path.isfile(file_):
                continue
            yield file_

def write_to_file(folder, all_files, out_file):
    '''Write to a text file the summary of all files present.'''
    if not out_file[-3:] == 'txt':
        out_file = out_file + '.txt'

    with open(os.path.join(folder, out_file), 'w', encoding='utf8') as text_file:
        text_file.write('List of all files inside folders and its subfolders in "{}".\n\n'.upper().format(folder))

        for size, file in all_files:
            size = format_size(size)
            line = '{} - {}\n'.format(size, file)
            text_file.write(line)

        summary = '*********** SUMMARY ***********'.center(53)
        text_file.write('\n\n' + summary + '\n\n')

        t_size = 'Total Size'.center(17).upper()
        n_files = 'Number of Files'.center(17).upper()
        file_ = 'File'.center(17).upper()
        text_file.write(t_size + '|' + n_files + '|' + file_ + '\n')

        file_dic = dict(sorted(File.file_dic.items(), key=operator.itemgetter(1), reverse=True))

        for file_type, value in file_dic.items():
            num, size = value
            size = format_size(str(size)).center(17)
            file_type = test_if_hash(file_type)
            file_type = file_type.center(17)
            num = str(num).center(17)
            text_file.write(size + '|' + num + '|' + file_type + '\n')

def test_if_hash(file_type):
    '''Test if if the seen file is a hash.'''
    if len(file_type) == 0:
        return 'hash'
    else:
        return file_type

def main(folder, file_ext, del_flag, org_flag, out_file):

    all_files = []

    for file in file_gen(folder):
        # print(file)
        file_obj = File(file)

        # Checks if file extension is given and then compares it with the searched file extension
        if file_ext and not file_obj._extension == file_ext:
            continue
        all_files.append((str(file_obj.size()).zfill(12), file_obj.file()))
        if del_flag:
            os.remove(file)

    if org_flag:
        all_files.sort(key=lambda x: x[1])
        write_to_file(folder, all_files, out_file)
    else:
        all_files.sort(reverse=True)
        write_to_file(folder, all_files, out_file)

if __name__ == '__main__':

    my_parser = ArgumentParser(prog='Search files',
                               usage='Search files in a given directory. May choose specific filetypes. utilities -path "D:\Documents" -ext pdf -o pdf.txt -org',
                               description='Command Line Application searches a given directory then list and ' \
                               'summarize and write on a text file.',)
    my_parser.add_argument('-path', type=str,
                                   default=os.curdir,
                                   help='dir/folder to be search of files')
    my_parser.add_argument('-ext', '--file_extension', type=str,
                                                   default=None,
                                                   help='Optional file extension that will be search')
    my_parser.add_argument('-del', '--delete',
                                         action='store_true',
                                         help='Delete all files with the specified file extension')
    my_parser.add_argument('-out', '--output', type=str,
                                            default='output.txt',
                                            help= 'Output file where reports will be saved')
    my_parser.add_argument('-org', '--organize', action='store_true',
                                                 help= 'Organize files and put different files into each folders in the summary, default is organized by size.')
    args = my_parser.parse_args()

    main(args.path, args.file_extension, args.delete, args.organize, args.output)


import math
epoch_num = 3
batch_size = 1 # for scheduler, ignore.
batches = 32 #16 or 1
SGD = True #when true, ignores batches
slot_dim = 2 + 2 + 4 #TODO automate this # init + final + ruleOHE
slot_len = 16
slot_hidden = 16
rule_len = 16
slot_num = 2
rule_num = 4 # 5 for 'no move' also
mlp_input = 4
mlp_hidden = 32
mlp_output = 2
dropout = 0.0 #lower is giving much better results, 0.0 is best i think
tau = 0.8
lr = (5e-4)/2 #5e-4
max_lr = (5e-3)/2 #5e-3
maxMSE = math.sqrt(2)
KLtype = 'output' #rulevec or output or slotvec or model
betaKL = -1e-1 #make this negative only
betaKLMulti = 1
maxBetaKL = 100
minBetaKL = 0.02
betaRuleCE = 0 #1e-1
epsKL = 1e-7
betaCE = 1 #1e-1
betaMSE = 5 #1
endstate = (100.0,100.0)
seed = 0
test_train_split = 0.8
version = 3

import gc
gc.set_threshold(0)

globalBetaCE = betaCE
globalBetaMSE = betaMSE
globalBetaKL = betaKL

import numpy as np
import matplotlib.pyplot as plt
data = np.load(f'demo_data/datasetv{version}.npy', allow_pickle = True)
import torch.nn as nn
import torch
import torch.nn.functional as F
import sklearn.ensemble as se
from sklearn.model_selection import train_test_split
from tqdm import tqdm
import os
import pickle
import random
os.environ['PYTHONHASHSEED'] = str(seed)
# Torch RNG
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
# Python RNG
np.random.seed(seed)
random.seed(seed)

data, datatest = train_test_split(data, test_size = 1-test_train_split, random_state = seed)

class MLP(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim = 32, dropout = 0.0):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.dropout = dropout
        self.linear = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(hidden_dim, output_dim)
    def forward(self, x):
        x = self.linear(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.linear2(x)
        return x
oldrules = torch.load('Rule-Learning/Rule-Learning/learnt/1Object/4rules_4_2.pth')

class entity():
    def __init__(self, data):
        super().__init__()
        self.init = data[0]
        self.fin = data[1]
        self.size = data[2]
        self.action = data[3]
        ruleOHE = torch.zeros(rule_num)
        ruleOHE[int(self.action)] = 1.0
        self.data = (*self.init, *self.fin, *ruleOHE)

class slot(nn.Module):
    def __init__(self, entity, slot_len = slot_len):
        super().__init__()
        self.data = torch.tensor(entity.data)
    def query_gen(self, W):
        if type(W) == torch.nn.Parameter:
            return torch.matmul(self.data, W)
        else:
            return W(self.data)
    def key_gen(self, W):
        if type(W) == torch.nn.Parameter:
            return torch.matmul(self.data, W)
        else:
            return W(self.data)

class rule(nn.Module):
    def __init__(self, rule_len, rule_mlp):
        super().__init__()
        self.rule_vec = nn.Parameter(torch.randn(rule_len))
        self.rule_mlp = rule_mlp
    def key_gen(self, W):
        if type(W) == torch.nn.Parameter:
            return torch.matmul(self.rule_vec, W)
        else:
            return W(self.rule_vec)
    def rule_output(self, x):
        return self.rule_mlp(x)

class model(nn.Module):
    def __init__(self):
        super().__init__()
        if slot_hidden == -1:
            self.Wq = nn.Parameter(torch.randn(slot_dim, slot_len))
            self.Wk = nn.Parameter(torch.randn(rule_len, slot_len))
            self.Wqn = nn.Parameter(torch.randn(slot_dim, slot_len))
            self.Wkn = nn.Parameter(torch.randn(slot_dim, slot_len))
        else:
            self.MLPq = MLP(slot_dim, slot_len, slot_hidden)
            self.MLPk = MLP(rule_len, slot_len, slot_hidden)
            self.MLPqn = MLP(slot_dim, slot_len, slot_hidden)
            self.MLPkn = MLP(slot_dim, slot_len, slot_hidden)
        self.rule_pred = MLP(mlp_input+rule_num, 1, mlp_hidden)

    def forward(self, primarySlots, secondarySlots, rules):
        '''
        Part 1: Primary slot and Rule selection
        Part 1a: Key and Query generation
        '''
        if slot_hidden == -1:
            self.slot_query = torch.stack([slot.query_gen(self.Wq) for slot in primarySlots])
            self.rule_key = torch.stack([rule.key_gen(self.Wk) for rule in rules])
        else:
            self.slot_query = torch.stack([slot.query_gen(self.MLPq) for slot in primarySlots])
            self.rule_key = torch.stack([rule.key_gen(self.MLPk) for rule in rules])
        self.rule_mlp = [rule.rule_mlp for rule in rules]
        '''
        Part 1b: Choosing Primary slot and Rule
        '''
        self.slot_rule_matrix = torch.stack([torch.stack([torch.dot(self.slot_query[i], self.rule_key[j]) for j in range(len(self.rule_key))]) for i in range(len(self.slot_query))])
        self.best_slot_rule_mask = nn.functional.gumbel_softmax(self.slot_rule_matrix.view(-1), tau = tau, hard = True).view(len(self.slot_query), len(self.rule_key))
        self.chosen_primary_slot = torch.mm(self.slot_query.T,self.best_slot_rule_mask).sum(dim=1)
        self.primary_slot_mask = self.best_slot_rule_mask.sum(dim=1)
        self.chosen_rule = torch.mm(self.rule_key.T,self.best_slot_rule_mask.T).sum(dim=1)
        self.rule_mask = self.best_slot_rule_mask.sum(dim=0)
        '''
        Part 2: Contextual slot selection
        Part 2a: Key and Query generation
        '''
        if slot_hidden == -1:
            self.slot_query_2 = torch.stack([slot.query_gen(self.Wqn) for slot in primarySlots])
            self.slot_key = torch.stack([slot.key_gen(self.Wkn) for slot in secondarySlots])
        else:
            self.slot_query_2 = torch.stack([slot.query_gen(self.MLPqn) for slot in primarySlots])
            self.slot_key = torch.stack([slot.key_gen(self.MLPkn) for slot in secondarySlots])
        '''
        Part 2b: Choosing Contextual slot
        '''
        self.primary_slot_query = torch.mm(self.slot_query_2.T, self.primary_slot_mask.view(-1,1)).sum(dim=1)
        self.slot_slot_matrix = torch.stack([torch.dot(self.primary_slot_query, self.slot_key[j]) for j in range(len(self.slot_key))])
        self.best_slot_mask = nn.functional.gumbel_softmax(self.slot_slot_matrix.view(-1), tau = tau, hard = True).view(-1,1)
        self.chosen_secondary_slot = torch.mm(self.slot_key.T,self.best_slot_mask).sum(dim=1)
        self.secondary_slot_mask = self.best_slot_mask.sum(dim=1)
        '''
        Part 3: Applying MLP
        '''
        self.primary_slot = torch.matmul(torch.stack([slot.data for slot in primarySlots]).T, self.primary_slot_mask).T
#         print(primary_slot[:2])
        self.secondary_slot = torch.matmul(torch.stack([slot.data for slot in secondarySlots]).T, self.secondary_slot_mask).T
#         print(secondary_slot[:2])
        self.all_predicted_outputs = torch.stack([rule_mlp(torch.cat((self.primary_slot[:2], self.primary_slot[:2]))) for rule_mlp in self.rule_mlp]) # check if the :2 works or not
        self.predicted_output = torch.matmul(self.all_predicted_outputs.T, self.rule_mask)
#         print(self.rule_mask)
        #TODO !!!! maybe add 'with torch.no_grad:' here
#         print(torch.cat((self.primary_slot[:2], self.secondary_slot[:2], self.rule_mask)))
        self.no_grad_rule_mask = self.rule_mask.detach().clone()
#         print(self.no_grad_rule_mask.requires_grad)
#         print('----')
#         print(self.rule_mask.requires_grad)
        self.predicted_collision = self.rule_pred(torch.cat((self.primary_slot[:2], self.secondary_slot[:2], self.no_grad_rule_mask))) # check if the :2 works or not
#         print(self.predicted_collision)
        return self.primary_slot, self.secondary_slot, self.rule_mask, self.predicted_output, self.all_predicted_outputs, self.predicted_collision[0]

class JSD(nn.Module):
    def __init__(self):
        super(JSD, self).__init__()
        self.kl = nn.KLDivLoss(reduction='batchmean', log_target=True)

    def forward(self, p: torch.tensor, q: torch.tensor):
        p, q = p.view(-1, p.size(-1)), q.view(-1, q.size(-1))
        m = (0.5 * (p + q)).log()
        return 0.5 * (self.kl(m, p.log()) + self.kl(m, q.log()))

rules = [rule(rule_len, oldrules[_]) for _ in range(rule_num)]
m = model()
mse = nn.MSELoss()
cross = nn.CrossEntropyLoss()
jsdloss = JSD()
params = list(m.parameters())
for i in range(len(rules)):
    params += list(rules[i].parameters())
optimiser = torch.optim.Adam(params, lr = lr)
######## #TODO we have a lot of 'constants' here, make them options at the beginning of the notebook
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimiser, max_lr = max_lr, total_steps = (2*epoch_num+10)*batch_size)
mselosses = []
belosses = []
kllosses = []
rcelosses = []
XColl = []
yColl = []
XColltest = []
yColltest = []

for epoch in tqdm(range(epoch_num)):
    m.train()
    operation_to_rule = {k:{j:0 for j in range(rule_num)} for k in range(rule_num)}
    pAccuracy = 0
    pTotal = 0
    pFalsePos = 0
    for k in range(rule_num):
        rules[k].train()
        for params in rules[k].rule_mlp.parameters():
            params.requires_grad = False
        rules[k].rule_vec.requires_grad = True

    betaCE = globalBetaCE/1e5
    betaMSE = globalBetaMSE
    betaKL = globalBetaKL
    for params in m.rule_pred.parameters():
        params.requires_grad = False
    if slot_hidden == -1:
        m.Wq.requires_grad = True
        m.Wk.requires_grad = True
        m.Wqn.requires_grad = True
        m.Wkn.requires_grad = True
    else:
        for mlpmodel in [m.MLPq, m.MLPk, m.MLPqn, m.MLPkn]:
            for params in mlpmodel.parameters():
                params.requires_grad = True

    summselosses = 0.0
    sumbelosses = 0.0
    sumkllosses = 0.0
    sumrcelosses = 0.0
    cnt = 0
    shuffledData = list(range(len(data)))
    random.shuffle(shuffledData)
    for i in tqdm(shuffledData):
        cnt+=1
        inits, finals, object_size, initsobs, finalsobs, obstacle_size, obj, action1 = data[i]
#         print(inits, finals, object_size, initsobs, finalsobs, obstacle_size, obj, action1)
#         print("wowow")
        primarySlots = [slot(entity([inits[j],finals[j],object_size[j], action1]), slot_len = slot_len) for j in range(len(object_size))]
        secondarySlots = [slot(entity([initsobs[j],finalsobs[j],obstacle_size[j], action1]), slot_len = slot_len) for j in range(len(obstacle_size))]
        primary_slot, secondary_slot, rule_mask, predicted_output, all_predicted_outputs, predicted_collision = m(primarySlots, secondarySlots, rules)
#         print("start")
#         print(predicted_output)
#         print("========================\n",all_predicted_outputs)
#         print("--------")

#         print(primary_slot, secondary_slot, rule_mask, predicted_output, all_predicted_outputs )
#         print("--------")
#         print(predicted_output, predicted_collision[0], all_predicted_collision,torch.tensor(float(int(p))))
        KLdivs = torch.tensor(0.0)
        _n = 0
        if KLtype == 'output':
            for i in range(len(all_predicted_outputs)):
                for j in range(i):
                    _n+=1
                    KLdivs += (jsdloss(F.softmax(all_predicted_outputs[i]),F.softmax(all_predicted_outputs[j])))
        if KLtype == 'rulevec':
            for i in range(len(rules)):
                for j in range(i):
                    _n+=1
                    KLdivs += (jsdloss(F.softmax(rules[j].rule_vec),F.softmax(rules[i].rule_vec)))
        if KLtype == 'slotvec':
            for i in range(len(primarySlots)):
                for j in range(i):
                    _n+=1
                    KLdivs += (jsdloss(F.softmax(primarySlots[j].data),F.softmax(primarySlots[i].data)))
            for i in range(len(secondarySlots)):
                for j in range(i):
                    _n+=1
                    KLdivs += (jsdloss(F.softmax(secondarySlots[j].data),F.softmax(secondarySlots[i].data)))
        if KLtype == 'model':
            for i in range(len(rules)):
                for j in range(i):
                    _n+=1
                    #TODO made into torch.cat instead of torch.concat due to old torch version?
                    KLdivs += (jsdloss(F.softmax(torch.cat((rules[j].rule_mlp.linear.weight.flatten(),rules[j].rule_mlp.linear2.weight.flatten(),rules[j].rule_mlp.linear.bias.flatten(),rules[j].rule_mlp.linear2.bias.flatten()))),F.softmax(torch.cat((rules[i].rule_mlp.linear.weight.flatten(),rules[i].rule_mlp.linear2.weight.flatten(),rules[i].rule_mlp.linear.bias.flatten(),rules[i].rule_mlp.linear2.bias.flatten())))))
        try:
            KLdivs /= (_n*(_n-1)/2)
        except:
            pass
        p = (finals[obj] == endstate) #TODO currently this is just for 1 datapoint at a time.
        if (epoch == epoch_num - 1) and (len(XColl) < len(data)):
            XColl.append([*(primary_slot.detach().clone().numpy()[:2]), *(secondary_slot.detach().clone().numpy()[:2]), *(rule_mask.detach().clone().numpy())])
            yColl.append(int(p))
        loss = betaMSE*(1-int(p))*mse(predicted_output, torch.tensor(finals[obj]))/(maxMSE)
#         print("-----")
#         print(finals[obj])
        loss1 = betaCE*F.binary_cross_entropy(F.sigmoid(predicted_collision),torch.tensor(float(int(p))))
#         print(F.sigmoid(predicted_collision))
        loss2 = betaKL*KLdivs
        #TODO added a new loss function here
        # print(rule_mask)
        # print(torch.tensor(action1))
        #TODO loss3 = betaRuleCE*F.cross_entropy(rule_mask, torch.tensor(action1))
        loss3 = (1-int(p))*betaRuleCE*nn.CrossEntropyLoss()(rule_mask.view(1,-1), torch.tensor([action1]))
#         print(F.cross_entropy(rule_mask, torch.tensor(action1)), torch.argmax(rule_mask.view(1,-1)), action1)
        loss += loss3 + loss1
        summselosses += (1-int(p))*mse(predicted_output, torch.tensor(finals[obj]))
        sumbelosses += F.binary_cross_entropy(F.sigmoid(predicted_collision),torch.tensor(float(int(p))))
        sumkllosses += KLdivs
        sumrcelosses += (1-int(p))*F.cross_entropy(rule_mask.view(1,-1), torch.tensor([action1]))
        loss.backward()
#         print(m.rule_pred.linear.bias.grad)
        if (SGD is True) or (cnt % (len(data)//batches) == 0) or (cnt == len(data) - 1):
            optimiser.step()
            optimiser.zero_grad()
            # gc.collect()
            # print("yay")
    betaKL *= betaKLMulti
    if -1*betaKL > maxBetaKL:
        betaKL = -1*maxBetaKL
    if -1*betaKL < minBetaKL:
        betaKL = -1*minBetaKL
    scheduler.step()
    mselosses.append(summselosses)
    belosses.append(sumbelosses)
    kllosses.append(sumkllosses)
    rcelosses.append(sumrcelosses)
    print(f"MSE Loss = {loss - loss1 - loss2 - loss3}, BCE Loss = {loss1}, KL Loss = {loss2}, RuleCE Loss = {loss3}")
    m.eval()
    for k in range(rule_num):
        rules[k].eval()
#         print(k,rules[k].rule_vec)
    totalmae = 0
    for i in tqdm(range(len(datatest))):
        inits, finals, object_size, initsobs, finalsobs, obstacle_size, obj, action1 = datatest[i]
        primarySlots = [slot(entity([inits[j],finals[j],object_size[j],action1]), slot_len = slot_len) for j in range(len(object_size))]
        secondarySlots = [slot(entity([initsobs[j],finalsobs[j],obstacle_size[j],action1]), slot_len = slot_len) for j in range(len(obstacle_size))]
#         slots = [slot(entity(datatest[i]), slot_len = slot_len) for _ in range(slot_num)]
        primary_slot, secondary_slot, rule_mask, predicted_output, all_predicted_outputs, predicted_collision = m(primarySlots, secondarySlots, rules)
        rule_index = torch.argmax(rule_mask)
        p = (finals[obj] == endstate)
        if (epoch == epoch_num - 1) and (len(XColltest) < len(datatest)):
            XColltest.append([*(primary_slot.detach().clone().numpy()[:2]), *(secondary_slot.detach().clone().numpy()[:2]), *(rule_mask.detach().clone().numpy())])
            yColltest.append(int(p))
            if int(p) == 1:
                # print(int(rule_index), action1)
                pass
        if finals[obj] == endstate:
            if predicted_collision > 0:
                pAccuracy += 1
            pTotal += 1
        else:
            if predicted_collision > 0:
                pFalsePos += 1
            operation_to_rule[action1][int(rule_index)] += 1
        if predicted_collision > 0:
            totalmae += 100*(1-int((finals[obj] == endstate)))
        else:
            totalmae += np.sum(np.abs(predicted_output.detach().numpy()-np.array(finals[obj]))) #TODO made this :-1, shouldn't hard code it like this
    totalmae /= len(datatest)
    print(f"Total MAE = {totalmae}")
    # print('------------------')
    # print(rules[0].rule_mlp)
    # print('------------------')
    for k in operation_to_rule:
        print(k, end = ' ')
        print(operation_to_rule[k])
    print('Accuracy: {}/{} = {}'.format(np.sum([operation_to_rule[k][k] for k in range(len(operation_to_rule))]),np.sum(np.sum([operation_to_rule[k][j] for j in range(len(operation_to_rule)) for k in range(len(operation_to_rule))])),np.sum([operation_to_rule[k][k] for k in range(len(operation_to_rule))])/np.sum(np.sum([operation_to_rule[k][j] for j in range(len(operation_to_rule)) for k in range(len(operation_to_rule))]))))
    print('Recall of unreachable: {}/{} = {} with {} false positives'.format(pAccuracy, pTotal, pAccuracy/pTotal, pFalsePos))
    print('========================================================')

# 0 {0: 0, 1: 0, 2: 264, 3: 0}
# 1 {0: 0, 1: 0, 2: 0, 3: 261}
# 2 {0: 0, 1: 228, 2: 0, 3: 0}
# 3 {0: 247, 1: 0, 2: 0, 3: 0}
# 97.5 accuracy

# exit the python script

pickle.dump((XColl,yColl,XColltest,yColltest),open(f"demo_data/XYXtYtCollv{version}.pkl","wb"))

allrules = nn.ModuleList()
for k in range(rule_num):
    allrules.append(rules[k].rule_mlp)
torch.save(allrules, f'models/RulesMLPv{version}.pth')
allrules = []
for k in range(rule_num):
    allrules.append(rules[k].rule_vec)
f = open(f'models/RulesVecv{version}.pkl', 'wb')
pickle.dump(allrules, f)
torch.save(m, f'models/Modelv{version}.pth')

exit()

pmodel = se.RandomForestClassifier(n_estimators=100, max_depth=2,random_state=0)
pmodel.fit(XColl,yColl)

print("pmodel")
print(pmodel.score(XColl,yColl),pmodel.score(XColltest,yColltest))
# find f1 score as pmodel on xcoll and xcolltest
from sklearn.metrics import classification_report
target_names = ['feasible', 'not feasible']
print(classification_report(yColl, pmodel.predict(XColl),target_names = target_names))
print(classification_report(yColltest, pmodel.predict(XColltest),target_names = target_names))
exit()
# # import joblib
# # joblib.dump(pmodel, "../learnt/pmodelv15.2.joblib")
# import pickle
# import joblib

# In[77]:

# tmselosses = [a.item() for a in mselosses]
# tbelosses = [a.item() for a in belosses]
# tkllosses = [a.item() for a in kllosses]
# # trcelosses = [a.item() for a in rcelosses]
# plt.plot(range(len(tmselosses)+1)[1:],[20*a for a in tmselosses],label="mean squared error loss")
# plt.plot(range(len(tbelosses)+1)[1:], tbelosses, label = "binary cross entropy loss")
# plt.plot(range(len(tkllosses)+1)[1:], tkllosses, label = "KL divergence loss")
# # plt.plot(range(len(trcelosses)+1)[1:], trcelosses, label = "rule cross entropy loss")
# plt.title("Losses vs Epoch")
# plt.legend()
# pass

# In[78]:

# datacheck = datatest[:20]
# # rules = [rule(rule_len, mlp_input, mlp_output, mlp_hidden, dropout) for _ in range(rule_num)]
# m.eval()
# for k in range(rule_num):
#     rules[k].eval()
# for i in tqdm(range(len(datacheck))):
# #     slots = [slot(entity(datacheck[i]), slot_len = slot_len) for _ in range(slot_num)]
#     inits, finals, object_size, initsobs, finalsobs, obstacle_size, obj, action1 = datacheck[i]
#     primarySlots = [slot(entity([inits[j],finals[j],object_size[j],action1]), slot_len = slot_len) for j in range(len(object_size))]
#     secondarySlots = [slot(entity([initsobs[j],finalsobs[j],obstacle_size[j],action1]), slot_len = slot_len) for j in range(len(obstacle_size))]
#     primary_slot, secondary_slot, rule_mask, predicted_output, all_predicted_outputs, predicted_collision = m(primarySlots, secondarySlots, rules)
# #     primary_slot, secondary_slot, rule_mask, predicted_output = m(slots, rules)
#     print(predicted_output.detach().numpy(), inits[obj], finals[obj], action1, round(np.sum(np.abs(predicted_output.detach().numpy()-np.array(finals[obj]))),3)) #TODO made this :-1, shouldn't hard code it like this

# In[86]:

#TODO
# Maybe start saving whenever MAE is lowest ngl
allrules = nn.ModuleList()
for k in range(rule_num):
    allrules.append(rules[k].rule_mlp)
torch.save(allrules, '../learnt/RulesMLPv15.3.pth')
allrules = []
for k in range(rule_num):
    allrules.append(rules[k].rule_vec)
import pickle
f = open('../learnt/RulesVecv15.3.pkl', 'wb')
pickle.dump(allrules, f)
torch.save(m, '../learnt/Modelv15.3.pth')
import joblib
joblib.dump(pmodel, "../learnt/pmodelv15.3.joblib")

# In[ ]:


"""
Provides a function for splitting videos into consecutive chunks at provided intervals or at exact locations.
"""
from typing import Union, Tuple, List, Optional, Iterable

import diplomat.processing.type_casters as tc
from diplomat.utils.pretty_printer import printer
import cv2
from pathlib import Path
from os import PathLike
import tqdm
import math

FALLBACK_CODEC = "mp4v"
FALLBACK_EXT = ".mp4"

@tc.typecaster_function
def split_videos(
    videos: tc.Union[tc.Path, tc.List[tc.Path]],
    seconds_per_segment: tc.Union[tc.List[int], int] = 300,
    output_fourcc_string: tc.Optional[str] = None,
    output_extension: tc.Optional[str] = None
) -> tc.List[tc.List[Path]]:
    """
    Split a video into even length segments. This will produce a list of videos with "-part{number}" appended to the
    original video name in the same directory as the original video.

    :param videos: Either a single path-like object (string or Path) or a list of path-like objects, being the
                   paths to the videos to split into several segments.
    :param seconds_per_segment: An integer or a list of integers. If a single integer, represents the length of each
                                split segment in seconds (Ex. if 30, split the clip every 30 seconds). If a list of
                                integers, represents the locations to split the video at in seconds. The list can be
                                out of order, as it will be automatically sorted. Also, values that are out of range
                                will be ignored.
                                (Ex. if [10, 400, 100, 30], split the video at 10s, 400s, 100s, and 30s].
    :param output_fourcc_string: Optional, the fourcc string to use for output videos. If not specified uses input video fourcc code.
    :param output_extension: Optional, the file extension to use for output videos (including the dot).
                             Defaults to the input file extension if not specified.

    :returns: A list of lists of Path objects, being the new split video paths for each and every video...
    """
    videos = _sanitize_path_arg(videos)

    if(videos is None):
        raise ValueError("No videos provided!!!")

    return [_split_single_video(video, seconds_per_segment, output_fourcc_string, output_extension) for video in videos]

def _split_single_video(
    video_path: Path,
    seconds_per_segment: Union[int, List[int]],
    output_fourcc_string: Optional[str] = None,
    output_extension: Optional[str] = None
) -> List[Path]:
    """
    PRIVATE: Splits a single video.

    :param video_path: The path of the original video to split.
    :param seconds_per_segment: The duration of each segment, in seconds.
    :param output_fourcc_string: Optional, the output fourcc string. If not specified uses input fourcc code.
    :param output_extension: Optional, the output file extension. Defaults to the input file extension if not specified.

    :returns: The paths of the newly split videos, as a list of "Path"s...
    """
    printer(f"Processing video: {video_path}")
    vid = cv2.VideoCapture(str(video_path))

    width, height = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH)), int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = vid.get(cv2.CAP_PROP_FPS)
    total_frames = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))

    four_cc = int(vid.get(cv2.CAP_PROP_FOURCC))
    if(four_cc == 0):
        four_cc = -1
    else:
        four_cc = "".join([chr((four_cc >> i) & 255) for i in range(0, 32, 8)])
        four_cc = cv2.VideoWriter_fourcc(*four_cc)

    if(output_fourcc_string is not None):
        four_cc = cv2.VideoWriter_fourcc(*output_fourcc_string)

    extension = video_path.suffix if(output_extension is None) else output_extension

    writer = None
    segment = 0
    frame = 0

    if(isinstance(seconds_per_segment, int)):
        # Yes, a range can be indexed just like an array... :)
        split_loc = range(0, total_frames, int(seconds_per_segment * fps))
    else:
        split_loc = sorted(set([0] + [int(seg * fps) for seg in seconds_per_segment]))

    bar = tqdm.tqdm(total=total_frames)

    zero_pad_amt = int(math.ceil(math.log10(total_frames + 1)))
    paths = []

    while(vid.isOpened()):
        if(writer is None):
            try:
                start = _list_access(split_loc, [total_frames], segment)
                end = _list_access(split_loc, [total_frames], segment + 1) - 1
                writer, p = _new_video_writer(video_path, (start, end), zero_pad_amt, four_cc, fps, (width, height), extension)
                paths.append(p)
            except OSError:
                vid.release()
                bar.close()
                raise
            segment += 1

        res, frm = vid.read()

        if(not res):
            break

        if(writer.isOpened()):
            writer.write(frm)

        if((segment < len(split_loc)) and (frame >= split_loc[segment])):
            writer.release()
            writer = None

        frame += 1
        bar.update(1)

    if(writer is not None):
        writer.release()

    bar.close()
    vid.release()

    return paths

def _list_access(list1, list2, index):
    comb = (list1, list2)
    return comb[index // len(list1)][index % len(list1)]

def _sanitize_path_arg(
    paths: Union[None, Iterable[PathLike], PathLike]
) -> Optional[List[Path]]:
    """
    Sanitizes a pathlike or list of pathlike argument and returns a list of Path, or None if rogue data was passed...
    """
    if isinstance(paths, (PathLike, str)):
        return [Path(str(paths)).resolve()]
    elif isinstance(paths, Iterable):
        paths = list(paths)
        if len(paths) > 0:
            return [Path(str(path)).resolve() for path in paths]
        else:
            return None
    else:
        return None

def _new_video_writer(
    video_path: Path,
    segment: Tuple[int, int],
    padding: int,
    four_cc: int,
    fps: float,
    size: Tuple[int, int],
    ext: str
) -> Tuple[cv2.VideoWriter, Path]:
    """
    PRIVATE: Construct a new video writer. Will try to use the passed codec if possible, otherwise uses a fallback codec and
    format.
    """
    suffix = f"_part{segment[0]:0{padding}d}-{segment[1]:0{padding}d}"
    preferred_path = video_path.parent / f"{video_path.stem}{suffix}{ext}"
    writer = cv2.VideoWriter(str(preferred_path), four_cc, fps, size)
    if(writer.isOpened()):
        return (writer, preferred_path)

    writer.release()
    fallback_path = video_path.parent / f"{video_path.stem}{suffix}{FALLBACK_EXT}"
    writer = cv2.VideoWriter(str(fallback_path), cv2.VideoWriter_fourcc(*FALLBACK_CODEC), fps, size)
    if(writer.isOpened()):
        return (writer, fallback_path)

    writer.release()
    raise OSError("Can't find a codec to write with!!!")


import torch
import numpy as np
import scipy.sparse as sp
from attackers.base import AttackABC
from torch_sparse import SparseTensor

class RandomAttack(AttackABC):

    def __init__(self,
                 attack_config, pyg_data,
                 model, device, logger, **kwargs):
        super(RandomAttack, self).__init__(attack_config, pyg_data, model, device, logger)
        self.att_type = attack_config['att_type']

    def _attack(self, n_perturbations):
        adj = self.pyg_data.adj_t.to_scipy(layout='csr')
        modified_adj = adj.tolil()
        type = self.att_type.lower()
        assert type in ['add', 'remove', 'flip']

        if type == 'flip':
            # sample edges to flip
            edges = self.random_sample_edges(adj, n_perturbations, exclude=set())
            for n1, n2 in edges:
                modified_adj[n1, n2] = 1 - modified_adj[n1, n2]
                modified_adj[n2, n1] = 1 - modified_adj[n2, n1]

        if type == 'add':
            # sample edges to add
            nonzero = set(zip(*adj.nonzero()))
            edges = self.random_sample_edges(adj, n_perturbations, exclude=nonzero)
            for n1, n2 in edges:
                modified_adj[n1, n2] = 1
                modified_adj[n2, n1] = 1

        if type == 'remove':
            # sample edges to remove
            nonzero = np.array(sp.triu(adj, k=1).nonzero()).T
            indices = np.random.permutation(nonzero)[: n_perturbations].T
            modified_adj[indices[0], indices[1]] = 0
            modified_adj[indices[1], indices[0]] = 0

        self.adj_adversary = SparseTensor.from_dense(torch.tensor(modified_adj.todense()))

    def random_sample_edges(self, adj, n, exclude):
        itr = self.sample_forever(adj, exclude=exclude)
        return [next(itr) for _ in range(n)]

    def sample_forever(self, adj, exclude):
        """Randomly random sample edges from adjacency matrix, `exclude` is a set
        which contains the edges we do not want to sample and the ones already sampled
        """
        while True:
            # t = tuple(np.random.randint(0, adj.shape[0], 2))
            # t = tuple(random.sample(range(0, adj.shape[0]), 2))
            t = tuple(np.random.choice(adj.shape[0], 2, replace=False))
            if t not in exclude:
                yield t
                exclude.add(t)
                exclude.add((t[1], t[0]))