import math
import random
import sys

import numpy as np
import pygame

# RGB
RED = (255, 0, 0)
BLUE = (0, 0, 255)
YELLOW = (255, 255, 0)

BLACK = (0, 0, 0)

ROW_COUNT = 6
COLUMN_COUNT = 6

PLAYER = 0
SKY_NET = 1  # AI actually XD,just had fun

PLAYER_PIECE = 1
AI_PIECE = 2

WINDOW_LENGTH = 4
EMPTY = 0

def create_board():
    board = np.zeros((ROW_COUNT, COLUMN_COUNT))
    return board

def drop_piece(board, row, col, piece):
    board[row][col] = piece

def is_valid_location(board, col):
    return board[ROW_COUNT - 1][col] == 0

def get_valid_row(board, col):
    for r in range(ROW_COUNT):
        if board[r][col] == 0:
            return r

def print_board(board):
    print(np.flip(board, 0))

pygame.init()

SQUARE_SIZE = 50

width = COLUMN_COUNT * SQUARE_SIZE
height = (ROW_COUNT + 1) * SQUARE_SIZE

size = (width, height)

RADIUS = int(SQUARE_SIZE / 2 - 5)

screen = pygame.display.set_mode(size)

def wining_move(board, piece):
    # Check horizontal location to win
    for c in range(COLUMN_COUNT - 3):
        for r in range(ROW_COUNT):
            if board[r][c] == piece and board[r][c + 1] == piece and \
                    board[r][c + 2] == piece and board[r][c + 3] == piece:
                return True

    # Check vertical location to win
    for c in range(COLUMN_COUNT):
        for r in range(ROW_COUNT - 3):
            if board[r][c] == piece and board[r + 1][c] == piece and \
                    board[r + 2][c] == piece and board[r + 3][c] == piece:
                return True
    # Check for Positive diagonal
    for c in range(COLUMN_COUNT - 3):
        for r in range(ROW_COUNT - 3):
            if board[r][c] == piece and board[r + 1][c + 1] == piece and \
                    board[r + 2][c + 2] == piece and board[r + 3][c + 3] == piece:
                return True

    # Check for Negative diagonal
    for c in range(COLUMN_COUNT - 3):
        for r in range(3, ROW_COUNT):
            if board[r][c] == piece and board[r - 1][c + 1] == piece and \
                    board[r - 2][c + 2] == piece and board[r - 3][c + 3] == piece:
                return True

def evaluate(window, piece):
    score = 0
    opp_piece = PLAYER_PIECE
    if piece == PLAYER_PIECE:
        opp_piece = AI_PIECE
    # 4 in a row
    if window.count(piece) == 4:
        score += 100
    # 3 in a row
    elif window.count(piece) == 3 and window.count(EMPTY) == 1:
        score += 5
    elif window.count(piece) == 2 and window.count(EMPTY) == 2:
        score += 2

    if window.count(opp_piece) == 3 and window.count(EMPTY) == 1:
        score -= 4

    return score

