# Implemented by: Hayk Aleksanyan
# based on the approach by Jonathan Feinberg, see http://static.mrfeinberg.com/bv_ch03.pdf

from PIL import Image
from PIL import ImageColor
from PIL import ImageFont
from PIL import ImageDraw

import math     #for sin, cos
import random   #used for random placement of the initial position of a word

import timeit   #for calculating running time (TESTING purposes only)

import numpy as np
import sys
import re

import fileReader as FR

class Node:
    """
      used to model a rectangle in quad-tree partition of the 2d plane;
      each Node models a rectangle in the partition and can have at most 4 sub-rectanlges
      representing the members of quad-tree construction
      the value of the node is a 4-tuple of integers representing a rectangle
      in "left upper-coord, and right-buttom" form
    """
    def __init__(self, val, p):
        self.value = val
        self.parent = p  # is a Node  or None if this is the root

        self.isFull = False # shows if the node has maximum number of children
                            # i.e. has reached its full capacity (2 or 4 in our case)

        self.child1 = None
        self.child2 = None
        self.child3 = None
        self.child4 = None

    def isLeaf(self):
        # true, if this node is a leaf, i.e. has no children
        return ( (self.child1 is None) and (self.child2 is None) and (self.child3 is None) and (self.child4 is None) )

    def Children(self):
        # return the list of children, if any
        c = []
        if self.child1 is not None:
            c.append(self.child1)
        if self.child2 is not None:
            c.append(self.child2)
        if self.child3 is not None:
            c.append(self.child3)
        if self.child4 is not None:
            c.append(self.child4)

        return c

    def hasLeaf_ChildrenOnly(self):
        # True, if the all children of this node (if any) are leaves
        if (self.child1 is not None):
            if self.child1.isLeaf() == False:
                return False

        if (self.child2 is not None):
            if self.child2.isLeaf() == False:
                return False

        if (self.child3 is not None):
            if self.child3.isLeaf() == False:
                return False

        if (self.child4 is not None):
            if self.child4.isLeaf() == False:
                return False

        return True

class Tree:
    """
      stores the entire quad-tree partition, where each member of the partition is a Node class
    """

    def __init__(self, root):
        # root is a Node; serves as the root of this tree
        self.root = root

    def getLeafs(self):
        #returns the leafs of the tree as a list

        if self.root == None:
            return []

        res = []

        c = self.root.Children()
        while c:
            c1 = []
            for x in c:
                #print(' '*i + 'Level ' + str(i) + ' : ' + str(x.value) )

                if x.isLeaf():
                    res.append(x)
                else:
                    for u in x.Children():
                        c1.append(u)

            c = c1

        return res

    def nodeCount(self):
        # returns the leafs of the tree as a list

        if self.root == None:
            return 0

        res = 1
        c = self.root.Children()

        while c:
            c1 = []
            res += len(c)

            for x in c:
                #print(' '*i + 'Level ' + str(i) + ' : ' + str(x.value) )
                if x.isLeaf() == False:
                    for u in x.Children():
                        c1.append(u)

            c = c1

        return res

def Tree_traverse(T):
    # traverses the tree T from the root to its leaves
    # used for testing purposes only

    if T.root == None:
        print('The tree is empty')
        return

    print('Level 0: ' + str(T.root.value) + '\n')

    c = T.root.Children()
    i = 0

    while len(c) > 0:
        i += 1
        c1 = []
        for x in c:
            print(' '*i + 'Level ' + str(i) + ' : ' + str(x.value) )

            if not x.isLeaf():
                for u in x.Children():
                    c1.append(u)

        c = c1

def get_Tree_values(T):
    # traverses the tree T from the root to its leaves and returns a list of all values of all nodes

    if T.root == None:
        return []

    res = [ T.root.value ]

    c = T.root.Children()
    i = 0

    while len(c) > 0:
        i += 1
        c1 = []
        for x in c:
            #print(' '*i + 'Level ' + str(i) + ' : ' + str(x.value) )
            res.append(x.value)

            if not x.isLeaf():
                for u in x.Children():
                    c1.append(u)

        c = c1

    return res

def Tree_compress(T):
    """
     compresses the tree T, by removing all leaves whose siblings
     are leafs and whose parents have reached their full (child) capacity,
     i.e. have MAX number of children (2 or 4 in our case)

     NOTE!# T is passed by a reference, and all changes to its nodes affect the actual T
    """

    if T.root == None:
        return T

    current_level = [ T.root ]
    all_nodes = [ [current_level] ]

    while True:
        c = []
        for i in range( len(current_level) ):
            x = current_level[i]
            if x.child1 != None:
                c.append(x.child1)
            if x.child2 != None:
                c.append(x.child2)
            if x.child3 != None:
                c.append(x.child3)
            if x.child4 != None:
                c.append(x.child4)

        if c == []:
            break

        all_nodes.append(c)
        current_level = c[:]

    for i in range( len( all_nodes ) -1, 0, -1 ):
        for j in range( len( all_nodes[i] ) ):
            n = all_nodes[i][j]

            if n != None:
                p = n.parent
                if p.isFull:
                    # p has all 4 or all 2 children
                    if p.hasLeaf_ChildrenOnly():
                        # compression means that we destroy all children of a node
                        # with only leaf children and where all child nodes are occupied
                        p.child1 = None
                        p.child2 = None
                        p.child3 = None
                        p.child4 = None

    return T

def colorBBoxesBorders(im, T, shift = (0,0) ):
    """
        takes an image @im and a Tree @T corresponding to the bounding box hierarchy (quad-tree partition)
        colors only the borders of bounding rectangles

        this function is for illustration/test purposes only
    """

    w = 1  # the (symmetric) width of the border of rectangles to be colored

    im_1 = im.copy()
    (W, H) = im_1.size

    Boxes = get_Tree_values(T)

    for i in range( len(Boxes) ):
        z = Boxes[i]

        for u in range(z[0] + shift[0] , z[2] + shift[0] ):
            for v in range( z[1] + shift[1] - w, z[1] + shift[1] + w +1 ):
                if ( (v >= 0 ) and (v < H) ):
                    im_1.putpixel( (u,v), (255,0,0,0) )

            for v in range( z[3] + shift[1] - w, z[3] + shift[1] + w +1 ):
                if ( (v >= 0) and ( v<H ) ):
                    im_1.putpixel( (u,v), (255,0,0,0) )

        for u in range( z[1] + shift[1], z[3] + shift[1] ):
            for v in range( z[0] + shift[0] - w, z[0] + shift[0] + w + 1 ):
                if ( ( v >= 0 ) and ( v<W ) ):
                    im_1.putpixel( (v,u), (255,0,0,0) )

            for v in range( z[2] + shift[0] - w, z[2] + shift[0] + w + 1 ):
                if ( ( v >= 0 ) and ( v<W ) ):
                    im_1.putpixel( (v,u), (255,0,0,0) )

    return im_1

def rectArea(r):
    #returns the area of the rectangle given in min-max coordinates (top-left bottom-right)
    return abs( (r[0] - r[2])*(r[1] - r[3]) )

def intersectionRect(r1, r2, shift1 = (0,0), shift2 = (0,0), extraSize = 3 ):
    """
     gets two 4-tuples of integers representing a rectangle in min,max coord-s
     optional params. @shifts can be used to move boxes on a large canvas (2d plane)
                     @extraSize, force the rectangles to stay away from each other by the given size

     returns true, if the rectangles intersect
    """

    if ((min(r1[0] - extraSize + shift1[0] , r1[2] + extraSize + shift1[0] ) > max( r2[0] - extraSize + shift2[0], r2[2] + extraSize + shift2[0] ) )
        or ( max(r1[0] - extraSize + shift1[0], r1[2] + extraSize + shift1[0] ) < min( r2[0] - extraSize + shift2[0], r2[2] + extraSize + shift2[0] ) ) ):
        return False

    if ((min(r1[1] - extraSize + shift1[1] , r1[3] + extraSize + shift1[1] ) > max( r2[1] - extraSize + shift2[1], r2[3] + extraSize + shift2[1]) )
        or ( max(r1[1] - extraSize + shift1[1], r1[3]  + extraSize + shift1[1] ) < min( r2[1] - extraSize + shift2[1], r2[3]+ extraSize + shift2[1] ) ) ):
        return False

    return True

def Arch_spiral(a):
    # generator for the Archimedian spiral
    r = 0
    step_size = 0.5

    u, v = 0, 0

    yield(0,0)

    while True:
        r += step_size
        x, y = a*r*math.cos(r), a*r*math.sin(r)

        if (( int(x - u) == 0 ) and ( int(y - v) == 0 ) ):
            continue
        else:
            u, v = int(x), int(y)
            yield ( u, v )

def drawArchSpiral( spiral_param, canvas_width, canvas_height, N_of_iter):
    # draw a spiral on a canvas, for testing purposes only

    N = 0

    (c_x, c_y) = (  int( 0.5*canvas_width), int( 0.5*canvas_height ) )

    im_canvas = Image.new('RGBA', (canvas_width, canvas_height), color=None)

    M = Arch_spiral( spiral_param )

    for dx, dy in M:
        u, v = c_x + dx , c_y + dy

        if ((u<0)or(v<0)or(u>canvas_width-1)or(v>canvas_height-1)):
            # out of borders
            continue

        im_canvas.putpixel( (u,v), (255,0,0,0) )

        N += 1
        if N == N_of_iter:
            break

    im_canvas.show()

def Rect_spiral(a, reverse = 1):
    """
       generator for rectangular spiral
       directions = [ (0,-1), (1,0), (0, 1), (-1, 0) ]

       if reverse = 1, then we make the normal way,
       otherwise, if reverse = -1, then we do mirror reflection in (x,y)
    """

    x, y = 0, 0
    yield (x,y)
    m = a
    i = 0

    while True:
        i = 0
        while ( i < m ):
            y -= 1
            i += 1

            if reverse == 1:
                yield (x,y)
            else:
                yield (-x, -y)

        i = 0
        while ( i < m ):
            x += 1
            i += 1

            if reverse == 1:
                yield (x,y)
            else:
                yield (-x, -y)

        i = 0
        m = m + a

        while (i < m):
            y += 1
            i += 1

            if reverse == 1:
                yield (x,y)
            else:
                yield (-x, -y)

        i = 0
        while (i<m):
            x -= 1
            i += 1

            if reverse == 1:
                yield (x,y)
            else:
                yield (-x, -y)

        m = m + a

def drawRectSpiral( spiral_param, canvas_width, canvas_height, N_of_iter):
    # draws a rectangular spiral on 2d plane, for test only

    N = 0

    (c_x, c_y) = (  int( 0.5*canvas_width), int( 0.5*canvas_height ) )

    im_canvas = Image.new('RGBA', (canvas_width, canvas_height), color=None)

    M = Rect_spiral( spiral_param )

    for dx, dy in M:

        u, v = c_x + dx , c_y + dy

        if ((u<0)or(v<0)or(u>canvas_width-1)or(v>canvas_height-1)):
            # out of borders
            continue

        im_canvas.putpixel( (u,v), (255,0,0,0) )

        N += 1
        if N == N_of_iter:
            break

    im_canvas.show()

def computeFontSize(words):
    #gets the words_ dictionary, i.e. the list of words, and their frequencies
    #returns the list of suggested fonts for each of them

    w_Font = []

    keys = list(words.keys())
    #the 1st key is the word, the 2nd is the frequency

    #we simply make the font size = count of the words

    for i in range( len(words[keys[1]])):
        w_Font.append(words[keys[1]][i])

    return w_Font

