"""TensorFlow utilities and data structures for rapid prototyping.

This module is inspired by:
.. _Structuring Your TensorFlow Models by Danijar Hafner / CC BY-SA 3.0:
   https://danijar.com/structuring-your-tensorflow-models/
.. _Patterns for Fast Prototyping with TensorFlow by Danijar Hafner / CC BY-SA 3.0:
   https://danijar.com/patterns-for-fast-prototyping-with-tensorflow/
"""

__license__ = "MIT License"
__copyright__ = "Copyright (c) 2019 Piotr Januszewski"

import functools

class AttrDict(dict):
    """Defines an attribute dictionary that allows the user to address keys as
    if they were attributes.

    `AttrDict` inherits from `dict`, so it can be used everywhere where `dict` is expected.
    The __init__ constructs an empty `AttrDict` or recursively initialises an `AttrDict`
    from a (maybe nested) `dict`.

    Args:
        other (dict): Either dictionary (can be nested) to initialise an `AttrDict` from
        or None to construct an empty `AttrDict`.
    """

    def __init__(self, other=None):
        super().__init__()

        if other is not None:
            for key, value in other.items():
                # It is easier to ask for forgiveness, than to ask for permission
                try:
                    self[key] = AttrDict(value)
                except AttributeError:
                    self[key] = value

    def nested_update(self, other):
        """Updates the values in this and the nested `AttrDict`s.

        Args:
            other (dict): Dictionary (can be nested) to update the `AttrDict` from.
        """

        for key, value in other.items():
            # It is easier to ask for forgiveness, than to ask for permission
            try:
                self[key].nested_update(value)
            except AttributeError:
                self[key] = value

        return self

    __getattr__ = dict.__getitem__
    __setattr__ = dict.__setitem__

def attrdict_from_json(data):
    """Constructs an AttrDict from JSON formatted string.

    Args:
        data (str): JSON formatted string.
    """

    from json import loads
    return AttrDict(loads(data))

def attrdict_from_yaml(data):
    """Constructs an AttrDict from YAML formatted string.

    Args:
        data (str): YAML formatted string.
    """

    from yaml import safe_load
    return AttrDict(safe_load(data))

def lazy_property_with_scope(*args, scope_name=None, **kwargs):
    """Property decorator on steroids.

    Property decorator which additionally defines TensorFlow variable scope
    and adds lazy loading.

    This decorator can be used without parentheses if no arguments are provided.

    Args:
        scope_name (str): The scope name. If `None` decorated function's name is
                          used as the scope name.
        *args: Passed to `tf.variable_scope`.
        **kwargs: Passed to `tf.variable_scope`.
    """

    def define_scope(function, scope_name, *args, **kwargs):
        """Decorator for properties that define TensorFlow operations.

        The decorated function will only be executed once. Subsequent calls to
        it will directly return the result so that operations are added to
        the graph only once. The operations added by the function live within
        a tf.variable_scope(). If this decorator is used with arguments, they
        will be forwarded to the variable scope. The scope name defaults to
        the name of the decorated function.

        Args:
            scope_name (str): The scope name. If `None` decorated function's
                              name is used as the scope name.
            *args: Passed to `tf.variable_scope`.
            **kwargs: Passed to `tf.variable_scope`.
        """

        from tensorflow import variable_scope

        attribute = '_cache_' + function.__name__
        name = scope_name or function.__name__

        @property
        @functools.wraps(function)
        def decorator(self):
            if not hasattr(self, attribute):
                with variable_scope(name, *args, **kwargs):
                    setattr(self, attribute, function(self))
            return getattr(self, attribute)
        return decorator

    if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
        return define_scope(args[0], scope_name=scope_name)
    return lambda function: define_scope(function, scope_name=scope_name,
                                         *args, **kwargs)

def share_variables(function):
    """Decorator for methods that define TensorFlow operations.

    The decorated method is wrapped with tf.make_template() so that it does
    variable sharing. The tf.make_template() will only be executed once.
    Subsequent calls to the decorated function will reuse the template so that
    operations are added to the graph only once.
    """

    from tensorflow import make_template

    attribute = '_cache_' + function.__name__

    @functools.wraps(function)
    def decorator(self, *args, **kwargs):
        if not hasattr(self, attribute):
            setattr(self, attribute,
                    make_template(function.__name__,
                                  function,
                                  create_scope_now_=True))
        return getattr(self, attribute)(self, *args, **kwargs)
    return decorator


#load the libraries
import pandas as pd
import re
import numpy as np
from progressbar import ProgressBar,Bar,Percentage
from scanpy import AnnData
import itertools
from lxml import etree as ET
from colour import Color
from utils import ttest,mtest
from IPython.display import Image

"""
Class to compute the RAS values

"""

class RAS_computation:

    def __init__(self,adata,model):

        self._logic_operators = ['and', 'or', '(', ')']
        self.val_nan = np.nan

        # Build the dictionary for the GPRs
        df_reactions = pd.DataFrame(index=[reaction.id for reaction in model.reactions])
        gene_rules=[reaction.gene_reaction_rule for reaction in model.reactions]

        gene_rules=[el.replace("OR","or").replace("AND","and").replace("(","( ").replace(")"," )") for el in gene_rules]
        df_reactions['rule'] = gene_rules
        df_reactions = df_reactions.reset_index()
        df_reactions = df_reactions.groupby('rule').agg(lambda x: sorted(list(x)))

        self.dict_rule_reactions = df_reactions.to_dict()['index']

        # build useful structures for RAS computation
        self.model = model
        self.count_adata = adata.copy()
        self.genes = self.count_adata.var.index.intersection([gene.id for gene in model.genes])

        #check if there is one gene at least
        if len(self.genes)==0:
            print("ERROR: no gene of the count matrix is in the metabolic model!")
            print(" are you sure that the gene annotation is the same for the model and the count matrix?")
            return -1

        self.cell_ids = list(self.count_adata.obs.index.values)
        self.count_df_filtered = self.count_adata.to_df().T.loc[self.genes]

    def compute(self,
                or_expression=np.nansum,   # type of operation to do in case of an or expression (max, sum, mean)
                and_expression=np.nanmin,  # type of operation to do in case of an and expression(min, sum)
                drop_na_rows=True,         # if True remove the nan rows of the ras  matrix
                drop_duplicates=False,     # if true, remove duplicates rows
                regexp=re.compile(r"\([a-zA-Z0-9-.:\s]+\)"),  # regular expression inside a parenthesis
                print_progressbar=True,    # if True, print the progress bar
                add_count_metadata=True,   # if True add metadata of cells in the ras adata
                add_met_metadata=True,      # if True add metadata from the metabolic model (gpr and compartments of reactions)
                tag_counmatrix="countmatrix_"  #suffix for information coming from original countmatrix
                ):

        self.or_function = or_expression
        self.and_function = and_expression

        ras_df = pd.DataFrame(index=range(len(self.dict_rule_reactions)), columns=self.cell_ids)
        ras_df[:][:] = self.val_nan

        if print_progressbar:
            pbar = ProgressBar(widgets=[Percentage(), Bar()], maxval=len(self.dict_rule_reactions)).start()
            i = 0

        # for loop on reactions
        ind = 0
        for rule, reaction_ids in self.dict_rule_reactions.items():
            if len(rule) != 0:
                # there is one gene at least in the formula
                rule_split = rule.split()
                rule_split_elements = list(filter(lambda x: x not in self._logic_operators, rule_split))  # remove of all logical operators
                rule_split_elements = list(np.unique(rule_split_elements))                                # genes in formula

                # which genes are in the count matrix?
                genes_in_count_matrix = list(set([el for el in rule_split_elements if el in self.genes]))
                genes_notin_count_matrix = list(set([el for el in rule_split_elements if el not in self.genes]))

                if len(genes_in_count_matrix) > 0: #there is at least one gene in the count matrix
                     if len(rule_split) == 1:
                         #one gene --> one reaction
                         ras_df.iloc[ind] = self.count_df_filtered.loc[genes_in_count_matrix]
                     else:
                        # more genes in the formula
                        lista = re.findall(regexp, rule)

                        if len(lista) == 0:
                             #or/and sequence
                             matrix = self.count_df_filtered.loc[genes_in_count_matrix].values
                             if len(genes_notin_count_matrix) > 0:
                                matrix = np.vstack([matrix, [self.val_nan for el in self.cell_ids]])

                             if 'or' in rule_split:
                                ras_df.iloc[ind] = self.or_function(matrix, axis=0)
                             else:
                                ras_df.iloc[ind] = self.and_function(matrix, axis=0)
                        else:
                            # ho almeno una tonda
                            data = self.count_df_filtered.loc[genes_in_count_matrix]  # dataframe of genes in the GPRs
                            genes = data.index
                            j = 0

                            for cellid in self.cell_ids:    #for loop on the cells
                                lista_cell = lista.copy()
                                rule_cell = rule

                                while len(lista_cell) > 0:
                                    #
                                    for el in lista_cell:
                                        #print(el[1:-1])
                                        value = self._evaluate_expression(el[1:-1].split(), data[cellid], genes)
                                        rule_cell = rule_cell.replace(el, str(value))
                                    lista_cell = re.findall(regexp, rule_cell)

                                ras_df.iloc[ind, j] = self._evaluate_expression(rule_cell.split(), data[cellid], genes)
                                j=j+1

            ind = ind+1
            #update percentage
            if print_progressbar:
                pbar.update(i+1)
                i = i+1

        if print_progressbar:
            pbar.finish()

        ras_df=ras_df.astype("float")
        ras_df['REACTIONS'] = [reaction_ids for rule,reaction_ids in self.dict_rule_reactions.items()]

        reactions_common = pd.DataFrame()
        reactions_common["REACTIONS"] = ras_df['REACTIONS']
        reactions_common["proof2"] = ras_df['REACTIONS']
        reactions_common = reactions_common.explode('REACTIONS')
        reactions_common = reactions_common.set_index("REACTIONS")

        ras_df = ras_df.explode("REACTIONS")
        ras_df = ras_df.set_index("REACTIONS")

        if drop_na_rows:
            ras_df = ras_df.dropna(how="all")

        if drop_duplicates:
            ras_df = ras_df.drop_duplicates()

        #create AnnData structure for RAS
        ras_adata = AnnData(ras_df.T)

        #add metadata
        if add_count_metadata:
            ras_adata.var["common_gprs"] = reactions_common.loc[ras_df.index]
            ras_adata.var["common_gprs"] = ras_adata.var["common_gprs"].apply(lambda x: ",".join(x))
            for el in self.count_adata.obs.columns:
                ras_adata.obs[tag_countmatrix+el]=self.count_adata.obs[el]

        if add_met_metadata:
            if len(self.model.compartments)>0:
                  ras_adata.var['compartments']=[list(self.model.reactions.get_by_id(reaction).compartments) for reaction in ras_adata.var.index]
                  ras_adata.var['compartments']=ras_adata.var["compartments"].apply(lambda x: ",".join(x))

            ras_adata.var['GPR rule'] = [self.model.reactions.get_by_id(reaction).gene_reaction_rule for reaction in ras_adata.var.index]

        return ras_adata

    def _check_number(self,value):
      try:
        float(value)
        return True
      except ValueError:
        return False

    def _evaluate_expression(self, rule_split, values_cellid, genes):

        #ci sono per forza solo or
        rule_split2 = list(filter(lambda x: x != "or" and x!="and", rule_split))

        values = list()

        for el in rule_split2:
             if self._check_number(el):
                 values.append(float(el))
             elif el in genes:
                 values.append(values_cellid[el])
             else:
                 values.append(self.val_nan)

        if "or" in rule_split:
            #or sequence
            return self.or_function(values)
        else:
            #and sequence
            return self.and_function(values)