def score_position(board, piece):
    score = 0
    # Score center column
    center_array = [int(i) for i in list(board[:, COLUMN_COUNT // 2])]
    center_count = center_array.count(piece)
    score += center_count * 3

    # Score Horizontal
    for r in range(ROW_COUNT):
        row_array = [int(i) for i in list(board[r, :])]
        for c in range(COLUMN_COUNT - 3):
            window = row_array[c:c + WINDOW_LENGTH]
            # 4 in a row
            score += evaluate(window, piece)
    # Score Vertical
    for c in range(COLUMN_COUNT):
        column_array = [int(i) for i in list(board[:, c])]
        for r in range(ROW_COUNT - 3):
            window = column_array[r:r + WINDOW_LENGTH]
            score += evaluate(window, piece)
    # Score positive diagonal
    for r in range(ROW_COUNT - 3):
        for c in range(COLUMN_COUNT - 3):
            window = [board[r + i][c + i] for i in range(WINDOW_LENGTH)]
            score += evaluate(window, piece)
    # Score negative diagonal
    for r in range(ROW_COUNT - 3):
        for c in range(COLUMN_COUNT - 3):
            window = [board[r + 3 - i][c + i] for i in range(WINDOW_LENGTH)]
            score += evaluate(window, piece)

    return score

def get_valid_locations(board):
    valid_locations = []
    for col in range(COLUMN_COUNT):
        if is_valid_location(board, col):
            valid_locations.append(col)
    return valid_locations

def pick_best_move(board, piece):
    valid_locations = get_valid_locations(board)
    best_score = -10000
    best_col = random.choice(valid_locations)
    for column in valid_locations:
        row = get_valid_row(board, column)
        # we make copy to take new memory location, to any modifications do not change main board
        temp_board = board.copy()
        drop_piece(temp_board, row, column, piece)
        score = score_position(temp_board, piece)
        if score > best_score:
            best_score = score
            best_col = column

    return best_col

def is_terminal_node(board):
    return wining_move(board, PLAYER_PIECE) or wining_move(board, AI_PIECE) or len(get_valid_locations(board)) == 0

def minimax(board, depth, maximizing_player, alpha, beta):
    valid_locations = get_valid_locations(board)
    is_terminal = is_terminal_node(board)
    if depth == 0 or is_terminal:
        if is_terminal:
            if wining_move(board, AI_PIECE):
                return None, 9999999999999999
            elif wining_move(board, PLAYER_PIECE):
                return None, -9999999999999999
            else:  # no more valid moves
                return None, 0
        else:
            return None, score_position(board, AI_PIECE)
    if maximizing_player:
        value = -math.inf
        column = random.choice(valid_locations)
        for col in valid_locations:
            row = get_valid_row(board, col)
            b_copy = board.copy()
            drop_piece(b_copy, row, col, AI_PIECE)
            new_score = minimax(b_copy, depth - 1, False, alpha, beta)[1]
            if new_score > value:
                value = new_score
                column = col
            alpha = max(alpha, value)
            if alpha > beta:
                break
        return column, value
    else:  # minimizing player
        value = math.inf
        column = random.choice(valid_locations)
        for col in valid_locations:
            row = get_valid_row(board, col)
            b_copy = board.copy()
            drop_piece(b_copy, row, col, PLAYER_PIECE)
            new_score = minimax(b_copy, depth - 1, True, alpha, beta)[1]
            if new_score < value:
                value = new_score
                column = col
            beta = min(beta,value)
            if beta <= alpha:
                break
        return column, value

def draw_board(board):
    for c in range(COLUMN_COUNT):
        for r in range(ROW_COUNT):
            # rect(Surface, color, Rect,width = 0)
            pygame.draw.rect(screen, BLUE, (c * SQUARE_SIZE, r * SQUARE_SIZE + SQUARE_SIZE, SQUARE_SIZE, SQUARE_SIZE))
            # circle(Surface, color, pos, radius, width=0)
            pygame.draw.circle(screen,
                               BLACK,
                               (int(c * SQUARE_SIZE + SQUARE_SIZE / 2),
                                int(r * SQUARE_SIZE + SQUARE_SIZE + SQUARE_SIZE / 2)),
                               RADIUS)
    for c in range(COLUMN_COUNT):
        for r in range(ROW_COUNT):
            if board[r][c] == PLAYER_PIECE:
                pygame.draw.circle(
                    screen,
                    RED,
                    (int(c * SQUARE_SIZE + SQUARE_SIZE / 2),
                     height - int(r * SQUARE_SIZE + SQUARE_SIZE / 2)),
                    RADIUS)
            elif board[r][c] == AI_PIECE:
                pygame.draw.circle(
                    screen,
                    YELLOW,
                    (int(c * SQUARE_SIZE + SQUARE_SIZE / 2),
                     height - int(r * SQUARE_SIZE + SQUARE_SIZE / 2)),
                    RADIUS)
    pygame.display.update()

board = create_board()
print_board(board)
draw_board(board)
pygame.display.update()

my_font = pygame.font.SysFont("monospace", 75)

game_over = False
turn = random.randint(PLAYER, SKY_NET)

while not game_over:

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            sys.exit()

        if event.type == pygame.MOUSEMOTION:
            pygame.draw.rect(screen, BLACK, (0, 0, width, SQUARE_SIZE))
            pos_x = event.pos[0]
            if turn == PLAYER:
                pygame.draw.circle(screen, RED, (pos_x, int(SQUARE_SIZE / 2)), RADIUS)
        pygame.display.update()

        # MouseButtonDown give us position
        if event.type == pygame.MOUSEBUTTONDOWN:
            pygame.draw.rect(screen, BLACK, (0, 0, width, SQUARE_SIZE))
            # Ask for PLayer 1 input
            if turn == PLAYER:
                pos_x = event.pos[0]
                col = int(math.floor(pos_x / SQUARE_SIZE))

                if is_valid_location(board, col):
                    row = get_valid_row(board, col)
                    drop_piece(board, row, col, PLAYER_PIECE)

                    if wining_move(board, PLAYER_PIECE):
                        label = my_font.render("Player 1 wins", 1, RED)
                        screen.blit(label, (40, 10))
                        print("Player 1 wins")
                        game_over = True

                    turn += 1
                    turn = turn % 2

    # Ask for PLayer 2 input
    if turn == SKY_NET and not game_over:

        #col = random.randint(0, COLUMN_COUNT - 1)
        #col = pick_best_move(board, AI_PIECE)
        col, minimax_score = minimax(board, 4, True, -math.inf, math.inf)
        if is_valid_location(board, col):
            # pygame.time.wait(500)
            row = get_valid_row(board, col)
            drop_piece(board, row, col, AI_PIECE)

            if wining_move(board, AI_PIECE):
                label = my_font.render("Player 2 wins", 7, YELLOW)
                screen.blit(label, (40, 10))
                print("Player 2 wins")
                game_over = True

            turn += 1
            turn = turn % 2

            print_board(board)
            draw_board(board)

    if game_over:
        pygame.time.wait(3000)


import os
from os import path, listdir
from subprocess import Popen, PIPE
import sys
from datetime import datetime
import yaml
import re
from timeout import timeout

stream = open("data.yaml", "r")
doc = yaml.load(stream)
directory = doc["directory_of_homework"] + sys.argv[1] + "/" + sys.argv[2] + "/" + sys.argv[3]

if sys.argv[4] == "all":
	stream = open("scenarios_" + sys.argv[3] + ".yaml", "r")
	scenarios = yaml.load(stream)
	tasks = scenarios["tasks"]
else:
	tasks = ["task{0}.c".format(sys.argv[4])]
	number_of_task = re.findall(r'\d', sys.argv[4])[0]
	scenarios = {"task{0}".format(sys.argv[4]): { "number_of_tests": 1, "input{0}".format(number_of_task): sys.argv[5], "output{0}".format(number_of_task): sys.argv[6] } }

summary = {}

@timeout(3)
def check_homework(tasks, scenarios):
	num_of_hw = 0
	print("Evaluating homework {0} on student {1} from {2} class\n".format(sys.argv[3], sys.argv[2], sys.argv[1]))

	files = []
	result = {}
	for root, dirs, files in os.walk(directory, topdown=False):
		files = [f for f in listdir(root) if (path.isfile(path.join(root, f))
				 and f in tasks)]
		if not files:
			continue

		for current_file in files:
			print("\nEvalutaing {0}".format(current_file))
			abs_path = path.abspath(path.join(directory, current_file))
			exec_path = path.abspath(path.join(directory, "a.out"))

			gpp_invoke = "gcc -Wall -std=c11 {0} -o {1} 2> homeworks_result.txt".format(abs_path, exec_path)

			result = os.system(gpp_invoke)

			f = open('homeworks_result.txt', 'r')
			result_of_compilation = f.read()

			successful = True
			scenario = current_file.split('.')[0]

			if result is not 0:
				print("\tFiles can't compile (gcc err code:{0})".format(result))
				summary[scenario] = "0/{0} - error".format(scenarios[scenario]["number_of_tests"])
				continue

			if result_of_compilation != "":
				print("\tWarning:\n \t\t{0}".format(result_of_compilation))
				summary[scenario] = "0/{0} - warning".format(scenarios[scenario]["number_of_tests"])
				continue

			if scenario in scenarios.keys():

				actual_successful_test = 0

				for i in range(scenarios[scenario]["number_of_tests"]):
					p = Popen([exec_path], stdout=PIPE, stderr=PIPE, stdin=PIPE)
					input_data = scenarios[scenario]["input{0}".format(i+1)]
					std_out_data, _ = p.communicate(input_data.encode())

					output = std_out_data.decode('latin-1').rstrip('\n').replace('\0', '').strip()
					output = output.replace('\n', ' ')
					sys.stdout.write("\t")

					if output == scenarios[scenario]["output{0}".format(i+1)]:
						sys.stdout.write(" ".join("Test {0} passed".format(i+1)
						   .split()))
						actual_successful_test += 1
					else:
						sys.stdout.write("Test {0} failed:\n".format(i+1))
						sys.stdout.write("\t\tExpected:\n\t\t\t{0}\n".format(scenarios[scenario]["output{0}".format(i+1)]))
						sys.stdout.write("\t\tGot: \n \t\t\t{0}\n".format(output))
						successful = False
					sys.stdout.write("\n")

					summary[scenario] = "{0}/{1}".format(actual_successful_test, scenarios[scenario]["number_of_tests"])
				if successful:
					num_of_hw += 1

	print("\nSummary\n")
	if summary.keys():
		for task in summary.keys():
			print("\t{0} - {1}".format(task, summary[task]))
	else:
		print("\tMissing homework\n")
	return num_of_hw

check_homework(tasks, scenarios)


from django.conf.urls.defaults import *
from django.conf import settings
from django.contrib import admin
from django.http import HttpResponse
admin.autodiscover()

urlpatterns = patterns('',
    url(r'^admin/', include(admin.site.urls)),
    url(r'^$', 'tips.views.index'),
    url(r'^robots\.txt$', lambda r: HttpResponse("User-agent: *\nAllow: /", mimetype="text/plain")),
)

if settings.DEBUG:
	urlpatterns += patterns('',
		(r'^static/(?P<path>.*)$', 'django.views.static.serve', {
			'document_root': settings.STATIC_ROOT,
		}),
		url(r'^media/(?P<path>.*)$', 'django.views.static.serve', {
            'document_root': settings.MEDIA_ROOT,
        }),
	)


import yaml
import os

class AbsoluteConfigFactory(object):
    def __init__(self, *secret_files, keep_secrets=False):
        data = ""
        for file in secret_files:
            with open(file, 'r') as fin:
                    data += fin.read() + '\n'
        self.secrets = data
        loaded = yaml.load(data)
        self.secret_keys = [] if keep_secrets or loaded is None else loaded.keys()

    def get(self, file):
        if not hasattr(self, file):
            with open(file, 'r') as fin:
                data = fin.read()
            decoded_yaml = yaml.load(self.secrets + data)
            for key in self.secret_keys:
                decoded_yaml.pop(key, None)
            setattr(self, file, decoded_yaml)
        return getattr(self, file)

class ConfigFactory(AbsoluteConfigFactory):
    def __init__(self, basedir, *secret_files, keep_secrets=False, secrets_absolute_path=False):
        self.basedir = basedir
        if not secrets_absolute_path:
            secret_files = [self.format_path(file) for file in secret_files]
        super().__init__(*secret_files, keep_secrets=keep_secrets)

    def get(self, file):
        file = self.format_path(file)
        return super().get(file)

    def format_path(self, file):
        return os.path.join(self.basedir, file if os.path.splitext(file)[1] is not '' else os.path.splitext(file)[0] + ".yaml")


import requests
import re
import json
import pickle
import gc
import time

src = 'http://www.baidu.com/s?wd='

pattern1 = re.compile('style=\"text-decoration:none;\">(.*?) </a>')
pattern2 = re.compile('<!--STATUS (.*?)-->')

s_t = 60001
e_d = 80000

# with open('../src/test.json', 'r') as f:
with open('../src/train.json', 'r') as f:
    json_line_t = 0
    urls = []
    for line in f:
        text = json.loads(line)

        json_line_t += 1

        if json_line_t < s_t:
            continue
        print(json_line_t)

        search_content = text['title']

        tmp_src = src + "\"" + search_content + "\""

        test_flag = True
        while True:
            try:
                content = requests.get(url=tmp_src).text
            except:
                print('contect error')
                time.sleep(60)
                content = requests.get(url=tmp_src).text
            status_pd = re.findall(pattern2, content)
            if status_pd[0] == 'OK':
                tmp_ans = re.findall(pattern1, content)
                if len(tmp_ans) == 0:
                    if test_flag == True:
                        print('get none')
                        time.sleep(10)
                        test_flag = False
                        continue
                if test_flag == False:
                    print('new get', len(tmp_ans))
                urls.append(tmp_ans)
                break
            else:
                print(json_line_t, 'Failed!')
                time.sleep(60)

        if json_line_t % 10000 == 0:
            with open('../src/train_urls/'+str(int(json_line_t/10000))+'.pkl', 'wb') as fw:
                pickle.dump(urls, fw, -1)
            del urls
            gc.collect()
            urls = []

        if json_line_t == e_d:
            break


import numpy as np
import sklearn.metrics.pairwise

import evm_enabler
import general_model as gm

class EVMOptions():
    """Options class for an EVM classifier

    The options class allows us to set custom options for each classifier
    and still pass these instructions to a General Model.  For more information
    about the EVM visit the following: https://arxiv.org/pdf/1506.06112.pdf,
    https://github.com/EMRResearch/ExtremeValueMachine

    Attributes:
        threshold (double): threshold used to call something novel
        tailsize (int): the number of extreme vectors allowed for the weibull distributions
        cover_threshold (double): probabilistic threshold to designate redundancy between points. Paper sets at 0.5
        for all experiments
        dist_func (string): name of function used to calculate distance (value: 'cosine' or 'euclidean')
        num_to_fuse (int): when picking y* need to find max probability.  Num_to_fuse looks at average max prob for
        top 'num_to_fuse'. Paper found increase of 1-2 percentage points by tuning this.
        margin_scale (double): Default 0.5 where to put margin between points

    """

    def __init__(self, tailsize, cover_threshold, dist_func, num_to_fuse, margin_scale=0.5):
        self.tailsize = tailsize
        self.cover_threshold = cover_threshold
        self.dist_func = dist_func
        self.num_to_fuse = num_to_fuse
        self.margin_scale = margin_scale

class ExtremeValueMachine(gm.General_Model):
    """Class for the flat EVM classifier

    This class implements the General Model Abstract class.
    The result is an EVM classifier that
    has no hierarchy.  It has the options as defined in
    the options object.  The init builds and compiles
    the model.

    Attributes:
        name (string): the name of the classifier
        options (object): options object for the classifier
    """

    def __init__(self, name, options):

        self.name = name
        self.trained = False

        ##
        ## EVM parameters
        ##
        self.tailsize = options.tailsize  ## <- number of extreme vectors allowed
        self.cover_threshold = options.cover_threshold
        self.dist_func = options.dist_func
        self.num_to_fuse = options.num_to_fuse
        self.margin_scale = options.margin_scale

        ##
        ## Items for the model
        ##
        self.points = None
        self.weibulls = None
        self.labels = None

        ##
        ## Have to handle distances somewhere
        ##
        def euclidean_cdist(X, Y):
            return sklearn.metrics.pairwise.pairwise_distances(X, Y, metric="euclidean", n_jobs=1)

        def euclidean_pdist(X):
            return sklearn.metrics.pairwise.pairwise_distances(X, metric="euclidean", n_jobs=1)

        def cosine_cdist(X, Y):
            return sklearn.metrics.pairwise.pairwise_distances(X, Y, metric="cosine", n_jobs=1)

        def cosine_pdist(X):
            return sklearn.metrics.pairwise.pairwise_distances(X, metric="cosine", n_jobs=1)

        dist_func_lookup = {
            "cosine": {"cdist": cosine_cdist,
                       "pdist": cosine_pdist},

            "euclidean": {"cdist": euclidean_cdist,
                          "pdist": euclidean_pdist}
        }

        self.cdist_func = dist_func_lookup[self.dist_func]["cdist"]
        self.pdist_func = dist_func_lookup[self.dist_func]["pdist"]

    def train(self, x, y):

        """Train the the EVM classifier

        This function trains the EVM classifier.

        Args:
            train (list): [x_train, y_train] numpy arrays (y_train NOT one-hot)
            val (list): [x_val, y_val] numpy arrays (y_val NOT one-hot)
            options (object): options object for the classifier

        """
        ##
        ## Separate train and test
        ##

        x_train = x.astype('double')
        y_train = np.argmax(y, axis = 1)

        ##
        ## Fit the model (the model ends up being points, weibulls and labels)
        ## See https://arxiv.org/pdf/1506.06112.pdf for details
        ##
        self.weibulls = evm_enabler.fit(x_train, y_train, self.tailsize, self.pdist_func, self.margin_scale)
        self.points, self.weibulls, self.labels = evm_enabler.reduce_model(x_train, self.weibulls, y_train,
                                                                           self.cover_threshold, self.pdist_func)
        self.trained = True

    def confidence(self, pool):

        """Report confidence of predictions of the pool

        Args:
            pool (numpy.ndarray): instances to predict

        Returns
            results: array of prediction confidences for the pool
        """
        if (self.trained):
            predictions, probs = evm_enabler.predict_evm(pool, self.points, self.weibulls, self.labels, self.cdist_func,
                                                         self.num_to_fuse)  ## Predict the pool

            return np.transpose(probs)
        else:
            print('Not trained! (prediction)')


import numpy as np
from sklearn.neighbors import KNeighborsClassifier
from sklearn.cross_validation import train_test_split, cross_val_score
from sklearn import svm
from sklearn.decomposition import PCA

def load_data(filename):
    return np.genfromtxt(filename, delimiter=',')

def perform_pca(Xs, n_components):
    # FIXME: TOTALLY WRONG! use fit() to train and then
    # transform
    print("Performing PCA with %s components" % (str(n_components),))
    pca = PCA(n_components=n_components)
    return [pca.fit_transform(X) for X in Xs]

def find_best_svm(X, y):
    print("Searching SVMs")
    scores = []
    for kernel in ['linear', 'rbf', 'sigmoid']:
        print(kernel)
        # For some reason, poly needed really long or
        # did even get stuck, so removed it
        clf = svm.SVC(kernel=kernel, C=1.0)
        score = cross_val_score(clf, X, y, cv=5)
        print("%s: %0.2f" % (kernel, score.mean()))

        scores.append((score.mean(), kernel))

    return max(scores)

def semi_supervised_set(clf, X, y, X_test, thresh=0.9):
    """
    Create a semi-supervised training set
    using estimations from classifier clf which will be
    trained on X and y before, but only take
    guesses where the probability is at least thresh
    """
    clf.fit(X, y)
    y_test = clf.predict(X_test)

    new_X = np.zeros((X.shape[0] + X_test.shape[0], X.shape[1]))
    new_y = np.zeros(X.shape[0] + X_test.shape[0])
    new_X[0:X.shape[0], :] = X
    new_y[0:X.shape[0]] = y

    added_count = 0
    for i, probs in enumerate(clf.predict_proba(X_test)):
        if probs[y_test[i]] > thresh:
            new_X[X.shape[0] + added_count, :] = X_test[i, :]
            new_y[X.shape[0] + added_count] = y_test[i]

            added_count += 1

    new_X.resize((X.shape[0] + added_count, X.shape[1]))
    new_y.resize(X.shape[0] + added_count)

    return new_X, new_y

def use_classifier(clf, X, y, X_test, filename):
    """
    Perform classification onto test data and save the
    result into filename
    """
    clf.fit(X, y)

    res = clf.predict(X_test)
    output = np.zeros((res.shape[0], 2), dtype='int')
    output[:, 0] = np.arange(1, res.shape[0] + 1)
    output[:, 1] = res
    np.savetxt(filename, output, fmt='%d', delimiter=',',
               header='Id,Solution', comments='')

X = load_data('train.csv')
y = load_data('trainLabels.csv')
X_test = load_data('test.csv')

scores_svm = []
for components in range(1, 2):  # removed mle, due to bad result
    #X_pca, X_test_pca = perform_pca((X, X_test), components)

    #score, kernel = find_best_svm(X_pca, y)
    score, kernel = find_best_svm(X, y)
    scores_svm.append((round(score, 2), components, kernel))

score, components, kernel = max(scores_svm, key=lambda x: (x[0], -x[1]))
print("Using %s kernel with %s components due to score %0.2f" %
      (kernel, str(components), score))
#X_pca, X_test_pca = perform_pca((X, X_test), components)
clf = svm.SVC(kernel=kernel, C=1.0, probability=True)
for i in range(2):
    X, y = semi_supervised_set(clf, X, y, X_test, 0.9)
clf = svm.SVC(kernel=kernel, C=1.0, probability=True)
#use_classifier(clf, X_pca, y, X_test_pca, 'submission_svm.csv')
use_classifier(clf, X, y, X_test, 'submission_svm_semi.csv')


"""
Interactive syllable segmentation

1) Relevant audio?
2) Boundary?
3) Good clusters?
"""

__author__ = "Jack Goffinet"
__date__ = "March-April 2019"

import os
import numpy as np
import h5py
import joblib
from tqdm import tqdm
from skimage.transform import resize, downscale_local_mean
from scipy.io import wavfile, loadmat
from scipy.signal import stft, welch

from time import sleep

# import matplotlib
# matplotlib.use('qt4agg')
# matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
from matplotlib.widgets import RectangleSelector
# plt.switch_backend('agg')

import umap

from bokeh.plotting import figure, output_file, show, ColumnDataSource
from bokeh.models import HoverTool
from bokeh.models.glyphs import ImageURL

from .preprocessing import time_from_filename

FS = 44100
SPEC_THRESH = 1.0 # -8.5 homer, 4.0 mouse, -9 marmoset, -3.5 zebra finch
EPSILON = 1e-9

DOWNSAMPLE_TUPLE = (1,1) # marmoset (1,2)

MIN_FREQ, MAX_FREQ = 300, 12e3 # 40k-110k mouse, 350-13k marmoset

PSD_DUR, SPEC_DUR = 0.5, 0.5 # 0.1 mouse, 0.5 marmoset, 0.5 zebra finch

CENTER_T = 0.5 # how far into the song are we looking?
SONG_DURATION = 2.0 # how long is the song?

X1, Y1, X2, Y2 = 0, 0, 0, 0
LINES = []

def segment(load_dirs):
	""" Main method """
	global X1, Y1, X2, Y2
	# Collect filenames.
	filenames = []
	for load_dir in load_dirs:
		filenames += [os.path.join(load_dir,i) for i in os.listdir(load_dir) if i[-4:] in ['.wav', '.mat']]

	"""
	run_num = 0
	while True:
		run_num += 1
		# Collect PSDs and spectrograms.
		psds, specs = [], []
		np.random.seed(42)
		psd_f_ind, spec_f_ind = None, None
		print("Collecting data...")

		# from multiprocessing import Pool
		# from itertools import repeat
		# with Pool(min(3, os.cpu_count()-1)) as pool:
		# 	pool.starmap(get_rand_psds_specs, zip())
		# # quit()

		for i in tqdm(range(5000)): #10000

			flag = True
			while flag:
				filename = filenames[np.random.randint(len(filenames))]
				psd, spec, psd_f_ind, spec_f_ind = get_rand_psd_spec(filename, psd_dur=PSD_DUR, spec_dur=SPEC_DUR, psd_f_ind=psd_f_ind, spec_f_ind=spec_f_ind)
				if run_num > 1:
					flag = not psd_in_region(spec, transform, boundaries)
				else:
					flag = False
			# print("psd", psd.shape)
			# print("spec", spec.shape)
			# print(np.min(spec), np.max(spec))
			# plt.imshow(spec, aspect='auto', origin='lower')
			# plt.savefig('temp.pdf')
			# plt.close('all')
			# sleep(2)
			# if i == 10:
			# 	quit()
			psds.append(psd)
			specs.append(spec)
		psds = np.array(psds)
		np.save('psds.npy', psds)
		specs = np.array(specs)
		np.save('specs.npy', specs)
		# UMAP the PSDs.
		print("Dimensionality reduction...")
		transform = umap.UMAP(n_components=2, n_neighbors=20, min_dist=0.1, metric='euclidean', random_state=42)
		embedding = transform.fit_transform(specs.reshape(len(specs), -1))
		joblib.dump(transform, 'transform_'+str(run_num)+'.sav')
		perm = np.random.permutation(len(embedding))
		embedding, specs = embedding[perm], specs[perm]
		np.save('embedding.npy', embedding)

		# Make an HTML mouse-over plot.
		print("Making HTML...")
		make_html_plot(embedding, specs, output_dir='temp2', num_imgs=400)
		embedding = np.load('embedding.npy')

		# Non-interactive:
		plt.title("Select relevant audio")
		plt.scatter(embedding[:,0], embedding[:,1], c='b', alpha=0.1, s=0.7)
		plt.savefig('temp.pdf')
		plt.close('all')
		x1 = float(input("x1: "))
		x2 = float(input("x2: "))
		y1 = float(input("y1: "))
		y2 = float(input("y2: "))
		x1s, x2s, y1s, y2s = [x1], [x2], [y1], [y2]

		boundaries = {'x1s':x1s, 'y1s':y1s, 'x2s':x2s, 'y2s':y2s}
		np.save('boundaries.npy', boundaries)

		if input("q for quit ") == 'q':
			quit()

	"""

	"""
	# Plot an interactive selcting tool.
	print("Interactive plot...")
	x1s, y1s, x2s, y2s = [], [], [], []
	colors = np.array(['b'] * len(embedding))
	X1 = 0
	while X1 is not None:
		print("here")
		X1 = None
		for i, embed in enumerate(embedding):
			if colors[i] == 'r':
				continue
			for x1, y1, x2, y2 in zip(x1s, y1s, x2s, y2s):
				if embed[0] < max(x1, x2) and embed[0] > min(x1, x2) and \
						embed[1] < max(y1, y2) and embed[1] > min(y1, y2):
					colors[i] = 'r'
		fig, current_ax = plt.subplots()
		toggle_selector.RS = RectangleSelector(current_ax, line_select_callback,
				drawtype='box', useblit=True,
				button=[1, 3],  # don't use middle button
				minspanx=5, minspany=5,
				spancoords='pixels',
				interactive=True)
		plt.title("Select relevant audio")
		plt.connect('key_press_event', toggle_selector)
		print("making plot 1")
		plt.scatter(embedding[:,0], embedding[:,1], c=colors, alpha=0.1, s=0.7)
		plt.show()
		sleep(20)
		print("making plot 2")
		if X1 is not None:
			x1s.append(X1)
			y1s.append(Y1)
			x2s.append(X2)
			y2s.append(Y2)
		"""

	"""
	boundaries = {'x1s':x1s, 'y1s':y1s, 'x2s':x2s, 'y2s':y2s}
	np.save('boundaries.npy', boundaries)
	"""

	"""
	PART II
	-------
	"""

	"""
	transform = joblib.load('transform_'+str(run_num)+'.sav')
	boundaries = np.load('boundaries.npy').item()

	# Find segmentation boundaries.
	specs = []
	np.random.seed(42)
	psd_f_ind, spec_f_ind = None, None
	print("Collecting data...")

	# # Collect non-random spectrograms.
	# delta_t = 0.025
	# fs, audio = wavfile.read('001.wav')
	# t = 0.0
	# f_ind = None
	# specs = []
	# print("len(audio)", len(audio)/fs)
	# while t + SPEC_DUR + 0.01 < min(len(audio)/fs, 2000):
	# 	i1, i2 = int(fs*t), int(fs*(t+SPEC_DUR))
	# 	spec, f_ind = get_spec(audio[i1:i2], f_ind=f_ind, min_freq=MIN_FREQ, max_freq=MAX_FREQ)
	# 	specs.append(spec)
	# 	t += delta_t
	# print("f_ind", f_ind)

	# Collect random spectrograms.
	for i in tqdm(range(5000)): # 100000
		flag = True
		while flag:
			filename = filenames[np.random.randint(len(filenames))]
			psd, spec, psd_f_ind, spec_f_ind = get_rand_psd_spec(filename, psd_dur=PSD_DUR, spec_dur=SPEC_DUR, psd_f_ind=psd_f_ind, spec_f_ind=spec_f_ind)
			flag = not psd_in_region(spec, transform, boundaries)
		specs.append(spec)
		if (i + 1) % 500 == 0:
			np.save('specs.npy', np.array(specs))
	print("spec_f_ind:", spec_f_ind)

	specs = np.array(specs)
	np.save('specs.npy', specs)

	print('specs', specs.shape)
	# weighted_specs = np.copy(specs)
	# for i in range(len(specs)):
	# 	for j in range(len(specs[i])):
	# 		weighted_specs[i,j,:specs.shape[2]//2] *= np.linspace(0,1,specs.shape[2]//2)
	# 		weighted_specs[i,j,specs.shape[2]//2:] *= np.linspace(1,0,specs.shape[2]-specs.shape[2]//2)
	print("Dimensionality reduction...")
	transform_final = umap.UMAP(n_components=2, n_neighbors=20, min_dist=0.1, metric='euclidean', random_state=42)
	print("specs shape:", specs.shape)
	embedding = transform_final.fit_transform(specs.reshape(len(specs), -1))
	joblib.dump(transform_final, 'transform_2.sav')
	perm = np.random.permutation(len(embedding))
	embedding, specs = embedding[perm], specs[perm]
	np.save('embedding_2.npy', embedding)

	# Make an HTML mouse-over plot.
	print("Making HTML...")
	marked_specs = get_marked_specs(specs)
	make_html_plot(embedding, marked_specs, output_dir='temp3')

	# Non-interactive:
	plt.scatter(embedding[:,0], embedding[:,1], c='b', alpha=0.1, s=0.7)
	plt.savefig('temp.pdf')
	plt.close('all')
	center_x = float(input("center x: "))
	center_y = float(input("center y: "))
	radius = float(input("radius: "))
	radii = [radius]
	centers = [[center_x, center_y]]
	"""

	"""
	# Plot an interactive selcting tool.
	print("Interactive plot...")
	x1s, y1s, x2s, y2s = [], [], [], []
	radii = []
	centers = []
	X1 = 0
	while X1 is not None:
		X1 = None
		fig, current_ax = plt.subplots()
		toggle_selector.RS = RectangleSelector(current_ax, line_select_callback,
				drawtype='line', useblit=True,
				button=[1, 3],  # don't use middle button
				minspanx=0, minspany=0,
				spancoords='pixels',
				interactive=True)
		plt.title("Select relevant audio")
		plt.connect('key_press_event', toggle_selector)
		plt.scatter(embedding[:,0], embedding[:,1], c='b', alpha=0.1, s=0.7)
		for x1, y1, x2, y2, center, radius in zip(x1s, y1s, x2s, y2s, centers, radii):
			plt.plot([x1,x2], [y1,y2], c='r')
			circle = plt.Circle(tuple(center), radius, color='r', fill=False)
			current_ax.add_artist(circle)
		plt.show()
		if X1 is not None:
			x1s.append(X1)
			y1s.append(Y1)
			x2s.append(X2)
			y2s.append(Y2)
			centers.append([(X1+X2)/2,(Y1+Y2)/2])
			radii.append(0.5 * ((X1-X2)**2+(Y1-Y2)**2)**0.5)
	"""

	# TEMP!
	centers = [[12.5,7]]
	radii = [2]

	seg_stuff = {'radii':radii, 'centers':np.array(centers)}
	np.save('seg_stuff.npy', seg_stuff)

	# Parse through all the files, collecting
	seg_stuff = np.load('seg_stuff.npy').item()
	radius = seg_stuff['radii'][0]
	center = seg_stuff['centers'][0]

	transform = joblib.load('transform_2.sav')

	# Segment a file.
	segment_audio('data/raw/bird_data/79', 'temp_save', transform, center, radius)
	quit()
	fs, audio = wavfile.read('data/raw/bird_data/79/red291_43566.27672737_4_11_7_41_12.wav')
	print(filenames[0])
	temp = segment_audio_chunk(audio, transform, center, radius)
	print(temp)
	quit()

def segment_audio(load_dir, save_dir, transform, center, radius):
	"""Save all the song bouts."""
	load_filenames = [os.path.join(load_dir,i) for i in os.listdir(load_dir) if i[-4:] in ['.wav', '.mat']]
	# load_filenames = ['data/raw/bird_data/79/red291_43566.27672737_4_11_7_41_12.wav']
	audio_segs, filenames, file_times, times = [], [], [], []
	num_audio_frames = int(SONG_DURATION*FS)
	write_file_num = 0
	songs_per_file = 10
	max_num_songs = 2000
	for load_filename in load_filenames:
		print(load_filename, len(audio_segs))
		filename_time = time_from_filename(load_filename)
		start_t = 0.0
		center_ts = []
		audio = get_audio(load_filename)
		while True:
			result = segment_audio_chunk(audio[int(FS*start_t):], transform, center, radius)
			if result == -1:
				continue
			if result is None:
				break
			center_ts.append(result + start_t)
			start_t += result + SONG_DURATION / 2
		for center_t in center_ts:
			t1 = int((center_t - CENTER_T)*FS)
			t2 = t1 + num_audio_frames
			audio_segs.append(audio[t1:t2])
			filenames.append(load_filename)
			file_times.append(center_t - CENTER_T)
			times.append(filename_time + center_t - CENTER_T)
		while len(audio_segs) >= songs_per_file:
			save_filename = os.path.join(save_dir, str(write_file_num).zfill(2)+'.hdf5')
			wavfile.write(os.path.join(save_dir, 'temp.wav'), FS, np.concatenate(audio_segs, axis=0))
			with h5py.File(save_filename, "w") as f:
				f.create_dataset('audio', data=np.array(audio_segs[:songs_per_file]))
				f.create_dataset('filename', data=np.array(filenames[:songs_per_file]).astype('S'))
				f.create_dataset('file_time', data=np.array(file_times[:songs_per_file]))
				f.create_dataset('time', data=np.array(times[:songs_per_file]))
			audio_segs = audio_segs[songs_per_file:]
			filenames = filenames[songs_per_file:]
			file_times = file_times[songs_per_file:]
			times = times[songs_per_file:]
			write_file_num += 1
			if write_file_num * songs_per_file >= max_num_songs:
				return

def segment_audio_chunk(audio, transform, center, radius):
	r2 = radius*radius
	f_ind = None
	STEP = 0.01
	SMALL_STEP = 0.002
	t = CENTER_T
	while t + SPEC_DUR + 0.01 < len(audio)/FS:
		embed, f_ind = get_embed_from_time(audio, t, f_ind, transform)
		current_r2 = np.sum(np.power(embed-center,2))
		# print(t, current_r2)
		if current_r2 > r2:
			t += STEP
			continue
		# Step until radius increases.
		prev_r2 = current_r2
		while True:
			t += SMALL_STEP
			embed, f_ind = get_embed_from_time(audio, t, f_ind, transform)
			current_r2 = np.sum(np.power(embed-center,2))
			# print(t, current_r2, "inside loop")
			if current_r2 > prev_r2:
				t -= SMALL_STEP
				# print(t)
				# print(CENTER_T)
				# print(t-CENTER_T+SONG_DURATION)
				# print(len(audio)/FS)
				if t - CENTER_T + SONG_DURATION < len(audio)/FS:
					return t
				return None
			prev_r2 = current_r2
	return None

def get_embed_from_time(audio, t, f_ind, transform):
	i1, i2 = int(t*FS), int((t+SPEC_DUR)*FS)
	spec, f_ind = get_spec(audio[i1:i2], f_ind=f_ind)
	return transform.transform(spec.reshape(1,-1)).flatten(), f_ind

def get_rand_psd_spec(filename, psd_dur=0.2, spec_dur=0.3, psd_f_ind=None, \
		spec_f_ind=None, return_middle_sample=False):
	audio = get_audio(filename)
	psd_samples, spec_samples = int(psd_dur * FS), int(spec_dur * FS)
	max_samples = max(psd_samples, spec_samples)
	start = np.random.randint(len(audio) - max_samples)
	delta = (max_samples - psd_samples) // 2
	psd, psd_f_ind = get_psd(audio[start+delta:start+delta+psd_samples], f_ind=psd_f_ind)
	delta = (max_samples - spec_samples) // 2
	spec, spec_f_ind = get_spec(audio[start+delta:start+delta+spec_samples], f_ind=spec_f_ind)
	if return_middle_sample:
		return psd, spec, psd_f_ind, spec_f_ind, start+psd_samples//2
	return psd, spec, psd_f_ind, spec_f_ind

def get_rand_segmented_spec(filename, lines, transform, boundaries, \
			psd_dur=0.2, spec_dur=0.3, psd_spec_ind=None, spec_f_ind=None, \
			max_syll_dur=1.5):
	"""

	Returns
	-------
	spec : numpy array or None
		spectrogram of a single syllable or None if none is found.
	"""
	# Get a psd
	psd, _, psd_f_ind, spec_d_ind, mid = get_rand_psd_spec(filename, \
			psd_dur=psd_dur, spec_dur=spec_dur, psd_spec_ind=psd_spec_ind, \
			spec_f_ind=spec_f_ind, return_middle_sample=True)
	# If there's no nearby birdsong, give up.
	if not psd_in_region(psd, transform, boundaries):
		return None
	# Otherwise, find the nearest syllable.
	max_syll_samples = int(max_syll_dur * FS)
	# First search left, looking for local minima in distances to segmenting lines.
	distances = {} # NOTE: TEMP to see how stable these are.
	for i in range(max(0,mid-max_syll_samples), min(len(audio)-1, mid+max_syll_samples), 1):
			distances[i] = get_distance(filename) # NOTE: HERE

def get_spec(audio, min_freq=MIN_FREQ, max_freq=MAX_FREQ, f_ind=None, fs=FS):
	"""Get a spectrogram."""
	# Get the spectrogram.
	f, t, Zxx = stft(audio, fs=fs)
	# Cut out the frequencies we don't care about.
	if f_ind is None:
		f_ind = np.searchsorted(f, [min_freq, max_freq])
	f = f[f_ind[0]:f_ind[1]]
	Zxx = Zxx[f_ind[0]:f_ind[1]]
	# Set a threshold & normalize.
	Zxx = np.log((np.abs(Zxx) + EPSILON))
	Zxx -= SPEC_THRESH
	Zxx[Zxx < 0.0] = 0.0
	Zxx /= np.max(Zxx) + EPSILON
	Zxx = downscale_local_mean(Zxx, DOWNSAMPLE_TUPLE)
	return Zxx, f_ind

def get_psd(audio, min_freq=MIN_FREQ, max_freq=MAX_FREQ, f_ind=None):
	"""Get a power spectral density."""
	# Get the PSD.
	f, Pxx = welch(audio, fs=FS)
	# Cut out the frequencies we don't care about.
	if f_ind is None:
		f_ind = np.searchsorted(f, [min_freq, max_freq])
	f = f[f_ind[0]:f_ind[1]]
	Pxx = Pxx[f_ind[0]:f_ind[1]]
	return Pxx, f_ind

def make_html_plot(embedding, specs, output_dir='temp', num_imgs=2000, title=""):
	"""Make an HTML tooltip mouse-over plot."""
	# Write the tooltip images.
	if not os.path.exists(output_dir):
		os.makedirs(output_dir)
	num_imgs = min(len(specs), num_imgs)
	for i in tqdm(range(num_imgs)):
		save_image(specs[i], os.path.join(output_dir, str(i)+'.jpg'))
	output_file(os.path.join(output_dir,"main.html"))
	source = ColumnDataSource(
			data=dict(
				x=embedding[:num_imgs,0],
				y=embedding[:num_imgs,1],
				imgs = ['./'+str(i)+'.jpg' for i in range(num_imgs)],
			)
		)
	source2 = ColumnDataSource(
			data=dict(
				x=embedding[num_imgs:,0],
				y=embedding[num_imgs:,1],
			)
		)
	p = figure(plot_width=800, plot_height=600, title=title)
	p.scatter('x', 'y', size=3, fill_color='blue', fill_alpha=0.1, source=source2)
	tooltip_points = p.scatter('x', 'y', size=5, fill_color='red', source=source)
	hover = HoverTool(
			renderers=[tooltip_points],
			tooltips="""

			"""
		)
	p.add_tools(hover)
	p.title.align = "center"
	p.title.text_font_size = "25px"
	p.axis.visible = False
	p.xgrid.visible = False
	p.ygrid.visible = False
	show(p)

def save_image(data, filename):
	"""https://fengl.org/2014/07/09/matplotlib-savefig-without-borderframe/"""
	sizes = np.shape(data)
	height = float(sizes[0])
	width = float(sizes[1])
	fig = plt.figure()
	fig.set_size_inches(width/height, 1, forward=False)
	ax = plt.Axes(fig, [0., 0., 1., 1.])
	ax.set_axis_off()
	fig.add_axes(ax)
	ax.imshow(data, cmap='viridis', origin='lower')
	plt.savefig(filename, dpi=height)
	plt.close('all')

def line_select_callback(eclick, erelease):
	"""https://matplotlib.org/examples/widgets/rectangle_selector.html"""
	# eclick and erelease are the press and release events
	global X1, Y1, X2, Y2
	X1, Y1 = eclick.xdata, eclick.ydata
	X2, Y2 = erelease.xdata, erelease.ydata

def toggle_selector(event):
	"""https://matplotlib.org/examples/widgets/rectangle_selector.html"""
	print(' Key pressed.')
	if event.key in ['Q', 'q'] and toggle_selector.RS.active:
		print(' RectangleSelector deactivated.')
		toggle_selector.RS.set_active(False)
	if event.key in ['A', 'a'] and not toggle_selector.RS.active:
		print(' RectangleSelector activated.')
		toggle_selector.RS.set_active(True)

def get_audio(filename):
	"""Retrieve audio from a given file."""
	if filename[-4:] == '.wav':
		fs, audio = wavfile.read(filename)
	elif filename[-4:] == '.mat':
		temp = loadmat(filename)
		audio = temp['spike2Chunk'].reshape(-1)
		fs = temp['fs']
	else:
		raise NotImplementedError
	assert fs == FS, "found "+str(fs)+", required " + str(FS)
	return audio

def psd_in_region(psd, transform, b):
	"""Does this PSD contain relevant audio?"""
	embed = transform.transform(psd.reshape(1,-1)).flatten()
	for x1, y1, x2, y2 in zip(b['x1s'], b['y1s'], b['x2s'], b['y2s']):
		if embed[0]>min(x1,x2) and embed[0]<max(x1,x2) and embed[1]>min(y1,y2) and embed[1]<max(y1,y2):
			return True
	return False

def get_marked_specs(specs):
	"""Mark the center in time."""
	marked_specs = np.zeros(shape=(specs.shape[0], specs.shape[1]+4, specs.shape[2]))
	marked_specs[:,2:-2,:] = specs[:,:,:]
	marked_specs[:,:2,specs.shape[2]//2] = 1.2
	marked_specs[:,-2:,specs.shape[2]//2] = 1.2
	return marked_specs

def onclick(event):
	if event.button == 3:
		global LINES
		x = event.xdata
		y = event.ydata
		if len(LINES) == 0 or len(LINES[-1]) % 4 == 0:
			LINES.append([x,y])
		else:
			LINES[-1].append(x,y)
			x1, y1, x2, y2 = LINES[-1]
			plt.plot([x1,x1], [y1,y2], '-')
			radius = 0.5 * ((x1-x2)**2 + (y1-y2)**2)**0.5
			circle = plt.Circle((0.5*(x1+x2), 0.5*(y1+y2)), radius, color='r', alpha=0.3)
			plt.gca().add_artist(circle)
			plt.show()

if __name__ == '__main__':
	# load_dirs = ['human/']
	load_dirs = ['data/raw/bird_data/79']
	segment(load_dirs)

###


import cgi
from flask import render_template, abort, request
from jinja2 import TemplateNotFound
from twilio import twiml
from twilio.rest import TwilioRestClient
from app import app, redis_db, socketio

TWILIO_ACCOUNT_SID = app.config['TWILIO_ACCOUNT_SID']
TWILIO_AUTH_TOKEN = app.config['TWILIO_AUTH_TOKEN']
TWILIO_NUMBER = app.config['TWILIO_NUMBER']

client = TwilioRestClient(account=TWILIO_ACCOUNT_SID,
						  token=TWILIO_AUTH_TOKEN)

@app.route('/<presentation_name>/', methods=['GET'])
def landing(presentation_name):
    try:
        return render_template(presentation_name + '.html', telephone=TWILIO_NUMBER)
    except TemplateNotFound:
        abort(404)

@app.route('/presentation/twilio/webhook/', methods=['POST'])
def twilio_callback():
    to = request.form.get('To', '')
    from_ = request.form.get('From', '')
    message = request.form.get('Body', '').lower()
    if to == TWILIO_NUMBER:
        redis_db.incr(cgi.escape(message))
        socketio.emit('msg', {'div': cgi.escape(message),
                              'val': redis_db.get(message)},
                               namespace='/presentation')
    #resp = twiml.Response()
    #resp.message('Merci/Cheers.')
    return str(resp)


import logging
import sys

from modules import Configuration

logging.basicConfig(level=logging.DEBUG)

class SubstitutionDecipher:
    """
    Decipher the file being passed through.
    """

    def __init__(self, file):
        self.logger = logging.getLogger(__name__)

        self.configuration = Configuration()
        self.config = self.configuration.load()

        if file:
            self.file = file
        else:
            self.logger.critical(f'You have not provided a file to decipher.')
            sys.exit(1)

    def _get_unfixed_characters(self):
        """
        We need to know how many unfixed characters there are,
        so this function will load the config and generate a figure
        based on the content of the config file key 'fixed-cipher-text'.
        :return:
        """

        self.logger.debug('Calculating Amount Of Unfixed Characters')
        self.un_fixed_characters = []

        # As the key 'fixed-cipher-text' is a list containing dictionaries,
        # we need to loop over it to enable us to check each item and run a
        # comparative check.
        for item in self.config['fixed-cipher-text']:
            # Let's grab the index of the item in this list,
            # incase we need to update or change the item later.
            index = self.config['fixed-cipher-text'].index(item)
            # This next line is rather complex, but it
            # basically just outputs the key in the dictionary.
            key = [key for key in item.keys()][0]
            # We are specifically looking for a False value in the list,
            # hence the negative check.
            if not item[str(key)]:
                # If a match is found, stick it in a new list inside of a
                # dictionary with it's index, so we can adjust it later on.
                self.un_fixed_characters.append(dict(item=item, index=index))

        # Output how many items are in the list of
        # unfixed characters.
        self.logger.info(f'Found {len(self.un_fixed_characters)} Unfixed Characters.')
        return(len(self.un_fixed_characters))


from setuptools import setup, find_packages
import httpscreenshot

install_requires = [
    'selenium',
    'requests',
    'm2crypto',
    'pillow']

setup(name='httpscreenshot',
      version=httpscreenshot.__version__,
      description=httpscreenshot.__doc__.strip(),
      url='https://github.com/breenmachine/httpscreenshot',
      author=httpscreenshot.__author__,
      author_email='unknown',
      license=httpscreenshot.__license__,
      packages=find_packages(),
      install_requires=install_requires,
      entry_points={
          'console_scripts': [
              'httpscreenshot = httpscreenshot.__main__:main',
          ],
      },
      zip_safe=False)


# -*- coding: utf-8 -*-
import pdb
import time
import json
import sys
import os
import glob

# Author: Preston Law
# 12/01/2018

# This script counts the number of tweets in all_tweets.json and prints
# out the range of tweet dates.

# This looks for the last tweet collection file, finds the lastest tweet id, and finds all tweets after that.

def custom_load_tweets(fp, date_range=None):
    tweets = []
    for line in fp:
        tweet = json.loads(line)
        tweets.append(tweet)
        date_str = my_tweet['tweet_date']
        date_list = date_str.split('-')
        year = int(date_list[0])
        month = int(date_list[1])
        date_list = date_list[2].split('T')
        day = int(date_list[0])
        tmp = month * 100 + day
        if date_range is None:
            date_range = [tmp,tmp]
        else:
            date_range[0] = min(date_range[0], tmp)
            date_range[1] = max(date_range[1], tmp)

    print(date_range)
    return tweets, date_range

def get_number_of_tweets_in_file(filename, date_range):
    month_dict = {'Jan': 1, 'Feb':2, 'Mar':3, 'Apr':4, 'May':5, 'Jun':6, 'Jul':7, 'Aug':8, 'Sep':9, 'Oct':10, 'Nov':11, 'Dec':12}
    print(filename)

    with open(filename, 'r') as fp:
        tweets = json.load(fp)

        for tweet in tweets:
            date_str = tweet['tweet_date']
            date_list = date_str.split('-')
            year = int(date_list[0])
            month = int(date_list[1])
            date_list = date_list[2].split('T')
            day = int(date_list[0])
            tmp = month * 100 + day

            if date_range is None:
                date_range = [tmp,tmp]
            else:
                date_range[0] = min(date_range[0], tmp)
                date_range[1] = max(date_range[1], tmp)

    return len(tweets), date_range

date_range = None
total = 0
filename = 'all_tweets.json'
count, date_range = get_number_of_tweets_in_file(filename, date_range)
total += count

print(total)
print(date_range)


# PROTOSPACER PRECISION LASSO
# Tristan Henser-Brownhill, September, 2018
# For Chakrabarti et al.

# IMPORTS
import os
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from sklearn.linear_model import LassoCV, Lasso
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import mean_squared_error, mean_absolute_error, median_absolute_error
from sklearn.externals import joblib
from scipy.stats.stats import pearsonr, spearmanr, linregress, kendalltau # scipy v. 1.0.0
from sklearn.preprocessing import LabelEncoder, OneHotEncoder
from time import time
from scipy.stats import randint as sp_randint

#DATA PREPROCESSING

# import
data = pd.read_csv(os.path.join('data','random_train_set80pc.csv'))
data_test = pd.read_csv(os.path.join('data','random_test_set20pc.csv'))
#data_ob = pd.read_csv(os.path.join('data','ob_all48h.csv'))
data_ob = pd.read_csv(os.path.join('data','ob_HCT116_48h.csv'))
#data_ob = pd.read_csv(os.path.join('data','ob_HEK293_48h.csv'))
#data_ob = pd.read_csv(os.path.join('data','ob_K562_48h.csv'))

# prepare the data
df_train_x = data.iloc[:,2:] #extract pixels; all rows, all except sgRna name (0) and label (1)
y_train = data.iloc[:,1].values #extract labels
df_test_x = data_test.iloc[:,2:]
y_test = data_test.iloc[:,1].values
df_ob_x = data_ob.iloc[:,2:]
y_ob = data_ob.iloc[:,1].values

# one hot encoding

# Note label order encoded ['A','C','G','T']
df_train_x = df_train_x.apply(LabelEncoder().fit_transform)
df_test_x = df_test_x.apply(LabelEncoder().fit_transform)
df_ob_x = df_ob_x.apply(LabelEncoder().fit_transform)
enc = OneHotEncoder(sparse = False)
enc.fit(df_train_x)
x_train = enc.transform(df_train_x)
x_test = enc.transform(df_test_x)
x_ob = enc.transform(df_ob_x)
'''
#MAIN MODEL SETUP

# parameter setup and model fitting
las = LassoCV(alphas = np.logspace(-5, 2, 30), cv = 10, verbose = 1, n_jobs = 4).fit(x_train, y_train) #fit the pixel values (x) to the labels (y)

print("Chosen alpha:", las.alpha_)

# save model
joblib.dump(las, os.path.join('saved_models','lasso_model.pkl'))
'''

las = joblib.load(os.path.join('saved_models','lasso_model.pkl'))
# MAKE PREDICTIONS AND CHECK ACCURACY

# get predictions
train_labels = y_train
train_pred = las.predict(x_train)
test_labels = y_test
test_pred = las.predict(x_test)
ob_labels = y_ob
ob_pred = las.predict(x_ob)

train_slope, train_intercept, train_r_value, train_p_value, train_std = linregress(train_labels, train_pred)
test_slope, test_intercept, test_r_value, test_p_value, test_std = linregress(test_labels, test_pred)
ob_slope, ob_intercept, ob_r_value, ob_p_value, ob_std = linregress(ob_labels, ob_pred)

# print metrics

with open('lasso_stats.csv', "w") as csv_file:

    csv_file.write('{0}{1}{2}{3}{4}{5}{6}{7}{8}'.format("SET,","MSE,","MAE,","MAD,","LR_slope,","LR_intercept,","LR_r_val,","LR_wald_p_val,","LR_stdev\n"))
    csv_file.write('{0}{1}{2}{3}{4}{5}{6}{7}{8}{9}{10}{11}{12}{13}{14}{15}{16}'.format("TRAIN,",mean_squared_error(train_labels, train_pred),",",mean_absolute_error(train_labels, train_pred),",",median_absolute_error(train_labels, train_pred),",",train_slope,",",train_intercept,",",train_r_value,",",train_p_value,",",train_std,"\n"))
    csv_file.write('{0}{1}{2}{3}{4}{5}{6}{7}{8}{9}{10}{11}{12}{13}{14}{15}{16}'.format("TEST,",mean_squared_error(test_labels, test_pred),",",mean_absolute_error(test_labels, test_pred),",",median_absolute_error(test_labels, test_pred),",",test_slope,",",test_intercept,",",test_r_value,",",test_p_value,",",test_std,"\n"))
    csv_file.write('{0}{1}{2}{3}{4}{5}{6}{7}{8}{9}{10}{11}{12}{13}{14}{15}'.format("OVERBEEK,",mean_squared_error(ob_labels, ob_pred),",",mean_absolute_error(ob_labels, ob_pred),",",median_absolute_error(ob_labels, ob_pred),",",ob_slope,",",ob_intercept,",",ob_r_value,",",ob_p_value,",",ob_std))

print('TRAINING SET RESULTS:')
print('MSE: ', mean_squared_error(train_labels, train_pred))
print('MAE: ', mean_absolute_error(train_labels, train_pred))
print('Median AE: ', median_absolute_error(train_labels, train_pred))
print('LR:', linregress(train_labels, train_pred))
print('Pearson:', pearsonr(train_labels, train_pred))
print('Spearman:', spearmanr(train_labels, train_pred))
print('Kendall:', kendalltau(train_labels, train_pred))

print('\nTEST SET RESULTS:')
print('MSE: ', mean_squared_error(test_labels, test_pred))
print('MAE: ', mean_absolute_error(test_labels, test_pred))
print('Median AE: ', median_absolute_error(test_labels, test_pred))
print('LR:', linregress(test_labels, test_pred))
print('Pearson:', pearsonr(test_labels, test_pred))
print('Spearman:', spearmanr(test_labels, test_pred))
print('Kendall:', kendalltau(test_labels, test_pred))

print('\nOVERBEEK RESULTS:')
print('MSE: ', mean_squared_error(ob_labels, ob_pred))
print('MAE: ', mean_absolute_error(ob_labels, ob_pred))
print('Median AE: ', median_absolute_error(ob_labels, ob_pred))
print('LR:', linregress(ob_labels, ob_pred))
print('Pearson:', pearsonr(ob_labels, ob_pred))
print('Spearman:', spearmanr(ob_labels, ob_pred))
print('Kendall:', kendalltau(ob_labels, ob_pred))

# FEATURE IMPORTANCES

importances = las.coef_
nt_importance = []
allAs = []
allCs = []
allGs = []
allTs = []
As = []
Cs = []
Gs = []
Ts = []
negAs = []
negCs = []
negGs = []
negTs = []
# output |summed| coeffs
for i in range(80):
    if ((i+1)%4 == 0):
        # sum of the abs of coeffs from the individual one-hot encoded features
        nt_importance.append(abs(importances[i-3])+abs(importances[i-2])+abs(importances[i-1])+abs(importances[i]))

        allAs.append(importances[i-3])
        if importances[i-3] < 0:
            As.append(0)
            negAs.append(importances[i-3])
        else:
            As.append(importances[i-3])
            negAs.append(0)

        allCs.append(importances[i-2])
        if importances[i-2] < 0:
            Cs.append(0)
            negCs.append(importances[i-2])
        else:
            Cs.append(importances[i-2])
            negCs.append(0)

        allGs.append(importances[i-1])
        if importances[i-1] < 0:
            Gs.append(0)
            negGs.append(importances[i-1])
        else:
            Gs.append(importances[i-1])
            negGs.append(0)

        allTs.append(importances[i])
        if importances[i] < 0:
            Ts.append(0)
            negTs.append(importances[i])
        else:
            Ts.append(importances[i])
            negTs.append(0)

        # output table of individual nt letter at nt pos contributions
nt_importance.append(abs(importances[80])+abs(importances[81])+abs(importances[82])+abs(importances[83]))

allAs.append(importances[80])
if importances[80] < 0:
    As.append(0)
    negAs.append(importances[80])
else:
    As.append(importances[80])
    negAs.append(0)

allCs.append(importances[81])
if importances[81] < 0:
    Cs.append(0)
    negCs.append(importances[81])
else:
    Cs.append(importances[81])
    negCs.append(0)

allGs.append(importances[82])
if importances[82] < 0:
    Gs.append(0)
    negGs.append(importances[82])
else:
    Gs.append(importances[82])
    negGs.append(0)

allTs.append(importances[83])
if importances[83] < 0:
    Ts.append(0)
    negTs.append(importances[83])
else:
    Ts.append(importances[83])
    negTs.append(0)

nt_importance.append(abs(importances[84]))
allAs.append(0), allCs.append(0), allGs.append(0), allTs.append(0)
nt_importance.append(abs(importances[85]))
allAs.append(0), allCs.append(0), allGs.append(0), allTs.append(0)

# DRAW PLOTS

#feature coefficients vs alphas
alphas = las.alphas_
coefs = []
for a in alphas:
    check_las = Lasso(alpha=a)
    check_las.fit(x_train, y_train)
    coefs.append(check_las.coef_)

ax = plt.gca()
ax.plot(alphas, coefs)
ax.set_xscale('log')
ax.set_xlim(ax.get_xlim()[::-1])  # reverse axis
plt.xlabel('alpha', fontsize = 18)
plt.ylabel('coefficients',fontsize = 18)
plt.yticks(fontsize = 12)
plt.xticks(fontsize = 12)
plt.title('LASSO coefficients as a function alpha', fontsize = 18)
plt.axis('tight')
plt.tight_layout()
plt.savefig('lasso_coeff_lambda.pdf', dpi = 300)

# Lasso model selection
m_log_alphas = -np.log10(las.alphas_)
plt.figure()
plt.plot(m_log_alphas, las.mse_path_, ':')
plt.plot(m_log_alphas, las.mse_path_.mean(axis=-1), 'k',
         label='Average across the folds', linewidth=2)
plt.axvline(-np.log10(las.alpha_), linestyle='--', color='k',
            label='alpha: CV estimate')
plt.legend()
plt.xlabel('-log(alpha)', fontsize = 18)
plt.ylabel('MSE', fontsize = 18)
plt.yticks(fontsize = 12)
plt.xticks(fontsize = 12)
plt.title('LASSO model selection', fontsize = 18)
plt.axis('tight')
plt.tight_layout()
plt.savefig('lasso_model_selection_cv.pdf', dpi = 300)

# plot the feature importance of LASSO (all)

plt.figure()
ind = np.arange(21)    # the x locations for the groups
width = 0.8       # the width of the bars: can also be len(x) sequence
p1 = plt.bar(ind, As, width, color = 'green')
p2 = plt.bar(ind, Ts, width, bottom = As, color = 'red')
p3 = plt.bar(ind, Cs, width, bottom = np.add(As, Ts), color = 'blue')
p4 = plt.bar(ind, Gs, width, bottom = np.add(Cs, np.add(As, Ts)), color = 'orange')
p5 = plt.bar(ind, negAs, width, color = 'green')
p6 = plt.bar(ind, negTs, width, bottom = negAs, color = 'red')
p7 = plt.bar(ind, negCs, width, bottom = np.add(negAs, negTs), color = 'blue')
p8 = plt.bar(ind, negGs, width, bottom = np.add(negCs, np.add(negAs, negTs)), color = 'orange')
plt.legend((p1[0], p2[0], p3[0], p4[0]), ('A', 'T', 'C', 'G'), fontsize  = 14, loc = 2)
plt.title('LASSO nucleotide importances', fontsize = 18)
plt.ylabel('coefficients', fontsize = 18)
plt.ylim([-0.1,0.2])
plt.yticks(fontsize = 14)
plt.xticks(range(23), ['-20', '-19', '-18', '-17', '-16', '-15', '-14', '-13', '-12', '-11', '-10', '-9', '-8', '-7', '-6', '-5', '-4', '-3', '-2', '-1', 'N', 'G', 'G'], rotation = 60, fontsize = 14)
plt.tight_layout()
plt.savefig('lasso_coeff.pdf', dpi = 300)

# plot the feature importances of LASSO (|summed| coeff)
plt.figure()
plt.title('LASSO nucleotide importances', fontsize = 18)
plt.bar(range(23), nt_importance, width = 0.8, color='gray')
plt.ylabel('absolute coefficients', fontsize = 18)
plt.yticks(fontsize = 12)
plt.xticks(range(23), ['-20', '-19', '-18', '-17', '-16', '-15', '-14', '-13', '-12', '-11', '-10', '-9', '-8', '-7', '-6', '-5', '-4', '-3', '-2', '-1', 'N', 'G', 'G'], rotation = 60, fontsize = 14)
plt.tight_layout()
plt.savefig('summed_lasso_coeff.pdf', dpi = 300)

##################

# Setting the positions and width for the bars
pos = list(range(len(allAs)))
width = 0.25
x_labels = ['-20', '-19', '-18', '-17', '-16', '-15', '-14', '-13', '-12', '-11', '-10', '-9', '-8', '-7', '-6', '-5', '-4', '-3', '-2', '-1', 'N', 'G', 'G']

# Plotting the bars
fig, ax = plt.subplots(figsize=(10,5))

# Create a bar with pre_score data,
# in position pos,
plt.bar(pos,
        #using df['pre_score'] data,
        allAs,
        # of width
        width,
        # with color
        color='green')

# Create a bar with mid_score data,
# in position pos + some width buffer,
plt.bar([p + width for p in pos],
        #using df['mid_score'] data,
        allTs,
        # of width
        width,
        # with color
        color='red')

# Create a bar with post_score data,
# in position pos + some width buffer,
plt.bar([p + width*2 for p in pos],
        #using df['post_score'] data,
        allCs,
        # of width
        width,
        # with color
        color='blue')

plt.bar([p + width*3 for p in pos],
        #using df['post_score'] data,
        allGs,
        # of width
        width,
        # with color
        color='orange')

# Set the y axis label
ax.set_ylabel('coefficient', fontsize = 20)

# Set the chart's title
ax.set_title('Nucleotide importance (LASSO)', fontsize = 20)

# Set the labels for the x ticks
plt.xticks(range(23), x_labels, rotation = 60, fontsize = 14)

# Set the position of the x ticks
ax.set_xticks([p + 2 * width for p in pos])

# Setting the x-axis and y-axis limits
plt.xlim(min(pos)-width, max(pos)+width*4)

# Adding the legend and showing the plot
plt.legend(['A', 'T', 'C', 'G'], loc='upper right', framealpha = 0, fontsize = 14)

for p in range(22):
    plt.axvline(p + width*3.5, color = 'black', linewidth = 0.3)

plt.savefig('trail_lasso_coeff.pdf', dpi = 600)

##################

# Plot pred vs observed (test)

plot_text = '{0}{1:.2f}{2}{3}{4:.2E}'.format('R-value: ', test_r_value,'\n','P-value: ', test_p_value)

fig2, ax2 = plt.subplots()
ax2.plot(test_labels, test_pred,'o')
ax2.plot(test_labels, test_intercept + test_slope*test_labels, 'r')
ax2.text(0.3,0.85,plot_text,fontsize=18)
ax2.set_title('Unseen test_data\n', fontsize = 18)
ax2.set_ylim([0.0,1.0])
ax2.set_xlim([0.0,1.0])
ax2.tick_params('both', labelsize = 18)
ax2.set_ylabel('Predicted precision', fontsize = 18)
ax2.set_xlabel('Observed precision', fontsize = 18)
ax2.set_aspect('equal', 'box')
fig2.tight_layout()
plt.savefig('lasso_predictions_vs_labels_test.pdf', dpi = 300)

# Plot pred vs observed (Overbeek)

plot_text = '{0}{1:.2f}{2}{3}{4:.2E}'.format('R-value: ', ob_r_value,'\n','P-value: ', ob_p_value)

fig2, ax2 = plt.subplots()
ax2.plot(ob_labels, ob_pred,'o', color = 'y')
ax2.plot(ob_labels, ob_intercept + ob_slope*ob_labels, 'r')
ax2.text(0.3,0.85,plot_text,fontsize=18)
ax2.set_title('HCT116 (van Overbeek et al)\n', fontsize = 18)
ax2.set_ylim([0.0,1.0])
ax2.set_xlim([0.0,1.0])
ax2.tick_params('both', labelsize = 18)
ax2.set_ylabel('Predicted precision', fontsize = 18)
ax2.set_xlabel('Observed precision', fontsize = 18)
ax2.set_aspect('equal', 'box')
fig2.tight_layout()
plt.savefig('lasso_predictions_vs_labels_overbeek.pdf', dpi = 300)


#BBN_LICENSE_START -- DO NOT MODIFY BETWEEN LICENSE_{START,END} Lines
# Copyright (c) <2017,2018,2019,2020,2021>, <Raytheon BBN Technologies>
# To be applied to the DCOMP/MAP Public Source Code Release dated 2018-04-19, with
# the exception of the dcop implementation identified below (see notes).
#
# Dispersed Computing (DCOMP)
# Mission-oriented Adaptive Placement of Task and Data (MAP)
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
# Redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
# TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
# TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#BBN_LICENSE_END
#!/usr/bin/env python3

import warnings
with warnings.catch_warnings():
    import re
    import sys
    import argparse
    import os
    import os.path
    import logging
    import logging.config
    import json
    from pathlib import Path
    import pandas as pd
    import csv
    import matplotlib.pyplot as plt
    import map_utils
    import multiprocessing

script_dir=os.path.abspath(os.path.dirname(__file__))

"""
Statistics for the whole system have a pid of -1
"""
system_stats_pid = -1

def get_logger():
    return logging.getLogger(__name__)

def bytes_to_gigabytes(b):
    return b / 1024 / 1024 / 1024

def pid_label(pid, pid_names):
    return f"{pid_names[pid]}_{pid}"

def make_pid_key(pid_names):
    """
    Create a key function that calls process_key with the result of pid_label.
    """
    def make_key(x_pid):
        x_label = pid_label(x_pid, pid_names)
        return process_key(x_label)
    return make_key

def process_key(x):
    """
    Compare process names preferring map-agent.
    """

    if x.startswith('map-agent'):
        return f"0_{x}"
    else:
        return f"9_{x}"

def is_interesting_process(unfiltered, pid, row):
    """
    Check if this is a graph that we should graph.

    Arguments:
        unfiltered(boolean): if true, return all processes
        pid(int): the process ID parsed as an int from 'row'
        row(dict): dictionary from csv.DictReader
    Returns:
        boolean: true if the process is "interesting"
    """
    if unfiltered:
        return True
    elif pid == system_stats_pid:
        # totals
        return True
    elif row['user'] == 'map':
        # map service
        return True
    elif re.search(r'iftop', row['cmdline']):
        return True
    elif row['name'] == 'fprobe':
        return True
    elif row['name'] == 'java' and re.search(r'map-dns', row['cmdline']):
        # DNS server
        return True
    elif row['docker'] == 'True':
        # docker container or docker system process
        return True
    else:
        return False

def organize_node_data(unfiltered, stats_dir):
    """
    Returned timestamps are unix epoch seconds.

    Arguments:
        unfiltered(boolean): if true, return all processes
        stats_dir(Path): base directory to read stats data from
    Returns:
        data_cpu(dict): ts -> pid -> cpu
        data_memory(dict): ts -> pid -> memory
        all_pids(set): all process ids
        pid_names(dict): pid -> display name
        network_data(dict): attr -> interface -> ts -> value
    """
    # timestamps are unix epoch seconds
    data_cpu = dict() # ts -> pid -> cpu
    data_memory = dict() # ts -> pid -> memory
    all_pids = set()
    pid_names = dict() # pid -> display name

    network_data = dict() # attr -> interface -> ts -> value

    for file in stats_dir.iterdir():
        match = re.match(r'^(\d+)\.csv$', file.name)
        net_match = re.match(r'^(\d+)_net\.csv$', file.name)
        if file.is_file() and match:
            timestamp = int(match.group(1))
            ts_data_cpu = dict() # pid -> cpu
            ts_data_memory = dict() # pid -> memory
            with open(file) as f:
                reader = csv.DictReader(f)
                for row in reader:
                    pid = int(row['pid'])
                    if is_interesting_process(unfiltered, pid, row):
                        cpu = float(row['cpu_user'])
                        # vms is not useful, it includes things such as file backed memory
                        memory = float(row['memory_rss'])
                        if pid != system_stats_pid:
                            cpu = cpu + float(row['cpu_sys'])

                            if row['name'] == 'java':
                                match = re.search(r'([^/\s]+)\.jar', row['cmdline'])
                                if match:
                                    name = match.group(1)
                                else:
                                    name = row['name']
                            else:
                                name = row['name']
                            pid_names[pid] = name

                        ts_data_cpu[pid] = cpu
                        ts_data_memory[pid] = memory
                        if pid != system_stats_pid:
                            all_pids.add(pid)

            # verify that we have a system entry
            if system_stats_pid not in ts_data_cpu:
                get_logger().warning("Missing system entry in %s, skipping", file)
                continue

            data_cpu[timestamp] = ts_data_cpu
            data_memory[timestamp] = ts_data_memory

        elif net_match:
            timestamp = int(net_match.group(1))
            with open(file) as f:
                reader = csv.DictReader(f)
                for row in reader:
                    interface = row['interface']
                    interface = interface.lower()
                    for attr, value in row.items():
                        if attr == 'interface':
                            continue
                        attr_data = network_data.get(attr, dict())
                        interface_data = attr_data.get(interface, dict())

                        try:
                            interface_data[timestamp] = int(value)
                        except ValueError:
                            # missing column
                            interface_data[timestamp] = 0

                        attr_data[interface] = interface_data
                        network_data[attr] = attr_data

    return data_cpu, data_memory, all_pids, pid_names, network_data

def compute_node_differential_percent(all_pids, data):
    """Compute percentage of the system used at each timestamp as a
    difference from the previous timestamp.  This assumes that the
    data measured is a counter that increases over time and the
    difference between 2 timestamps is the useful information.

    """
    data_percent = dict() # timestamp -> pid -> data

    prev_values = None
    for timestamp, ts_data in sorted(data.items()):
        ts_percent = dict() # pid -> data
        if prev_values is not None:
            system_cur = float(ts_data[system_stats_pid])
            system_prev = (prev_values[system_stats_pid])
            system_diff = system_cur - system_prev

            for pid in all_pids:
                cur = ts_data.get(pid, None)
                prev = prev_values.get(pid, None)
                if cur is None or prev is None:
                    percent = 0
                else:
                    diff = float(cur - prev)
                    if diff < 0:
                        get_logger().warn("Found negative usage difference at %d for %d diff %f", timestamp, pid, diff)
                        percent = 0
                    else:
                        percent = diff / system_diff
                        if percent > 1:
                            get_logger().warn("Found diff (%f) that is greater than system diff (%f)", diff, system_diff)

                ts_percent[pid] = percent * 100

        data_percent[timestamp] = ts_percent
        prev_values = ts_data

    return data_percent

def compute_node_absolute_percent(all_pids, data):
    """Compute percentage of the system that is used at each timestamp.
    This assumes that the values at each timestamp are absolute
    measurements.

    """
    data_percent = dict() # timestamp -> pid -> data

    for timestamp, ts_data in sorted(data.items()):
        ts_percent = dict() # pid -> data
        for pid in all_pids:
            system_cur = float(ts_data[system_stats_pid])
            cur = ts_data.get(pid, None)
            if cur is None:
                percent = 0
            else:
                percent = float(cur) / system_cur
            ts_percent[pid] = percent * 100
        data_percent[timestamp] = ts_percent

    return data_percent

def write_percents(all_pids, pid_names, filename, data_percent):
    with open(filename, 'w') as f:
        writer = csv.writer(f)
        header = ['timestamp']
        for pid in sorted(all_pids):
            label = pid_label(pid, pid_names)
            header.append(label)
        writer.writerow(header)

        for timestamp, ts_percent in sorted(data_percent.items()):
            row = [timestamp]
            for pid in sorted(all_pids):
                if pid in ts_percent:
                    row.append(ts_percent[pid])
                else:
                    row.append(0)
            writer.writerow(row)

def write_cpu_plot(node_name, file_cpu, output, first_timestamp_sec):
    get_logger().debug("Writing plot for %s with first timestamp of %d from %s", node_name, first_timestamp_sec, file_cpu)

    df = pd.read_csv(file_cpu)
    df['time_minutes'] = (df['timestamp'] - first_timestamp_sec) / 60
    max_minutes = df['time_minutes'].max()

    fig, ax = map_utils.subplots()
    ax.set_title(f"{node_name} Process CPU load")
    ax.set_ylabel('% CPU')
    ax.set_xlabel('time (minutes)')
    ax.set_xlim(left=0, right=max_minutes)

    ydata = list()
    labels = list()
    for process in sorted(df.columns, key=process_key):
        if process != 'timestamp' and process != 'time_minutes':
            ydata.append(df[process])
            labels.append(process)

    ax.stackplot(df['time_minutes'], ydata, labels=labels)
    handles, labels = ax.get_legend_handles_labels()
    lgd = ax.legend(handles, labels, bbox_to_anchor=(1.04, 1), loc="upper left")

    output_name = output / f"{file_cpu.stem}.png"
    fig.savefig(output_name.as_posix(), format='png', bbox_extra_artists=(lgd,), bbox_inches='tight')
    plt.close(fig)

def write_memory_plot(node_name, file_memory, output, first_timestamp_sec):
    df = pd.read_csv(file_memory)
    df['time_minutes'] = (df['timestamp'] - first_timestamp_sec) / 60
    max_minutes = df['time_minutes'].max()

    fig, ax = map_utils.subplots()
    ax.set_title(f"{node_name} Process memory load")
    ax.set_ylabel('% memory')
    ax.set_xlabel('time (minutes)')
    ax.set_xlim(left=0, right=max_minutes)

    ydata = list()
    labels = list()
    for process in sorted(df.columns, key=process_key):
        if process != 'timestamp' and process != 'time_minutes':
            ydata.append(df[process])
            labels.append(process)

    ax.stackplot(df['time_minutes'], ydata, labels=labels)

    handles, labels = ax.get_legend_handles_labels()
    lgd = ax.legend(handles, labels, bbox_to_anchor=(1.04, 1), loc="upper left")

    output_name = output / f"{file_memory.stem}.png"
    fig.savefig(output_name.as_posix(), format='png', bbox_extra_artists=(lgd,), bbox_inches='tight')
    plt.close(fig)

def write_memory_absolute_plot(node_name, data_memory, all_pids, pid_names, output, first_timestamp_sec):
    """
    Write out the graph of absolute memory.

    Arguments:
        node_name(str): name of the node
        data_memory(dict): ts(int) -> pid(int) -> memory(float)
        all_pids(set): all process names
        pid_names(dict): pid(int) -> pid name(str)
        output(Path): where the graphs are being written
    """

    timestamps = sorted(data_memory.keys())
    minutes = [ (ts - first_timestamp_sec) / 60 for ts in timestamps ]
    max_minutes = max(minutes)

    fig, ax = map_utils.subplots()
    ax.set_title(f"{node_name} Process memory")
    ax.set_ylabel('GB memory')
    ax.set_xlabel('time (minutes)')
    ax.set_xlim(left=0, right=max_minutes)

    ydata = list()
    labels = list()
    for pid in sorted(all_pids, key=make_pid_key(pid_names)):
        plot_data = list()
        for timestamp in timestamps:
            ts_data = data_memory[timestamp]
            if pid in ts_data:
                memory = bytes_to_gigabytes(ts_data[pid])
            else:
                memory = 0
            plot_data.append(memory)
        label = pid_label(pid, pid_names)
        ydata.append(plot_data)
        labels.append(label)

    ax.stackplot(minutes, ydata, labels=labels)
    handles, labels = ax.get_legend_handles_labels()
    lgd = ax.legend(handles, labels, bbox_to_anchor=(1.04, 1), loc="upper left")

    output_name = output / f"{node_name}_stats_memory_absolute.png"
    fig.savefig(output_name.as_posix(), format='png', bbox_extra_artists=(lgd,), bbox_inches='tight')
    plt.close(fig)

def compute_network_differentials(network_data):
    """
    Arguments:
        network_data(dict): attr -> interface -> ts -> value
    Returns:
        network_data(dict): attr -> interface -> ts -> diff from previous per second
    """
    retval = dict()

    for attr, attr_data in network_data.items():
        ret_attr_data = retval.get(attr, dict())

        for interface, interface_data in attr_data.items():
            ret_interface_data = ret_attr_data.get(interface, dict())

            prev_ts = None
            prev_value = None
            for ts, value in sorted(interface_data.items()):
                if prev_value is not None:
                    diff_value = value - prev_value
                    diff_ts = ts - prev_ts
                    diff_per_second = diff_value / diff_ts
                    ret_interface_data[ts] = diff_per_second
                prev_value = value
                prev_ts = ts

            ret_attr_data[interface] = ret_interface_data
        retval[attr] = ret_attr_data
    return retval

def write_network_bandwidth_plot(output_dir, first_timestamp_sec, node_name, graph_label, file_label, bandwidth_data, ip_info):
    """
    Arguments:
        output_dir(Path): base directory for output
        label(str): label for the data
        bandwidth_data(dict): interface -> ts -> value
        first_timestamp_sec(int): the first timestamp for relative computation
        ip_info(dict): node -> ifce -> ip
    """
    node_ip_info = ip_info.get(node_name, dict())
    get_logger().debug("Got ip info for %s -> %s", node_name, node_ip_info)

    fig, ax = map_utils.subplots()
    ax.set_title(graph_label)
    ax.set_ylabel('mbps')
    ax.set_xlabel('time (minutes)')

    max_minutes=0
    for interface, interface_data in bandwidth_data.items():
        xdata = list()
        ydata = list()
        for ts, value in interface_data.items():
            rel_sec = ts - first_timestamp_sec
            rel_min = rel_sec / 60
            mbps = value / 1024 / 1024 * 8 # 1024 kb / byte, 1024 kb / mb, 8 bits per byte
            xdata.append(rel_min)
            ydata.append(mbps)
            max_minutes = max(max_minutes, rel_min)

        if interface in node_ip_info:
            ip = f" - {node_ip_info[interface]}"
        else:
            ip = ""
        ax.plot(xdata, ydata, label=f"{interface}{ip}")

    ax.set_xlim(left=0, right=max_minutes)

    handles, labels = ax.get_legend_handles_labels()
    lgd = ax.legend(handles, labels, bbox_to_anchor=(1.04, 1), loc="upper left")

    output_name = output_dir / f"{node_name}_stats_network-{file_label}.png"
    fig.savefig(output_name.as_posix(), format='png', bbox_extra_artists=(lgd,), bbox_inches='tight')
    plt.close(fig)

def gather_network_ips(sim_output):
    """
    Determine the IP address for each interface on each node.
    If the information isn't available, then the result will be an empty dictionary.
    Use of the dictionary should always check for the information existing

    Arguments:
        sim_output(Path): directory that the sim output is in

    Returns:
        dict: node_name -> ifce_name -> ip
    """

    network_info = sim_output / "inputs/network-information.txt"
    if not network_info.exists():
        # no information
        return dict()

    info = dict()
    node_name = None
    ifce_name = None
    node_name_re = re.compile(r'^(\S+)\s+\|\s+CHANGED')
    ifce_name_re = re.compile(r'^\d+:\s+([^@]+)(@\S+)?: \<')
    ip_re = re.compile(r'^\s+inet ([0-9\.]+)/')
    with open(network_info) as f:
        for line in f:
            node_name_match = node_name_re.match(line)
            if node_name_match:
                node_name = node_name_match.group(1).split('.')[0]
                ifce_name = None
                get_logger().debug("Found node %s", node_name)
                continue

            ifce_name_match = ifce_name_re.match(line)
            if ifce_name_match:
                ifce_name = ifce_name_match.group(1)
                get_logger().debug("Found interface %s", ifce_name)
                continue

            ip_match = ip_re.match(line)
            if ip_match:
                ip = ip_match.group(1)
                if node_name is not None and ifce_name is not None:
                    node_name = node_name.lower()
                    ifce_name = ifce_name.lower()
                    get_logger().debug("Found ip %s on %s", ip, node_name)
                    node_info = info.get(node_name, dict())
                    node_info[ifce_name] = ip
                    info[node_name] = node_info
                else:
                    get_logger().warning("Found network interface name %s and node_name (%s) or ifce_name (%s) is missing", ip, node_name, ifce_name)

                ifce_name = None

    return info

def process_node_dir(unfiltered, output, node_dir, first_timestamp_sec, ip_info):
    try:
        stats_dir = node_dir / 'system_stats'
        if not stats_dir.exists():
            # not a directory we care about
            return

        node_name = node_dir.stem.lower()
        get_logger().info("Processing node %s", node_name)

        (data_cpu, data_memory, all_pids, pid_names, network_data) = organize_node_data(unfiltered, stats_dir)
        percent_cpu = compute_node_differential_percent(all_pids, data_cpu)
        percent_memory = compute_node_absolute_percent(all_pids, data_memory)

        file_cpu = output / f"{node_name}_stats_cpu.csv"
        file_memory = output / f"{node_name}_stats_memory.csv"
        write_percents(all_pids, pid_names, file_cpu, percent_cpu)
        write_percents(all_pids, pid_names, file_memory, percent_memory)
        write_cpu_plot(node_name, file_cpu, output, first_timestamp_sec)
        write_memory_plot(node_name, file_memory, output, first_timestamp_sec)

        write_memory_absolute_plot(node_name, data_memory, all_pids, pid_names, output, first_timestamp_sec)

        network_per_second = compute_network_differentials(network_data)

        if 'rx_bytes' in network_per_second:
            write_network_bandwidth_plot(output, first_timestamp_sec, node_name, 'RX mbps', 'rx', network_per_second['rx_bytes'], ip_info)

        if 'tx_bytes' in network_per_second:
            write_network_bandwidth_plot(output, first_timestamp_sec, node_name, 'TX mbps', 'tx', network_per_second['tx_bytes'], ip_info)

    except:
        get_logger().exception("Unexpected error")

def main_method(args):
    sim_output = Path(args.sim_output)
    if not sim_output.exists():
        get_logger().error("%s does not exist", sim_output)
        return 1

    with open(args.first_timestamp_file) as f:
        ts_str = f.readline().strip()
        first_timestamp = map_utils.log_timestamp_to_datetime(ts_str)
    get_logger().info("Simulation started at %s", first_timestamp)

    ip_info = gather_network_ips(sim_output)
    get_logger().debug("ip_info: %s", ip_info)

    first_timestamp_secs = first_timestamp.timestamp()

    output = Path(args.output)
    output.mkdir(parents=True, exist_ok=True)

    with multiprocessing.Pool(processes=os.cpu_count()) as pool:
        results = list()

        for node_dir in sim_output.iterdir():
            if node_dir.is_dir():
                results.append(pool.apply_async(func=process_node_dir, args=[args.unfiltered, output, node_dir, first_timestamp_secs, ip_info]))

        for result in results:
            result.wait()

def main(argv=None):
    if argv is None:
        argv = sys.argv[1:]

    class ArgumentParserWithDefaults(argparse.ArgumentParser):
        '''
        From https://stackoverflow.com/questions/12151306/argparse-way-to-include-default-values-in-help
        '''
        def add_argument(self, *args, help=None, default=None, **kwargs):
            if help is not None:
                kwargs['help'] = help
            if default is not None and args[0] != '-h':
                kwargs['default'] = default
                if help is not None:
                    kwargs['help'] += ' (default: {})'.format(default)
            super().add_argument(*args, **kwargs)

    parser = ArgumentParserWithDefaults(formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("-l", "--logconfig", dest="logconfig", help="logging configuration (default: logging.json)", default='logging.json')
    parser.add_argument("--debug", dest="debug", help="Enable interactive debugger on error", action='store_true')
    parser.add_argument("-o", "--output", dest="output", help="Output directory (Required)", required=True)
    parser.add_argument("-s", "--sim-output", dest="sim_output", help="Chart output directory (Required)", required=True)
    parser.add_argument("--first-timestamp-file", dest="first_timestamp_file", help="Path to file containing the log timestamp that the simulation started", required=True)
    parser.add_argument("--unfiltered", dest="unfiltered", help="Show all processes, not just MAP processes", action='store_true')

    args = parser.parse_args(argv)

    map_utils.setup_logging(default_path=args.logconfig)
    if 'multiprocessing' in sys.modules:
        import multiprocessing_logging
        multiprocessing_logging.install_mp_handler()

    if args.debug:
        import pdb, traceback
        try:
            return main_method(args)
        except:
            extype, value, tb = sys.exc_info()
            traceback.print_exc()
            pdb.post_mortem(tb)
    else:
        return main_method(args)

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


import baseline as bl
import baseline.tf.embeddings
import baseline.tf.classify
import time
import tensorflow as tf
import numpy as np
import os
import argparse
import logging
log = logging.getLogger('baseline.timing')
NUM_PREFETCH = 2
SHUF_BUF_SZ = 5000

def get_logging_level(ll):
    ll = ll.lower()
    if ll == 'debug':
        return logging.DEBUG
    if ll == 'info':
        return logging.INFO
    return logging.WARNING

def get_tf_logging_level(ll):
    ll = ll.lower()
    if ll == 'debug':
        return tf.logging.DEBUG
    if ll == 'info':
        return logging.INFO
    return tf.logging.WARN

def to_tensors(ts):
    X = []
    y = []
    for sample in ts:
        X.append(sample['word'].squeeze())
        y.append(sample['y'].squeeze())
    return np.stack(X), np.stack(y)

parser = argparse.ArgumentParser(description='Train a Baseline model with TensorFlow Estimator API')
parser.add_argument('--export_dir', help='Directory for TF export (for serving)', default='./models', type=str)
parser.add_argument('--checkpoint_dir', help='Directory for model checkpoints', default='./checkpoints', type=str)
parser.add_argument('--model_type', help='What type of model to build', type=str, default='default')
parser.add_argument('--poolsz', help='How many hidden units for pooling', type=int, default=100)
parser.add_argument('--stacksz', help='How many hidden units for stacking', type=int, nargs='+')
parser.add_argument('--text', help='raw value', type=str)
parser.add_argument('--backend', help='backend', default='tf')
parser.add_argument('--remote', help='(optional) remote endpoint', type=str) # localhost:8500
parser.add_argument('--name', help='(optional) signature name', type=str)
parser.add_argument('--epochs', help='Number of epochs to train', type=int, default=2)
parser.add_argument('--batchsz', help='Batch size', type=int, default=50)
parser.add_argument('--filts', help='Parallel convolution filter widths (if default model)', type=int, default=[3, 4, 5], nargs='+')
parser.add_argument('--mxlen', help='Maximum post length (number of words) during training', type=int, default=100)
parser.add_argument('--train', help='Training file', default='../data/stsa.binary.phrases.train')
parser.add_argument('--valid', help='Validation file', default='../data/stsa.binary.dev')
parser.add_argument('--test', help='Testing file', default='../data/stsa.binary.test')
parser.add_argument('--embeddings', help='Pretrained embeddings file', default='/data/embeddings/GoogleNews-vectors-negative300.bin')
parser.add_argument('--ll', help='Log level', type=str, default='info')
parser.add_argument('--tf_ll', help='TensorFlow Log level', type=str, default='warning')
parser.add_argument('--lr', help='Learning rate', type=float, default=0.001)
parser.add_argument('--optim', help='Optimizer (sgd, adam) (default is adam)', type=str, default='adam')
parser.add_argument('--gpus', help='Number of GPUs to use', default=1, type=int)
args = parser.parse_known_args()[0]

logging.basicConfig(level=get_logging_level(args.ll))
tf.logging.set_verbosity(get_tf_logging_level(args.tf_ll))

pool_field = 'cmotsz' if args.model_type == 'default' else 'rnnsz'

model_params = {
    'model_type': args.model_type,
    'filtsz': args.filts,
    pool_field: args.poolsz
}

if args.stacksz is not None:
    model_params['hsz'] = args.stacksz

feature_desc = {
    'word': {
        'vectorizer': bl.Token1DVectorizer(mxlen=args.mxlen, transform_fn=bl.lowercase),
        'embed': {'file': args.embeddings, 'type': 'default', 'unif': 0.25}
    }
}
# Create a reader that is using our vectorizers to parse a TSV file
# with rows like:
# <label>\t<sentence>\n

vectorizers = {k: v['vectorizer'] for k, v in feature_desc.items()}
reader = bl.TSVSeqLabelReader(vectorizers,
                              clean_fn=bl.TSVSeqLabelReader.do_clean)

train_file = args.train
valid_file = args.valid
test_file = args.test

# This builds a set of counters
vocabs, labels = reader.build_vocab([train_file,
                                     valid_file,
                                     test_file])

# This builds a set of embeddings objects, these are typically not DL-specific
# but if they happen to be addons, they can be
embeddings = dict()
for k, v in feature_desc.items():
    embed_config = v['embed']
    embeddings_for_k = bl.load_embeddings('word',
                                          embed_file=embed_config['file'],
                                          known_vocab=vocabs[k],
                                          embed_type=embed_config.get('type', 'default'),
                                          unif=embed_config.get('unif', 0.),
                                          use_mmap=True)

    embeddings[k] = embeddings_for_k['embeddings']
    # Reset the vocab to the embeddings one
    vocabs[k] = embeddings_for_k['vocab']

X_train, y_train = to_tensors(reader.load(train_file, vocabs=vocabs, batchsz=1))
X_valid, y_valid = to_tensors(reader.load(valid_file, vocabs=vocabs, batchsz=1))
X_test, y_test = to_tensors(reader.load(test_file, vocabs=vocabs, batchsz=1))

def train_input_fn():
    dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train))
    dataset = dataset.shuffle(buffer_size=SHUF_BUF_SZ)
    # https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/distribute/README.md
    # effective_batch_sz = args.batchsz*args.gpus
    dataset = dataset.batch(args.batchsz//args.gpus)
    dataset = dataset.map(lambda x, y: ({'word': x}, y))
    dataset = dataset.repeat(args.epochs)
    dataset = dataset.prefetch(NUM_PREFETCH)
    _ = dataset.make_one_shot_iterator()
    return dataset

def eval_input_fn():
    dataset = tf.data.Dataset.from_tensor_slices((X_valid, y_valid))
    dataset = dataset.batch(args.batchsz)
    dataset = dataset.map(lambda x, y: ({'word': x}, y))
    _ = dataset.make_one_shot_iterator()
    return dataset

def predict_input_fn():
    dataset = tf.data.Dataset.from_tensor_slices((X_test, y_test))
    dataset = dataset.batch(args.batchsz)
    dataset = dataset.map(lambda x, y: ({'word': x}, y))
    _ = dataset.make_one_shot_iterator()
    return dataset

def server_input_fn():
    tensors = {
        'word': tf.placeholder(tf.int64, [None, None])
    }
    features = {
        'word': tensors['word']
    }
    return tf.estimator.export.ServingInputReceiver(receiver_tensors=tensors, features=features)

def one_epoch(X_train):
    """

    effective_batch_sz = args.batchsz = gpu_batchsz*args.gpus

    :param X_train:
    :return:
    """
    return len(X_train)//args.batchsz

def fit(estimator):

    for i in range(args.epochs):
        start_time = time.time()
        estimator.train(input_fn=train_input_fn, steps=one_epoch(X_train))
        elapsed = time.time() - start_time
        log.info('EPOCH {} {}s'.format(i+1, elapsed))
        eval_results = estimator.evaluate(input_fn=eval_input_fn)
        print(eval_results)
    predictions = np.array([p['classes'] for p in estimator.predict(input_fn=predict_input_fn)])

    cm = bl.ConfusionMatrix(labels)
    for truth, guess in zip(y_test, predictions):
        cm.add(truth, guess)

    print(cm.get_all_metrics())

def model_fn(features, labels, mode, params):

    if labels is not None:
        y = tf.one_hot(tf.reshape(labels, [-1, 1]), len(params['labels']))

    if mode == tf.estimator.ModeKeys.PREDICT:
        bl.tf.SET_TRAIN_FLAG(False)
        model = bl.model.create_model(embeddings, labels=params['labels'], word=features['word'], y=None, sess=None, **model_params)
        predictions = {
            'classes': model.best,
            'probabilities': model.probs,
            'logits': model.logits,
        }
        outputs = tf.estimator.export.PredictOutput(predictions['classes'])
        return tf.estimator.EstimatorSpec(mode, predictions=predictions, export_outputs={'classes': outputs})

    elif mode == tf.estimator.ModeKeys.EVAL:
        bl.tf.SET_TRAIN_FLAG(False)
        model = bl.model.create_model(embeddings, labels=params['labels'], word=features['word'], y=y, sess=None, **model_params)
        loss = model.create_loss()
        predictions = {
            'classes': model.best,
            'probabilities': model.probs,
            'logits': model.logits,
        }
        eval_metric_ops = {
            'accuracy': tf.metrics.accuracy(
                labels=labels, predictions=predictions['classes'])}
        return tf.estimator.EstimatorSpec(mode=mode, predictions=model.logits, loss=loss, eval_metric_ops=eval_metric_ops)

    bl.tf.SET_TRAIN_FLAG(True)
    model = bl.model.create_model(embeddings, labels=params['labels'], word=features['word'], y=y, sess=None, **model_params)
    loss = model.create_loss()
    #Optimizer = tf.train.GradientDescentOptimizer if args.optim == 'sgd' else tf.train.AdamOptimizer
    #print(Optimizer)
    #optimizer = Optimizer(learning_rate=args.lr)
    #train_op = optimizer.minimize(loss=loss, global_step=tf.train.get_global_step())
    colocate = True if args.gpus > 1 else False
    global_step, train_op = bl.tf.optz.optimizer(loss, optim=args.optim, eta=args.lr, colocate_gradients_with_ops=colocate)

    return tf.estimator.EstimatorSpec(mode=mode, predictions=model.logits,
                                      loss=loss,
                                      train_op=train_op)
params = {'labels': labels}

checkpoint_dir = '{}-{}'.format(args.checkpoint_dir, os.getpid())
if args.gpus > 1:
    config = tf.estimator.RunConfig(model_dir=checkpoint_dir, train_distribute=tf.contrib.distribute.MirroredStrategy(num_gpus=args.gpus))
else:
    config = tf.estimator.RunConfig(model_dir=checkpoint_dir)

estimator = tf.estimator.Estimator(model_fn=model_fn, config=config, params=params)

fit(estimator)

export_dir = '{}-{}'.format(args.export_dir, os.getpid())
estimator.export_savedmodel(export_dir_base=export_dir, serving_input_receiver_fn=server_input_fn)