def drawWord(word_, fsize):

    #returns the cropped image of word_ in Font.size = fsize

    font = ImageFont.truetype("arial.ttf", fsize)
    w, h = font.getsize(word_)

    im = Image.new('RGBA', (w,h), color=None)
    draw = ImageDraw.Draw(im)
    draw.text((0, 0),word_, font=font)

    return im

def computeCanvasSize(words_):
    #gets the dictionary of words (with keys: ['words', 'freq', 'size'])
    #returns the proposed (width, height) of the canvas

    W, H = 0, 0

    keys_  = list( words_.keys() )

    for i in range(len(words_[keys_[2]])):
        fsize = words_[keys_[2]][i]
        font = ImageFont.truetype("arial.ttf", fsize)
        w, h = font.getsize(words_[keys_[0]][i])

        W += w
        H += h

    return  int(0.5*W), H

def colorBoxes( im, Boxes, shift = (0,0) ):
    #colors boxes in image in red

    im_1 = im.copy()

    for i in range( len(Boxes) ):
        z = Boxes[i].value
        for u in range(z[0] + shift[0] , z[2] + shift[0] ):
            for v in range(z[1] +shift[1] , z[3] + shift[1] ):
                im_1.putpixel( (u,v), (255,0,0,0) )

    return im_1

def getBoxes_Nested(im, minW, minH):
    #returns the quad-tree of the image im, i.e.
    #a Tree, where the value of each node is a 4-tuple of ints (can be 2 for some of the leafs),
    #representing min-max of hierarchic boxes
    #here minW and minH are the width, height of the minimal box

    box_0 = im.getbbox() #the initial box

    p = Node(box_0, None) # the root of the tree

    T = Tree( p )

    stack = [ p ] # (4 ints, a tuple)

    W, H = abs(box_0[0] - box_0[2]), abs(box_0[1] - box_0[3])

    if ( (H <= minH) and (W <= minW) ):
        return T

    full_node = True

    while (stack):

        x = stack.pop()
        x_box = x.value #the box coordinates
        full_node = True #shows if the x as a parent is full

        W, H = abs(x_box[0] - x_box[2]), abs(x_box[1] - x_box[3])

        if ( (H <= minH) and (W <= minW)):
            continue

        #consider the 4 sub-boxes:

        if (x_box[0] + x_box[2])%2 == 0:
            d1 = (x_box[0] + x_box[2])>>1
        else:
            d1 = (x_box[0] + x_box[2] + 1)>>1

        if (x_box[1] + x_box[3])%2 == 0:
            d2 = (x_box[1] + x_box[3])>>1
        else:
            d2 = (x_box[1] + x_box[3] + 1)>>1

        # (x0, x1, d1, x3), (d1+1, x1, x2, d2), (x0, d2+1, d1, x3), (d1+1, d2+1, x2, x3)

        #we now set the children of x
        if ((H > minH) and (W > minW) ): #we need 4 sub-rectangles

            x_child = (x_box[0], x_box[1], d1, d2 )
            if im.crop(x_child).getbbox() is not None:
                x.child1 = Node( x_child, x )
                stack.append(x.child1)
            else:
                full_node = False

            x_child = (d1 , x_box[1], x_box[2], d2)
            if im.crop(x_child).getbbox() is not None:
                x.child2 = Node( x_child,  x )
                stack.append(x.child2)
            else:
                full_node = False

            x_child = (x_box[0], d2, d1, x_box[3])
            if im.crop(x_child).getbbox() is not None:
                x.child3 = Node( x_child,  x )
                stack.append(x.child3)
            else:
                full_node = False

            x_child = (d1 , d2 , x_box[2], x_box[3])
            if im.crop(x_child).getbbox() is not None:
                x.child4 = Node( x_child,  x )
                stack.append(x.child4)
            else:
                full_node = False

            x.isFull = full_node

        else:
            if (( H <= minH  ) and ( W > minW ) ): # don't split the y-coord, but only x

                x_child = (x_box[0], x_box[1], d1, x_box[3] )
                if im.crop(x_child).getbbox() is not None:
                    x.child1 = Node( x_child, x )
                    stack.append(x.child1)
                else:
                    full_node = False

                x_child = (d1, x_box[1], x_box[2], x_box[3] )
                if im.crop(x_child).getbbox() is not None:
                    x.child2 = Node( x_child, x )
                    stack.append(x.child2)
                else:
                    full_node = False

                x.isFull = full_node

            else: #we're in a position that we only split H

                x_child = (x_box[0], x_box[1], x_box[2], d2 )
                if im.crop(x_child).getbbox() is not None:
                    x.child1 = Node( x_child, x )
                    stack.append(x.child1)
                else:
                    full_node = False

                x_child = (x_box[0], d2, x_box[2], x_box[3] )
                if im.crop(x_child).getbbox() is not None:
                    x.child2 = Node( x_child, x )
                    stack.append(x.child2)
                else:
                    full_node = False

                x.isFull = full_node

    return T

def computeWordArea( T ):
    #given the quadtreee of a word, we compute the area occupied by the word's shape

    a = 0
    c = T.getLeafs()

    for r in c:
        a += rectArea(r.value)

    return a

def collision_test(T1, T2, shift1, shift2):
    # the input is a pair of trees, representing the objects as a quad-tree
    # shift_1 = (a1, b1) and shift2 = (a2, b2) are the left-top coordinats of the boxes on the large canvas
    # this means that all boxes in T1, T2 must be shifted by (a_i, b_i)

    r1 = T1.root
    r2 = T2.root

    if ((not r1)or(not r2)):
        return False

    if intersectionRect(r1.value, r2.value, shift1, shift2 ) == False: #need to shift
        return False

    stack = [ (r1,r2) ]

    while stack:

        pair = stack.pop()

        p1, p2 = pair[0], pair[1]  #the node of the 1st and the 2nd trees

        if intersectionRect(p1.value, p2.value, shift1, shift2) == False: #needs shift
            continue #no need to go for sub-nodes

        if ( p1.isLeaf() and p2.isLeaf() ):
            return True #we've found a collision in leafs

        else:
            c1 = p1.Children()
            c2 = p2.Children()

            if (c1 == []):
                for x in c2:
                    stack.append( (p1,x) )
            else:
                if (c2 == []):
                    for x in c1:
                        stack.append( (x, p2) )
                else: #none are empty, i.e. the nodes are not leafs
                    for x in c1:
                        for y in c2:
                            stack.append((x,y))

    return False

def insideCanvas( T, shift, canvas_size ):
    #T is the tree of the word, shift is its upper-left corner coord in canvas
    # and canvas_size is a tuple (width, height) of the canvas

    #this function checks, if the word's leafs stay inside the canvas

    r = T.root

    stack = [r]
    W, H = canvas_size[0], canvas_size[1]
    sh_w, sh_h = shift[0], shift[1]

    while stack:
        v = stack.pop()

        coord = v.value #is a 4 tuple, left upper-coord, and right-buttom
        if ( (coord[0] + sh_w < 0)or( coord[2] + sh_w > W )or(coord[1] + sh_h <0)or(coord[3] + sh_h>H) ) == False:
            continue
        else:
            if v.isLeaf():
                return False
            else:
                c = v.Children()

                for x in c:
                    stack.append(x)

    return True

def copyTokens(tokens_with_freq, N_of_tokens_to_use):
    res = ''

    words = tokens_with_freq[0][0:N_of_tokens_to_use]
    sizes = tokens_with_freq[1][0:N_of_tokens_to_use]

    for i in range(0, len(words)):
        res += sizes[i]*(words[i] + ' ')

    return res

def drawOnCanvas(tokens_with_freq,   placeInfo ):

    words = tokens_with_freq[0]
    sizes = tokens_with_freq[1]

    (c_W,c_H) = placeInfo[0]     # the suggested canvas size, might change here
    place = placeInfo[1]         # list of tuples, each for each word, showing the coordinate of the upper left corner
    word_img_size = placeInfo[2] # lisf of tuples, for each word, showing the size of the word's image

    word_img_path = placeInfo[3]

    #there might be some positions of words which fell out of the canvas
    # we first need to go through these exceptions, and expand the canvas and (or) shift the coordinate's origin.

    X_min, Y_min = 0, 0

    for i in range(0, len(place)):

        if place[i] is None:
            continue

        if X_min > place[i][0]:
            X_min = place[i][0]

        if Y_min > place[i][1]:
            Y_min = place[i][1]

    x_shift, y_shift = 0, 0
    if X_min < 0:
        x_shift = -X_min
    if Y_min < 0:
        y_shift = -Y_min

    X_max , Y_max = 0, 0
    for i in range( len(place) ):
        if place[i] is None:
            continue

        place[i] = ( place[i][0] + x_shift, place[i][1] + y_shift )
        if X_max < place[i][0] + word_img_size[i][0]:
            X_max = place[i][0] + word_img_size[i][0]
        if Y_max < place[i][1] + word_img_size[i][1]:
            Y_max = place[i][1] + word_img_size[i][1]

        word_img_path[i] = ( word_img_path[i][0] + x_shift, word_img_path[i][1] )

    if c_W < X_max:
        c_W = X_max
    if c_H < Y_max:
        c_H = Y_max

    im_canvas = Image.new('RGBA', (c_W + 10 ,c_H + 10 ), color = None )
    im_canvas_white = Image.new('RGBA', (c_W + 10 ,c_H + 10 ), color = (255,255,255,255) )

    dd = ImageDraw.Draw(im_canvas)
    dd_white = ImageDraw.Draw(im_canvas_white)

    color_schemes = [
        [ (89, 97, 113), (115, 124, 140), (141, 150, 168), (179, 188, 204), (219, 227, 240) ],
        [ (89, 97, 113), (115, 124, 140), (141, 150, 168) , (89, 97, 113), (89, 97, 113)],
        [ (120, 120, 120), (0,100,149),  (242, 99, 95), (0, 76, 112), (244,208,12) ],
        [ (14, 38, 50), (1,70,99),  (35, 118, 150), (180, 200, 207), (159,195,185) ],
        [ (3, 113, 146), (99,167,190),  (10, 31, 78), (252, 105, 53), (252,105,53) ],
        [ (12, 6, 54), (9,81,105),  (5, 155, 154), (83, 186, 131), (159,217,107) ] ,
        [ (100, 101, 165), (105,117,167),  (244, 233, 109), (242, 138, 49), (241,88,56) ],
        [ (171, 165, 191), (90,87,118),  (88, 62, 47), (241, 224, 214), (191,153,144) ],
        [ (25, 46, 91), (30,101,167),  (115, 162, 192), (0, 116, 63), (241,161,4) ]
    ]

    word_colors = color_schemes[ random.randint(0, len(color_schemes) - 1) ]

    max_size = max(sizes)

    for i in range( len(place) ):
        if place[i] is None:
            print('the word <' + words[i] + '> was skipped' )
        else:
            font1 = ImageFont.truetype("arial.ttf", sizes[i])

            if sizes[i] >= 0.7*max_size:
                c = word_colors[-1]
            elif ((sizes[i] >= 0.5*max_size)and(sizes[i]<0.7*max_size) ):
                c = word_colors[-2]
            elif ((sizes[i] >= 0.35*max_size)and(sizes[i]<0.5*max_size) ):
                c = word_colors[-3]
            elif ((sizes[i] >= 0.15*max_size)and(sizes[i]<0.35*max_size) ):
                c = word_colors[-4]
            else:
                c = word_colors[-5]

            dd.text( place[i], words[i], fill = c,  font = font1 )
            dd_white.text( place[i], words[i], fill = c,  font = font1 )

    enlarge_ = 5
    margin_size = 10

    box = im_canvas.getbbox()

    im_canvas_1 = Image.new('RGBA', ( box[2] - box[0] + 2*margin_size, box[3] - box[1] + 2*margin_size ), color = (100,100,100,100)  )
    im_canvas_1.paste( im_canvas_white.crop(box), ( margin_size, margin_size, margin_size + box[2] - box[0], margin_size + box[3] - box[1] ) )

    return im_canvas_1