"""
Class to color a metabolic map

"""

class RAS_map():

    def __init__(self):
        return None

    def colorMap(self,
                 fileMap,
                 fileMapColor,
                 df,                            #dataset of reaction to color the map
                 colors=["blue", "red"],        #color scale              #
                 nosignificant_color="grey",    #color of reaction whose differences results not significant
                 ref_width_value=10,            #dimension of reaction rows
                 nosignificant_width=5,         #dimension of the row of reaction whose differences results not significant
                 min_val=0,                     #minimum width value. The minimum width of the row is min_val*ref_width_value
                 max_val=10,                    #maximum width value. The maximum width of the row is max_val*ref_width_value
                 unconfined=True,               #Set unconfined=True to disable max-width confinement of the image.
                 width_image=1000               #dimension of the image
                 ):

        df=df[df[df.columns[0]].isna() == False]
        df_significant = df[np.abs(df[df.columns[0]]) > 0]
        df_significant = df_significant.sort_values(by=df_significant.columns[0])
        df_nosignificant = df[np.abs(df[df.columns[0]]) == 0]

        #create a vector
        color_range = [Color(colors[0]),Color(colors[1])]

        dict_colors = dict()
        dict_widths = dict()

        for el in df_significant.index:
            dict_colors[el] = str(color_range[0]) if df_significant.loc[el,df.columns[0]]<0 else str(color_range[1])

            width=np.max([np.abs(df_significant.loc[el,df.columns[0]]),min_val])
            width=np.min([width,max_val])

            dict_widths[el] = str(np.round(ref_width_value*width,2))

        for el in df_nosignificant.index:
            dict_colors[el]=nosignificant_color
            dict_widths[el]=str(nosignificant_width)

        root = ET.parse(open(fileMap, 'r'))

        for element in root.iter():
            eTag = element.tag.split("}")[1]
            if (eTag=="g") and element.get('class')=="reaction":
                for key in dict_colors.keys():
                    if str(ET.tostring(element,pretty_print=True)).find(">"+key+"<")!=-1:
                        for el in element.iter():
                            if el.tag.split("}")[1]=="path":
                                el.set("style", "stroke:"+dict_colors[key]+";"+"stroke-width:"+dict_widths[key]+";")

        with open(fileMapColor, 'w') as newMapFile:
            string=ET.tostring(root, pretty_print=True)
            string=self._to_utf8(string)
            newMapFile.write(string)

        if width_image:
            image = Image(url=fileMapColor, unconfined =unconfined, width=width_image)
        else:
            image = Image(url=fileMapColor, unconfined =unconfined)

        return image

    def _to_utf8(self,s):
        return s if isinstance(s, str) else s.decode('utf-8')

"""
Function to evalute possible significant change between the RAS of two groups
"""

def computeRAS_diff(ras_adata,        #expect normalize and logarithmized  data
                name_feature,         #name of the feature for which the RAS differences are computed
                func=np.mean,         #method to aggregate ras values to perform group differences
                function=lambda x: np.log2(x[0]/x[1]), #funcion used to test how much RAS changes between two groups
                threshold_min=np.log2(1.2),    #minimum threshold below which there is no significant change between two groups
                threshold_max=np.log2(3),      #maximum threshold
                which_test="mtest",   #which statistical test used to test highly differentiated reactions
                p_min=0.05            #pvalue for statistical t-test
                ):

    if which_test not in [None,"mtest","ttest"]:
        return print("statistical test must be ttest, mtest or None!")

    ras_adata2 = ras_adata.copy()

    groups = ras_adata2.obs.groupby(name_feature).indices
    names = list(groups.keys())

    df = pd.DataFrame(index=ras_adata2.var.index)
    df_list=pd.DataFrame(index=ras_adata2.var.index)
    for inds,name in zip(groups.values(),names):
         adata=ras_adata2[ras_adata2.obs.iloc[inds].index,:]
         df[name]=adata.to_df().apply(lambda x:func(x),axis=0)
         df_list[name]=adata.to_df().T.apply(lambda x:list(x),axis=1)

    tests=itertools.combinations_with_replacement(names,2)
    tests=[(test[0],test[1]) for test in tests if test[0]!=test[1]]

    df_comparison=pd.DataFrame(index=df.index)
    for test in tests:
        val=[df[test[0]],df[test[1]]]
        name=str(test[0]+" vs "+test[1])
        df_comparison[name]=function(val)
        df_comparison[name][np.abs(df_comparison[name])<=threshold_min]=0
        df_comparison[name][df_comparison[name]>=threshold_max]=threshold_max
        df_comparison[name][df_comparison[name]<=-threshold_max]=-threshold_max

        if which_test:
            if which_test=="ttest":
                pvalues = df_list.apply(lambda x: ttest(x[test[0]], x[test[1]]), axis=1)
                df_comparison[name][pvalues>p_min]=0
            elif which_test=="mtest":
                pvalues = df_list.apply(lambda x: mtest(x[test[0]], x[test[1]]), axis=1)
                df_comparison[name][pvalues>p_min]=0

    return df_comparison


import os
import asyncio
import aiohttp
from a5_4_downloadimage import BASE_URL, save_image, show, main

@asyncio.coroutine
def get_image(suffix):
    url = os.path.join(BASE_URL, suffix)
    # response = yield from aiohttp.request('GET', url)
    response = yield from aiohttp.ClientSession().get(url)
    image = yield from response.read()
    return image

@asyncio.coroutine
def download_one(img):
    image = yield from get_image(img)
    show(img)
    save_image(image, filename=img)
    return img

def download_all(image_list):
    loop = asyncio.get_event_loop()
    to_do = [download_one(img) for img in sorted(image_list)]
    wait_coro = asyncio.wait(to_do)
    res, _ = loop.run_until_complete(wait_coro)
    loop.close()
    return len(res)

if __name__ == '__main__':
    main(download_all)
    #  download 80 images in 1.0823283195495605s


from pixel import pixels

class Frames:
    def __init__(self, num):
        """
        hold an array of frame buffers.
        the Buffur first has a boolean. If True, the buffer is displayed and can be reused.
        If false, it is waiting to be displayed.

        :param num: number of frames
        """
        self.fps = 24
        self.num = num
        self.buffer = [bytearray(pixels.num_leds*3)] * num
        self.shown = [True] * num
        self.display = 0

    def store(self, frame) -> bool:
        """
        Store the frame in the first available buffer.
        Return False is all buffers are filled.
        :param frame:
        :return:
        """
        for i in range(self.num):
            num = i + self.display
            if num >= self.num:
                num -= self.num
            if self.shown[num]:
                self.buffer[num] = frame
                self.shown[num] = False
                return True

        return False

    def run_next(self):
        if self.shown[self.display]:
            return False

        for i in range(pixels.num_leds):
            data = self.buffer[self.display][3*i:3*i+3]
            pixels[i] = (data[1], data[0], data[2])
        pixels.show()

        self.shown[self.display] = True
        self.display += 1
        if self.display == self.num:
            self.display = 0

        return True

frames = Frames(10)


import sys
sys.path.append("..")
import logging
import psutil
import torch
import torch.distributed.autograd as dist_autograd
from torch.distributed import rpc
from torch.distributed.optim import DistributedOptimizer

from utils.model_utils import construct_model

class Node:
    def __init__(self,
        server_rref:rpc.RRef,
        model_type:str,
        lr:float,
        cut_point:int
    ) -> None:
        self.__server_rref = server_rref
        self.__lr = lr
        self.__model = construct_model(model_type).get_splited_module(cut_point)[0]
        self.__p = psutil.Process()

    def set_next_node(self,
        next_node_rref:rpc.RRef
    ) -> None:
        logging.info(f"setting up next node: {next_node_rref.owner_name()}")
        self.__next_node = next_node_rref

    def set_trainloader(self,
        trainloader: torch.utils.data.DataLoader
    ) -> None:
        self.__trainloader = trainloader
        self.__trainloader_iter = iter(self.__trainloader)

    def train(self,
        client_model:dict,
        step_counter:int,
        eval_step:int,
    ) -> None:
        logging.info(f"training {step_counter} ...")
        self.__model.load_state_dict(client_model)
        try:
            inputs, labels = self.__trainloader_iter.next()
        except StopIteration:
            self.__trainloader_iter = iter(self.__trainloader)
            inputs, labels = self.__trainloader_iter.next()

        with dist_autograd.context() as context_id:
            fm = self.__model(inputs)
            self.__server_rref.rpc_sync().relay_forward(context_id, fm, labels)
            self.__dist_optim.step(context_id)

        if step_counter == (eval_step-1):
            flag = self.__server_rref.rpc_sync().eval(self.__model.state_dict())
            if flag == True:
                step_counter = 0
                self.__next_node.rpc_async().train(self.__model.state_dict(), step_counter, eval_step)
        else:
            step_counter += 1
            self.__next_node.rpc_async().train(self.__model.state_dict(), step_counter, eval_step)
        print(self.__p.cpu_times().user, self.__p.cpu_times().system)

    def start_init(self) -> None:
        self.__rrefs = [rpc.RRef(param) for param in self.__model.parameters()] + self.__server_rref.rpc_sync().relay_init()
        self.__dist_optim = DistributedOptimizer(
            torch.optim.SGD,
            self.__rrefs,
            lr=self.__lr
        )
        logging.info("client initialized ...")


from __future__ import print_function
import unittest