def normalizeWordSize(tokens, N_of_tokens_to_use, max_size, min_size):

    words = tokens[0][0:N_of_tokens_to_use]
    sizes = tokens[1][0:N_of_tokens_to_use]

    sizes = [int(max_size*x/max(sizes)) for x in sizes]

    i = 0
    while i < len(sizes):
        if sizes[i] < min_size:
            del sizes[i]
            del words[i]
        else:
            i += 1

    return [words, sizes]

def placeWords(tokens_with_freq ):
    #gets a list of tokens and their frequencies
    #returns canvas size, locations of upper-left corner of words and words' sizes

    #1. we first create the QuadTrees for words, and determine a size for the canvas

    words = tokens_with_freq[0]
    sizes = tokens_with_freq[1]

    word_img_size = [] #list of tuples, collecting the drawn image sizes (width, height)

    word_img_path = [] #shows the path passed through the spiral before hitting a free space

    print('Number of words is ' + str(len(words)) + '\n')

    word_Trees = [] #list of trees
    H_W_quotient = [] #shows the quotient of Height/Width

    T_start = timeit.default_timer()

    c_W, c_H = 0, 0     #the proposed size of the canvas
    ensure_first = 2    #shows for how many words we need to reserve space on canvas

    total_area = 0
    avg_quotient = 0

    for i in range( len(words) ):

        im_tmp = drawWord(words[i], sizes[i] )
        T = getBoxes_Nested( im_tmp , 7, 7 )

        im_tmp = im_tmp.crop(im_tmp.getbbox())
        (a,b) = im_tmp.size
        H_W_quotient.append( b/a )

        Tree_compress(T) # we merge some leafs here, making the tree smaller
        word_Trees.append( T )

        word_img_size.append(im_tmp.size)

        if i < ensure_first:
            c_W += im_tmp.size[0]
            c_H += im_tmp.size[1]
        else:
            avg_quotient += H_W_quotient[i]
            total_area += computeWordArea( word_Trees[i] )

    if len(words)>ensure_first:
        avg_quotient = avg_quotient/( len(words) - ensure_first  )
    else:
        avg_quotient = 1

    T_stop = timeit.default_timer()

    #2. We now find places for the words on our canvas

    print('(i)  QuadTrees have been made for all words in ' + str( T_stop - T_start ) + ' seconds.\n')

    c_W, c_H = 2000, 1200

    print('(ii) Now trying to place the words.\n')

    T_start = timeit.default_timer()

    #3a. we start with the 1st word

    places = [ ]  #returned value; list of tuples representing places of words: if no suggested place, we put NONE
    ups_and_downs = [ random.randint(0,20)%2  for i in range(0, len(words) )]

    strLog = '' #the log

    for i in range(0, len(words)):

        print(  words[i], end = ', ' )

        strLog_word = '<' + words[i] + '>\n'

        a = 0.1 #the parameter of the spiral
        place_found = False

        if ups_and_downs[i] == 1:
            a = -a

        no_collision_place = [] #in case we don't get a place inside canvas, we collect legal (i.e. collision-free) places for further use

        #get some starting position on the canvas, in a strip near half of the width of canvas
        w, h = 0, 0
        if i == 0:
            w, h =  np.random.randint( int(0.3*c_W), int(0.5*c_W) ),   int(0.5*c_H) + np.random.randint( -50,50 )
        else:

            if sizes[i] < 0.3*max( sizes ):
                w, h =  np.random.randint( int(0.3*c_W), int(0.7*c_W) ),   int(0.5*c_H) + np.random.randint( -50,50 )
            else:
                w, h =  np.random.randint( int(0.3*c_W), int(0.7*c_W) ),   int(0.5*c_H) + np.random.randint( -50,50 )

        strLog_word += '   Starting position: (' + str(w) + ',' +str(h) + ')\n'

        if ups_and_downs[i] == 1:
            A = Arch_spiral(a)
        else:
            A = Rect_spiral(2, ups_and_downs[i])

        dx0, dy0 = 0, 0

        place1 = (w,h)

        word_img_path.append( (w,h) )

        last_hit_index = 0 #we cache the index of last hit

        iter_ = 0

        start_countdown = False
        max_iter = 0

        for dx, dy in A:

            w, h = place1[0] + dx, place1[1] + dy

            if start_countdown == True:
                max_iter -= 1
                if max_iter == 0:
                    break
            else:
                iter_ += 1

            if ( (w<0)or(w>c_W)or(h<0)or(h>c_H) ): #fell outside the canvas area
                if start_countdown == False:
                    start_countdown = True
                    max_iter  = 1 + 10*iter_

                    strLog_word += '    Exited the canvas on step =' + str(iter_) + ' at (' + str(w) + ','  + str(h) +')\n'
                    strLog_word += '    Max_iter to complete =' + str(max_iter) +'\n'

            place1 = ( w, h )

            collision = False

            if last_hit_index < i:
                j = last_hit_index
                if places[j] is not None:
                    if collision_test(word_Trees[i], word_Trees[j], place1, places[j]) == True:
                        collision = True

            if collision == False:
                for j in range(0, i ): #check for collisions with the rest
                    if ((j != last_hit_index) and (places[j] is not None)):
                        if collision_test(word_Trees[i], word_Trees[j], place1, places[j]) == True:
                            collision = True
                            last_hit_index = j
                            break

            if collision == False:
                if insideCanvas( word_Trees[i] , place1, (c_W, c_H) ) == True:
                    places.append( place1 )
                    place_found = True
                    strLog_word += '    Place was found\n'
                    break # breaks the move in the Archimedian spiral
                else:
                    no_collision_place.append(  place1  )

        if place_found == False:

            print('no place was found')

            strLog_word += '    N of collision_free_outside canvas places =' + str(len(no_collision_place)) +'\n'

            if len(no_collision_place) == 0:
                places.append(None)
            else:
                if  i == 0 : # the first word
                    places.append( no_collision_place[0]  )
                else:
                    # print('BUT there is some outside the canvas')

                    #k1, k2 = int(0.5*c_W), int(0.5*c_H) #the centre of the canvas
                    k1, k2 =  int(0.5*(places[0][0] + word_img_size[0][0])), int(0.5*(places[0][1] + word_img_size[0][1]))

                    place1 = no_collision_place[0]
                    c_x, c_y = int(0.5*(place1[0] + word_img_size[i][0])),  int(0.5*(place1[1] + word_img_size[i][1]))
                    ds = (k1 - c_x)**2 + (k2 - c_y)**2
                    d_index = 0
                    for p in range(1, len(no_collision_place)):
                        place1 = no_collision_place[p]
                        c_x, c_y = int(0.5*(place1[0] + word_img_size[i][0])),  int(0.5*(place1[1] + word_img_size[i][1]))
                        ds_1 = (k1 - c_x)**2 + (k2 - c_y)**2
                        if ds_1<ds:
                            ds, d_index = ds_1 , p

                    places.append( no_collision_place[d_index] )
                    #NOTE!!! - - -  - - - TODO -  - - - - - -- -
                    #if the no collision place has negative (x or y), we need
                    #to to replace the origin of the canvas to accommodate those words

        strLog += strLog_word + '\n' + 'New size of the canvas = (' + str(c_W) +',' +  str(c_H)  +  ')\n\n'

    T_stop = timeit.default_timer()

    print('\n Words have been placed in ' + str( T_stop - T_start ) + ' seconds.\n')

    with open('LogFile.txt', 'w+') as f:
        f.write(strLog)

    return [ (c_W, c_H),  places , word_img_size, word_img_path ]

def createWordle_fromFile( fName ):

    tokens = FR.tokenize_file_IntoWords(fName)
    tokens_with_freq = FR.tokenGroup(tokens)

    print('\n')

    for i in range(  min(10, len(tokens_with_freq) ) ):
        s = tokens_with_freq[1][i]
        print( str(s) +  (7-len(str(s)))*' ' + ':  ' + tokens_with_freq[0][i]  )

    N_of_words = 150
    tokens =  normalizeWordSize(tokens_with_freq, N_of_words, 250, 15)
    finalPlaces = placeWords(tokens)

    wordle = drawOnCanvas(tokens, finalPlaces)

    wordle.save( fName[0:-4] + '_wordle.png')

    print( 'the wordle image was sucessfully saved on the disc as <' + fName[0:-4]  + '_wordle.png >' )

if __name__ == "__main__":

    createWordle_fromFile( sys.argv[1])


#
# Format CSS/SASS/SCSS/LESS code according to LaterPay styleguide.
#   written by Dominik Rodler <drodler@laterpay.net>
#
# usage:
#   format_code(code, action)
#

"""Format CSS/SASS/SCSS/LESS code according to LaterPay styleguide."""

import re

def format_code(code):
    code = normalize_code(code)
    code = apply_LaterPay_style(code)
    code = trim(code)

    return code

def normalize_code(code):
    code = re.sub(r'\s*([\{\}:;,])\s*', r'\1', code)                        # remove \s before and after characters {}:;,
    code = re.sub(r'\n', r'', code)                                         # remove \n
    code = re.sub(r',[\d\s\.\#\+>:]*\{', '{', code)                         # remove invalid selectors
    code = re.sub(r';\s*;', ';', code)                                      # remove superfluous ;

    code = re.sub(r'\/\*\s*([\s\S]+?)\s*\*\/', r'/* \1 */', code)           # add space before and after comment content
    code = re.sub(r'\}\s*(\/\*[\s\S]+?\*\/)\s*', r'}\n\1\n', code)          # add \n before and after outside comment
    code = re.sub(r'(http[s]?:) \/\/', r'\1//', code)                       # fix space after http[s]:
    code = re.sub(r'\s*!important', ' !important', code)                    # add space before !important

    return code

def apply_LaterPay_style(code):
    code = fix_0_values(code)

    code = re.sub(r'(\S)\{', r'\1 {', code)                                 # add space before {
    code = re.sub(r'{(\s)(\S)', r'{\2', code)                               # remove space after {
    code = re.sub(r'((@media|@[\w-]*keyframes)[^\{]+\{)\s*', r'\1\n', code) # add \n after @media {

    code = re.sub(r'\s*([\:])\s*', r'\1', code)                             # remove space after :
    code = re.sub(r'(\S);([^\}])', r'\1; \2', code)                         # add space after ;
    code = re.sub(r'\;\s*(\/\*[^\n]*\*\/)\s*', r'; \1\n', code)             # fix comment after ;
    code = re.sub(r'((?:@charset|@import)[^;]+;)\s*', r'\1\n', code)        # add \n after @charset & @import
    code = re.sub(r';\s*([^\};]+?\{)', r';\n\1', code)                      # add \n before included selector

    # TODO: add space around > + /
    # TODO: add space around *, if at least one adjacent character is a number (in order to preserve hacks)
    # TODO: remove " or ' with url(...)
    # TODO: compress hex color values

    code = re.sub(r'(\/\*[^\n]*\*\/)\s+\}', r'\1}', code)                   # remove \n between comment and }
    code = re.sub(r'(\s)\}', r'\1}', code)                                  # remove space before }
    code = re.sub(r'([^;])\}', r'\1;}', code)                               # add missing ; before }
    code = re.sub(r'\}\s*', r'}\n', code)                                   # add \n after }

    code = comma_rules(code)                                                # add space or \n after ,
    code = sort_properties(code)
    code = expand_long_rules(code)
    code = indent_rules(code)

    return code

def comma_rules(code):
    block = code.split('}')

    for i in range(len(block)):
        b = block[i].split('{')
        for j in range(len(b)):
            if b[j].count('@import'):
                s = b[j].split(';')
                for k in range(len(s)):
                    if not s[k].count('@import'):                           # ignore @import
                        s[k] = re.sub(r',(\S)', r',\n\1', s[k])
                b[j] = ';'.join(s)
            else:
                if j == len(b) - 1 or b[j].count('@media'):
                    b[j] = re.sub(r',(\S)', r', \1', b[j])                  # add space after properties' or @media's ,
                else:
                    b[j] = re.sub(r',(\S)', r',\n\1', b[j])                 # add \n after selectors' ,
        block[i] = '{'.join(b)
    code = '}'.join(block)

    return code

def sort_properties(code):
    lines = code.split('\n')

    for i in range(len(lines)):
        # get CSS rule by trimming current line
        rule = trim(lines[i])

        # extract CSS selector from rule
        selectorCharacters = rule.find('{')
        if selectorCharacters > 0:
            selector = rule[:selectorCharacters]
            # delete selector from rule
            rule = re.sub(r'(\S*\s*\{)', r'', rule)
            rule = re.sub(r'\}', r'', rule)

            # sort CSS properties alphabetically
            properties = re.sub(r';\s', r';', rule).split(';')
            properties = filter(None, properties)
            properties = sorted(properties)

            rule = selector + '{' + '; '.join(properties) + ';}'            # reconstruct rule

        lines[i] = rule

    code = '\n'.join(lines)

    return code

def expand_long_rules(code):
    expand_threshold = 120  # expand rules with > 120 charactes
# TODO: make this a setting
    lines = code.split('\n')

    for i in range(len(lines)):
        if (expand_threshold < len(lines[i])):
            lines[i] = re.sub(r'\{(\S)', r'{\n\1', lines[i])                # add \n after {
            lines[i] = re.sub(r'(\S);([^\}])', r'\1;\n    \2', lines[i])    # add \n and indentation after ; except for ; before }
            lines[i] = re.sub(r'([^\}])\s*\}', r'\1\n}', lines[i])          # add \n before }

    code = '\n'.join(lines)

    return code

def indent_rules(code):
    lines = code.split('\n')
    level = 0
    tabSize = 4
    spaces = ' ' * tabSize

    for i in range(len(lines)):
        increment = lines[i].count('{') - lines[i].count('}')
        level = level + increment
        thisLevel = level - increment if increment > 0 else level
        lines[i] = trim(lines[i])
        lines[i] = spaces * thisLevel + lines[i]

    code = '\n'.join(lines)

    return code

def trim(code):
    code = re.sub(r'^\s*(\S+(\s+\S+)*)\s*$', r'\1', code)

    return code

def fix_0_values(code):
    code = re.sub(r'\s0[emprtx%]+', r' 0', code)                            # remove unit from 0 values after \s
    code = re.sub(r':0[emprtx%]+', r':0', code)                             # remove unit from 0 values after :

    return code


import torch
import numpy as np
import argparse
import os
import nrrd
from tqdm import tqdm

from data.dataloader2d import create_loader_2d
from data.dataloader3d import create_loader_3d
from models.backbone import Unet
from util.yaml_util import load_config_yaml
from util.model_util import load_checkpoint_model

def pred_with_model(model, ckpt_dir, ckpt_fn, pred_save_dir, config, data_config, roi_names, device):
    model = load_checkpoint_model(model, ckpt_dir, ckpt_fn, device)

    ###################################################################################
    # Dataset
    ###################################################################################
    dataloader_3d = create_loader_3d(data_config, 'pred')
    n_batch_3d = len(dataloader_3d)
    with torch.no_grad():
        with tqdm(total=n_batch_3d, desc=f"Predicting", unit='batch') as pbar:
            for batch_3d in dataloader_3d:
                assert len(batch_3d) == 1, 'len(batch_3d) in pred.py must be set to 1.'
                pid = batch_3d[0]['pid']
                mask_pred_all = torch.zeros_like(batch_3d[0]['mask'], dtype=torch.bool).to(device)  # shape -> (N_roi, D, H, W)
                dataloader_2d = create_loader_2d(batch_3d, data_config, 'pred')
                for idx, batch_2d in enumerate(dataloader_2d):
                    img = batch_2d['img'].to(device=device, dtype=torch.float32)  # [N, n_channel, H, W]
                    target_slice = batch_2d['target_slice']  # shape -> [N]
                    mask_pred = model(img)  # shape -> (N, C, H, W), C is N_roi
                    mask_pred = mask_pred[0]
                    mask_pred = torch.transpose(mask_pred, 0, 1)  # shape -> (C, N, H, W)
                    mask_pred_all[:, target_slice] = mask_pred.gt(0)

                os.makedirs(config['pred_save_dir'], exist_ok=True)
                for i, roi in enumerate(roi_names):
                    nrrd.write(os.path.join(pred_save_dir, f'{pid}_{roi}.nrrd'), mask_pred_all[i].numpy().astype(np.uint8))
                pbar.update()

def main():
    parser = argparse.ArgumentParser(description='Predicting...')
    parser.add_argument('--config', type=str, required=True)
    args = parser.parse_args()

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

    config = load_config_yaml(args.config)
    data_config = load_config_yaml(config['data_config'])

    ###################################################################################
    # construct net
    ###################################################################################
    roi_names = data_config['dataset']['3d']['roi_names']
    n_channel = data_config['dataset']['2d']['n_slice']
    n_class = len(roi_names)
    start_channel = int(config['start_channel'])
    model = Unet(n_channel=n_channel, start_channel=start_channel, n_class=n_class).to(device)

    ckpt_dir = config['ckpt_dir']
    ckpt_fn = config['ckpt_fn']
    pred_save_dir = config['pred_save_dir']
    pred_with_model(model, ckpt_dir, ckpt_fn, pred_save_dir, config, data_config, roi_names, device)

if __name__ == '__main__':
    main()


#Copyright 2021, Battelle Energy Alliance, LLC
## Logic:
# We want to be able to iterate through the entire pcap once (this will allow the speed to be O(n) for any operations we do).
# We need to be able to record what ips/port/protocol combinations have already been created/exist in a local manner (I don't trust the speed of the memorystore).
# Two-three modes may be required, one that summarized conversations, one that displays each packet, and one that summarizes all conversations as one or two objects.

## what data do we need to collect:
# start (time of first packet seen - if possible)
# end (time of last packet seen - if possible)
# is_active (should prob just not include)
# src_ref, dst_ref (should reference ipv4,ipv6, domain name, mac)
# src_port, dst_port
# protocols (if possible)
# src_byte_count
# dst_byte_count
# src_packets
# dst_packets
# ipfix?
# src_payload_ref
# dst_payload_ref
# encapsulation(maybe)

# conversation will be identified by proto (TCP/UDP), dport, sport (need to be able to flip these), and highest level src/dst (need to be able to flip those)
# lookup table for converstation, if it doesn't already exist, need to try to lookup hosts on either side.
# try by higest level proto and descend from there, if not found create.

from scapy.all import rdpcap, sniff
from scapy.utils import PcapReader
from scapy.layers.inet import IP, ICMP, UDP, TCP
from scapy.layers.inet6 import IPv6
from scapy.layers.l2 import Ether, Dot3
from inspect import getmembers
import sys
from pprint import pprint
from collections import Counter
from tqdm import tqdm
from autodiscover.util.helper import get_object_or_create
from uuid import uuid4, uuid5
from stix2 import Infrastructure, IPv4Address, IPv6Address, MACAddress, Relationship, NetworkTraffic
from datetime import datetime, timedelta
import logging

# from scapy import packet,

class Host:
    # key represents whether this host is referenced as the src or dst in the packet.
    def __init__(self, pkt, top_addr_proto, src=True):

        addr_key = 'src' if src else 'dst'
        self.addrs = {}
        self.stix_addrs = {}
        self.top_addr_proto = top_addr_proto
        # self.top_addr = f'{top_addr_proto}-{addr_key}'
        self.top_addr = None
        for _, proto in PacketProcessor.src_dst_protos:
            try:
                # self.addrs[proto] = pkt[f'{proto}_{addr_key}']
                self.addrs[pkt[f'{proto}_{addr_key}']] = proto
                if self.top_addr is None:
                    self.top_addr = pkt[f'{proto}_{addr_key}']

            except KeyError:
                pass

    def get_addrs(self):
        return self.addrs.keys()

    def to_stix(self):
        l = []
        infra = Infrastructure(name=self.top_addr)
        l.append(infra)
        for addr, proto in self.addrs.items():
            if proto == 'IPv4':
                ipv4 = IPv4Address(value=addr)
                self.stix_addrs[addr] = ipv4
                l.append(ipv4)
                rel = Relationship(source_ref=infra, target_ref=ipv4,  relationship_type='has')
                l.append(rel)
            elif proto == 'IPv6':
                ipv6 = IPv6Address(value=addr)
                self.stix_addrs[addr] = ipv6
                l.append(ipv6)
                rel = Relationship(source_ref=infra, target_ref=ipv6,  relationship_type='has')
                l.append(rel)
            elif proto == 'MAC':
                mac = MACAddress(value=addr)
                self.stix_addrs[addr] = mac
                l.append(mac)
                rel = Relationship(source_ref=infra, target_ref=mac,  relationship_type='has')
                l.append(rel)
            elif proto == 'Dot3':
                mac = MACAddress(value=addr)
                self.stix_addrs[addr] = mac
                l.append(mac)
                rel = Relationship(source_ref=infra, target_ref=mac,  relationship_type='has')
                l.append(rel)
        return l

class HostHolder:
    def __init__(self):
        self.d = {}

    def add_hosts(self, pkt, top_addr_proto):
        # we need to store the addresses here
        h_src = None
        h_dst = None

        _h_src = Host(pkt, top_addr_proto)
        for addr in _h_src.get_addrs():
            if addr in self.d:
                # Currently causing badness with external ips have the routers mac.
                if set(_h_src.get_addrs()) != set(self.d[addr].get_addrs()):
                    self.d[addr].addrs.update(_h_src.addrs)
                h_src = self.d[addr]
                break

        if h_src is None:
            h_src = _h_src
            for addr in h_src.get_addrs():
                self.d[addr] = h_src

        _h_dst = Host(pkt, top_addr_proto, src=False)
        for addr in _h_dst.get_addrs():
            if addr in self.d:
                if set(_h_dst.get_addrs()) != set(self.d[addr].get_addrs()):
                    self.d[addr].addrs.update(_h_dst.addrs)
                h_dst = self.d[addr]
                break

        if h_dst is None:
            h_dst = _h_dst
            for addr in h_dst.get_addrs():
                self.d[addr] = h_dst

        return h_src, h_dst

    def to_stix(self):
        stix_objs = []
        reduced_hosts = set(self.d.values())
        for host in reduced_hosts:
            stix_objs.extend(host.to_stix())
        return stix_objs