class PriorityQueue(object):
    class Node(object):
        def __init__(self, value, next_node):
            self.value = value
            self.next_node = next_node

    def __init__(self, initial=None):
        self.front = self.back = self.current = None
        self.__length = 0
        if initial != None:
            for element in tuple(initial):
                self.enqueue(element[0], element[1])

    def empty(self):
        return self.front == self.back == None

    def __len__(self):
        return self.__length

    def __iter__(self):
        self.current = self.front
        return self

    def __next__(self):
        if self.current:
            tmp = self.current.value
            self.current = self.current.next_node
            return tmp
        else:
            raise StopIteration()

    def __str__(self):
        if self.empty():
            return str(None)
        if self.front == self.back:
            return str(self.front.value)
        vals = ''
        for item in self:
            if item == self.back.value:
                vals += str(item)
                break
            vals += str(item) + ', '
        return vals

    def __repr__(self):
        return 'PriorityQueue((' + str(self) + '))'

    def __push_front(self, value):
        new = self.Node(value, None)
        if self.empty():
            self.front = self.back = new
            self.__length += 1
        else:
            new.next_node = self.front
            self.front = new
            self.__length += 1

    def __pop_front(self):
        tmp = self.front
        if self.empty():
            raise RuntimeError("PriorityQueue is empty")
        if self.front == self.back:
            value = self.front.value
            self.front = self.back = None
            self.__length -= 1
            return value
        else:
            tmp = self.front.next_node
            value = self.front.value
            self.front.next_node = None
            self.front = tmp
            self.__length -= 1
            return value

    def __push_back(self, value):
        new = self.Node(value, None)
        if self.empty():
            self.front = self.back = new
            self.__length += 1
            return
        self.back.next_node = new
        self.back = new
        self.__length += 1

    def __pop_back(self):
        if self.empty():
            raise RuntimeError("PriorityQueue is empty")
        if self.front is self.back:
            value = self.front.value
            self.front = self.back = None
            self.__length -= 1
            return value
        else:
            current = self.front
            while current.next_node is not self.back:
                current = current.next_node
            value = current.next_node.value
            current.next_node = None
            self.back = current
            self.__length -= 1
            return value

    def delete(self, value=None):
        if value == None:
            raise RuntimeError("PriorityQueue's delete method must be called with a value")
        if self.empty():
            raise RuntimeError("PriorityQueue is empty")
        if self.front.value["name"] == value:
            self.__pop_front()
            return
        if self.back.value["name"] == value:
            self.__pop_back()
            return
        else:
            prev = self.front
            current = self.front.next_node
            nxt = self.front.next_node.next_node
            if current.value["name"] == value:
                prev.next_node = nxt
                current.next_node = None
                current = None
                return
            while current is not self.back:
                if current.value["name"] == value:
                    prev.next_node = nxt
                    current.next_node = None
                    current = None
                    return
                else:
                    prev = current
                    current = nxt
                    nxt = nxt.next_node
            return
        raise RuntimeError("Delete method is not working correctly")

    def middle_from_node(self, start=None):
        if self.empty():
            raise RuntimeError("PriorityQueue is empty")
        if self.front == self.back:
            return start
        if start == None:
            return start
        fast = start.next_node.next_node
        slow = start
        while fast is not None:
            fast = fast.next_node
            if fast is not None:
                fast = fast.next_node
            slow = slow.next_node
        return slow

    def names(self):
        name_list = list()
        for item in self:
            name_list.append(item["name"])
        return name_list

    def enqueue(self, priority, name, time=0):
        if priority > 5:
            priority = 5
        elif priority < 0:
            priority = 0.0
        item = {"priority": priority,
                "name": name,
                "time": time}

        # Uncomment print statement for debugging
        #print(f"Queueing item: {item}")
        needsSort = False
        if self.empty() or (item["priority"] <= self.back.value["priority"]):
            self.__push_back(item)
        elif self.front.value["priority"] < item["priority"]:
            self.__push_front(item)
            nxt = self.front.next_node
            while nxt != None:
                nxt.value["priority"] = round((nxt.value["priority"] + 0.4), 1)
                if nxt.value["priority"] > 5:
                    nxt.value["priority"] = 5
                nxt = nxt.next_node
            needsSort = True
        else:
            inserted = False
            nxt = self.front
            while nxt != self.back.next_node:
                if (inserted == False) and (nxt.next_node.value["priority"] < item["priority"]):
                    temp = nxt.next_node
                    nxt.next_node = self.Node(item, temp)
                    inserted = True
                    nxt = nxt.next_node
                    self.__length += 1
                    continue
                if inserted == True and nxt.value != item:
                    nxt.value["priority"] = round((nxt.value["priority"] + 0.4), 1)
                    if nxt.value["priority"] > 5:
                         nxt.value["priority"] = int(5)
                nxt = nxt.next_node
            needsSort = True
        if needsSort == True:
            self.front = self.mergeSort(self.front)

    def mergeSort(self, noderino):
        if (noderino == None) or (noderino.next_node == None):
            return noderino
        mid = self.middle_from_node(noderino)
        mid_next = mid.next_node
        mid.next_node = None

        left_side = self.mergeSort(noderino)
        right_side = self.mergeSort(mid_next)
        return self.mergeBack(left_side, right_side)

    def mergeBack(self, left, right):
        sortedList = None
        if left == None:
            return right
        if right == None:
            return left
        if left.value["priority"] >= right.value["priority"]:
            sortedList = left
            sortedList.next_node = self.mergeBack(left.next_node, right)
        else:
            sortedList = right
            sortedList.next_node = self.mergeBack(left, right.next_node)
        return sortedList

    def dequeue(self):
        item = self.__pop_front()
        # Uncomment below for debugging
        # print(f"Dequeueing item: {item}")
        return item

''''''''''''''''''''
''''''''''''''''''''
'''''Unit Tests'''''
''''''''''''''''''''
''''''''''''''''''''
# Run unittests with:
##  OSX:      'python3 -m unittest PriorityQueue.py'
##  Windows:  'python  -m unittest PriorityQueue.py'

class BasicTests(unittest.TestCase):
    def test_empty(self):
        self.assertTrue(PriorityQueue().empty())

    def test_enqueue_simple(self):
        priority_queue = PriorityQueue()
        priority = 5
        name = "test"
        priority_queue.enqueue(priority, name)
        self.assertFalse(priority_queue.empty())

    def test_length(self):
        priority_queue = PriorityQueue()
        priority = 5
        top = 11
        for task_id in range(1, top):
            priority_queue.enqueue(priority, task_id)
            priority -= .5
        self.assertEqual(len(priority_queue), top-1)

    def test_enqueue_same_priority_simple(self):
        priority_queue = PriorityQueue()
        priority = 5
        top = 5
        for task_id in range(1, top):
            priority_queue.enqueue(priority, task_id)
        self.assertEqual(len(priority_queue), top-1)
        self.assertEqual(priority_queue.names(), [1, 2, 3, 4])

    def test_enqueue_increasing_priority(self):
        priority_queue = PriorityQueue()
        priority = 1
        top = 11
        for task_id in range(1, top):
            priority_queue.enqueue(priority, task_id)
            priority += .5
        self.assertEqual(len(priority_queue), top-1)
        self.assertEqual(priority_queue.names(), [9, 10, 8, 7, 6, 5, 4, 3, 2, 1])

    def test_init(self):
        students = [[4, "John"], [5,  "Bill"], [2, "Joe"], [3, "Janet"]]
        priority_queue = PriorityQueue(students)
        self.assertEqual(len(priority_queue), len(students))

    def test_deqeue_empty(self):
        self.assertRaises(RuntimeError, lambda: PriorityQueue().dequeue())

    def test_dequeue(self):
        students = [[4, "John"], [5,  "Bill"], [2, "Joe"], [3, "Janet"]]
        priority_queue = PriorityQueue(students)
        self.assertEqual(priority_queue.dequeue()["name"], "Bill")
        self.assertEqual(priority_queue.dequeue()["name"], "John")
        self.assertEqual(priority_queue.dequeue()["name"], "Janet")
        self.assertEqual(priority_queue.dequeue()["name"], "Joe")
        self.assertTrue(priority_queue.empty())

    def test_str(self):
        students = [[4, "John"], [5,  "Bill"], [2, "Joe"]]
        priority_queue = PriorityQueue(students)
        self.assertEqual(priority_queue.__str__(), "{'priority': 5, 'name': 'Bill', 'time': 0}, " \
                                                   "{'priority': 4.4, 'name': 'John', 'time': 0}, {'priority': 2, 'name': 'Joe', 'time': 0}")

    def test_repr(self):
        students = [[4, "John"], [5,  "Bill"], [2, "Joe"]]
        priority_queue = PriorityQueue(students)
        self.assertEqual(priority_queue.__repr__(), "PriorityQueue(({'priority': 5, 'name': 'Bill', 'time': 0}, " \
                                                    "{'priority': 4.4, 'name': 'John', 'time': 0}, {'priority': 2, 'name': 'Joe', 'time': 0}))")

class TestDelete(unittest.TestCase):
    def test_delete_empty(self):
        self.assertRaises(RuntimeError, lambda: PriorityQueue().delete(3))

    def test_delete_bad_call(self):
        self.assertRaises(RuntimeError, lambda: PriorityQueue().delete())

    def test_delete_first_item(self):
        students = [[4, "John"], [5,  "Bill"], [2, "Joe"]]
        priority_queue = PriorityQueue(students)
        priority_queue.delete("Bill")
        self.assertEqual(priority_queue.__str__(), "{'priority': 4.4, 'name': 'John', 'time': 0}, " \
                                                   "{'priority': 2, 'name': 'Joe', 'time': 0}")

    def test_delete_last_item(self):
        students = [[4, "John"], [5,  "Bill"], [2, "Joe"]]
        priority_queue = PriorityQueue(students)
        priority_queue.delete("Joe")
        self.assertEqual(priority_queue.__str__(), "{'priority': 5, 'name': 'Bill', 'time': 0}, " \
                                                   "{'priority': 4.4, 'name': 'John', 'time': 0}")

    def test_delete_middle_item(self):
        students = [[4, "John"], [5,  "Bill"], [2, "Joe"]]
        priority_queue = PriorityQueue(students)
        priority_queue.delete("John")
        self.assertEqual(priority_queue.__str__(), "{'priority': 5, 'name': 'Bill', 'time': 0}, " \
                                                   "{'priority': 2, 'name': 'Joe', 'time': 0}")

    def test_first_edge(self):
        students = [[4, "John"], [5,  "Bill"], [2, "Joe"], [3, "Janet"]]
        priority_queue = PriorityQueue(students)
        priority_queue.delete("John")
        self.assertEqual(priority_queue.__str__(), "{'priority': 5, 'name': 'Bill', 'time': 0}, "  \
                                                   "{'priority': 3, 'name': 'Janet', 'time': 0}, {'priority': 2.4, 'name': 'Joe', 'time': 0}")

    def test_last_edge(self):
        students = [[4, "John"], [5,  "Bill"], [2, "Joe"], [3, "Janet"]]
        priority_queue = PriorityQueue(students)
        priority_queue.delete("Janet")
        self.assertEqual(priority_queue.__str__(), "{'priority': 5, 'name': 'Bill', 'time': 0}, " \
                                                   "{'priority': 4.4, 'name': 'John', 'time': 0}, {'priority': 2.4, 'name': 'Joe', 'time': 0}")

    def test_delete_all(self):
        students = [[4, "John"], [5,  "Bill"], [2, "Joe"], [3, "Janet"]]
        priority_queue = PriorityQueue(students)
        for student in students:
            priority_queue.delete(student[1])
        self.assertTrue(priority_queue.empty())