class ConversationHolder:
    def __init__(self, host_holder):
        self.hh = host_holder
        self.lookup_dict = {}

    #def add_pkt(self, port_proto, src_port, dst_port, src_dst_proto, src_address, dst_address, pkt)

    def add_pkt(self, port_proto, src_dst_proto, pkt):
        # This is ugly...
        # key = (f"{port_proto}-{pkt[f'{port_proto}_sport']}-{pkt[f'{port_proto}_dport']}-"
        # "{src_dst_proto}-{pkt[f'{src_dst_proto}_src']}-{pkt[f'{src_dst_proto}_dst']}")
        # rev_key = (f"{port_proto}-{pkt[f'{port_proto}_dport']}-{pkt[f'{port_proto}_sport']}-"
        # "{src_dst_proto}-{pkt[f'{src_dst_proto}_dst']}-{pkt[f'{src_dst_proto}_src']}")

        if port_proto is not None:
            key = (port_proto, pkt[f'{port_proto}_sport'], pkt[f'{port_proto}_dport'],
            src_dst_proto, pkt[f'{src_dst_proto}_src'], pkt[f'{src_dst_proto}_dst'])
            rev_key = (port_proto, pkt[f'{port_proto}_dport'], pkt[f'{port_proto}_sport'],
            src_dst_proto, pkt[f'{src_dst_proto}_dst'], pkt[f'{src_dst_proto}_src'])

        else:
            key = (None, None, None,
            src_dst_proto, pkt[f'{src_dst_proto}_src'], pkt[f'{src_dst_proto}_dst'])
            rev_key = (None, None, None,
            src_dst_proto, pkt[f'{src_dst_proto}_dst'], pkt[f'{src_dst_proto}_src'])

        if key in self.lookup_dict:
            self.add_to_conv(key, pkt)
        elif rev_key in self.lookup_dict:
            self.add_to_conv(rev_key, pkt, rev=True)
        else:
            self.make_conversation(key, pkt, port_proto, src_dst_proto)

    def add_to_conv(self, key, pkt, rev=False):
        if pkt['arr_time'] < self.lookup_dict[key]['start']:
            self.lookup_dict[key]['start'] = pkt['arr_time']
        if pkt['arr_time'] > self.lookup_dict[key]['end']:
            self.lookup_dict[key]['end'] = pkt['arr_time']

        self.lookup_dict[key]['protocols'].union(pkt['protocols'])

        if rev == True:
            self.lookup_dict[key]['dst_byte_count'] += pkt['size']
            self.lookup_dict[key]['dst_packets'] += 1
        else:
            self.lookup_dict[key]['src_byte_count'] += pkt['size']
            self.lookup_dict[key]['src_packets'] += 1
        # print(self.lookup_dict[key])

    def make_conversation(self, key, pkt, p_proto, sd_proto):
        # start (time of first packet seen - if possible)
        # end (time of last packet seen - if possible)
        # is_active (should prob just not include)
        # src_ref, dst_ref (should reference ipv4,ipv6, domain name, mac)
        # src_port, dst_port
        # protocols (if possible)
        # src_byte_count
        # dst_byte_count
        # src_packets
        # dst_packets
        # ipfix?
        # src_payload_ref
        # dst_payload_ref
        # encapsulation(maybe)
        src_object, dst_object = self.hh.add_hosts(pkt, sd_proto)
        addr_proto = pkt['addr_proto']
        self.lookup_dict[key] = {
            'start': pkt['arr_time'],
            # the stix2 library refuses to have converstation in which end == start, however in 1 packet converstations that occurs... So we do hacky stuff
            'end': pkt['arr_time'] + timedelta(microseconds=1),
            'is_active': False,
            'src_ref': src_object,
            'dst_ref': dst_object,
            'src_addr' : pkt[f'{addr_proto}_src'],
            'dst_addr' : pkt[f'{addr_proto}_dst'],
            'src_port': None if p_proto is None else pkt[f'{p_proto}_sport'],
            'dst_port': None if p_proto is None else pkt[f'{p_proto}_dport'],
            'protocols': pkt['protocols'],
            'src_byte_count': pkt['size'],
            'dst_byte_count':0,
            'src_packets': 1,
            'addr_proto': addr_proto,
            'dst_packets': 0
            # 'src_object':1,
            # 'dst_object':1
        }

        #   'IPv4_dst': '8.8.8.8', 'IPv4_src': '192.168.60.20', 'UDP_dport': 53, 'UDP_sport': 63000, 'MAC_dst': 'b4:fb:e4:8d:fb:76', 'MAC_src': 'cc:48:3a:5a:b3:3a'}

    def to_stix(self):
        l = []
        for key, conversation in self.lookup_dict.items():
            addr_proto = conversation.pop('addr_proto')
            # if addr_proto not in conversation['src_ref'].stix_addrs:
            #     logging.warning(f"Proto not in object: {conversation['src_ref'].stix_addrs}")
            # if addr_proto not in conversation['dst_ref'].stix_addrs:
            #     logging.warning(f"Proto not in object: {conversation['dst_ref'].stix_addrs}")
            conversation['src_ref'] = conversation['src_ref'].stix_addrs[conversation.pop('src_addr')]
            conversation['dst_ref'] = conversation['dst_ref'].stix_addrs[conversation.pop('dst_addr')]
            #REMOVE CUSTOM
            netraffic = NetworkTraffic(**conversation, allow_custom=True)
            src_rel = Relationship(source_ref=conversation['src_ref'], target_ref=netraffic, relationship_type='communicates-with')
            dst_rel = Relationship(source_ref=netraffic, target_ref=conversation['dst_ref'], relationship_type='communicates-with')
            l.extend([netraffic, src_rel, dst_rel])

        return l

class PacketProcessor:
    # The tuple exists for display purposes, might be better to just do it later on
    src_dst_protos = [(IPv6, 'IPv6'), (IP, 'IPv4'), (Ether, 'MAC'), (Dot3, 'MAC')]
    # This one is just to we're not at the mercy of others tostrings
    port_protos = [(TCP, 'TCP'), (UDP, 'UDP')]
    def __init__(self, args, stix_loader, pcap=None):
        self.args = args
        self.sl = stix_loader

        self.hh = HostHolder()
        self.ch = ConversationHolder(self.hh)
        if pcap is not None:
            self.open_pcap(pcap)

    @staticmethod
    def get_numbered_protocols(packet):
        counter = 0
        yield counter, packet.name
        while packet.payload:
            packet = packet.payload
            counter += 1
            yield counter, packet.name

    @staticmethod
    def get_protocols(packet):
        yield packet.name
        while packet.payload:
            packet = packet.payload
            yield packet.name

    @staticmethod
    def get_clean_protocols(packet):
        for proto in PacketProcessor.get_protocols(packet):
            if proto not in ['Raw', 'Padding']:
                yield proto

    def get_pkt_info(self, pkt):
        d = {}
        d['size'] = len(pkt)
        d['arr_time'] = datetime.fromtimestamp(pkt.time)
        d['protocols'] = set(self.get_clean_protocols(pkt))

        top_sd_proto = None
        top_p_proto = None

        for src_dst_proto in self.src_dst_protos:
            sd_proto, sd_proto_name = src_dst_proto
            #trying make sure we don't duplicate a single system for each address type
            if sd_proto in pkt:
                if top_sd_proto is None:
                    top_sd_proto = sd_proto_name
                d[f'{sd_proto_name}_dst'] = pkt[sd_proto].dst
                d[f'{sd_proto_name}_src'] = pkt[sd_proto].src
                # may want to remove this later
                d['addr_proto'] = top_sd_proto
                for port_proto in self.port_protos:
                    p_proto, p_proto_name = port_proto
                    if p_proto in pkt:
                        if top_p_proto is None:
                            top_p_proto = p_proto_name

                        d[f'{p_proto_name}_dport'] = pkt[p_proto].dport
                        d[f'{p_proto_name}_sport'] = pkt[p_proto].sport

        if ('IPv6_src' in d or 'IPv4_src' in d) and ('addr_proto' in d and d['addr_proto'] == 'MAC'):
            logging.warning(f'MAC is top proto when it shouldn\'t be : {d} {pkt.summary()}')

        if 'addr_proto' not in d:
            logging.warning(f'Packet has no address protocol: {d} {pkt.summary()}')
        self.ch.add_pkt(top_p_proto, top_sd_proto, d)
        return d

    def open_pcap(self, path):
        # The below (will speedup) disables any sort of processing past TCP/UDP, keep protocol list clean as well, will need to be disabled if we wish to get more in-depth info
        # Ether.payload_guess = [({"type": 0x800}, IP)]
        # IP.payload_guess = [({"frag": 0, "proto": 0x11}, UDP), ({"frag": 0, "proto": 0x6}, TCP)]
        # TCP.payload_guess = []
        # UDP.payload_guess = []
        self.packets = PcapReader(path)

    def run(self):
        # print(packets.summary())
        # l = []
        for packet in tqdm(self.packets):
            self.get_pkt_info(packet)
        self.sl.merge(self.to_stix())
            # self.get_pkt_info(packet)

    def to_stix(self):
        objs = []
        objs.extend(self.hh.to_stix())
        objs.extend(self.ch.to_stix())
        return objs
        # objs.extend()

if __name__ == "__main__":
    packetprocessor = PacketProcessor(pcap='/home/jack/code/autodiscover/top_4000.pcapng')
    packetprocessor.run()
    print(packetprocessor.to_stix())
    # pprint(Counter(loop_data(p)))


# This source file is part of DiffAI
# Copyright (c) 2018 Secure, Reliable, and Intelligent Systems Lab (SRI), ETH Zurich
# This software is distributed under the MIT License: https://opensource.org/licenses/MIT
# SPDX-License-Identifier: MIT
# For more information see https://github.com/eth-sri/diffai

# THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT ANY WARRANTY OF ANY KIND, EITHER
# EXPRESS, IMPLIED OR STATUTORY, INCLUDING BUT NOT LIMITED TO ANY WARRANTY
# THAT THE SOFTWARE WILL CONFORM TO SPECIFICATIONS OR BE ERROR-FREE AND ANY
# IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE,
# TITLE, OR NON-INFRINGEMENT.  IN NO EVENT SHALL ETH ZURICH BE LIABLE FOR ANY
#  DAMAGES, INCLUDING BUT NOT LIMITED TO DIRECT, INDIRECT,
# SPECIAL OR CONSEQUENTIAL DAMAGES, ARISING OUT OF, RESULTING FROM, OR IN
# ANY WAY CONNECTED WITH THIS SOFTWARE (WHETHER OR NOT BASED UPON WARRANTY,
# CONTRACT, TORT OR OTHERWISE).

import future
import builtins
import past
import six

from timeit import default_timer as timer
from datetime import datetime
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import Dataset

import inspect
from inspect import getargspec
import os
import helpers as h
from helpers import Timer
import copy
import random

from components import *
import models

import domains
from domains import *
import math

from losses import *

def chunks(l, n):
    """Yield successive n-sized chunks from l."""
    for i in range(0, len(l), n):
        yield l[i:i + n]

class Top(nn.Module):
    def __init__(self, args, net, ty = Point):
        super(Top, self).__init__()
        self.net = net
        self.ty = ty
        self.cw = args.width_weight

        self.tw = args.tot_weight
        self.sw = args.std_weight
        self.time_mult = args.time_mult * 100000
        self.mix_mult = args.mix_mult * 30000
        self.w = args.width
        self.global_num = 0
        self.getSpec = getattr(self, args.spec)
        self.sub_batch_size = args.sub_batch_size
        self.curve_width = args.curve_width
        self.use_log = args.use_log
        self.regularize = args.regularize
        self.log_ai_loss = args.log_ai_loss

        self.min_tw = args.tot_weight_min
        self.tw_time_pow = args.tot_weight_time_pow
        self.tw_time_falloff = pow(float(args.tot_weight_time), self.tw_time_pow)

        self.speedCount = 0
        self.speed = 0.0

    def addSpeed(self, s):
        self.speed = (s + self.speed * self.speedCount) / (self.speedCount + 1)
        self.speedCount += 1

    def forward(self, x):
        return self.net(x)

    def clip_norm(self):
        self.net.clip_norm()

    def getPred(self, x):
        return self(x).log_softmax()

    def widthL(self, outDom):
        return outDom.diameter()

    def widthDomL(self, dom):
        return self.widthL(self(dom))

    def boxSpec(self, x, target):
        if h.hasMethod(self.ty, "attack"):
            x = self.ty.attack(self , self.w, x, target)
            return [(x,None,target)]
        return [(x, self.ty.box(x, self.w), target)]

    def lineSpec(self, x, target):
        eps = h.ones(x.size()) * self.w
        return [(x, self.ty.line(x - eps, x + eps, None), target)]

    def mixSpec(self, x, target):
        if self.ty in SYMETRIC_DOMAINS:
            return self.boxSpec(x, target)
        tm = float(self.global_num) / float(self.time_mult)
        lg = math.log10(tm * tm + 10)
        pr_hbox = 1 - 1. / (lg * lg)
        s = random.uniform(0,1)
        if s < pr_hbox:
            return [(x, domains.HBox.box(x,self.w), target)]
        else:
            return [(x, domains.Box.box(x,self.w), target)]

    def curveSpec(self, x, target):
        if self.ty in SYMETRIC_DOMAINS:
            return self.boxSpec(x,target)

        batch_size = x.size()[0]

        newTargs = [ None for i in range(batch_size) ]
        newSpecs = [ None for i in range(batch_size) ]
        bestSpecs = [ None for i in range(batch_size) ]

        for i in range(batch_size):
            newTarg = target[i]
            newTargs[i] = newTarg
            newSpec = x[i]

            best_x = newSpec
            best_dist = float("inf")
            for j in range(batch_size):
                potTarg = target[j]
                potSpec = x[j]
                if (not newTarg.data.equal(potTarg.data)) or i == j:
                    continue
                curr_dist = (newSpec - potSpec).norm(1).item()  # must experiment with the type of norm here
                if curr_dist <= best_dist:
                    best_x = potSpec

            newSpecs[i] = newSpec
            bestSpecs[i] = best_x

        new_batch_size = self.sub_batch_size
        batchedTargs = chunks(newTargs, new_batch_size)
        batchedSpecs = chunks(newSpecs, new_batch_size)
        batchedBest = chunks(bestSpecs, new_batch_size)

        def batch(t,s,b):
            t = h.ltype(t)
            s = torch.stack(s)
            b = torch.stack(b)

            if h.use_cuda:
                t.cuda()
                s.cuda()
                b.cuda()

            m = self.ty.line(s, b, self.curve_width)
            return (s, m , t)

        return [ batch(t,s,b) for t,s,b in zip(batchedTargs, batchedSpecs, batchedBest)]

    def stdLoss(self, x, dom, target, **args):
        return F.nll_loss(self.getPred(x), target, **args, size_average = False, reduce = False) # definitely includes averaging

    def preDomRes(self, outDom, target, **args): # TODO: make faster again by keeping sparse tensors sparse
        t = h.one_hot(target.data.long(), outDom.size()[1]).to_dense()
        tmat = t.unsqueeze(2).matmul(t.unsqueeze(1))

        tl = t.unsqueeze(2).expand(-1, -1, tmat.size()[1])

        inv_t = h.eye(tmat.size()[1]).expand(tmat.size()[0], -1, -1)
        inv_t = inv_t - tmat

        tl = tl.bmm(inv_t)

        fst = outDom.bmm(tl)
        snd = outDom.bmm(inv_t)

        return (fst - snd) + t
    def domRes(self, outDom, target, **args):
        return self.preDomRes(outDom, target, **args).lb()

    def outDomL(self, outDom, target, **args):
        r = -self.domRes(outDom,target, **args)
        return torch.exp( torch.clamp(r, -1000, 25) )

    def outRDomL(self, outDom, target, **args):
        r = -self.domRes(outDom,target, **args)
        out = F.softplus(r.max(1)[0])      # kinda works
        return (out).pow(self.log_ai_loss) if not self.log_ai_loss is None and self.log_ai_loss > 0 else out

    def isSafeDom(self, outDom, target, **args):
        od,_ = torch.min(self.domRes(outDom,target, **args), 1)
        return od.gt(0.0).long()

    def regLoss(self):
        if self.regularize is None:
            return 0
        reg_loss = 0
        for param in self.parameters():
            reg_loss += param.norm(2)
        return self.regularize * reg_loss

    def domLoss(self, x, dom, target, **args):
        if self.ty in POINT_DOMAINS:
            return self.regLoss() + self.stdLoss(x,dom, target,**args)
        outDom = self(dom)
        return self.outDomL(outDom, target, **args) + self.regLoss()

    def domRLoss(self, x, dom, target, **args):
        if self.ty in POINT_DOMAINS:
            return self.regLoss() + self.stdLoss(x,dom, target,**args)
        outDom = self(dom)
        return self.outRDomL(outDom, target, **args) + self.regLoss()

    def getMult(self):
        time_passage = float(self.global_num) / float(self.time_mult)
        m = torch.log(h.pyval(time_passage + 2)) if self.use_log else h.pyval(1)
        return m * m * m * m * m * m

    def getTW(self):
        tm = pow(float(self.global_num), self.tw_time_pow)

        if tm >= self.tw_time_falloff:
            return self.tw
        return (self.tw - self.min_tw) * tm / self.tw_time_falloff + self.min_tw

    def totalLoss(self, x, dom, target, **args):
        if self.ty in POINT_DOMAINS:
            return self.regLoss() + self.stdLoss(x,dom,target,**args)
        outDom = self(dom)
        tw = self.tw * self.getMult()
        cw = self.cw
        sw = self.sw
        return self.regLoss() + ((self.outRDomL(outDom, target, **args) + 1 / tw) * ((self.widthL(outDom) * cw + 1) * self.stdLoss(x, dom, target, **args) + 1 / sw) - 1 / ( sw * tw) )

    def aiLoss(self, x, dom, target, **args):
        if self.ty in POINT_DOMAINS:
            return self.regLoss()  +  self.stdLoss(x,dom,target,**args)
        outDom = self(dom)
        cw = self.cw / self.getMult()
        sw = self.sw
        tw = self.getTW() * self.getMult()

        if cw > 0:
            ww = self.widthL(outDom) * cw
        else:
            ww = 0

        if tw > 0:
            tt = self.outRDomL(outDom, target, **args) * tw
        else:
            tt = 0

        return self.regLoss() + (tt + ww + self.stdLoss(x, dom, target, **args) * sw) / (tw +  sw + cw)

    def fgsmLoss(self, x, dom, target, **args):
        if self.ty in POINT_DOMAINS:
            return self.regLoss()  +  self.stdLoss(x,dom,target,**args)
        atk = domains.FGSM.attack(self, self.w, x, target)
        outDom = self(dom)
        sw = self.sw
        tw = self.tw * self.getMult()
        return (self.outRDomL(outDom, target, **args) * tw
                + self.stdLoss(atk, None, target, **args) * sw) / (tw + sw)

    def printNet(self, f):
        self.net.printNet(f)

# Training settings
parser = argparse.ArgumentParser(description='PyTorch DiffAI Example',  formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--batch-size', type=int, default=10, metavar='N', help='input batch size for training')
parser.add_argument('--test-freq', type=int, default=1, metavar='N', help='number of epochs to skip before testing')
parser.add_argument('--test-batch-size', type=int, default=10, metavar='N', help='input batch size for testing')
parser.add_argument('--sub-batch-size', type=int, default=3, metavar='N', help='input batch size for curve specs')

parser.add_argument('--test', type=str, default=None, metavar='net', help='Saved net to use, in addition to any other nets you specify with -n')

parser.add_argument('--epochs', type=int, default=1000, metavar='N', help='number of epochs to train')
parser.add_argument('--log-freq', type=int, default=10, metavar='N', help='The frequency with which log statistics are printed')
parser.add_argument('--lr', type=float, default=0.001, metavar='LR', help='learning rate')
parser.add_argument('--threshold', type=float, default=-0.01, metavar='TH', help='threshold for lr schedule')
parser.add_argument('--patience', type=int, default=0, metavar='PT', help='patience for lr schedule')
parser.add_argument('--factor', type=float, default=0.5, metavar='R', help='reduction multiplier for lr schedule')
parser.add_argument('--max-norm', type=float, default=10000, metavar='MN', help='the maximum norm allowed in weight distribution')
parser.add_argument('--time-mult', type=float, default=1, metavar='MN', help='the time falloff for standard training')
parser.add_argument('--mix-mult', type=float, default=1, metavar='MN', help='the time falloff for mix-spec training')
parser.add_argument('--width-weight', type=float, default=0.0, metavar='CW', help='the weight of width in a combined loss')
parser.add_argument('--curve-width', type=float, default=None, metavar='CW', help='the width of the curve spec')

parser.add_argument('--tot-weight', type=float, default=0.001, metavar='TW', help='the weight of domain total in a combined loss')

parser.add_argument('--tot-weight-min', type=float, default=0, metavar='TW', help='the minimum weight of domain total in a combined loss')
parser.add_argument('--tot-weight-time', type=int, default=0, metavar='TW', help='the number of timesteps to use for domain weight scheduling')
parser.add_argument('--tot-weight-time-pow', type=float, default=1, metavar='TW', help='the time-power to use for domain weight scheduling')

parser.add_argument('--std-weight', type=float, default=1, metavar='TW', help='the weight of standard loss in a combined loss')
parser.add_argument('--width', type=float, default=0.01, metavar='CW', help='the width of either the line or box')
parser.add_argument('--loss', choices = [ x for x in dir(Top) if x[-4:] == "Loss" and len(getargspec(getattr(Top, x)).args) == 4]
                    , default="aiLoss", help='picks which loss function to use for training')
parser.add_argument('--spec', choices = [ x for x in dir(Top) if x[-4:] == "Spec" and len(getargspec(getattr(Top, x)).args) == 3]
                    , default="boxSpec", help='picks which spec builder function to use for training')

parser.add_argument('--seed', type=int, default=1, metavar='S', help='random seed')
parser.add_argument("--use-schedule", type=h.str2bool, nargs='?',
                    const=True, default=False,
                    help="activate learning rate schedule")
parser.add_argument("--use-log", type=h.str2bool, nargs='?',
                    const=True, default=False,
                    help="Activate log falloff of time")

parser.add_argument('-d', '--domain', choices = h.getMethodNames(domains), action = 'append'
                    , default=[], help='picks which abstract domains to use for training', required=True)

parser.add_argument('-t', '--test-domain', choices = h.getMethodNames(domains), action = 'append'
                    , default=[], help='picks which abstract domains to use for testing', required=True)

parser.add_argument('-n', '--net', choices = h.getMethodNames(models), action = 'append'
                    , default=[], help='picks which net to use for training')  # one net for now

parser.add_argument('-D', '--dataset', choices = [n for (n,k) in inspect.getmembers(datasets, inspect.isclass) if issubclass(k, Dataset)]
                    , default="MNIST", help='picks which dataset to use.')

parser.add_argument('-o', '--out', default="out/", help='picks which net to use for training')
parser.add_argument('--test-size', type=int, default=2000, help='number of examples to test with')

parser.add_argument('-r', '--regularize', type=float, default=None, help='use regularization')
parser.add_argument('--log-ai-loss', type=float, default=None, help='use the log of the output. Must be non-negative. ')

args = parser.parse_args()

args.log_interval = int(50000 / (args.batch_size * args.log_freq))

h.max_c_for_norm = args.max_norm

if h.use_cuda:
    torch.cuda.manual_seed(1 + args.seed)
else:
    torch.manual_seed(args.seed)

kwargs = {'num_workers': 1, 'pin_memory': True} if h.use_cuda else {}

def loadDataset(train):
    oargs = {}
    if args.dataset in ["MNIST", "CIFAR10", "CIFAR100", "FashionMNIST", "PhotoTour"]:
        oargs['train'] = train
    elif args.dataset in ["STL10", "SVHN"] :
        oargs['split'] = 'train' if train else 'test'
    elif args.dataset in ["LSUN"]:
        oargs['classes'] = 'train' if train else 'test'
    else:
        raise Exception(args.dataset + " is not yet supported")

    if args.dataset in ["MNIST"]:
        transformer = transforms.Compose([ transforms.ToTensor(),
                             transforms.Normalize((0.1307,), (0.3081,))])
    elif args.dataset in ["CIFAR10", "CIFAR100"]:
        transformer = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])
    elif args.dataset in ["SVHN"]:
        transformer = transforms.Compose([
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5,0.5,0.5), (0.2,0.2,0.2))])
    else:
        transformer = transforms.ToTensor()

    return torch.utils.data.DataLoader(
        getattr(datasets, args.dataset)('../data', download=True,
                                        transform=transformer, **oargs)
        , batch_size=args.batch_size if train else args.test_batch_size
        , shuffle=True, **kwargs)