class TestMiddleFromNode(unittest.TestCase):
    def test_empty(self):
        self.assertRaises(RuntimeError, lambda: PriorityQueue().middle_from_node())

    def test_find_middle_single(self):
        students = [[4, "John"]]
        priority_queue = PriorityQueue(students)
        self.assertEqual(priority_queue.middle_from_node(priority_queue.front).value["name"], "John")

    def test_find_middle_value_of_two(self):
        students = [[4, "John"], [5,  "Bill"]]
        priority_queue = PriorityQueue(students)
        self.assertEqual(priority_queue.middle_from_node(priority_queue.front).value["name"], "Bill")

    def test_find_middle_value_of_three(self):
        students = [[4, "John"], [5,  "Bill"], [2, "Joe"]]
        priority_queue = PriorityQueue(students)
        self.assertEqual(priority_queue.middle_from_node(priority_queue.front).value["name"], "John")

    def test_find_middle_value_of_four(self):
        students = [[4, "John"], [5,  "Bill"], [2, "Joe"], [3, "Janet"]]
        priority_queue = PriorityQueue(students)
        self.assertEqual(priority_queue.middle_from_node(priority_queue.front).value["name"], "John")

    def test_find_middle_value_of_five(self):
        students = [[4, "John"], [5,  "Bill"], [2, "Joe"], [3, "Janet"], [2, "Jimothy"]]
        priority_queue = PriorityQueue(students)
        self.assertEqual(priority_queue.middle_from_node(priority_queue.front).value["name"], "Janet")

'''''''''''''''''''''''''''
''''Main Runtime Script''''
'''''''''''''''''''''''''''

if '__main__' == __name__:
    def print_queue(priority_queue):
        num = 1
        print("List of Students")
        print("----------------")
        for item in priority_queue:
            print(f"{num} : {item}")
            num += 1
        print("\n\n")

    def file_read(file_name, mode=""):
        priority_queue = PriorityQueue()
        print_queue(priority_queue)
        with open(file_name) as inputFile:
            for line in inputFile:
                temp = line.split(" ")
                priority_queue.enqueue(int(temp[0]), temp[1].rstrip())
                print()
                print_queue(priority_queue)
        if inputFile.closed == False:
            raise RuntimeError("Input file did not close correctly")
        if mode == "dequeue":
            for i in range(0, 3):
                priority_queue.dequeue()
                print()
                print_queue(priority_queue)
        print("Final:")
        print_queue(priority_queue)

    input_type = input("Input from a file or from user input? ('file' or 'user')")

    while input_type.lower() not in ('user', 'file'):
        print(f"'{input_type}' is not a valid option")
        input_type = input("Input from a file or from user input? ('file' or 'user')  ")
    if input_type.lower() == 'file':
        print("\n-----------------------------------")
        print("Test Cases with Detailed Event Logs")
        print("-----------------------------------")
        print("\n\n---------------------------------------------\n")
        print("File: test_data1.dat\n")
        file_read("test_data1.dat")
        print("\n\n---------------------------------------------\n")
        print("File: test_data2.dat\n")
        file_read("test_data2.dat")
        print("\n\n---------------------------------------------\n")
        print("File: test_data3.dat\n")
        file_read("test_data3.dat")
        print("\n\n---------------------------------------------\n")
        print("File: test_data3.dat w/ dequeue\n")
        file_read("test_data3.dat", "dequeue")
    else:
        done = False
        priority_queue = PriorityQueue()
        print()
        while(done != True):
            name = input("Student name? ")
            priority = input("Task priority? ")
            priority_queue.enqueue(int(priority), name)
            print()
            print_queue(priority_queue)
            dequeue = input("Dequeue a student? (yes or no)")
            if dequeue.lower() in ('yes'):
                priority_queue.dequeue()
                print()
                print_queue(priority_queue)
            exit = input("Add another student? (yes or no)")
            if exit.lower() in ('no'):
                done = True
        print("Final:")
        print_queue(priority_queue)


"""Utils for Zinnia's tests"""
import cStringIO
from xmlrpclib import Transport

from django.test.client import Client

class TestTransport(Transport):
    """Handles connections to XML-RPC server
    through Django test client."""

    def __init__(self, *args, **kwargs):
        Transport.__init__(self, *args, **kwargs)
        self.client = Client()

    def request(self, host, handler, request_body, verbose=0):
        self.verbose = verbose
        response = self.client.post(handler,
                                    request_body,
                                    content_type="text/xml")
        res = cStringIO.StringIO(response.content)
        res.seek(0)
        return self.parse_response(res)


import json
import sys

from sklearn.model_selection import StratifiedKFold

from db import table_load, table_write
from utils import timer

if __name__ == "__main__":

    if len(sys.argv) == 2:
        config_file_name = sys.argv[1]
    else:
        config_file_name = "lightgbm_0"

    print("Config file Name: ", config_file_name)

    with timer("kfold"):
        config: dict = json.load(open("./configs/{}.json".format(config_file_name)))

        train = table_load("train")

        folds = StratifiedKFold(
            n_splits=config["cv"]["n_splits"],
            shuffle=True,
            random_state=config["cv"]["random_state"],
        ).split(train, train[config["features"]["target"]])

        for n_fold, (train_index, valid_index) in enumerate(folds):
            cv_train_df = train.loc[train_index]
            cv_test_df = train.loc[valid_index]
            table_write(table_name="cv_train_{}".format(n_fold), df=cv_train_df)
            table_write(table_name="cv_test_{}".format(n_fold), df=cv_test_df)


from datetime import datetime
from loguru import logger
from transformers import TrainingArguments, Trainer

from settings import LogDir, LogVar
from utils.log_helper import get_log_name

logger.add(
    get_log_name(LogDir.model, datetime.now()),
    level=LogVar.level,
    format=LogVar.format,
    enqueue=LogVar.enqueue,
    diagnose=LogVar.diagnose,
    catch=LogVar.catch,
    serialize=LogVar.serialize,
    backtrace=LogVar.backtrace,
    colorize=LogVar.color
)

def setup_trainer(model, compute_metrics, tokenizer,
        train_dataset, test_dataset, epoch=3, batch_size=32, weight_decay=0.01):
    logging_steps = len(train_dataset) // batch_size
    args = TrainingArguments(
        output_dir="models",
        evaluation_strategy='epoch',
        save_strategy='epoch',
        learning_rate=2e-5,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size,
        num_train_epochs=epoch,
        weight_decay=weight_decay,
        logging_steps=logging_steps
    )

    logger.info(f"{args}")

    trainer = Trainer(
        model,
        args,
        train_dataset=train_dataset,
        eval_dataset=test_dataset,
        compute_metrics=compute_metrics,
        tokenizer=tokenizer)

    return trainer, args


#!/usr/bin/env python
# coding: utf-8

# ### Fine-tuning 6-Billion GPT-J (& other models) in colab with LoRA and 8-bit compression
#
# This notebook is a simple example for fine-tuning [GPT-J-6B](https://huggingface.co/EleutherAI/gpt-j-6B) with limited memory. A detailed explanation of how it works can be found in [this model card](https://huggingface.co/hivemind/gpt-j-6B-8bit). It is heavily based on [this Colab](https://colab.research.google.com/drive/1ft6wQU0BhqG5PRlwgaZJv2VukKKjU4Es#scrollTo=vfdLQHOuEU7h). Huge thanks to Hivemind!
#
# You can also finetune [GPT-Neo-2.7B](https://huggingface.co/gustavecortal/gpt-neo-2.7B-8bit), [French GPT-J (Cedille's Boris)](https://huggingface.co/gustavecortal/fr-boris-8bit) and [T0-3B](https://huggingface.co/gustavecortal/T0_3B-8bit) with limited memory.
#
# Twitter: [@gustavecortal](https://twitter.com/gustavecortal)

# In[1]:

from bitsandbytes.optim import Adam8bit
from torch.utils.data import DataLoader
from datasets import load_dataset
from sklearn.model_selection import train_test_split
import transformers

import pandas as pd

import torch
import torch.nn.functional as F
from torch import nn
from torch.cuda.amp import custom_fwd, custom_bwd
from bitsandbytes.functional import quantize_blockwise, dequantize_blockwise

from tqdm.auto import tqdm

# ## Converting the model to 8 bits

# In[2]:

class FrozenBNBLinear(nn.Module):
    def __init__(self, weight, absmax, code, bias=None):
        assert isinstance(bias, nn.Parameter) or bias is None
        super().__init__()
        self.out_features, self.in_features = weight.shape
        self.register_buffer("weight", weight.requires_grad_(False))
        self.register_buffer("absmax", absmax.requires_grad_(False))
        self.register_buffer("code", code.requires_grad_(False))
        self.adapter = None
        self.bias = bias

    def forward(self, input):
        output = DequantizeAndLinear.apply(
            input, self.weight, self.absmax, self.code, self.bias)
        if self.adapter:
            output += self.adapter(input)
        return output

    @classmethod
    def from_linear(cls, linear: nn.Linear) -> "FrozenBNBLinear":
        weights_int8, state = quantize_blockise_lowmemory(linear.weight)
        return cls(weights_int8, *state, linear.bias)

    def __repr__(self):
        return f"{self.__class__.__name__}({self.in_features}, {self.out_features})"

class DequantizeAndLinear(torch.autograd.Function):
    @staticmethod
    @custom_fwd
    def forward(ctx, input: torch.Tensor, weights_quantized: torch.ByteTensor,
                absmax: torch.FloatTensor, code: torch.FloatTensor, bias: torch.FloatTensor):
        weights_deq = dequantize_blockwise(
            weights_quantized, absmax=absmax, code=code)
        ctx.save_for_backward(input, weights_quantized, absmax, code)
        ctx._has_bias = bias is not None
        return F.linear(input, weights_deq, bias).clone()

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_output: torch.Tensor):
        assert not ctx.needs_input_grad[1] and not ctx.needs_input_grad[2] and not ctx.needs_input_grad[3]
        input, weights_quantized, absmax, code = ctx.saved_tensors
        # grad_output: [*batch, out_features]
        weights_deq = dequantize_blockwise(
            weights_quantized, absmax=absmax, code=code)
        grad_input = grad_output @ weights_deq
        grad_bias = grad_output.flatten(
            0, -2).sum(dim=0) if ctx._has_bias else None
        return grad_input, None, None, None, grad_bias

class FrozenBNBEmbedding(nn.Module):
    def __init__(self, weight, absmax, code):
        super().__init__()
        self.num_embeddings, self.embedding_dim = weight.shape
        self.register_buffer("weight", weight.requires_grad_(False))
        self.register_buffer("absmax", absmax.requires_grad_(False))
        self.register_buffer("code", code.requires_grad_(False))
        self.adapter = None

    def forward(self, input, **kwargs):
        with torch.no_grad():
            # note: both quantuized weights and input indices are *not* differentiable
            weight_deq = dequantize_blockwise(
                self.weight, absmax=self.absmax, code=self.code)
            output = F.embedding(input, weight_deq, **kwargs)
        if self.adapter:
            output += self.adapter(input)
        return output

    @classmethod
    def from_embedding(cls, embedding: nn.Embedding) -> "FrozenBNBEmbedding":
        weights_int8, state = quantize_blockise_lowmemory(embedding.weight)
        return cls(weights_int8, *state)

    def __repr__(self):
        return f"{self.__class__.__name__}({self.num_embeddings}, {self.embedding_dim})"

def quantize_blockise_lowmemory(matrix: torch.Tensor, chunk_size: int = 2 ** 20):
    assert chunk_size % 4096 == 0
    code = None
    chunks = []
    absmaxes = []
    flat_tensor = matrix.view(-1)
    for i in range((matrix.numel() - 1) // chunk_size + 1):
        input_chunk = flat_tensor[i * chunk_size: (i + 1) * chunk_size].clone()
        quantized_chunk, (absmax_chunk, code) = quantize_blockwise(
            input_chunk, code=code)
        chunks.append(quantized_chunk)
        absmaxes.append(absmax_chunk)

    matrix_i8 = torch.cat(chunks).reshape_as(matrix)
    absmax = torch.cat(absmaxes)
    return matrix_i8, (absmax, code)

def convert_to_int8(model):
    """Convert linear and embedding modules to 8-bit with optional adapters"""
    for module in list(model.modules()):
        for name, child in module.named_children():
            if isinstance(child, nn.Linear):
                print(name, child)
                setattr(
                    module,
                    name,
                    FrozenBNBLinear(
                        weight=torch.zeros(
                            child.out_features, child.in_features, dtype=torch.uint8),
                        absmax=torch.zeros(
                            (child.weight.numel() - 1) // 4096 + 1),
                        code=torch.zeros(256),
                        bias=child.bias,
                    ),
                )
            elif isinstance(child, nn.Embedding):
                setattr(
                    module,
                    name,
                    FrozenBNBEmbedding(
                        weight=torch.zeros(
                            child.num_embeddings, child.embedding_dim, dtype=torch.uint8),
                        absmax=torch.zeros(
                            (child.weight.numel() - 1) // 4096 + 1),
                        code=torch.zeros(256),
                    )
                )

# You have to Monkey-Patch GPT-J before loading:

# In[3]:

class GPTJBlock(transformers.models.gptj.modeling_gptj.GPTJBlock):
    def __init__(self, config):
        super().__init__(config)

        convert_to_int8(self.attn)
        convert_to_int8(self.mlp)

class GPTJModel(transformers.models.gptj.modeling_gptj.GPTJModel):
    def __init__(self, config):
        super().__init__(config)
        convert_to_int8(self)

class GPTJForCausalLM(transformers.models.gptj.modeling_gptj.GPTJForCausalLM):
    def __init__(self, config):
        super().__init__(config)
        convert_to_int8(self)

transformers.models.gptj.modeling_gptj.GPTJBlock = GPTJBlock

# If you're using another 8-bit quantized model (e.g. T0-3B), remember to Monkey-Patch the model using convert_to_int8()

# In[4]:

class T5ForConditionalGeneration(transformers.models.t5.modeling_t5.T5ForConditionalGeneration):
    def __init__(self, config):
        super().__init__(config)
        convert_to_int8(self)

transformers.models.t5.modeling_t5.T5ForConditionalGeneration = T5ForConditionalGeneration

# In[5]:

config = transformers.GPTJConfig.from_pretrained("NbAiLab/nb-gpt-j-6B")
tokenizer = transformers.AutoTokenizer.from_pretrained("NbAiLab/nb-gpt-j-6B")

# In[6]:

config.pad_token_id = config.eos_token_id
tokenizer.pad_token = config.pad_token_id

# In[7]:

gpt = GPTJForCausalLM.from_pretrained(
    "NbAiLab/nb-gpt-j-6B", low_cpu_mem_usage=True)
# gpt = GPTJForCausalLM.from_pretrained("gustavecortal/fr-boris-8bit", low_cpu_mem_usage=True) French GPT-J Cedille's Boris

# ## LoRA fine-tuning example

# You can load my very small dataset composed of philosophical sentences:

# In[8]:

dataset = load_dataset('MasterThesisCBS/Rettsavgjoerelser')

# In[9]:

def tokenize_function(examples):
    return tokenizer(examples["Prompt w/suffix"], padding=True, truncation=True, max_length=2048)

tokenized_datasets = dataset.map(tokenize_function, batched=True)
# print(tokenized_datasets)
tokenized_datasets = tokenized_datasets.remove_columns(
    ["Prompt w/suffix", 'Prompt w/o suffix', '__index_level_0__', 'prompt length', 'sak', 'stikkord', 'summary', 'url'])
# print(tokenized_datasets)
tokenized_datasets.set_format("torch")

# In[10]:

full_train_dataset = tokenized_datasets["train"]
train_dataloader = DataLoader(full_train_dataset, shuffle=False, batch_size=8)

# Add adapters to Embedding/MLP/Attention/LMHead layers

# In[11]:

def add_adapters(model, adapter_dim=4, p=0.1):
    assert adapter_dim > 0

    for name, module in model.named_modules():
        if isinstance(module, FrozenBNBLinear):
            if "attn" in name or "mlp" in name or "head" in name:
                print("Adding adapter to", name)
                module.adapter = nn.Sequential(
                    nn.Linear(module.in_features, adapter_dim, bias=False),
                    nn.Dropout(p=p),
                    nn.Linear(adapter_dim, module.out_features, bias=False),
                )
                print("Initializing", name)
                nn.init.zeros_(module.adapter[2].weight)

            else:
                print("Not adding adapter to", name)
        elif isinstance(module, FrozenBNBEmbedding):
            print("Adding adapter to", name)
            module.adapter = nn.Sequential(
                nn.Embedding(module.num_embeddings, adapter_dim),
                nn.Dropout(p=p),
                nn.Linear(adapter_dim, module.embedding_dim, bias=False),
            )
            print("Initializing", name)
            nn.init.zeros_(module.adapter[2].weight)

add_adapters(gpt)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
gpt.to(device)

# In[12]:

gpt.gradient_checkpointing_enable()
optimizer = Adam8bit(gpt.parameters(), lr=1e-5, weight_decay=0.01)

# In[13]:

num_epochs = 5
num_training_steps = num_epochs * len(train_dataloader)

# In[14]:

lr_scheduler = transformers.get_linear_schedule_with_warmup(
    optimizer, int(num_training_steps*0.1), num_training_steps
)

# In[15]:

filepath = '/content/model.pt'

# In[16]:

# get_ipython().system(
#    'export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/cuda/lib64')
scaler = torch.cuda.amp.GradScaler()
progress_bar = tqdm(range(num_training_steps))
gpt.train()
gpt.gradient_checkpointing_enable()
k = 0
for epoch in range(num_epochs):
    for batch in train_dataloader:
        # print(batch)
        k = k + 1
        if k % 500 == 0:
            print(k)
            state = {'k': k, 'epoch': num_epochs, 'lr_scheduler': lr_scheduler.state_dict(
            ), 'state_dict': gpt.state_dict(), 'optimizer': optimizer.state_dict()}
            torch.save(state, filepath)

        batch = {k: v.to(device) for k, v in batch.items()}

        optimizer.zero_grad()

        with torch.cuda.amp.autocast():
            out = gpt.forward(**batch,)

            loss = F.cross_entropy(out.logits[:, :-1, :].flatten(0, -2), batch['input_ids'][:, 1:].flatten(),
                                   reduction='mean', label_smoothing=0.1)

        print(loss)

        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(gpt.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()

        lr_scheduler.step()
        progress_bar.update(1)

# ## Text generation example

# In[ ]:

gpt.eval()
with torch.no_grad():
    prompt = tokenizer("Quote:", truncation=True, padding=True,
                       max_length=128, return_tensors='pt')
    prompt = {key: value.to(device) for key, value in prompt.items()}
    out = gpt.generate(**prompt, max_length=128, top_k=50, top_p=0.9,
                       temperature=1.0, do_sample=True, repetition_penalty=1.2, num_beams=1)
    print(tokenizer.decode(out[0]))


import numpy as np
from complex_cell import *
from gabor_kernel import *
import shutil
from numpy.linalg import norm
from multiprocessing import Pool
from functools import partial

CPU_NUMBER = os.cpu_count()

class MT_cell(object):
    """
        Creates a bank of MT neurons, which respond to stimuli moving at a certain velocity in a certain direction.
        The response is obtained from combining the responses of complex cells (in our case, we compute the response
        directly from simple cells, as complex cell only combine the responses of simple cells with 4 different
        phases spaced in steps of pi/2).

        The attributes of the class are detailled in the __init__ method.
    """

    def __init__(self, im_shape, v_sensitivity, wavelengths, phases, c, alpha,  comp_mt_params,
                 deactivation_rate = 0.0, filt_width = 0.56,  causal_time = False,
                 mt_directory = 'mt_cell', mt_file = 'response', comp_directory = 'complex_cell',
                 comp_file = 'response', verbose = False):
        """
            Detail of the inputs to the class:
            - im_shape: shape of the stimulus to which the MT_cell will react.
            - v_senstivity: an array of the shape 3 x N_MT_neurons, the first and second line give the
                            direction of the stimulus to which the neuron is sensitive (not normalised)
                            and the third line gives the magnitude of the velocity to which the neuron will react.
            - wavelengths: an array containing the different wavelengths to which the cells are sensitive.
            - phases: same principle than the wavelengths.
            - c: useful when time is converted to space and vice versa in the rotations, to account for the
                 different resolutions.
            - alpha: spread of the filter in time.
            - comp_mt_params: a set of 6 parameters, accounting for the spontaneous activations, semi-
                              saturation constants and maximal responses of both complex and MT cells.
            - deactivation rate: percentage of the repsonses that will be deactivated after computation
                                 in order to simulate ageing of the brain.
            - filt_width: ratio of the width of the filter with respect to the wavelength
                          (cf. the computation of sigma further down in the compute_responses method)
            - directory and file: where the results will be stored
            - verbose: to get some extra information at execution.

            Detail of the attributes derivated from those:
            - mt_v: array of velocity sensitivities formatted into the correct way: i.e. 2D vector with
                    the correctly expected magnitude. (e.g. Move from [ [1] [1] [2] ] to [sqrt 2, sqrt 2])
            - N_cells: number of mt cells that are simualted
            - mt_map: list which contains the angles theta and phi of the four complex cells that have to be
                      combined into one mt cell response.
            - thetas and velocities: parameters that will be passed to the complex cell declaration.
            - unique_pos: For each MT cell, tells whether the spatiotemporal parameters of the complex
                          cells appear for the first time or where was the first apparition. If it appears
                          for the first time, unique_pos will contain a None at the entry and if it already
                          appeared previously in mt_map, it will contain the location of the first apparition
                          of this spatiotemporal orientation. This is especially useful for inhibition and
                          deactivation of cells.
            - deactivation_map: for each MT cell, mask that will be multiplied with the response to deactivate
                                a certain percentage of the reponses.
        """

        ## 1. INITIALISE DIVERSE ATTRIBUTES OF THE CLASS
        self.shape = np.array(im_shape)

        self.wavelengths = wavelengths
        self.phases = phases
        self.c = c
        self.alpha = alpha
        self.a1, self.s1, self.K1, self.a2, self.s2, self.K2, self.inhib_resp, self.inhib_weight = comp_mt_params
        self.directory = mt_directory
        self.file = mt_file
        self.stimulus = np.zeros(im_shape)
        self.verbose = verbose
        self.filt_width = filt_width
        self.deactivation_rate = deactivation_rate

        ## 2. COMPUTE THE SET OF ANGLES AND VELOCITIES NEEDED

        directions = v_sensitivity[:2, :].T
        if v_sensitivity.shape[0] > 2:
            magnitudes = v_sensitivity[2, :]
        else:
            magnitudes = np.array([np.sqrt(np.dot(v, v)) for v in directions])

        self. v_max = magnitudes.max()
        # Computing the 2D vector of velocity from the directions and magnitudes
        self.mt_v = np.array([magn*v/np.sqrt(np.dot(v, v)) if np.dot(v, v) != 0 else [0, 0]
                 for v, magn in zip(directions, magnitudes)])

        # Declaring a few more attributes of the class
        self.N_cells = len(magnitudes)
        self.mt_map = np.zeros((self.N_cells, 4, 2))
        self.thetas = np.zeros((self.N_cells * 4))
        psis = np.zeros((self.N_cells* 4))

        for i, v in enumerate(self.mt_v):
            if (v == 0).all():
                angles_v = np.array([[0,0],[np.pi/4,0],[np.pi/2,0],[3*np.pi/4,0]])
            else:
                angles_v =  np.array(self.get_spacetime_angles(v))
            self.mt_map[i] = angles_v
            self.thetas[4*i:4*(i + 1)] = angles_v[:, 0]
            psis[4*i:4*(i + 1)] = angles_v[:, 1]

        self.velocities = self.c * np.tan(psis)

        ## 3. COMPUTE THE CROSS PRODUCT BETWEEN THE PLANE ORIENTATION AND THE DIRECTION OF THE DIFFERENT VECTORS.
        self.compute_cross_product()

        ## 4. COMPUTE THE MAP WHICH SAYS IF THE COMPLEX CELL COMPOSING AN MT CELL NEEDS TO BE RECOMPUTED
        self.unique_pos = list()

        for i, mt_cell in enumerate(self.mt_map):
            cell_pos_list = list()
            for j,orientation in enumerate(mt_cell):
                cell_pos, comp_pos = np.where(np.all(np.abs(self.mt_map-orientation)<1e-15,axis=2))
                cell_pos, comp_pos = cell_pos[0], comp_pos[0]
                if cell_pos == i and comp_pos == j:
                    cell_pos_list.append([cell_pos, comp_pos])
                else:
                    cell_pos_list.append(None)
            self.unique_pos.append(cell_pos_list)
        self.unique_pos = np.array(self.unique_pos)

        ## 5. GENERATE THE COMPLEX CELLS AND THE AGEING MAP
        self.comp_directory = comp_directory
        self.comp_file = comp_file
        self.causal_time = causal_time
        self.wavelengths = wavelengths

        if (self.shape != [1,1,1]).all():
            self.generate_comp_ageing()
        else:
            print('At least one of the components of the shape of the image is invalid. Please change the attribute shape and run the method generate_comp_ageing to terminate the initialisation of the class.')

    def generate_comp_ageing(self):
        """
            Method which generates the complex cells and the deactivation linked to ageing.
        """
        self.comp_cell = ComplexCell(self.shape, self.thetas, self.velocities, self.wavelengths,
                                     self.phases, self.c, self.alpha, self.a1, self.s1, self.K1, self.deactivation_rate,
                                     self.filt_width, directory_cell = self.comp_directory, file_cell = self.comp_file,
                                     filters_theta_v = True, causal_time = self.causal_time, verbose = self.verbose)

        if self.deactivation_rate > 0:
            self.deactivation_map = np.random.choice([0, 1], size=(int(self.N_cells),self.shape[1], self.shape[2]),
                                                  p=[ self.deactivation_rate, 1-self.deactivation_rate])

    def clear_directory(self):
        """
            Removes the folder in which the results are stored
        """
        self.comp_cell.clear_directory()

        if os.path.isdir(self.directory):
            shutil.rmtree(self.directory)

###################################################################################################
### 3 methods that compute 4 couples of angles from a velocity v = (v_x, v_y), as described in the
### appendix of the paper of Simoncelli and Heeger, 1996.
###     1. get_angles:           maps a space-time orientation vector in Fourier space to the corresponding
###                              angles (theta, psi) which generate a Gabor filter
###     2. get_orientations:     from a given velocity v, computes the 4 vectors u_1, u_2, u_3, u_4
###     3. get_spacetime_angles: calls both previous methods to return the 4 couples of angles (theta, psi)
###################################################################################################

    def get_angles(self, vect):
        """
            Transform a spatio-temporal unit vector for preferred orientation of
            the filter into the two rotation angles to which it corresponds.
        """
        x_1, x_2, x_3 = vect

        return  -np.arctan2(x_2,x_1),  np.arcsin(-x_3)

    def get_orientations(self, v):
        """
            Computes the 4 vectors required to get a ring in Fourier domain which
            will detect a certain velocity
        """
        v_x, v_y = v
        if v_x == 0 and v_y == 0:
            u_1 = np.array([1, 0, 0])
            u_2 = np.array([1/np.sqrt(2), 1/np.sqrt(2), 0])
            u_3 = np.array([0, 1, 0])
            u_4 = np.array([-1/np.sqrt(2), 1/np.sqrt(2), 0])

        else:
            s = np.sqrt(v_x**2 + v_y ** 2)

            u_1 = np.array([-v_x, -v_y, s**2])/np.sqrt(s**4 + s**2)
            u_2 = np.array([-v_y, v_x, 0])/s
            u_3 = (u_1 + u_2) / np.sqrt(2)
            u_4 = (u_1 - u_2) / np.sqrt(2)

        return u_1, u_2, u_3, u_4

    def get_spacetime_angles(self, v):
        """
            Extracts the 4 spatio-temporal unit vectors that we need in order
            to obtain a MT cell sensitive to the velocity v (2D vector with
            the preferred x and y velocities of the cell).
        """
        u_1, u_2, u_3, u_4 = self.get_orientations(v)

        return self.get_angles(u_1), self.get_angles(u_2), self.get_angles(u_3), self.get_angles(u_4)

###################################################################################################
### 4 methods to compute the MT cells responses to a given stimulus.
###
###     1. compute_responses_raw:    Computes the raw MT cell response for a subset of the total
###                                  MT cells that are simulated, and stores the response in the
###                                  self.directory directory, with file name raw_file
###     2. computes_responses_inhib: Based on the raw MT cell responses, computes the inhibition
###                                  from the V1 neurons which are not close to the preferred
###                                  velocity plane and stores it in the same directory, with
###                                  file name self.file.
###     3. compute_responses:        Combines the two previous methods, and splits the whole MT cell
###                                  set into different subsets, which are delegated to different
###                                  processes.
###     4. get_response:             Loads the response which was computed by a call at compute_responses
###                                  normalises it and returns it.
###################################################################################################

    def compute_responses_raw(self, cell_split, raw_file = 'mt_raw'):
        """
            Computes the raw response of the MT cells in the subset cell_split, and stores it
            with a file name starting with mt_raw. The computation is simply the sum of the responses
            of the complex cells with 4 different phases, all possible wavelengths along the 4 preferred
            directions which are given as 4 couples of (theta, psi) angles.

            The verbose mode will give details about the magnitude, mean, max, min of responses.
        """
        for num_cell in cell_split:

            if self.verbose:
                print(os.getpid(),'- Computing the raw response of the MT Cell n°',num_cell,
                      'sensitive to v = [', round(self.mt_v[num_cell, 0],3), ',',round(self.mt_v[num_cell,1],3),
                      '] made of the following responses.\n','->'.ljust(9),
                      'Theta'.ljust(5),'Psi'.ljust(5),
                      '|', 'sum'.ljust(15),'max'.ljust(5),'min'.ljust(5),'mean > 0'.ljust(10))

            response = np.zeros(self.shape)

            for i in range(4):
                theta, psi = self.mt_map[num_cell, i]
                v = self.c * np.tan(psi)
                response += self.comp_cell.get_response(theta, v)

                if self.verbose:
                    print((str(os.getpid())+' -').ljust(10), str(int(theta*180/np.pi)).ljust(5),
                        str(int(psi*180/np.pi)).ljust(5), '|',
                        str(round(np.sum(response),3)).ljust(15),
                        str(round(response.max(), 3)).ljust(5),
                        str(round(response.min(), 3)).ljust(5),
                        str(round(response[response>0].mean(),3)).ljust(10))

            response_path = self.get_path(num_cell, raw_file)
            if self.deactivation_rate > 0:
                np.save(response_path, self.deactivation_map[num_cell]*response)
            else:
                np.save(response_path, response)

    def compute_responses_inhib(self, cell_split):
        """
            Compute the third and last part of the response output from the MT cells to the stimulus given as
            input of the compute_responses_raw model. Has to be ran after compute_responses_raw

            Compute the model output, especially adding inhibition from the V1 cells
            which are not close to a particular MT cell preferred velocity plane.
            (for each mt cell, load and subtract the responses of the cells which are not near the
            preferred velocity plane)
        """
        # 3. MT CELL WITH INHIBITION FROM OTHER CELLS
        # After getting the raw MT cell responses, we compute the inhibition
        # from the v1 cells which are not close to the velocity plane
        #print(os.getpid(), 'second_part', cell_split, cell_split is None)

        normalisation_response = np.zeros(self.shape)

        for i in cell_split:
           # print(os.getpid(), i-cell_split[0],'/',cell_split[-1]-cell_split[0])

            response_path_load = self.get_path(i, self.raw_file)
            response = np.load(response_path_load)

            for j in range(self.N_cells):
                if j != i:
                    for k in range(4):
                        # We express the fact of not being close as the magnitude of the
                        # cross product between the normal vector to the velocity plane
                        # and the directional vector of the filter.
                        if self.cross_table[i,j,k] < self.inhib_resp and self.unique_pos[j,k] is not None:
                            #print(i,j,k)
                            # Retrieve and subtract the responses from filters
                            # with the same phase and wavelength
                            theta, psi = self.mt_map[j,k]
                            v = self.c * np.tan(psi)
                            response -= self.inhib_weight*self.comp_cell.get_response(theta, v)

            response += self.a2 # adding the spontaneous firing rate constant
            response = (response > 0) * response * response # half-squaring rectification
            if self.verbose:
                print(os.getpid(),'- Computing the model response of the MT Cell n°',i,
                      'sensitive to v = [', round(self.mt_v[i, 0],3), ',',round(self.mt_v[i,1],3),
                      '] rectified but not normalised.\n',os.getpid(),'- sum =', str(response.sum()).ljust(15),
                      'max =', str(round(response.max(), 3)).ljust(5),
                      'min =', str(round(response.min(), 3)).ljust(5))

            response_path_save = self.get_path(i)
            np.save(response_path_save, response)

            normalisation_response += response

        return normalisation_response

    def compute_responses(self, stimulus, raw_file = 'mt_raw'):
        """
            Computes the complex cell respones if the stimulus is different that the one stored
            (meaning that the previous iteration already worked on the same stimulus, and thus
            already computed the complex cell and raw responses). Then, splits the list of MT cells
            and delegates the computation to different processes, which will call
            compute_responses_raw and compute_responses_inhib.
        """

        if not os.path.isdir(self.directory):
            os.mkdir(self.directory)

        self.raw_file = raw_file
        cell_split = np.array_split(np.arange(self.N_cells), CPU_NUMBER)

        diff_stimulus = False

        # If the stimulus is the same as the one stored, then the stored responses
        # already were computed. It is hence not needed to recompute them.
        if (np.array(stimulus.shape) != self.shape).any():
            raise ValueError('Wrong shape of image for this instance of MT cell')

        if (stimulus != self.stimulus).any():
            diff_stimulus = True

            if self.verbose:
                print('Computing the complex cell responses')
            # 1. COMPLEX CELLS RESPONSES
            # This first line takes a long time to run, as we need to get all the possible responses
            # of the complex cell to the stimulus
            self.comp_cell.compute_responses(stimulus)
            self.stimulus = stimulus

        pool = Pool(CPU_NUMBER)

        # 2. MT CELLS RAW RESPONSES

        # Computing and saving the responses to all the filters, as well
        # as the normalisation (sum of all the responses), needed for further
        # steps of the model.
        if diff_stimulus:
            pool.map(partial(self.compute_responses_raw, raw_file = raw_file), cell_split)
        # computing normalisation and adding the semi-saturation constant
        normalisation_response = np.sum(pool.map(self.compute_responses_inhib, cell_split),axis=0) + self.s2
        pool.close()
        pool.join()

        norm_path = self.directory + '/normalisation.npy'
        np.save(norm_path, normalisation_response)

    def get_response(self, *args):
        """
            Retrieves the locally stored response. Has to be run after the
            `compute_responses` method, or the result will not be found.
        """
        #print(os.getpid(), 'get_response', args)
        path = self.get_path(*args)

        normalisation = np.load(self.directory + '/normalisation.npy')
        response = np.load(path)

        return self.K2*response/normalisation

###################################################################################################
### 2 methods to return the optic flow from the MT Cell responses
###
###     1. get_optic_flow:          Displays an arrow representing the preferred velocity of
###                                 the MT cell which fires the most, and returns the magnitude
###                                 of this firing.
###     2. get_optic_flow_combined: Combines the reponses of all the MT Cell at each location,
###                                 summing the preferred velocities according to the firing rates
###                                 of the cell at a given spatio-temporal location.
###
###################################################################################################

    def get_optic_flow(self):
        """
            Returns the optic flow arrows that can be drawn using a quiver plot.
            The returned objects are:
            1. response: the maximal response at each pixel in space and time
            2. v_x : the x-component of the velocity of the MT cell which responds more strongly to the stimulus
                        at each location.
            3. v_y : the y-component of the velocity of the MT cell which responds more strongly to the stimulus
                        at each location.
        """
        response = np.zeros(self.shape)
        v_x = np.zeros(self.shape)
        v_y = np.zeros(self.shape)

        for num_cell in range(self.N_cells):
            resp = self.get_response(num_cell)
            for t in range(self.shape[0]):
                for x in range(self.shape[1]):
                    for y in range(self.shape[2]):
                        if resp[t, x, y] > response[t, x, y]:
                            response[t, x, y] = resp[t, x, y]
                            v_x[t, x, y] = self.mt_v[num_cell,0]
                            v_y[t, x, y] = self.mt_v[num_cell,1]
        return response, v_x, v_y

    def get_optic_flow_combined(self):
        """
            Returns the optic flow arrows that can be drawn using a quiver plot.
            The returned objects are:
            1. response: the maximal response at each pixel in space and time
            2. v_x : the x-component of the v
        """
        response = np.zeros(self.shape)
        response_normalisation = np.zeros(self.shape)
        v_x = np.zeros(self.shape)
        v_y = np.zeros(self.shape)

        for num_cell in range(self.N_cells):
            resp = self.get_response(num_cell)
            v_x += resp*self.mt_v[num_cell,0]#/ (1.5 * self.v_max)
            v_y += resp*self.mt_v[num_cell,1]#/ (1.5 * self.v_max)
            response_normalisation += resp
        # Divide by the normalisation of the activity in the case where everything isn't zero
        # as it would otherwise return an error.
        v_x[response_normalisation != 0] /= response_normalisation[response_normalisation != 0]
        v_y[response_normalisation != 0] /= response_normalisation[response_normalisation != 0]
        return v_x, v_y

    def get_optic_flow_direction(self):
        """
            Returns the direction in which the mt_cells respond more strongly. It computes
            a weighted average of the direction of the MT cells responses.
            It is used for the MSTd cell heading direction detection.
        """
        #response = np.zeros(self.shape)
        response_normalisation = np.zeros(self.shape)
        v_x = np.zeros(self.shape)
        v_y = np.zeros(self.shape)

        for num_cell in range(self.N_cells):
            resp = self.get_response(num_cell)
            if (self.mt_v[num_cell,0] != 0 or self.mt_v[num_cell,1] != 0):
                v_x += resp*self.mt_v[num_cell,0]/norm(self.mt_v[num_cell])#/ (1.5 * self.v_max)
                v_y += resp*self.mt_v[num_cell,1]/norm(self.mt_v[num_cell])#/ (1.5 * self.v_max)
                response_normalisation += resp
        # Divide by the normalisation of the activity in the case where everything isn't zero
        # as it would otherwise return an error.
        v_x[response_normalisation != 0] /= response_normalisation[response_normalisation != 0]
        v_y[response_normalisation != 0] /= response_normalisation[response_normalisation != 0]
        return v_x, v_y#, response_normalisation

###################################################################################################
### 4 methods to compute the MT cells responses to a given stimulus.
###
###     1. int_to_str:            Converts an integer to its correct string version (1.0 -> '1')
###     2. get_path:              Constructs the path to the file from the parameters given as input
###     3. cross_plane_vect:      Computes the cross product between the normal vector of the velocity plane
###                               and a V1 neuron preferred orientation vector
###     4. compute_cross_product: Computes all possible combinations of plane normal vector and
###                               V1 neuron preferred orientation cross products. Delegates the computation
###                               to the compute_cross_parallel method.
###     5. compute_cross_product_parallel: Method which computes what is needed in compute_cross_product.
###################################################################################################

    def int_to_str(self, val):
        """
            Checks whether the given number is an integer, and if it is, cuts the decimal part.
            It is useful to make sure that having inputs like 2 and 2.0 will give filters which
            have the same file name, which would not be the case otherwise.
        """
        if isinstance(val, (np.int_, int)):
            return str(val)
        else:
            return str(int(val)) if val.is_integer() else str(val)

    def get_path(self, *args):
        """
            Returns the path to which a particular response is/will be stored from the
            velocity or position in the MT cell velocity array. The types of inputs are
            as follows:
                1. Position in the mt_v velocity array
                2. Velocity given as (v_x, v_y)
                3. Velocity given as (v_x, v_y, magnitude)
                4. Position in the mt_v velocity array and file name
                5. Velocity given as (v_x, v_y) and file name
                6. Velocity given as (v_x, v_y, magnitude) and file name
            There are 6 cases because we store the both the raw responses and the inhibited ones
            under different file names
        """

        file_str = self.file
        error_msg = ('Wrong type of input for the method get_path. Please make sure you either '+
                    'pass an integer or 2/3 parameters uniquely identifying the desired filter.')
        if len(args) == 1 and isinstance(args[0], (int, np.int_)):
            v = self.mt_v[args[0]]

        elif len(args) == 2:
            if all(isinstance(arg, (float, int, np.int_, np.float_)) for arg in args):
                v = np.array(args)

            elif isinstance(args[0], (int, np.int_)) and isinstance(args[1], str):
                v = self.mt_v[args[0]]
                file_str = args[1]

            else:
                raise TypeError(error_msg)

        elif len(args) == 3:
            if all(isinstance(arg, (float,int,np.int_,np.float_)) for arg in args):
                v = np.array(args[0:2])
                v = v/np.sqrt(np.dot(v,v)) if (v != 0).any() else 0
                v = args[2]*v

            elif all(isinstance(arg, (float, int, np.int_, np.float_)) for arg in args[:2]) and isinstance(args[2], str):
                v = np.array(args[:2])
                file_str = args[2]

            else:
                raise TypeError(error_msg)

        elif len(args) == 4 and all(isinstance(arg, (float,int,np.int_,np.float_)) for arg in args[:3]) and (
            isinstance(args[3], str)):
                v = np.array(args[0:2])
                v = v/np.sqrt(np.dot(v,v)) if (v != 0).any() else 0
                v = args[2]*v
                file_str = args[3]

        else:
            raise TypeError(error_msg)

        v = np.round(v,3)

        vx = self.int_to_str(v[0])
        vy = self.int_to_str(v[1])

        response_path = self.directory + '/' + file_str + '_' + vx + '_' + vy + '.npy'

        return response_path

    def cross_plane_vect(self, v_plane, vect):
        """
            Computes the cross product between the normal preferred velocity plane
            of a MT cell (computed as the cross product between the vectors u_1 and
            u_2) and the preferred velocity vector of a V1 neuron.

            v_plane is the velocity we consider (e.g. v = (2, 0) )
            vect is a velocity vector for a V1 neuron.
        """
        if norm(vect) == 0:
            raise ValueError('Vector with norm zero passed in the method.')

        v_normal = -np.array(np.r_[v_plane, 1])
        v_normal = v_normal/np.linalg.norm(v_normal)

        return np.linalg.norm(np.cross(v_normal, vect/ np.linalg.norm(vect)))

    def compute_cross_product(self):
        """
            Creates a pool of workers which computes and stores all possible
            combinations of cross products between the normals to the velocity
            planes in Fourier space and the V1 neuron orientations.
        """

        pool = Pool(CPU_NUMBER)

        cell_split = np.array_split(np.arange(self.N_cells), CPU_NUMBER)

        # Retrieving the pooling results as a list, the last step is to combine
        # it into the table we want.
        cross_list = pool.map(self.compute_cross_product_parallel, cell_split)

        pool.close()
        pool.join()

        self.cross_table = np.zeros((self.N_cells, self.N_cells, 4))

        for n_cells, cross in zip(cell_split, cross_list):
            self.cross_table[n_cells,:,:] = np.array(cross)

    def compute_cross_product_parallel(self, cell_split):
        """
            Computes and stores all possible combinations of cross products
            between the normals to the velocity planes in Fourier space and
            the V1 neuron orientations for a small part of the velocity planes.
        """

        cross_table = np.zeros((len(cell_split), self.N_cells, 4))

        for i in cell_split:
            v_plane = self.mt_v[i]
            for j in range(self.N_cells):
                u1j, u2j, u3j, u4j = self.get_orientations(self.mt_v[j])
                cross_table[i-cell_split[0], j, :] = [self.cross_plane_vect(v_plane, u1j),
                                                      self.cross_plane_vect(v_plane, u2j),
                                                      self.cross_plane_vect(v_plane, u3j),
                                                      self.cross_plane_vect(v_plane, u4j)]

        return cross_table


import random

def create_deck():
    suits = ['Hearts', 'Diamonds', 'Clubs', 'Spades']
    ranks = ['2', '3', '4', '5', '6', '7', '8', '9', '10', 'Jack', 'Queen', 'King', 'Ace']
    deck = [{'rank': rank, 'suit': suit} for suit in suits for rank in ranks]
    random.shuffle(deck)
    return deck

def play_war():
    print("Welcome to the card game 'War' against the computer!")
    deck = create_deck()

    player_deck = deck[:26]
    computer_deck = deck[26:]

    while True:
        input("Press Enter to draw a card...")

        player_card = player_deck.pop(0)
        computer_card = computer_deck.pop(0)

        print(f"Your card: {player_card['rank']} of {player_card['suit']}")
        print(f"Computer's card: {computer_card['rank']} of {computer_card['suit']}")

        player_rank_index = ranks.index(player_card['rank'])
        computer_rank_index = ranks.index(computer_card['rank'])

        if player_rank_index > computer_rank_index:
            print("You win this round!")
            player_deck.extend([player_card, computer_card])
        elif player_rank_index < computer_rank_index:
            print("Computer wins this round!")
            computer_deck.extend([player_card, computer_card])
        else:
            print("It's a tie! War begins...")

            if len(player_deck) < 4 or len(computer_deck) < 4:
                print("Insufficient cards for a war. The game is a draw.")
                break

            player_war_cards = player_deck[:4]
            computer_war_cards = computer_deck[:4]

            player_card = player_deck.pop(0)
            computer_card = computer_deck.pop(0)

            player_rank_index = ranks.index(player_card['rank'])
            computer_rank_index = ranks.index(computer_card['rank'])

            if player_rank_index > computer_rank_index:
                print("You win the war!")
                player_deck.extend(player_war_cards + [player_card, computer_war_cards, computer_card])
            elif player_rank_index < computer_rank_index:
                print("Computer wins the war!")
                computer_deck.extend(player_war_cards + [player_card, computer_war_cards, computer_card])
            else:
                print("It's another tie!")

        print(f"You have {len(player_deck)} cards. Computer has {len(computer_deck)} cards.")

        if not player_deck:
            print("You ran out of cards. Computer wins!")
            break
        elif not computer_deck:
            print("Computer ran out of cards. You win!")
            break

if __name__ == "__main__":
    ranks = ['2', '3', '4', '5', '6', '7', '8', '9', '10', 'Jack', 'Queen', 'King', 'Ace']
    play_war()


import socket
import Constants

localIP     = "127.0.0.1"

# Create a datagram socket: (ipv4, UDP)
UDPServerSocket = socket.socket(family=socket.AF_INET, type=socket.SOCK_DGRAM)

# Bind to address and ip
server_port = int(input(f"Please enter server port: "))
UDPServerSocket.bind((localIP, server_port))

print("UDP server up and listening")

# Listen for incoming datagrams
while(True):
    bytesAddressPair = UDPServerSocket.recvfrom(Constants.BUFFER_SIZE)

    message = bytesAddressPair[0]

    address = bytesAddressPair[1]

    print(f"Message from Client: {message.decode('ascii')}")
    print(f"Client IP Address: {address}")

    # Sending a reply to client
    UDPServerSocket.sendto(message, address)


from datetime import datetime, timedelta
from random import randint

import pyotp
from fastapi import Body, Depends, HTTPException
from starlette import status
from starlette.background import BackgroundTask
from starlette.requests import Request

from authentication import models
from authentication.crud import save_verification, get_verification_rec_by_user_id
from authentication.dependencies.auth import get_current_active_user
from authentication.dependencies.token import create_and_save_access_token, create_temp_token
from authentication.models import VerificationTypes
from common.constants import TwoFaCodesStorageTypes
from common.resources import response_strings
from common.schemas import ResponseSchema
from db.pgsql.engine import AsyncDBSession
from db.redis.handler import redis_save, redis_get
from db.redis.namespaces import auth_mail_codes_namespace
from logger.loggers import APP_LOGGER
from logger.tasks import write_log
from mail.handler import send_mail
from mail.shared import MailMessage
from authentication.background_tasks import new_login_happened
from settings import settings

async def create_access_token(request: Request, user: models.User, bg_tasks: BackgroundTask):
    auth_data = await create_and_save_access_token(user)
    bg_tasks.add_task(new_login_happened, request, user)

    return auth_data

async def send_2fa_code(user: models.User):
    """Returns user 2fa enabled types and send their codes"""
    # TODO (PHASE 2) check mobile and send its code
    temp_token = await create_temp_token(user)
    if user.two_fa_enabled:
        return ResponseSchema(data={'msg': response_strings.TWO_FA_ACTIVATED,
                                    'tmp_token': temp_token,
                                    'types': '2fa'})
    else:
        await send_verification_mail(user)
        return ResponseSchema(data={'msg': response_strings.TWO_FA_CODE_SENT,
                                    'tmp_token': temp_token,
                                    'types': 'email'})

async def send_verification_mail(user: models.User):
    """Send a verification email to user"""
    code = str(randint(100000, 999999))
    msg = MailMessage(subject="Email Verification",
                      recipients=[user.email],
                      template_body={"code": code},
                      subtype="html")
    await send_mail(msg, 'mail_verification_code.jinja2')
    # save code in db or redis by checking the settings
    if settings.TWO_FA_CODES_STORAGE == TwoFaCodesStorageTypes.DB:
        async with AsyncDBSession() as async_session:
            # save to the db
            await save_verification(async_session, user, VerificationTypes.email_code_verify, code)
    elif settings.TWO_FA_CODES_STORAGE == TwoFaCodesStorageTypes.REDIS:
        await redis_save(auth_mail_codes_namespace, str(user.id), code, settings.TWO_FA_CODES_TIMEOUT)
    else:
        write_log(APP_LOGGER, 'warning', 'mail_verification_sender',
                  f"verification storage type {settings.TWO_FA_CODES_STORAGE} is not implemented.")
        raise NotImplementedError

async def send_mobile_code():
    # TODO (PHASE 2) implement this
    raise NotImplementedError

async def get_email_verification_record_from_db(user: models.User, code) -> models.Verification:
    """Checks the verifications table and get the records for the current user
    with type of 'email_code_verify', iterate over them and check the enter code
     is okay or not, if yes returns that record."""
    async with AsyncDBSession() as async_session:
        user_verification = await get_verification_rec_by_user_id(async_session, user, code,
                                                                  VerificationTypes.email_code_verify)
        if user_verification.created_at <= datetime.now() + timedelta(seconds=settings.TWO_FA_CODES_TIMEOUT):
            return user_verification

async def validate_email_code_from_cache(user: models.User, code) -> bool:
    saved_code = await redis_get(auth_mail_codes_namespace, str(user.id))
    if saved_code is not None and saved_code == code:
        return True
    return False

def verify_totp_code(user, code):
    """Verifies user TOTP code"""
    totp = pyotp.TOTP(user.two_fa_secret_key)
    if not totp.verify(code):
        return False
    return True

def verify_totp_dep(code: str = Body(..., min_length=6, max_length=6, embed=True),
                    current_user: models.User = Depends(get_current_active_user)) -> models.User:
    """Verifies user TOTP code or raise 401 exception"""
    if not verify_totp_code(current_user, code):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=response_strings.WRONG_2FA_CODE)
    else:
        return current_user


from setuptools import setup

from pathlib import Path
dir = Path(__file__).parent
desc = (dir / 'README.md').read_text()

setup(
    name="colors",
    version="0.1",
    description="A [Unix] Colorscheme Generator",
    packages=["colors"],
    author="Jonathan Doyle",
    install_requires=["chevron", "pyyaml", "rich"],
    python_requires=">=3.5",
    entry_points={"console_scripts": ["colors = colors.main:run"]},
    long_description=desc
)


import random
rock = '''
    _______
---'   ____)
      (_____)
      (_____)
      (____)
---.__(___)
'''

paper = '''
    _______
---'   ____)____
          ______)
          _______)
         _______)
---.__________)
'''

scissors = '''
    _______
---'   ____)____
          ______)
       __________)
      (____)
---.__(___)
'''

#Write your code below this line ðŸ‘‡
userchoice = int(input("What do you choose? Type 0 for Rock, 1 for Paper or 2 for Scissors. "))
comp = random.randint(0,2)
i=0
if userchoice == 0:
  print(f"\n user chose {rock}")
elif userchoice == 1:
  print(f"\n user chose {paper}")
elif userchoice == 2:
  print(f"\n user chose {scissors}")
else:
  print("\n Wrong choice entered")
  i += 1
while(i == 0):
  if comp == 0:
    print(f"\n computer chose {rock}")
  elif comp == 1:
    print(f"\n computer chose {paper}")
  else:
    print(f"\n computer chose {scissors}")

  if ((userchoice == 0 and comp == 2) or (userchoice == 1 and comp == 0) or (userchoice == 2 and comp == 1) ):
    print("You won... Yay!!!!")
  elif(userchoice == comp):
    print("It's a draw..... No one won")
  else:
    print("You lost....oops you can try again")
  i += 1


# Python bytecode 2.7 (decompiled from Python 2.7)
# Embedded file name: scripts/client/gui/prb_control/entities/base/__init__.py
from adisp import adisp_process
from gui.shared.utils import functions
from helpers import dependency
from skeletons.gui.lobby_context import ILobbyContext

def vehicleAmmoCheck(func):
    from CurrentVehicle import g_currentVehicle

    @adisp_process
    def wrapper(*args, **kwargs):
        res = yield functions.checkAmmoLevel((g_currentVehicle.item,))
        if res:
            func(*args, **kwargs)
        elif kwargs.get('callback') is not None:
            kwargs.get('callback')(False)
        return

    return wrapper

def lobbyHeaderNavigationPossibleCheck(func):

    @adisp_process
    def wrapper(*args, **kwargs):
        lobbyContext = dependency.instance(ILobbyContext)
        res = yield lobbyContext.isHeaderNavigationPossible()
        if res:
            func(*args, **kwargs)
        elif kwargs.get('callback') is not None:
            kwargs.get('callback')(False)
        return

    return wrapper

@adisp_process
def checkVehicleAmmoFull(vehicle, callback=None):
    result = yield functions.checkAmmoLevel((vehicle,))
    callback(result)