train_loader = loadDataset(True)
test_loader = loadDataset(False)

input_dims = train_loader.dataset[0][0].size()
num_classes = int(max(getattr(train_loader.dataset, 'train_labels' if args.dataset != "SVHN" else 'labels'))) + 1

print("input_dims: ", input_dims)
print("Num classes: ", num_classes)
def train(epoch, models):
    for model in models:
        model.train()

    ep_tot = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        if h.use_cuda:
            data, target = data.cuda(), target.cuda()

        for model in models:
            model.global_num += data.size()[0]

            timer = Timer("train", "sample from " + model.name + " with " + model.ty.name, data.size()[0], False)
            lossy = 0
            with timer:
                for s in model.getSpec(data,target):
                    model.optimizer.zero_grad()
                    loss = getattr(model, args.loss)(*s).sum() / data.size()[0]
                    lossy += loss.item()
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(model.parameters(), 1)
                    model.optimizer.step()
                    model.clip_norm()

            model.addSpeed(timer.getUnitTime())

            if batch_idx % args.log_interval == 0:
                print('Train Epoch {:12} {:20}: {:3} [{:7}/{} ({:.0f}%)] \tAvg sec/ex {:1.8f}\tMult {:.6f}\tTW {:.12f}\tLoss: {:.6f}'.format(
                    model.name,  model.ty.name,
                    epoch,
                    batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader),
                    model.speed,
                    model.getMult().data[0],
                    model.getTW(),
                    lossy))

def test(models, epoch, f = None):
    class MStat:
        def __init__(self, model):
            model.eval()
            self.model = model
            self.correct = 0
            self.test_loss = 0
            class Stat:
                def __init__(self, d, dnm):
                    self.domain = d
                    self.name = dnm
                    self.width = 0
                    self.safe = 0
                    self.proved = 0
                    self.time = 0
            self.domains = [ Stat(getattr(domains,d), d) for d in args.test_domain ]
    model_stats = [ MStat(m) for m in models ]

    num_its = 0
    for data, target in test_loader:
        if num_its >= args.test_size:
            break
        num_its += data.size()[0]
        if h.use_cuda:
            data, target = data.cuda(), target.cuda()

        for m in model_stats:
            with torch.no_grad():
                m.test_loss += m.model.stdLoss(data, None, target).sum().item() # sum up batch loss

            tyorg = m.model.ty

            with torch.no_grad():
                pred = m.model(data).data.max(1, keepdim=True)[1] # get the index of the max log-probability
                m.correct += pred.eq(target.data.view_as(pred)).sum()

            for stat in m.domains:
                timer = Timer(shouldPrint = False)
                with timer:
                    m.model.ty = stat.domain
                    def calcData(data, target):
                        box = m.model.boxSpec(data, target)[0]
                        with torch.no_grad():
                            if m.model.ty in POINT_DOMAINS:
                                preder = m.model(box[0]).data
                                pred = preder.max(1, keepdim=True)[1] # get the index of the max log-probability
                                org = m.model(data).max(1,keepdim=True)[1]
                                stat.proved += float(org.eq(pred).sum())
                                stat.safe += float(pred.eq(target.data.view_as(pred)).sum())
                            else:
                                bs = m.model(box[1])
                                stat.width += m.model.widthL(bs).data[0] # sum up batch loss
                                stat.safe  += m.model.isSafeDom(bs, target).sum().item()
                                stat.proved += sum([ m.model.isSafeDom(bs, (h.ones(target.size()) * n).long() ).sum().item() for n in range(num_classes) ])
                    if m.model.net.neuronCount() < 5000 or stat.domain in SYMETRIC_DOMAINS:
                        calcData(data, target)
                    else:
                        for d,t in zip(data, target):
                            calcData(d.unsqueeze(0),t.unsqueeze(0))
                stat.time += timer.getUnitTime()
            m.model.ty = tyorg

    l = num_its # len(test_loader.dataset)
    for m in model_stats:

        pr_corr = float(m.correct) / float(l)
        if args.use_schedule:
            m.model.lrschedule.step(1 - pr_corr)

        h.printBoth('Test: {:12} trained with {:8} - Mult {:1.8f}, Avg sec/ex {:1.12f}, Average loss: {:8.4f}, Accuracy: {}/{} ({:3.1f}%)'.format(
            m.model.name, m.model.ty.name,
            m.model.getMult().data[0],
            m.model.speed,
            m.test_loss / l,
            m.correct, l, 100. * pr_corr), f)

        model_stat_rec = ""
        for stat in m.domains:
            pr_safe = stat.safe / l
            pr_proved = stat.proved / l
            pr_corr_given_proved = pr_safe / pr_proved if pr_proved > 0 else 0.0
            h.printBoth("\t{:10} - Width: {:<22.4f} Pr[Proved]={:<1.3f}  Pr[Corr and Proved]={:<1.3f}  Pr[Corr|Proved]={:<1.3f}    Time = {:<7.5f}".format(
                stat.name, stat.width / l, pr_proved, pr_safe, pr_corr_given_proved, stat.time), f)
            model_stat_rec += "{}_{:1.3f}_{:1.3f}_{:1.3f}__".format(stat.name, pr_proved, pr_safe, pr_corr_given_proved)
        net_file = os.path.join(out_dir, m.model.name + "_checkpoint_"+str(epoch)+"_with_{:1.3f}".format(pr_corr)+"__"+model_stat_rec+".net")

        h.printBoth("\tSaving netfile: {}\n".format(net_file), f)

        if epoch == 1 or epoch % 10 == 0:
            torch.save(m.model.net, net_file)

def createModel(net, domain, domain_name):
    net_weights, net_create = net
    domain.name = domain_name

    net = net_create(num_classes).infer(input_dims)
    net.load_state_dict(net_weights.state_dict())

    model = Top(args, net, domain)
    model.clip_norm()
    if h.use_cuda:
        model.cuda()

    model.optimizer = optim.Adam(model.parameters(), lr=args.lr if not domain in POINT_DOMAINS else 1e-4)
    model.lrschedule = optim.lr_scheduler.ReduceLROnPlateau(
        model.optimizer,
        'min',
        patience=args.patience,
        threshold= args.threshold,
        min_lr=0.000001,
        factor=args.factor,
        verbose=True)

    model.name = net_create.__name__

    return model

out_dir = os.path.join(args.out, args.dataset, str(args.net)[1:-1].replace(", ","_").replace("'",""),
                       args.spec, "width_"+str(args.width), str(datetime.now()).replace(":","").replace(" ", "") )

print("Saving to:", out_dir)

if not os.path.exists(out_dir):
    os.makedirs(out_dir)

print("Starting Training with:")
with open(os.path.join(out_dir, "config.txt"), "w") as f:
    for k in sorted(vars(args)):
        h.printBoth("\t"+k+": "+str(getattr(args,k)), f)
print("")

if not args.test is None:
    net = torch.load(args.test)
    nets = [ lambda *args, **kargs: net ]
elif args.net == []:
    raise Exception("Need to specify at least one net with either -n or --test")
else:
    nets = []

nets += [ getattr(models,n) for n in args.net ]

nets = [ (n(num_classes).infer(input_dims),n) for n in nets ]

for net, net_create in nets:
    print("Name: ", net_create.__name__)
    print("Number of Neurons (relus): ", net.neuronCount())
    print("Number of Parameters: ", sum([h.product(s.size()) for s in net.parameters()]))
    print()

if args.domain == []:
    models = [ createModel(net, domains.Box, "Box") for net in nets]
else:
    models = h.flat([[createModel(net, getattr(domains,d), d) for net in nets] for d in args.domain])

with open(os.path.join(out_dir, "log.txt"), "w") as f:

    startTime = timer()
    for epoch in range(1, args.epochs + 1):
        if (epoch - 1) % args.test_freq == 0:
            with Timer("test before epoch "+str(epoch),"sample", 10000):
                test(models, epoch, f)
        h.printBoth("Elapsed-Time: {:.2f}s\n".format(timer() - startTime), f)
        with Timer("train","sample", 60000):
            train(epoch, models)


# -*- coding: utf-8 -*-
"""
Created on Fri Jun  2 13:09:20 2017

@author: 凯风
"""

from sklearn.linear_model import ElasticNet
from sklearn.datasets import load_boston
from sklearn.model_selection import train_test_split

# 准备数据
boston = load_boston()
X,Y = boston.data,boston.target
X_train,X_test,Y_train,Y_test = train_test_split(X,Y,test_size=.3)
'''
    弹性网
        主要就是正则化项是L1+L2
        目标函数 = 模型损失 + αβ||w|| + α(1-β)/2 * ||w||^2
        优点：
            多个特征有关联关系的情况下，弹性网是有用的
            集成了ridge和lasso的一些好的特点
        问题：
            计算的复杂度可能高一些了就
        优化方法呢？
'''
rg = ElasticNet(alpha=1.0, l1_ratio=0.5, fit_intercept=True, normalize=False, precompute=False, max_iter=1000, copy_X=True, tol=0.0001, warm_start=False, positive=False, random_state=None, selection='cyclic')
rg.fit(X_train,Y_train)
Y_pre = rg.predict(X_test)
rg.score(X_test,Y_test)
rg.coef_
rg.intercept_
'''
    alpha                       两个惩罚项系数的相乘，如果是0就是最小二乘了
    l1_ratio                    混合参数，在[0,1]之间是弹性网的
    fit_intercept               是否训练截距
    normalize                   归一化否
    precompute                  是否使用Gram矩阵来加速
    max_iter                    最大迭代次数
    copy_X                      是否覆盖模型中的X
    tol                         精度
    warm_start                  是否使用上一次调用的解决方案作为初始化
    positive                    设置强制系数为正的嘛？
    random_state                随机器
    selection                   每次迭代的时候，是否更新随机系数
'''


# additional training with frozen encoder

import numpy as np
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data

import argparse
import time
import os, sys

from pysrc.cond_vae_ringdown import CondConvVAE, CondConvVAE_KLDiv, CondConvVAE_ReconstLoss
from pysrc.utils import noise_inject_1D, normalize_1D, detector_output

parser = argparse.ArgumentParser(description="parse args")
parser.add_argument('savedir', type=str)
parser.add_argument('rootmodeldir', type=str)
parser.add_argument('snr_min', type=float)
parser.add_argument('snr_max', type=float)
parser.add_argument('--datadir', default='data/ModifiedSXS0305_GW150914/', type=str)
args = parser.parse_args()

savedir = args.savedir
if not os.path.exists(savedir):
    os.makedirs(savedir)

modeldir = args.rootmodeldir
if not os.path.exists(modeldir):
    print("The modeldir {} does not exist.".format(modeldir))
    sys.exit()

#---------------------------------------------
# set the parameters
#---------------------------------------------
IN_FEATURES = 256
HIDDEN_FEATURES = 16
PHYS_FEATURES = 2
BATCH_SIZE_TRAIN = 32
lr = 1.0e-5
alpha = 1.0

#---------------------------------------------
# make the network instance
#---------------------------------------------
net = CondConvVAE(in_features = IN_FEATURES,
                  hidden_features = HIDDEN_FEATURES,
                  phys_features = PHYS_FEATURES,
                  cudaflg=True)
net.load_state_dict(torch.load(modeldir + 'vaenet_state.pth'))
net.cuda()

#---------------------------------------------
# freeze the encoder
#---------------------------------------------

net.freeze_encoder()
unfrozen_params = filter(lambda p: p.requires_grad, net.parameters())

#---------------------------------------------
# define the loss function and the optimizer
#---------------------------------------------
LOSS_MSE = CondConvVAE_ReconstLoss()
optimizer = optim.Adam(unfrozen_params, lr=lr)

#-------------------------------
# load the training data
#-------------------------------
datadir = args.datadir
kstart = 64
nlength = 256
trainwave_p = np.load(datadir+'QNM_Train_whitened_GW150914_hPlus.npy')[:,kstart:kstart+nlength]
trainwave_c = np.load(datadir+'QNM_Train_whitened_GW150914_hCross.npy')[:,kstart:kstart+nlength]
trainlabel = torch.tensor(np.genfromtxt(datadir+'Train_Label.dat', usecols=(2,3)), dtype=torch.float32)

#-------------------------------
# load the validate data
#-------------------------------
validatewave_p = np.load(datadir+'QNM_Validate_whitened_GW150914_hPlus.npy')[:,kstart:kstart+nlength]
validatewave_c = np.load(datadir+'QNM_Validate_whitened_GW150914_hCross.npy')[:,kstart:kstart+nlength]
validatelabel = torch.tensor(np.genfromtxt(datadir+'Validate_Label.dat', usecols=(2,3)), dtype=torch.float32)

#---------------------------------------------
# Training
#---------------------------------------------

pSNR_min = args.snr_min
pSNR_max = args.snr_max

pSNR = [pSNR_min, pSNR_max]
nepoch = 100

start_time = time.time()

for epoch in range(nepoch):

    #-------------------------------
    # train dataloader
    #-------------------------------
    trainsignal = detector_output(trainwave_p, trainwave_c).reshape(-1,1,nlength)
    trainsignal = normalize_1D(noise_inject_1D(trainsignal, pSNR))
    trainsignal = torch.tensor(trainsignal, dtype=torch.float32)
    traindata = torch.utils.data.TensorDataset(trainsignal,trainlabel)
    train_loader = torch.utils.data.DataLoader(traindata, batch_size=BATCH_SIZE_TRAIN, shuffle=True, num_workers=8, drop_last=True)

    #-------------------------------
    # validate dataloader
    #-------------------------------
    validatesignal = detector_output(validatewave_p, validatewave_c).reshape(-1,1,nlength)
    validatesignal = normalize_1D(noise_inject_1D(validatesignal, pSNR))
    validatesignal = torch.tensor(validatesignal, dtype=torch.float32)
    validatedata = torch.utils.data.TensorDataset(validatesignal, validatelabel)
    validate_loader = torch.utils.data.DataLoader(validatedata, batch_size=512, shuffle=True)

    #---------------------------------------------
    # Training iteration
    #---------------------------------------------
    train_klloss = 0.0
    train_mseloss = 0.0
    train_loss = 0.0
    for i, data in enumerate(train_loader, 0):

        inputs, labels = data
        inputs = inputs.cuda()
        labels = labels.cuda()

        optimizer.zero_grad()
        mu_y, logvar_y, mu1, logvar1, mu2, logvar2 = net.forward_training(inputs, labels)
        loss = LOSS_MSE(mu_y, logvar_y, labels)

        train_loss += loss.data

        loss.backward()
        optimizer.step()

    #---------------------------------------------
    # Validation iteration
    #---------------------------------------------
    val_klloss = 0.0
    val_mseloss = 0.0
    val_loss = 0.0
    with torch.no_grad():
        for j, data in enumerate(validate_loader, 0):
            inputs, labels = data
            inputs = inputs.cuda()
            labels = labels.cuda()

            mu_y, logvar_y, mu1, logvar1, mu2, logvar2 = net.forward_training(inputs, labels)
            loss = LOSS_MSE(mu_y, logvar_y, labels)
            val_loss += loss.data

    print("{:d} {:.5e} {:.5e} ".format(epoch+1, train_loss/(i+1), val_loss/(j+1)))

    with open(savedir+'log_train_{:.2f}-{:.2f}.txt'.format(pSNR_min, pSNR_max), 'a') as f:

        f.write("{:d} {:.5e} {:.5e}\n".format(epoch+1, train_loss/(i+1), val_loss/(j+1)))

if (epoch+1)%100==0:

    torch.save(net.state_dict(), savedir+'vaenet_state_{:d}.pth'.format(epoch+1))
    torch.save(optimizer.state_dict(), savedir+'optim_state_{:d}.pth'.format(epoch+1))

torch.save(net.state_dict(), savedir+'vaenet_state_{:.2f}-{:.2f}.pth'.format(pSNR_min, pSNR_max))
torch.save(optimizer.state_dict(), savedir+'optim_state_{:.2f}-{:.2f}.pth'.format(pSNR_min, pSNR_max))

end_time = time.time()

print("elapsed time:{0} [sec]".format(end_time - start_time))


import sqlite3
import datetime
import random

DataBase = sqlite3.connect('Sales.db')
cur = DataBase.cursor()
cur.execute(''' CREATE TABLE IF NOT EXISTS Sales (
    Identifier INTEGER PRIMARY KEY,
    Country TEXT,
    City TEXT,
    Shop TEXT,
    Date DATETIME,
    Summ INTEGER
);''')

Countries = ['Russia', 'Japan', 'German', 'USA', 'England']
Cities = {'Russia': ['Surgut', 'Kazan', 'Moscow'],
          'Japan': ['Tokio', 'Kioto', 'Kie'],
          'German': ['Berlin', 'Munchern', 'Hamburg'],
          'USA': ['New-York', 'Los-Angelas', 'Washington'],
          'England': ['London', 'Cambridge', 'Manchester']}
Shop = ['Channel', 'Dolce&Gabanna', 'Gucci', 'Zara']

i = 0
for i in range(1000000):
    country_ = random.choice(Countries)
    city_ = random.choice(City[country_])
    shop_ = random.choice(Shop)
    month = random.randint(1, 12)
    date_ = datetime.date(2019, month_, random.randint(1, 28))
    summ_ = random.randint(10000, 1000000)
    a.execute("INSERT INTO Sales VALUES (?,?,?,?,?,?);", (i, country_, city_, shop_, date_, summ_))

cur.execute(''' CREATE TABLE IF NOT EXISTS Shop (
    Shop_id INTEGER PRIMARY KEY,
    City TEXT,
    Country TEXT
);''')

Base.commit()
Base.close()


from django.db.models import Sum

import elasticutils
import pyes.exceptions as pyes

import amo
from stats.models import Contribution

def extract_installed_count(installed):
    date = installed.created.date()
    return {'date': date,
            'addon': installed.addon_id,
            'count': installed.__class__.objects.filter(
                created__year=date.year,
                created__month=date.month,
                created__day=date.day).count()}

def extract_contribution_counts(contribution):
    """
    number of contributions (sales) per app/day
    revenue per app/day
    """
    addon_id = contribution['addon']
    date = contribution['created']
    return {
        'date': date,
        'addon': addon_id,
        'count': Contribution.objects.filter(
            addon__id=addon_id,
            created__year=date.year,
            created__month=date.month,
            created__day=date.day).count() or 0,
        'revenue': Contribution.objects.filter(
            addon__id=addon_id,
            type=amo.CONTRIB_PURCHASE,
            created__year=date.year,
            created__month=date.month,
            created__day=date.day).aggregate(Sum('amount'))['amount__sum']
            or 0,
        'refunds': Contribution.objects.filter(
            addon__id=addon_id,
            refund__isnull=False,
            uuid__isnull=False,
            created__year=date.year,
            created__month=date.month,
            created__day=date.day).count() or 0
    }

def setup_mkt_indexes():
    es = elasticutils.get_es()
    for model in [Contribution]:
        index = model._get_index()
        try:
            es.create_index_if_missing(index)
        except pyes.ElasticSearchException:
            pass

        mapping = {
            'properties': {
                'id': {'type': 'long'},
                'count': {'type': 'long'},
                'data': {'dynamic': 'true',
                         'properties': {
                            'v': {'type': 'long'},
                            'k': {'type': 'string'}
                        }
                },
                'date': {'format': 'dateOptionalTime',
                         'type': 'date'}
            }
        }
        es.put_mapping(model._meta.db_table, mapping,
                       model._get_index())


from os import startfile, system
from nodoM import NodoM, NodoEncabezado
from encabezados import ListaEncabezado

class Matriz:

    def __init__(self):
        self.eRows = ListaEncabezado()
        self.eColumns = ListaEncabezado()

    def insertar(self, fila, columna, valor):
        nuevo = NodoM(fila,columna,valor)

        eRows = self.eRows.getEncabezado(fila)

        if eRows == None:
            eRows = NodoEncabezado(fila)
            eRows.acceso = nuevo
            self.eRows.setEncabezado(eRows)
        else:
            if nuevo.columna < eRows.acceso.columna:
                nuevo.derecha = eRows.acceso
                eRows.acceso.izquierda = nuevo
                eRows.acceso = nuevo
            else:
                actual = eRows.acceso
                while actual.derecha !=None:
                    if nuevo.columna < actual.derecha.columna:
                        nuevo.derecha = actual.derecha
                        actual.derecha.izquierda = nuevo
                        nuevo.izquierda = actual
                        actual.derecha = nuevo
                        break
                    actual = actual.derecha

                if actual.derecha == None:
                    actual.derecha = nuevo
                    nuevo.izquierda = actual

        eColumna = self.eColumns.getEncabezado(columna)
        if eColumna == None:
            eColumna = NodoEncabezado(columna)
            eColumna.acceso = nuevo
            self.eColumns.setEncabezado(eColumna)
        else:
            if nuevo.fila < eColumna.acceso.fila:
                nuevo.abajo = eColumna.acceso
                eColumna.acceso.arriba = nuevo
                eColumna.acceso = nuevo
            else:
                actual = eColumna.acceso
                while actual.abajo != None:
                    if nuevo.fila < actual.abajo.fila:
                        nuevo.abajo = actual.abajo
                        actual.abajo.arriba = nuevo
                        nuevo.arriba = actual
                        actual.abajo = nuevo
                        break
                    actual = actual.abajo

            if actual.abajo == None:
                actual.abajo = nuevo
                nuevo.arriba = actual

    def actualizar(self,fila,columna,valor):
        aux = self.eRows.primero
        while aux != None:
            actual = aux.acceso
            while actual != None:
                if actual.fila == fila and actual.columna == columna:
                    actual.valor = valor
                actual = actual.derecha
            aux = aux.siguiente

    def get(self,fila,columna):
        aux = self.eRows.primero
        while aux != None:
            actual = aux.acceso
            while actual != None:
                if actual.fila == fila and actual.columna == columna:
                    return actual.valor
                actual = actual.derecha
            aux = aux.siguiente

    def valores(self):
        eColumna = self.eColumns.primero
        fila=' \t'
        while eColumna != None:
            actual = eColumna.acceso
            while actual != None:
                fila+= '['+str(actual.fila)+str(actual.columna)+']'
                actual = actual.abajo

            print(fila+'')
            fila=' \t'
            eColumna = eColumna.siguiente