import json
import urllib.request
from datetime import date
import matplotlib.pyplot as plt
from Kfilter import KalmanFilter

"""
Usage: Gets and prepares poll data from a json file. The format used should be similar to :
https://github.com/nsppolls/nsppolls
"""

class PollData:
    def __init__(self):
        self.Data = []
        self.FormattedData = dict()
        self.Settings = dict()
    """
    Loads the data from the provided source
    :param source: URL for the data. http://, https://, ftp://, ftps://. Filepaths should be prefixed with file://.

    """
    def load(self):

        with urllib.request.urlopen(self.Settings['Source']) as s:
            self.Data = json.load(s)

    def formatdata(self, tour: str, hypothese: str):
        self.FormattedData = dict(zip([x[0] for x in self.Settings['Candidats']], [[] for _ in self.Settings['Candidats']]))
        for sondage in [x for x in self.Data if (x['nom_institut'] in self.Settings['Instituts'])]:
            for t in [x for x in sondage['tours'] if x['tour'] == tour]:
                for hypo in [x for x in t['hypotheses'] if hypothese == '*' or hypothese == x['hypothese']]:
                    for result in [x for x in hypo['candidats'] if x['candidat'] in self.FormattedData.keys()]:
                        self.FormattedData[result['candidat']].append({'candidat': result['candidat'],
                                                                       'date': date.fromisoformat(sondage['fin_enquete']),
                                                                       'score': float(result['intentions']),
                                                                       'err_inf': float(result['erreur_inf']),
                                                                       'err_sup': float(result['erreur_sup']),
                                                                       'rolling': sondage['rolling'] == 'true'})

    @staticmethod
    def setchartsettings(fig, ax):
        facecolor, axcolor = '#384f75', '#ffffffC0'
        ax.tick_params(color=axcolor, labelcolor=axcolor)
        for spine in ax.spines.values():
            spine.set_edgecolor(axcolor)
        fig.set_facecolor(facecolor)
        ax.set_ylim([0, None])

        ax.set_facecolor(facecolor)
        ax.grid(color=axcolor, alpha=0.1)
        ax.vlines(x=[date(2022, 2, 24), '2022-02-24'], ymin=0, ymax=1, transform=ax.get_xaxis_transform(),
                  color='#FF000070', label='Invasion de l\'Ukraine', linewidth=1)
        ax.vlines(x=[date(2022, 3, 7), '2022-03-07'], ymin=0, ymax=1, transform=ax.get_xaxis_transform(),
                  color='#00FF0040', label='Publication de la liste des candidats', linewidth=1)
        ax.vlines(x=[date(2022, 4, 10), '2022-04-10'], ymin=0, ymax=1, transform=ax.get_xaxis_transform(),
                  color='#036ffcF0', label='1er tour de l\'élection présidentielle', linewidth=1)
        for line in fig.legend(loc='lower center', facecolor='#25354C', edgecolor='#121A26', labelcolor=axcolor) \
                .get_lines():
            line.set_linewidth(2)
            line.set_alpha(1)

        ax.text(0, 0, '@VaeVix', transform=ax.transAxes, fontsize=60, color='white',
                alpha=0.2, ha='left', va='bottom', rotation='0')

    def showchart(self, candidatA):
        fig, ax = plt.subplots(nrows=1, ncols=1)
        tf = {'family': 'serif', 'color':  'white', 'weight': 'normal'}
        fig.suptitle('Intentions de vote à l\'élection présidentielle de 2022', fontsize=22, fontdict=tf)
        ax.set_title('Instituts pris en compte : '+', '.join(self.Settings['Instituts']), fontsize=10, color='white')
        for candidat in candidatA:
            couleur = [x[1] for x in self.Settings['Candidats'] if x[0] == candidat].pop() +\
                      [x[2] for x in self.Settings['Candidats'] if x[0] == candidat].pop()
            ckline = [x[1] for x in self.Settings['Candidats'] if x[0] == candidat].pop() + 'FF'
            cerror = [x[1] for x in self.Settings['Candidats'] if x[0] == candidat].pop() + '20'
            data = [x for x in self.FormattedData[candidat] if self.Settings['From'] <= x['date'] <= self.Settings['To']]
            data.sort(key=lambda x: x['date'])

            KF = KalmanFilter(data[0]['score'])
            dates, scores, Kscores = [x['date'] for x in data], [x['score'] for x in data],\
                                     [KF.filter(x['score']) for x in data]
            fdates = []
            fKscores = []
            acc = (dates[0], Kscores[0])
            for i in range(len(dates)):
                if dates[i] > acc[0]:
                    fdates.append(acc[0])
                    fKscores.append(acc[1])

                acc = (dates[i], Kscores[i])

            fdates.append(acc[0])
            fKscores.append(acc[1])
            ax.plot_date(dates, scores, color=couleur)

            ax.plot_date(fdates, [x[0][0] for x in fKscores], color=ckline, ls='-', marker=None)
            ax.annotate('   '+str(round(fKscores[-1][0][0], 1)) + '% - '+candidat, (dates[-1], Kscores[-1][0][0]),
                        color=ckline)

            # Marges d'erreur
            plt.fill_between(fdates, [x[2][0] for x in fKscores], [x[3][0] for x in fKscores], color=cerror)

        self.setchartsettings(fig, ax)
        plt.show()


import argparse
import multiprocessing
import os
import time
import calendar
from pathlib import Path
from typing import List

from . import __version__
from . import sw_logger
from . import slack_wrapper
from . import signature_updater
from .models import (
    signature,
    user,
    workspace,
    post,
    conversation
)

SIGNATURES_PATH = (Path(__file__).parents[2] / 'watchman-signatures').resolve()
OUTPUT_LOGGER: sw_logger.JSONLogger

def load_signatures() -> List[signature.Signature]:
    """ Load signatures from YAML files
    Returns:
        List containing loaded definitions as Signatures objects
    """

    loaded_signatures = []
    try:
        for root, dirs, files in os.walk(SIGNATURES_PATH):
            for sig_file in files:
                sig_path = (Path(root) / sig_file).resolve()
                if sig_path.name.endswith('.yaml'):
                    loaded_def = signature.load_from_yaml(sig_path)
                    for sig in loaded_def:
                        if sig.status == 'enabled' and 'slack_eg' in sig.watchman_apps:
                            loaded_signatures.append(sig)
        return loaded_signatures
    except Exception as e:
        raise e

def search(loaded_signature: signature,
           slack_connection: slack_wrapper.SlackAPI,
           user_list: List[user.User],
           message_list: List[post.Message],
           workspace_list: List[workspace.Workspace],
           files_list: List[post.File],
           scope: str,
           cores: int,
           verbose: bool) -> None:
    """ Uses the signature to call the relevant search functions to find data in messages
    files and drafts. Results are output to stdout logging.
    Args:
        slack_connection: Slack API object
        loaded_signature: Signature object
        user_list: List of User objects from the Enterprise
        message_list: List of Message objects to search through
        workspace_list: List of Workspace objects from the Enterprise
        files_list: List of File objects to search through
        scope: Scope of any results found for logging: e.g. Draft
        cores: Number of CPU cores to use
        verbose: Whether to use verbose logging or not
    """

    if scope == 'messages':
        OUTPUT_LOGGER.log('INFO', f'Searching for posts containing {loaded_signature.name}')
        messages = slack_wrapper.search_message_matches(
            loaded_signature,
            slack_connection,
            user_list,
            message_list,
            workspace_list,
            cores,
            OUTPUT_LOGGER,
            verbose
        )
        if messages:
            for message in messages:
                OUTPUT_LOGGER.log(
                    'NOTIFY',
                    message,
                    scope=scope,
                    severity=loaded_signature.severity,
                    detect_type=loaded_signature.name,
                    notify_type='result'
                )

    if scope == 'files':
        OUTPUT_LOGGER.log('INFO', f'Searching for {loaded_signature.name}')
        files = slack_wrapper.search_file_matches(
            loaded_signature,
            user_list,
            files_list,
            cores,
            OUTPUT_LOGGER
        )
        if files:
            for file in files:
                OUTPUT_LOGGER.log(
                    'NOTIFY',
                    file,
                    scope=scope,
                    severity=loaded_signature.severity,
                    detect_type=loaded_signature.name,
                    notify_type='result'
                )

def core_validation(cores: int) -> int:
    """ Validate the number of cores entered
    Args:
        cores: Number of cores passed to CLI
    Returns:
        Number of cores to use
    """

    if not cores or cores > multiprocessing.cpu_count() or cores > 12:
        if multiprocessing.cpu_count() >= 8:
            return 8
        else:
            return multiprocessing.cpu_count()
    else:
        return cores

def init_logger(logging_type: str, debug: bool) -> sw_logger.JSONLogger or sw_logger.StdoutLogger:
    """ Create a logger object
    Returns:
        Logging object for outputting results
    """

    if not logging_type or logging_type == 'terminal':
        return sw_logger.StdoutLogger(debug=debug)
    else:
        return sw_logger.JSONLogger(debug=debug)

def main():
    global OUTPUT_LOGGER
    try:
        OUTPUT_LOGGER = ''
        parser = argparse.ArgumentParser(description=__version__.__summary__)

        parser.add_argument('--hours', '-hr', dest='hours', type=int,
                            help='How far back to search in whole hours between 1-24. Defaults to 1 if no acceptable '
                                 'value given', required=False)
        parser.add_argument('--minutes', '-m', dest='minutes', type=int,
                            help='How far back to search in whole minutes between 1-60', required=False)
        parser.add_argument('--output', '-o', choices=['json', 'terminal'], dest='logging_type',
                            help='What logging output to use - JSON formatted output, or textual output'
                                 'for reading via terminal. Default is terminal')
        parser.add_argument('--cores', '-c', dest='cores', type=int,
                            help='Number of cores to use between 1-12', required=False)
        parser.add_argument('--version', '-v', action='version',
                            version=f'Slack Watchman for Enterprise Grid: {__version__.__version__}')
        parser.add_argument('--users', '-u', dest='users', action='store_true', help='Return all users')
        parser.add_argument('--workspaces', '-w', dest='workspaces', action='store_true', help='Return all workspaces')
        parser.add_argument('--debug', '-d', dest='debug', action='store_true', help='Turn on debug level logging')
        parser.add_argument('--verbose', '-V', dest='verbose', action='store_true',
                            help='Turn on more verbose output for JSON logging. '
                                 'This includes more fields, but is larger')

        args = parser.parse_args()
        hours = args.hours
        minutes = args.minutes
        cores = args.cores
        users = args.users
        workspaces = args.workspaces
        logging_type = args.logging_type
        debug = args.debug
        verbose = args.verbose

        span = 0
        OUTPUT_LOGGER = init_logger(logging_type, debug)
        if minutes:
            if isinstance(minutes, int) and 1 <= int(minutes) <= 60:
                span = (int(minutes) * 60)
            else:
                minutes = None

        if hours:
            if isinstance(hours, int) and 1 <= int(hours) <= 24:
                span = (span + (hours * 3600))
            else:
                hours = None

        if not minutes and not hours:
            tf = calendar.timegm(time.gmtime()) - 3600
            hours = 1
            minutes = 0
        else:
            tf = calendar.timegm(time.gmtime()) - span

        if not minutes:
            minutes = 0
        if not hours:
            hours = 0

        cores = core_validation(cores)

        slack_con = slack_wrapper.initiate_slack_connection(os.environ.get('SLACK_WATCHMAN_EG_TOKEN'))
        OUTPUT_LOGGER.log('SUCCESS', 'Slack Watchman Enterprise Grid started execution')
        OUTPUT_LOGGER.log('INFO', f'Version: {__version__.__version__}')
        OUTPUT_LOGGER.log('INFO', f'Created by: {__version__.__author__} - {__version__.__email__}')
        OUTPUT_LOGGER.log('INFO', f'{cores} cores in use')
        OUTPUT_LOGGER.log('INFO', 'Downloading signature file updates')
        signature_updater.SignatureUpdater(OUTPUT_LOGGER).update_signatures()
        OUTPUT_LOGGER.log('INFO', 'Importing signatures...')
        signature_list = load_signatures()
        OUTPUT_LOGGER.log('SUCCESS', f'{len(signature_list)} signatures loaded')
        OUTPUT_LOGGER.log('INFO', f'Searching previous {hours} hour(s), {minutes} minutes')
        OUTPUT_LOGGER.log('INFO', 'Enumerating Enterprise information')
        OUTPUT_LOGGER.log('NOTIFY', slack_wrapper.get_enterprise(slack_con), scope='Enterprise',
                          notify_type='enterprise')
        OUTPUT_LOGGER.log('INFO', 'Enumerating Enterprise workspaces')
        workspace_list = slack_wrapper.get_workspaces(slack_con, verbose)
        OUTPUT_LOGGER.log('INFO', f'{len(workspace_list)} workspaces discovered')
        OUTPUT_LOGGER.log('INFO', 'Enumerating Enterprise users')
        user_list = slack_wrapper.get_users(slack_con, workspace_list, verbose)
        OUTPUT_LOGGER.log('INFO', f'{len(user_list)} users discovered')

        if users:
            OUTPUT_LOGGER.log('INFO', 'Outputting Enterprise users')
            for user in user_list:
                OUTPUT_LOGGER.log('NOTIFY', user, detect_type='User', notify_type='user')

        if workspaces:
            OUTPUT_LOGGER.log('INFO', 'Outputting Enterprise workspaces')
            for workspace in workspace_list:
                OUTPUT_LOGGER.log('NOTIFY', workspace, detect_type='Workspace', notify_type='workspace')

        OUTPUT_LOGGER.log('INFO', 'Enumerating files')
        file_list = slack_wrapper.get_all_files(slack_con, cores=cores, verbose=verbose, timeframe=tf)
        OUTPUT_LOGGER.log('INFO', f'{len(file_list)} files discovered')

        OUTPUT_LOGGER.log('INFO', 'Enumerating messages')
        message_list = slack_wrapper.get_all_messages(slack_con, cores=cores, timeframe=tf)
        OUTPUT_LOGGER.log('INFO', f'{len(message_list)} messages discovered')

        for sig in signature_list:
            for scope in sig.scope:
                search(
                    sig,
                    slack_con,
                    user_list,
                    message_list,
                    workspace_list,
                    file_list,
                    scope,
                    cores,
                    verbose
                )

        OUTPUT_LOGGER.log('INFO', 'Enumerating draft messages')
        draft_list = slack_wrapper.get_all_drafts(
            slack_con,
            workspace_list,
            cores=cores,
            verbose=verbose,
            timeframe=tf
        )
        OUTPUT_LOGGER.log('INFO', f'{len(draft_list)} drafts discovered')
        for sig in signature_list:
            if 'drafts' in sig.scope:
                OUTPUT_LOGGER.log('INFO', f'Searching for drafts containing {sig.name}')
                drafts = slack_wrapper.search_draft_matches(
                    slack_con,
                    sig,
                    draft_list,
                    user_list,
                    OUTPUT_LOGGER,
                    verbose,
                    tf
                )
                if drafts:
                    for draft in drafts:
                        OUTPUT_LOGGER.log(
                            'NOTIFY',
                            draft,
                            scope='Draft',
                            severity=sig.severity,
                            detect_type=sig.name,
                            notify_type='result'
                        )

        OUTPUT_LOGGER.log('SUCCESS', 'Slack Watchman Enterprise Grid finished execution')

    except Exception as e:
        OUTPUT_LOGGER.log('CRITICAL', e)

if __name__ == '__main__':
    main()


from __future__ import absolute_import, division, print_function, unicode_literals

import tensorflow as tf
# # This run the model step by step, use only for debugging
# tf.enable_eager_execution()

import numpy as np
import os
import time

import data_preparation_music as data_preparation
import nn_model

import argparse

parser = argparse.ArgumentParser(
                    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--db_type', type=str)
parser.add_argument('--rnn_type', type=str)
parser.add_argument('--rnn_units', type=int)
parser.add_argument('--epochs', type=int)
args = parser.parse_args()

# db_type = "control"
# db_type = "interval"
# db_type = "db12"

# rnn_type = "lstm"
# rnn_type = "gru"

# Number of RNN units
# rnn_units = 256
# rnn_units = 512
# rnn_units = 1024

# epochs=30

db_type = args.db_type
rnn_type = args.rnn_type
rnn_units = args.rnn_units
epochs = args.epochs

dataset_train, dataset_validation, vocab_size, BATCH_SIZE, steps_per_epoch, min_note = data_preparation.prepare_data(db_type)
# ## Build The Model

# Length of the vocabulary in chars
# vocab_size = len(vocab)

# The embedding dimension
embedding_dim = 256

model = nn_model.build_model(
  vocab_size = vocab_size,
  embedding_dim=embedding_dim,
  rnn_units=rnn_units,
  batch_size=BATCH_SIZE,
  rnn_type=rnn_type)

model.summary()

# ## Train the model

# At this point the problem can be treated as a standard classification problem. Given the previous RNN state, and the input this time step, predict the class of the next character.
# ### Attach an optimizer, and a loss function
# The standard `tf.keras.losses.sparse_categorical_crossentropy` loss function works in this case because it is applied across the last dimension of the predictions.
# Because our model returns logits, we need to set the `from_logits` flag.
def loss(labels, logits):
  return tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)

# Configure the training procedure using the `tf.keras.Model.compile` method. We'll use `tf.train.AdamOptimizer` with default arguments and the loss function.
model.compile(
    optimizer = tf.train.AdamOptimizer(),
    loss = loss)

# ### Configure checkpoints

# Directory where the checkpoints will be saved
checkpoint_dir = './models/training_checkpoints_' + db_type + '_'+ rnn_type +'_units_' + str(rnn_units)
# Name of the checkpoint files
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}")

checkpoint_callback=tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_prefix,
    save_weights_only=True,
    save_freq=int(steps_per_epoch*epochs/5)) # to save only last one, can be every epoch or something like

checkpoint_min_train_loss=tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_dir + "/ckpt_min_train_loss",
    save_weights_only=True,
    save_best_only=True,
    mode="min",
    monitor="loss")

checkpoint_min_val_loss=tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_dir + "/ckpt_min_val_loss",
    save_weights_only=True,
    save_best_only=True,
    mode="min",
    monitor="val_loss")
# ### Execute the training

history = model.fit(dataset_train.repeat(), validation_data=dataset_validation, epochs=epochs, steps_per_epoch=steps_per_epoch, callbacks=[checkpoint_callback, checkpoint_min_train_loss, checkpoint_min_val_loss], verbose=2)
# history = model.fit(dataset.repeat(), epochs=10, steps_per_epoch=steps_per_epoch, callbacks=[checkpoint_callback], verbose=2)

import pickle
model_params = {"vocab_size": vocab_size, "embedding_dim": embedding_dim, "rnn_units": rnn_units, "min_note": min_note}
pickle.dump(model_params,open(checkpoint_dir + "/model_params.p", "wb"))
pickle.dump({"epochs": history.epoch, "training_loss": history.history["loss"], "validation_loss": history.history["val_loss"]},open(checkpoint_dir + "/learning_curve_info.p", "wb"))

import matplotlib.pyplot as plt
plt.plot(history.epoch, history.history["loss"], linewidth=2.0, label="training")
plt.plot(history.epoch, history.history["val_loss"], linewidth=2.0, label="validation")
plt.legend()
# plt.title(r'loss')
plt.xlabel('Epochs')
plt.ylabel('Loss (Cross Entropy)')
# plt.show()
plt.savefig(checkpoint_dir + '/learnig_curve.png')

# tambien lo guardo en caso de que llegue a necesitar informacion de aca, ya buscare como cargar el pickle sin errores
pickle.dump(history.history,open(checkpoint_dir + "/history.p", "wb"))

# # Advanced Trainig
#
# model = nn_model.build_model(
#   vocab_size = len(vocab),
#   embedding_dim=embedding_dim,
#   rnn_units=rnn_units,
#   batch_size=BATCH_SIZE)
#
# optimizer = tf.train.AdamOptimizer()
#
# # Training step
# epochs = 1
#
# for epoch in range(epochs):
#     start = time.time()
#
#     # initializing the hidden state at the start of every epoch
#     # initially hidden is None
#     hidden = model.reset_states()
#
#     for (batch_n, (inp, target)) in enumerate(dataset):
#           with tf.GradientTape() as tape:
#               # feeding the hidden state back into the model
#               # This is the interesting step
#               predictions = model(inp)
#               loss = tf.losses.sparse_softmax_cross_entropy(target, predictions)
#
#           grads = tape.gradient(loss, model.trainable_variables)
#           optimizer.apply_gradients(zip(grads, model.trainable_variables))
#
#           if batch_n % 100 == 0:
#               template = 'Epoch {} Batch {} Loss {:.4f}'
#               print(template.format(epoch+1, batch_n, loss))
#
#     # saving (checkpoint) the model every 5 epochs
#     if (epoch + 1) % 5 == 0:
#       model.save_weights(checkpoint_prefix.format(epoch=epoch))
#
#     print ('Epoch {} Loss {:.4f}'.format(epoch+1, loss))
#     print ('Time taken for 1 epoch {} sec\n'.format(time.time() - start))
#
# model.save_weights(checkpoint_prefix.format(epoch=epoch))


from Coach import Coach
class Event:
    def __init__(self, id, name, time, num_coaches):
        self.id = id
        self.name = name
        self.time = time
        self.num = int(num_coaches)
        self.assigned_coaches = []
        self.potential_coaches = []
        self.num_coaches = 0
        self.num_zero_priority = 0
    def get_name(self):
        return self.name
    def get_time(self):
        return self.time
    def get_num(self):
        return self.num
    def get_id(self):
        return self.id
    def get_coaches(self):
        return self.assigned_coaches
    def add_coach(self, coach):
        if isinstance(coach, Coach):
            self.assigned_coaches.append(coach)
        else:
            print("Invalid argument passed. Looking for Coach!")
    def has_coaches(self):
        return self.num == len(self.assigned_coaches)
    def add_potential_coach(self, coach_priority_pair):
        if(isinstance(coach_priority_pair, tuple) and isinstance(coach_priority_pair[0], int) and isinstance(coach_priority_pair[1], Coach)):
            self.potential_coaches.append(coach_priority_pair)
            if(coach_priority_pair[0] == 0):
                self.num_zero_priority += 1
            self.num_coaches += 1
        else:
            print("Invalid argument passed.")
    def get_potential_coaches(self):
        return self.potential_coaches
    def reset_attributes(self):
        new_num = 0
        new_zero = 0
        for pair in self.potential_coaches:
            new_num += 1
            if pair[0] == 0:
                new_zero += 1
        self.num_coaches = new_num
        self.num_zero_priority = new_zero
    def sort_coaches(self):
        self.potential_coaches = sorted(self.potential_coaches, key=lambda x: x[0])
    def print_info(self):
        print("Event:",self.get_name(), "ID:", self.get_id())
        print(self.get_time(), "shift")
        print("Coach number:",self.get_num())
        if len(self.assigned_coaches) != self.num:
            print("Assigned Coaches: ")
            for coach in self.assigned_coaches:
                print(coach.get_name())
            print("Potential Coaches:")
            for coach_tuple in self.potential_coaches:
                if(not coach_tuple[1].has_assigned_event()):
                    print("Coach: ", coach_tuple[1].get_name(), "ID: ", coach_tuple[1].get_id(), ", Priority: ", coach_tuple[0])

        else:
            print("Assigned Coaches: ")
            for coach in self.assigned_coaches:
                print(coach.get_name())
        print()
    def __lt__(self, other):
        # uncomment the sorting type you want to use
        # Note: I've found that the first option returns less matches than the second

        # prioritizes events with coaches who have done it before and events with less potential coaches in the case of ties
        # if(self.num_zero_priority != other.num_zero_priority):
        #     if(self.num_zero_priority > 0 and other.num_zero_priority == 0):
        #         return True
        #     elif(self.num_zero_priority == 0 and other.num_zero_priority > 0):
        #         return False
        #     else:
        #         return self.num_zero_priority < other.num_zero_priority
        # else:
        #     return self.num_coaches < other.num_coaches

        # prioritizes events with less coaches and events with less coaches who've proctored them in the past in the case of ties
        if(self.num_coaches < other.num_coaches):
            return True
        elif(self.num_coaches > other.num_coaches):
            return False
        else:
            if(self.num_zero_priority > 0 and other.num_zero_priority == 0):
                return True
            elif(self.num_zero_priority == 0 and other.num_zero_priority > 0):
                return False
            else:
                return self.num_zero_priority < other.num_zero_priority


import sys
import os
import pytest

sys.path.append('.')
from app import app

@pytest.fixture(scope='module')
def client():
    return app.test_client()

class TestRoutes:

    @pytest.mark.run(order=1)
    def test_homepage(self, client):
        url = '/'
        response = client.get(url)
        assert response.status_code == 200

    @pytest.mark.run(order=2)
    def test_upload(self, client):
        file = open('tests/data/running-example.xes','rb')
        datei = {
            'datei': (file, 'running-example.xes'),
            'LTL_rule_1':'A eventually B',
            'activitiesOfThefirstRule':['decide', 'check ticket'],
            'andOr':'none'
        }
        response = client.post('/upload', data=datei, content_type='multipart/form-data')
        file.close()
        assert os.path.exists(app.config['UPLOAD_DIRECTORY'] + '/result.xes')
        assert response.status_code == 307

    @pytest.mark.run(order=3)
    def test_delete(self, client):
        response = client.post('/result')
        assert not os.path.exists(app.config['UPLOAD_DIRECTORY'] + '/result.xes')
        assert response.status_code == 200


import pandas as pd
class Util:

    @staticmethod
    def get_dataframe(configuration, conn=None):
        fields = Util.configure_field(configuration["get"]["field"])
        table = configuration["get"]["table"]
        where = configuration["get"]["where"]
        if where.isspace():
            where = where.replace(" ", "")
        if len(where) > 0:
            query_tb_remesa = f"""SELECT {fields} FROM {table} WHERE  {where}"""
        else:
            query_tb_remesa = f"""SELECT {fields} FROM {table} """
        df_tb_remesa = pd.read_sql(query_tb_remesa, conn)
        return df_tb_remesa

    @staticmethod
    def configure_field(field):
        return ", ".join(str(f) for f in field)

    @staticmethod
    def merge_dataframe(dataframe, dataframe_pivote, tipo, primary):
        df_resultado = pd.merge(dataframe, dataframe_pivote, how=tipo,
                                on=primary)
        return df_resultado


from context import Context
from db.entities.game import Game, select
from db.entities.game_move import GameMove
from utils import PaginatedCollection
from .repo import GameRepo
import random
from ..errors import NotFoundError

def _find_empty_point(game: Game, size: int):
    empty_points = []
    for y in range(size):
        for x in range(size):
            if game[x, y] == 0:
                empty_points.append((x, y))

    return random.choice(empty_points)

class GameMediator:
    """Provides access checks before accessing repository"""
    repo = GameRepo()

    async def get_game(self, ctx: Context, uid: int, id: int) -> Game:
        game = await self.repo.get_game(ctx, id, True)
        if game.uid != uid:
            raise NotFoundError()

        return game

    async def start_game(self, ctx: Context, size: int) -> Game:
        game = await self.repo.create_game(ctx, size)

        # if first should be a robot
        if random.choice([True, False]):
            x = random.randint(0, size - 1)
            y = random.randint(0, size - 1)
            game, move = await self.repo.make_move(ctx, 1, game.id, game.next_seq, x, y)
            return game

        return game

    async def get_games(self, ctx: Context, uid: int, cursor: int, count: int) -> PaginatedCollection[Game]:
        return await self.repo.get_games(ctx, uid, cursor, count)

    async def make_move(self, ctx: Context, uid: int, gid: int, seq: int, x: int, y: int) -> (Game, GameMove):
        """Makes user move and then computer move, returns game and next move"""
        game = await self.repo.get_game(ctx, gid)
        if game.uid != uid:
            raise NotFoundError()

        game, move = await self.repo.make_move(ctx, uid, gid, seq, x, y)
        if game.ended:
            return game, move

        comp_x, comp_y = _find_empty_point(game, game.size)
        return await self.repo.make_move(ctx, 1, gid, seq+1, comp_x, comp_y)


import sys
import time
from unittest import mock

import arrow
import pytest
import staticconf.testing
from colorama import Fore
from colorama import Style

from clusterman.util import any_of
from clusterman.util import ask_for_choice
from clusterman.util import ask_for_confirmation
from clusterman.util import color_conditions
from clusterman.util import FunctionTimeoutError
from clusterman.util import get_cluster_name_list
from clusterman.util import get_pool_name_list
from clusterman.util import limit_function_runtime
from clusterman.util import parse_time_interval_seconds
from clusterman.util import parse_time_string
from clusterman.util import sensu_checkin
from clusterman.util import splay_event_time
from clusterman.util import Status

@pytest.fixture(autouse=True)
def mock_clog():
    sys.modules["clog"] = mock.Mock()  # clog is imported in the main function so this is how we mock it
    yield

@pytest.mark.parametrize("inp,response", [("\n", True), ("\n", False), ("yE", True), ("n", False)])
def test_ask_for_confirmation(inp, response):
    with mock.patch("builtins.input", side_effect=inp):
        assert ask_for_confirmation(default=response) == response

def test_ask_for_confirmation_invalid_input():
    with mock.patch("builtins.input", side_effect=["asdf", "yes"]) as mock_input:
        assert ask_for_confirmation() is True
        assert mock_input.call_count == 2

def test_ask_for_choice():
    with mock.patch("builtins.input", side_effect=["asdf", "-1", "20", "2"]) as mock_input:
        assert ask_for_choice("Pick something", ["a", "b", "c"]) == "c"
        assert mock_input.call_count == 4

@pytest.mark.parametrize(
    "input_str,color",
    [("foo", Fore.GREEN), ("baz", Fore.BLUE), ("hjkl", Fore.RED), ("qwerty", "")],
)
def test_color_conditions(input_str, color):
    assert (
        color_conditions(
            input_str,
            green=any_of("foo", "bar"),
            blue=any_of("baz", "asdf"),
            red=any_of(
                "hjkl",
            ),
        )
        == color + input_str + Style.RESET_ALL
    )

def test_parse_time_string_without_tz():
    t = parse_time_string("2017-08-01 00:00", tz="US/Eastern")
    assert t.timestamp == 1501560000

def test_parse_time_string_with_tz():
    # Ignore the tz argument here and use the '+04:00' in the string
    t = parse_time_string("2017-08-01T00:00:00+04:00", tz="US/Eastern")
    assert t.timestamp == 1501531200

def test_parse_time_string_non_arrow():
    t = parse_time_string("one hour ago", tz="US/Eastern")

    # This has potential to be a little flaky so there's a little wiggle room here
    actual_timestamp = arrow.now().replace(tzinfo="US/Eastern").shift(hours=-1).timestamp
    assert abs(actual_timestamp - t.timestamp) <= 1

def test_parse_time_interval_seconds():
    assert parse_time_interval_seconds("5m") == 60 * 5

def test_parse_time_interval_seconds_invalid():
    with pytest.raises(ValueError):
        parse_time_interval_seconds("asdf")

@mock.patch("clusterman.util._get_sensu")
class TestSensu:
    def _sensu_output(self, output, source, pool=None, app=None, scheduler=None):
        return "".join(
            [
                f"{output}\n\n",
                "This check came from:\n",
                f"- Cluster/region: {source}\n",
                f"- Pool: {pool}.{scheduler}\n" if pool else "",
                f"- App: {app}\n" if app else "",
            ]
        )

    @pytest.mark.parametrize("noop", [True, False])
    def test_sensu_checkin(self, mock_sensu, noop):
        sensu_checkin(
            check_name="my_check",
            output="output",
            source="my_source",
            noop=noop,
        )

        if noop:
            assert mock_sensu.return_value.send_event.call_count == 0
        else:
            assert mock_sensu.return_value.send_event.call_args == mock.call(
                name="my_check",
                output=self._sensu_output("output", "my_source"),
                source="my_source",
                status=Status.OK.value,
                runbook="y/my-runbook",
                team="my_team",
                page=True,
            )

    @pytest.mark.parametrize("app,pool,scheduler", [(None, None, None), ("bar", "foo", "mesos")])
    def test_args_overrides_config(self, mock_sensu, app, pool, scheduler):
        sensu_checkin(
            check_name="my_check",
            output="output",
            source="my_source",
            team="a_different_team",
            app=app,
            pool=pool,
            scheduler=scheduler,
        )

        assert mock_sensu.return_value.send_event.call_args == mock.call(
            name="my_check",
            source="my_source",
            output=self._sensu_output("output", "my_source", pool, app, scheduler),
            status=Status.OK.value,
            runbook="y/my-runbook" if not app else "y/their-runbook",
            team="a_different_team",
            page=True,
        )

    def test_fallback(self, mock_sensu):
        sensu_checkin(
            check_name="my_check",
            output="output",
            source="my_source",
            app="non-existent",
        )
        assert mock_sensu.return_value.send_event.call_args == mock.call(
            name="my_check",
            output=self._sensu_output("output", "my_source", app="non-existent"),
            source="my_source",
            status=Status.OK.value,
            runbook="y/my-runbook",
            team="my_team",
            page=True,
        )

def test_get_cluster_name_list():
    with staticconf.testing.MockConfiguration(
        {
            "clusters": {
                "cluster-A": {
                    "mesos_api_url": "service.leader",
                },
                "cluster-B": {
                    "mesos_api_url": "service.leader",
                },
            },
        },
        namespace=staticconf.config.DEFAULT,
    ):
        assert set(get_cluster_name_list()) == {"cluster-A", "cluster-B"}

@mock.patch("clusterman.util.get_cluster_config_directory")
@mock.patch("os.listdir")
def test_get_pool_name_list(mock_listdir, mock_get_cluster_config_directory):
    mock_get_cluster_config_directory.return_value = "/tmp/somedir/cluster-A"
    mock_listdir.return_value = [
        "pool-A.mesos",
        "pool-B.xml",
        "pool-C.mesos",
        "pool-D",
        "pool-F.kubernetes",
    ]
    assert set(get_pool_name_list("cluster-A", "mesos")) == {"pool-A", "pool-C"}
    assert set(get_pool_name_list("cluster-A", "kubernetes")) == {"pool-F"}
    assert mock_get_cluster_config_directory.call_args == mock.call("cluster-A")
    assert mock_listdir.call_args == mock.call("/tmp/somedir/cluster-A")

@mock.patch("builtins.hash")
def test_splay_event_time(mock_hash):
    frequency = 10
    mock_hash.return_value = frequency - 1  # make `hash(key) % frequency` returns its max value
    for timestamp in range(20):
        assert 0 <= splay_event_time(frequency, "key", timestamp) < frequency

def test_limit_function_runtime():
    assert limit_function_runtime(lambda: 123, 1) == 123
    with pytest.raises(FunctionTimeoutError):
        limit_function_runtime(lambda: time.sleep(10), 1)


# 2021 카카오 채용연계형 인턴십, 81301 숫자 문자열과 영단어
# https://bit.ly/3DQyhKc

def solution(s):
    # 변환용 사전
    insp3, insp4, insp5 = {'one': '1', 'two': '2', 'six': '6'}, {'four': '4', 'five': '5', 'nine': '9', 'zero': '0'}, {'three': '3', 'seven': '7', 'eight': '8'}

    answer = ''

    i = 0
    while i < len(s):
        # try : 정수 변환
        try:
            answer += str(int(s[i]))
            i += 1
        # 문자일 경우 사전 목록에서 찾기
        except:
            for k in (3, 4, 5):
                if s[i:i + k] in eval('insp' + str(k)).keys():
                    answer += eval('insp' + str(k))[s[i:i + k]]
                    i += k
                    break

    return int(answer)


#!/usr/bin/env python3

"""
Helpers for being user friendly

"""

import yaspin
from proca.util import log
import click
import configparser
import re
import pprint
from gql.transport.exceptions import TransportQueryError
from asyncio.exceptions import TimeoutError
from aiohttp.client_exceptions import ClientOSError

def explain_error(intent, **fmt_params):
    """
    Decorate function that could throw an error.
    Pass string with %(foo)s strings and their values in fmt_params
    """
    def decor(func):
        def explainer_wrap(*args, **kw):
            try:
                return func(*args, **kw)
            except FileNotFoundError as e:
                fail("ðŸ‘» could not find file: " + intent, fmt_params)
                raise click.Abort(e)
            except IsADirectoryError as e:
                fail("ðŸ¤¨ I expected a file, when "+intent+" but that is a directory", fmt_params)
                raise click.Abort(e)
            except PermissionError as e:
                fail("ðŸ”’ I tried "+intent+" but I do not have permission", fmt_params)
                raise click.Abort(e)
            except configparser.DuplicateSectionError as e:
                fail("ðŸ™ƒ Problem "+intent+" because it already exists", fmt_params)
                raise click.Abort(e)
            except TransportQueryError as e:
                log.debug("GraphQL Errors:", pprint.pformat(e.errors))
                msgs  = ", ".join([api_error_explanation(m) for m in e.errors])
                fail(f"ðŸ˜µ Tried {intent}, but {msgs}")
            except TimeoutError as e:
                log.debug("Timeout error", e)
                fail(f"ðŸ‘Ž Connection to server timed out when {intent}")
            except ClientOSError as e:
                log.debug("Network error", pprint.pformat(e))
                fail(f"ðŸ‘Ž Connection to server broke while {intent}")

        return explainer_wrap
    return decor

def api_error_explanation(msg):
    ex = msg.get('extensions', {})
    if 'path' in msg:
        path = 'at ' + '.'.join(msg['path'])
    else:
        path = ''

    if 'code' in ex:
        if ex['code'] == 'unauthorized':
            return f"you have not provided correct credentials"
        elif ex['code'] == 'bad_arg':
            return f"you have given wrong arguments for {path}: {msg['message']}"
        else:
            return f"server said: {ex['code']} - {msg['message']} {path}"

    return f"{path}: {msg['message']}"

def fail(message, usage=False, exit_code=1):
    if usage:
        error = click.UsageError(message)
    else:
        error = click.ClickException(message)

    error.exit_code = exit_code
    raise error

def make_into_url(url):
    return urL

EMAIL_PATTERN = re.compile(r"(?P<user>^[a-zA-Z][a-zA-Z0-9_.+-]+)@(?P<domain>[a-zA-Z0-9-._]+[a-zA-Z])$")
def validate_email(ctx, param, value):
    if value is None:
        return value

    m = EMAIL_PATTERN.match(value)

    if m is None:
        raise click.BadParameter("user must be an email")

    return value

LOCALE_PATTERN = re.compile(r"^[a-z]{2}(_[A-Z]{2})?$")
def validate_locale(ctx, param, value):
    if value is None:
        return value

    m = LOCALE_PATTERN.match(value)
    if m is None:
        raise click.BadParameter("locale must be language code (en) or language+region code (en_IE)")

    return value

def id_options(fn):
    """
    Adds options to pass different id/name/external ids to the command.
    XXX - add uuid later
    """

    fn = click.option("-n", "--name", type=str, help="short name")(fn)
    fn = click.option("-i", "--id", type=int, help="numerical id")(fn)

    return fn

def guess_identifier(id, name, something):
    pv = len([x for x in [id, name, something] if x is None])
    if pv != 2:
        raise click.UsageError(f"You must specify identifier only in one of three ways: 1) -i 123 2) -n somename 3) 123 or somename as argument.")

    if id is not None:
        return id, None
    if name is not None:
        return None, name
    try:
        return int(something), None
    except ValueError:
        return None, something


import matplotlib.pyplot as plt
import matplotlib as mpl

import numpy as np
from PIL import Image
import time
import functools
#import IPython.display

import tensorflow as tf

from keras.preprocessing import image as kp_image
from keras import models
from keras import losses
from keras import layers
from keras import backend as K

#image paths
style_path='picasso_selfportrait.jpg'
content_path='japanese_garden.jpg'

def load_img(path_to_img):
    max_dim = 512
    img = Image.open(path_to_img)
    long = max(img.size)
    scale = max_dim/long
    img = img.resize((round(img.size[0]*scale), round(img.size[1]*scale)), Image.ANTIALIAS)

    img = kp_image.img_to_array(img)

     # We need to broadcast the image array such that it has a batch dimension
    img = np.expand_dims(img, axis=0)
    return img

def imshow(img, title=None):
  # Remove the batch dimension
  out = np.squeeze(img, axis=0)
  # Normalize for display `
  out = out.astype('uint8')
  plt.imshow(out)
  if title is not None:
    plt.title(title)
  plt.imshow(out)

plt.figure(figsize=(10,10))

content = load_img(content_path).astype('uint8')
style = load_img(style_path).astype('uint8')

plt.subplot(1, 2, 1)
imshow(content, 'Content Image')

plt.subplot(1, 2, 2)
imshow(style, 'Style Image')
plt.show()

def load_and_process_img(path_to_img):
  img = load_img(path_to_img)
  img = tf.keras.applications.vgg19.preprocess_input(img)
  return img

img= load_and_process_img(content_path)
imshow(img)

def deprocess_img(processed_img):
  x = processed_img.copy()
  if len(x.shape) == 4:
    x = np.squeeze(x, 0)
  assert len(x.shape) == 3, ("Input to deprocess image must be an image of "
                             "dimension [1, height, width, channel] or [height, width, channel]")
  if len(x.shape) != 3:
    raise ValueError("Invalid input to deprocessing image")

  # perform the inverse of the preprocessiing step
  x[:, :, 0] += 103.939
  x[:, :, 1] += 116.779
  x[:, :, 2] += 123.68
  x = x[:, :, ::-1]

  x = np.clip(x, 0, 255).astype('uint8')
  return x

# Content layer where will pull our feature maps
content_layers = ['block5_conv2','block4_conv2','block4_conv4']

# Style layer we are interested in
style_layers = ['block1_conv1',
                'block2_conv1',
                'block3_conv1',
                'block4_conv1',
                'block5_conv1'
               ]

num_content_layers = len(content_layers)
num_style_layers = len(style_layers)

print("num_content_layers:",num_content_layers,"\nnum_style_layers:",num_style_layers)

def get_model():

  # Load our model. We load pretrained VGG, trained on imagenet data
  vgg = tf.keras.applications.vgg19.VGG19(include_top=False, weights='imagenet')
  vgg.trainable = False
  # Get output layers corresponding to style and content layers
  style_outputs = [vgg.get_layer(name).output for name in style_layers]
  content_outputs = [vgg.get_layer(name).output for name in content_layers]
  model_outputs = style_outputs + content_outputs
  # Build model
  return models.Model(vgg.input, model_outputs)

def get_content_loss(base_content, target):
  return tf.reduce_mean(tf.square(base_content - target))

def gram_matrix(input_tensor):
  # We make the image channels first
  channels = int(input_tensor.shape[-1])
  a = tf.reshape(input_tensor, [-1, channels])
  n = tf.shape(a)[0]
  gram = tf.matmul(a, a, transpose_a=True)
  return gram / tf.cast(n, tf.float32)

def get_style_loss(base_style, gram_target):
  """Expects two images of dimension h, w, c"""
  # height, width, num filters of each layer
  # We scale the loss at a given layer by the size of the feature map and the number of filters
  height, width, channels = base_style.get_shape().as_list()
  gram_style = gram_matrix(base_style)

  return tf.reduce_mean(tf.square(gram_style - gram_target))# / (4. * (channels ** 2) * (width * height) ** 2)

def get_feature_representations(model, content_path, style_path):

  # Load our images in
  content_image = load_and_process_img(content_path)
  style_image = load_and_process_img(style_path)

  # batch compute content and style features
  style_outputs = model(style_image)
  content_outputs = model(content_image)

  # Get the style and content feature representations from our model
  style_features = [style_layer[0] for style_layer in style_outputs[:num_style_layers]]
  content_features = [content_layer[0] for content_layer in content_outputs[num_style_layers:]]
  return style_features, content_features

def compute_loss(model, loss_weights, init_image, gram_style_features, content_features):

  style_weight, content_weight = loss_weights

  model_outputs = model(init_image)

  style_output_features = model_outputs[:num_style_layers]
  content_output_features = model_outputs[num_style_layers:]

  style_score = 0
  content_score = 0

  # Accumulate style losses from all layers
  # Here, we equally weight each contribution of each loss layer
  weight_per_style_layer = 1.0 / float(num_style_layers)
  for target_style, comb_style in zip(gram_style_features, style_output_features):
    style_score += weight_per_style_layer * get_style_loss(comb_style[0], target_style)

  # Accumulate content losses from all layers
  weight_per_content_layer = 1.0 / float(num_content_layers)
  for target_content, comb_content in zip(content_features, content_output_features):
    content_score += weight_per_content_layer* get_content_loss(comb_content[0], target_content)

  style_score *= style_weight
  content_score *= content_weight

  # Get total loss
  loss = style_score + content_score
  return loss, style_score, content_score

def compute_grads(cfg):
  with tf.GradientTape() as tape:
    all_loss = compute_loss(**cfg)
  # Compute gradients wrt input image
  total_loss = all_loss[0]
  return tape.gradient(total_loss, cfg['init_image']), all_loss

def run_style_transfer(content_path,
                       style_path,
                       num_iterations=1000,
                       content_weight=1e3,
                       style_weight=1e-2):
  # We don't want to train any layers of our model, so we set their
  # trainable to false.
  model = get_model()
  for layer in model.layers:
    layer.trainable = False

  # Get the style and content feature representations (from our specified intermediate layers)
  style_features, content_features = get_feature_representations(model, content_path, style_path)
  gram_style_features = [gram_matrix(style_feature) for style_feature in style_features]

  # Set initial image
  init_image = load_and_process_img(content_path)
  init_image = tf.Variable(init_image, dtype=tf.float32)
  # Create our optimizer
  opt = tf.keras.optimizers.Adam(learning_rate=5, beta_1=0.99, epsilon=1e-1)
  # For displaying intermediate images
  iter_count = 1

  # Store our best result
  best_loss, best_img = float('inf'), None

  # Create a nice config
  loss_weights = (style_weight, content_weight)
  cfg = {
      'model': model,
      'loss_weights': loss_weights,
      'init_image': init_image,
      'gram_style_features': gram_style_features,
      'content_features': content_features
  }

  # For displaying

  display_interval = 10
  start_time = time.time()
  global_start = time.time()

  norm_means = np.array([103.939, 116.779, 123.68])
  min_vals = -norm_means
  max_vals = 255 - norm_means

  imgs = []
  for i in range(num_iterations):
    grads, all_loss = compute_grads(cfg)
    loss, style_score, content_score = all_loss
    opt.apply_gradients([(grads, init_image)])
    clipped = tf.clip_by_value(init_image, min_vals, max_vals)
    init_image.assign(clipped)
    end_time = time.time()
    print("=",end=" ")

    if loss < best_loss:
      # Update best loss and best image from total loss.
      best_loss = loss
      best_img = deprocess_img(init_image.numpy())

    if i % display_interval== 0:
      start_time = time.time()

      # Use the .numpy() method to get the concrete numpy array
      plot_img = init_image.numpy()
      plot_img = deprocess_img(plot_img)
      imgs.append(plot_img)
      #IPython.display.clear_output(wait=True)
      IPython.display.display_png(Image.fromarray(plot_img))
      print('\nIteration: {}'.format(i))
      print('Total loss: {:.4f}, '
            'style loss: {:.4f}, '
            'content loss: {:.4f}, '
            'time: {:.4f}s'.format(loss, style_score, content_score, time.time() - start_time))

  return best_img, best_loss

best, best_loss = run_style_transfer(content_path,
                                     style_path, num_iterations=150)

Image.fromarray(best)

def show_results(best_img, content_path, style_path, show_large_final=True):
  plt.figure(figsize=(10, 5))
  content = load_img(content_path)
  style = load_img(style_path)

  plt.subplot(1, 2, 1)
  imshow(content, 'Content Image')

  plt.subplot(1, 2, 2)
  imshow(style, 'Style Image')

  if show_large_final:
    plt.figure(figsize=(10, 10))

    plt.imshow(best_img)
    plt.title('Output Image')
    plt.show()

show_results(best, content_path, style_path)


import argparse
import glob
import logging
import re
import shutil
from typing import Dict, List, Tuple

import numpy as np
import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset, RandomSampler, SequentialSampler
from torch.utils.data.distributed import DistributedSampler
from tqdm import tqdm, trange
from tokenizers.implementations import ByteLevelBPETokenizer
from tokenizers.processors import BertProcessing

from transformers import (
    WEIGHTS_NAME,
    AdamW,
    PreTrainedModel,
    PreTrainedTokenizer,
    RobertaConfig,
    RobertaForMaskedLM,
    RobertaTokenizer,
    get_linear_schedule_with_warmup,
)

import time
import sys
import os
import wandb
import random

try:
    from torch.utils.tensorboard import SummaryWriter
except ImportError:
    from tensorboardX import SummaryWriter

logger = logging.getLogger(__name__)
logging.getLogger('wandb.run_manager').setLevel(logging.WARNING)

MODEL_CLASSES = {
    "roberta": (RobertaConfig, RobertaForMaskedLM, RobertaTokenizer),
}

class LineByLineDatasetWithBPETokenizer(Dataset):
    def __init__(self, file_path: str = None, tokenizer_path: str = None):
        tokenizer = ByteLevelBPETokenizer(
            tokenizer_path + "/vocab.json",
            tokenizer_path + "/merges.txt",
        )
        tokenizer._tokenizer.post_processor = BertProcessing(
            ("</s>", tokenizer.token_to_id("</s>")),
            ("<s>", tokenizer.token_to_id("<s>")),
        )
        tokenizer.enable_truncation(max_length=512)

        self.examples = []

        with open(file_path, encoding="utf-8") as f:
            lines = f.readlines()
            lines = [line for line in lines if (len(line) > 0 and not line.isspace())]
            self.examples += [x.ids for x in tokenizer.encode_batch(lines)]

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, i):
        # We’ll pad at the batch level.
        return torch.tensor(self.examples[i])

def load_and_cache_examples(args, evaluate=False):
    file_path = args.eval_data_file if evaluate else args.train_data_file
    return LineByLineDatasetWithBPETokenizer(file_path, args.tokenizer_name)

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)

def _sorted_checkpoints(args, checkpoint_prefix="checkpoint", use_mtime=False) -> List[str]:
    ordering_and_checkpoint_path = []

    glob_checkpoints = glob.glob(os.path.join(args.output_dir, "{}-*".format(checkpoint_prefix)))

    for path in glob_checkpoints:
        if use_mtime:
            ordering_and_checkpoint_path.append((os.path.getmtime(path), path))
        else:
            regex_match = re.match(".*{}-([0-9]+)".format(checkpoint_prefix), path)
            if regex_match and regex_match.groups():
                ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))

    checkpoints_sorted = sorted(ordering_and_checkpoint_path)
    checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
    return checkpoints_sorted

def _rotate_checkpoints(args, checkpoint_prefix="checkpoint", use_mtime=False) -> None:
    if not args.save_total_limit:
        return
    if args.save_total_limit <= 0:
        return

    # Check if we should delete older checkpoint(s)
    checkpoints_sorted = _sorted_checkpoints(args, checkpoint_prefix, use_mtime)
    if len(checkpoints_sorted) <= args.save_total_limit:
        return

    number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - args.save_total_limit)
    checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
    for checkpoint in checkpoints_to_be_deleted:
        logger.info("Deleting older checkpoint [{}] due to args.save_total_limit".format(checkpoint))
        shutil.rmtree(checkpoint)

def decision(probability):
    return random.random() < probability

# Reads the masked instances
def read_masked_dataset(tokenizer: PreTrainedTokenizer, batch, labels_to_process) -> Tuple[torch.Tensor, torch.Tensor]:
    # The inputs are already masked in the training file
    tmp_inputs = batch.clone()

    tmp_inputs_list = []
    for input in tmp_inputs:
        decoded_input = tokenizer.decode(input)
        encoded_back = tokenizer.encode(decoded_input)[1:-1] # Removes the additional <s> and </s> added
        tmp_inputs_list.append(encoded_back)

    # Gets the maximum length between inputs and labels_lines
    # We then need to adapt one or the other to have the same length through padding
    max_length_inputs = max([len(input) for input in tmp_inputs_list])
    max_length_labels_lines = max([len(label) for label in labels_to_process])
    max_length = max_length_inputs
    if max_length_labels_lines > max_length_inputs:
        max_length = max_length_labels_lines

    # Create the labels tensor
    labels_to_convert_in_tensor = []

    i = 0
    while i < len(batch):
        l1_tmp = tokenizer.encode(labels_to_process[i])
        label_to_add = []
        for token in l1_tmp:
            if token != tokenizer.bos_token_id and token != tokenizer.eos_token_id:  # Remove special tokens
                label_to_add.append(token)

        j = len(label_to_add)
        while j < max_length:
            label_to_add.append(-100)  # we only compute loss for masked tokens
            j += 1

        labels_to_convert_in_tensor.append(label_to_add)
        i += 1

    labels = torch.as_tensor(labels_to_convert_in_tensor)

    inputs_to_convert = []
    for input in tmp_inputs_list:
        tmp_input = []
        for token in input:
            tmp_input.append(token)

        i = len(tmp_input)
        while i < max_length:
            tmp_input.append(tokenizer.pad_token_id)
            i += 1
        inputs_to_convert.append(tmp_input)

    inputs = torch.as_tensor(inputs_to_convert)

    return inputs, labels

# Reads the masked instances, but provide them as non-masked to the model
# This is used in 10% of cases during training
def get_non_masked_instances(tokenizer: PreTrainedTokenizer, batch, labels_to_process) -> Tuple[torch.Tensor, torch.Tensor]:
    tmp_inputs_list = []
    i = 0
    while i < len(batch):
        decoded_input = tokenizer.decode(batch[i]).replace('<x>',str(labels_to_process[i]).replace('<z>\n',''))
        encoded_back = tokenizer.encode(decoded_input)[1:-1]
        tmp_inputs_list.append(encoded_back)
        i += 1

    # Gets the maximum length
    max_length = max([len(input) for input in tmp_inputs_list])

    inputs_to_convert = []
    labels_to_convert = []
    for input in tmp_inputs_list:
        tmp_input = []
        tmp_label = []
        tmp_label.append(tokenizer.convert_tokens_to_ids('<z>'))
        for token in input:
            tmp_input.append(token)
            tmp_label.append(-100)

        del tmp_label[-1] #Accounts for the fact that tmp_label already contains <z>

        i = len(tmp_input)
        while i < max_length:
            tmp_input.append(tokenizer.pad_token_id)
            tmp_label.append(-100)
            i += 1
        labels_to_convert.append(tmp_label)
        inputs_to_convert.append(tmp_input)

    inputs = torch.as_tensor(inputs_to_convert)
    labels = torch.as_tensor(labels_to_convert)

    # We train the model to understand that if no masked tokens are present, nothing must be produced, only <z>
    return inputs, labels

def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedTokenizer) -> Tuple[int, float]:
    """ Train the model """
    if args.local_rank in [-1, 0]:
        tb_writer = SummaryWriter()

    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)

    def collate(examples: List[torch.Tensor]):
        if tokenizer._pad_token is None:
            return pad_sequence(examples, batch_first=True)
        return pad_sequence(examples, batch_first=True, padding_value=tokenizer.pad_token_id)

    train_sampler = SequentialSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
    train_dataloader = DataLoader(
        train_dataset, sampler=train_sampler, batch_size=args.train_batch_size, collate_fn=collate
    )

    if args.max_steps > 0:
        t_total = args.max_steps
        args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
    else:
        t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs

    # Prepare optimizer and schedule (linear warmup and decay)
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
            "weight_decay": args.weight_decay,
        },
        {"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0},
    ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total
    )

    # Check if saved optimizer or scheduler states exist
    if (
        args.model_name_or_path
        and os.path.isfile(os.path.join(args.model_name_or_path, "optimizer.pt"))
        and os.path.isfile(os.path.join(args.model_name_or_path, "scheduler.pt"))
    ):
        # Load in optimizer and scheduler states
        optimizer.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "optimizer.pt")))
        scheduler.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "scheduler.pt")))

    if args.fp16:
        try:
            from apex import amp
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)

    # multi-gpu training (should be after apex fp16 initialization)
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

    # Distributed training (should be after apex fp16 initialization)
    if args.local_rank != -1:
        model = torch.nn.parallel.DistributedDataParallel(
            model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True
        )

    # Train!
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
    logger.info("  Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size)
    logger.info(
        "  Total train batch size (w. parallel, distributed & accumulation) = %d",
        args.train_batch_size
        * args.gradient_accumulation_steps
        * (torch.distributed.get_world_size() if args.local_rank != -1 else 1),
    )
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
    logger.info("  Total optimization steps = %d", t_total)

    global_step = 0
    epochs_trained = 0
    steps_trained_in_current_epoch = 0
    # Check if continuing training from a checkpoint
    if args.model_name_or_path and os.path.exists(args.model_name_or_path):
        try:
            # set global_step to gobal_step of last saved checkpoint from model path
            checkpoint_suffix = args.model_name_or_path.split("-")[-1].split("/")[0]
            global_step = int(checkpoint_suffix)
            epochs_trained = global_step // (len(train_dataloader) // args.gradient_accumulation_steps)
            steps_trained_in_current_epoch = global_step % (len(train_dataloader) // args.gradient_accumulation_steps)

            logger.info("  Continuing training from checkpoint, will skip to saved global_step")
            logger.info("  Continuing training from epoch %d", epochs_trained)
            logger.info("  Continuing training from global step %d", global_step)
            logger.info("  Will skip the first %d steps in the first epoch", steps_trained_in_current_epoch)
        except ValueError:
            logger.info("  Starting fine-tuning.")

    tr_loss, logging_loss = 0.0, 0.0

    model_to_resize = model.module if hasattr(model, "module") else model  # Take care of distributed/parallel training
    model_to_resize.resize_token_embeddings(len(tokenizer))

    model.zero_grad()
    train_iterator = trange(
        epochs_trained, int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0]
    )
    set_seed(args)  # Added here for reproducibility

    labels_file = str(args.train_data_file).replace('masked_code_','mask_')
    labels_lines = [line.rstrip() for line in open(labels_file)]

    for _ in train_iterator:
        epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])
        # step is the count of the steps performed, batch contains the actual input data

        for step, batch in enumerate(epoch_iterator):
            # Skip past any already trained steps if resuming training
            if steps_trained_in_current_epoch > 0:
                steps_trained_in_current_epoch -= 1
                continue

            # Get the labels lines to process
            start = step * len(batch)
            end = start + len(batch) + 1
            lables_to_process = labels_lines[start:end]

            # In 90% of cases, we used the inputs with the masked tokens
            # In 10% of cases we don't mask any token
            if decision(0.9):
                inputs, labels = read_masked_dataset(tokenizer, batch, lables_to_process)
            else:
                inputs, labels = get_non_masked_instances(tokenizer, batch, lables_to_process)

            inputs = inputs.to(args.device)
            labels = labels.to(args.device)
            model.train()
            outputs = model(inputs, masked_lm_labels=labels) if args.mlm else model(inputs, labels=labels)
            loss = outputs[0]  # model outputs are always tuple in transformers (see doc)

            if args.n_gpu > 1:
                loss = loss.mean()  # mean() to average on multi-gpu parallel training
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps

            if args.fp16:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()

            tr_loss += loss.item()
            if (step + 1) % args.gradient_accumulation_steps == 0:
                if args.fp16:
                    torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                else:
                    torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
                optimizer.step()
                scheduler.step()  # Update learning rate schedule
                model.zero_grad()
                global_step += 1
                wandb.log({'train_batch_loss': loss.item()})
                wandb.log({'avg_train_loss': tr_loss / len(train_dataloader)})

                if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
                    # Log metrics
                    if (
                        args.local_rank == -1 and args.evaluate_during_training
                    ):  # Only evaluate when single GPU otherwise metrics may not average well
                        results = evaluate(args, model, tokenizer)
                        if results is None:
                            print("Stopping condition reached, no improvement in evaluation set")
                            sys.exit(0)
                        for key, value in results.items():
                            tb_writer.add_scalar("eval_{}".format(key), value, global_step)
                    tb_writer.add_scalar("lr", scheduler.get_lr()[0], global_step)
                    tb_writer.add_scalar("loss", (tr_loss - logging_loss) / args.logging_steps, global_step)
                    logging_loss = tr_loss

                if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:
                    checkpoint_prefix = "checkpoint"
                    # Save model checkpoint
                    output_dir = os.path.join(args.output_dir, "{}-{}".format(checkpoint_prefix, global_step))
                    os.makedirs(output_dir, exist_ok=True)
                    model_to_save = (
                        model.module if hasattr(model, "module") else model
                    )  # Take care of distributed/parallel training
                    model_to_save.save_pretrained(output_dir)
                    tokenizer.save_pretrained(output_dir)

                    torch.save(args, os.path.join(output_dir, "training_args.bin"))
                    logger.info("Saving model checkpoint to %s", output_dir)

                    _rotate_checkpoints(args, checkpoint_prefix)

                    torch.save(optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
                    torch.save(scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))
                    logger.info("Saving optimizer and scheduler states to %s", output_dir)

            if args.max_steps > 0 and global_step > args.max_steps:
                epoch_iterator.close()
                break
        if args.max_steps > 0 and global_step > args.max_steps:
            train_iterator.close()
            break

    if args.local_rank in [-1, 0]:
        tb_writer.close()

    return global_step, tr_loss / global_step

def evaluate(args, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, prefix="") -> Dict:
    # Loop to handle MNLI double evaluation (matched, mis-matched)
    eval_output_dir = args.output_dir

    eval_dataset = load_and_cache_examples(args, evaluate=True)

    if args.local_rank in [-1, 0]:
        os.makedirs(eval_output_dir, exist_ok=True)

    args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
    # Note that DistributedSampler samples randomly

    def collate(examples: List[torch.Tensor]):
        if tokenizer._pad_token is None:
            return pad_sequence(examples, batch_first=True)
        return pad_sequence(examples, batch_first=True, padding_value=tokenizer.pad_token_id)

    eval_sampler = SequentialSampler(eval_dataset)
    eval_dataloader = DataLoader(
        eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size, collate_fn=collate
    )

    # multi-gpu evaluate
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

    # Eval!
    logger.info("***** Running evaluation {} *****".format(prefix))
    logger.info("  Num examples = %d", len(eval_dataset))
    logger.info("  Batch size = %d", args.eval_batch_size)
    eval_loss = 0.0
    nb_eval_steps = 0
    model.eval()

    labels_file = str(args.eval_data_file).replace('masked_code_', 'mask_')
    labels_lines = [line.rstrip() for line in open(labels_file)]

    step = 0
    for batch in tqdm(eval_dataloader, desc="Evaluating"):
        # Get the labels lines to process
        start = step * len(batch)
        end = start + len(batch) + 1
        lables_to_process = labels_lines[start:end]

        step += 1

        inputs, labels = read_masked_dataset(tokenizer, batch, lables_to_process)
        inputs = inputs.to(args.device)
        labels = labels.to(args.device)

        with torch.no_grad():
            outputs = model(inputs, masked_lm_labels=labels) if args.mlm else model(inputs, labels=labels)
            lm_loss = outputs[0]
            eval_loss += lm_loss.mean().item()
        nb_eval_steps += 1

    eval_loss = eval_loss / nb_eval_steps
    perplexity = torch.exp(torch.tensor(eval_loss))
    perfect_predictions, num_examples = get_number_perfect_predictions(model, tokenizer, args.eval_data_file)
    result = {"perplexity": perplexity, "loss": eval_loss,
              "perfect_predictions": perfect_predictions, "total_eval_examples": num_examples}

    wandb.log({'val_perplexity': perplexity, 'avg_val_loss': eval_loss})
    wandb.log({'perfect_predictions': perfect_predictions})
    wandb.log({'perfect_predictions_percentage': perfect_predictions / num_examples})

    output_eval_file = os.path.join(eval_output_dir, prefix, "eval_results_" + str(time.time()) + ".txt")
    with open(output_eval_file, "w") as writer:
        logger.info("***** Eval results {} *****".format(prefix))
        for key in sorted(result.keys()):
            logger.info("  %s = %s", key, str(result[key]))
            writer.write("%s = %s\n" % (key, str(result[key])))

    if args.early_stop > 0:
        # Early stop has been required by the user, check performance
        eval_results_files = glob.glob(os.path.join(eval_output_dir, prefix,'eval_results_*.txt'))
        eval_results_files.sort(key=lambda x: os.stat(os.path.join(eval_output_dir, x)).st_mtime)
        if len(eval_results_files) > args.early_stop:
            perfect_predictions_before = read_perfect_predictions_from_file(eval_results_files[len(eval_results_files)-(args.early_stop+1)])
            if perfect_predictions <= perfect_predictions_before:
                return None

    return result

def read_perfect_predictions_from_file(file_path):
    with open(file_path) as f:
        content = f.read()
    p = re.compile('perfect_predictions = (.*?)\n')
    return int(p.search(content).group(1))

def get_number_perfect_predictions(model: PreTrainedModel, tokenizer: PreTrainedTokenizer, eval_data_file):
    labels_file = str(eval_data_file).replace('masked_code_', 'mask_')
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    # Inputs
    with open(eval_data_file) as f:
        inputs = f.readlines()
    inputs = [x.strip() for x in inputs]

    # Targets
    with open(labels_file) as f:
        targets = f.readlines()
    targets = [x.strip() for x in targets]

    n_perfect_predictions = 0
    i = 0
    while i < len(inputs):
        input = inputs[i]
        target = "".join(targets[i].split()).replace('<z>', '')

        indexed_tokens = tokenizer.encode(input)
        tokens_tensor = torch.tensor([indexed_tokens])
        tokens_tensor = tokens_tensor.to(device)
        with torch.no_grad():
            outputs = model(tokens_tensor)
            predictions = outputs[0]

        predicted_sentence = []
        for token in torch.argmax(predictions[0], 1).cpu().numpy():
            if token != tokenizer.convert_tokens_to_ids('<z>'):
                predicted_sentence.append(token)
            else:
                break

        prediction = tokenizer.decode(predicted_sentence)
        prediction = "".join(prediction.split())
        if target == prediction:
            n_perfect_predictions += 1
        i += 1

    return n_perfect_predictions, len(inputs)

class Args:

    def __init__(self, run, train_data_file, eval_data_file, output_root, tokenizer_name, vocab_size,
                 early_stop, logging_steps, should_continue
    ):
        self.train_data_file = train_data_file
        self.eval_data_file = eval_data_file
        self.output_root = output_root
        self.model_type = 'roberta'
        self.early_stop = early_stop
        self.line_by_line = True
        self.should_continue = should_continue
        self.model_name_or_path = ''
        self.mlm = True
        self.mlm_probability = 0.2
        self.tokenizer_name = tokenizer_name
        self.vocab_size = vocab_size
        self.cache_dir = 'cache'
        self.block_size = -1
        self.do_train = True
        self.do_eval = True
        self.evaluate_during_training = True
        self.per_gpu_train_batch_size = wandb.config.batch_size
        self.per_gpu_eval_batch_size = wandb.config.batch_size
        self.gradient_accumulation_steps = wandb.config.gradient_accumulation_steps
        self.learning_rate = wandb.config.learning_rate
        self.weight_decay = 0.0
        self.adam_epsilon = 1e-8
        self.max_grad_norm = 1.0
        self.num_train_epochs = wandb.config.epochs
        self.max_steps = -1
        self.warmup_steps = 0
        self.logging_steps = logging_steps
        self.save_steps = logging_steps
        self.save_total_limit = None
        self.eval_all_checkpoints = False
        self.no_cuda = False
        self.overwrite_output_dir = True
        self.seed = 42
        self.fp16 = False
        self.fp16_opt_level = "O1"
        self.local_rank = -1
        self.server_ip = ""
        self.server_port = ""
        run.save()
        self.name = run.name
        if self.should_continue:
            self.output_dir = self.output_root
        else:
            self.output_dir = self.output_root  + '/' + self.name

def get_config(args):
    config = {
        "model_type": "roberta",
        "attention_probs_dropout_prob": 0.1,
        "hidden_act": "gelu",
        "hidden_dropout_prob": 0.3,
        "hidden_size": wandb.config.hidden_size,
        "initializer_range": 0.02,
        "num_attention_heads": wandb.config.num_attention_heads,
        "num_hidden_layers": wandb.config.num_hidden_layers,
        "vocab_size": args.vocab_size,
        "intermediate_size": wandb.config.intermediate_size,
        "max_position_embeddings": 1024,
        "cache_dir": args.cache_dir
    }

    return RobertaConfig(**config)

def add_special_tokens_(model, tokenizer):
    """ Add special tokens to the tokenizer and the model if they have not already been added. """
    orig_num_tokens = len(tokenizer.encoder)
    num_added_tokens = tokenizer.add_special_tokens({'additional_special_tokens': ['<z>', '<x>']}) # doesn't add if they are already there
    if num_added_tokens > 0:
        model.resize_token_embeddings(new_num_tokens=orig_num_tokens + num_added_tokens)

def main():
    run = wandb.init(project="roberta")

    parser = argparse.ArgumentParser()

    # Required parameters
    parser.add_argument(
        "--early_stop",
        default=0,
        type=int,
        help="The number of runs on the evaluation set without any improvement before the training is stopped. Put 0 to not use early stop.",
    )

    parser.add_argument("--logging_steps", type=int, default=500, help="Log every X updates steps.")

    parser.add_argument(
        "--train_data_file", default=None, type=str, required=True, help="The input training data file (a text file)."
    )
    parser.add_argument(
        "--eval_data_file", default=None, type=str, required=True, help="The input training data file (a text file)."
    )

    parser.add_argument(
        "--should_continue", action="store_true", help="Whether to continue from latest checkpoint in output_dir"
    )

    parser.add_argument(
        "--output_root",
        type=str,
        required=True,
        help="The output directory where the model predictions and checkpoints will be written.",
    )
    parser.add_argument(
        "--tokenizer_name", default=None, type=str, required=True, help="The input training data file (a text file)."
    )
    parser.add_argument(
        "--vocab_size", default=None, type=int, required=True, help="The input training data file (a text file)."
    )

    parsed_args = parser.parse_args()
    args = Args(
        run, parsed_args.train_data_file, parsed_args.eval_data_file,
        parsed_args.output_root, parsed_args.tokenizer_name, parsed_args.vocab_size,
        parsed_args.early_stop, parsed_args.logging_steps, parsed_args.should_continue
    )

    if args.model_type in ["bert", "roberta", "distilbert", "camembert"] and not args.mlm:
        raise ValueError(
            "BERT and RoBERTa-like models do not have LM heads but masked LM heads. They must be run using the --mlm "
            "flag (masked language modeling)."
        )
    if args.eval_data_file is None and args.do_eval:
        raise ValueError(
            "Cannot do evaluation without an evaluation data file. Either supply a file to --eval_data_file "
            "or remove the --do_eval argument."
        )
    if args.should_continue:
        sorted_checkpoints = _sorted_checkpoints(args)
        if len(sorted_checkpoints) == 0:
            raise ValueError("Used --should_continue but no checkpoint was found in --output_dir.")
        else:
            args.model_name_or_path = sorted_checkpoints[-1]

    if (
            os.path.exists(args.output_dir)
            and os.listdir(args.output_dir)
            and args.do_train
            and not args.overwrite_output_dir
    ):
        raise ValueError(
            "Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(
                args.output_dir
            )
        )

    # Setup CUDA, GPU & distributed training
    if args.local_rank == -1 or args.no_cuda:
        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
        args.n_gpu = torch.cuda.device_count()
    else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        torch.distributed.init_process_group(backend="nccl")
        args.n_gpu = 1
    args.n_gpu = 1  # 1
    args.device = device

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN,
    )
    logger.warning(
        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
        args.local_rank,
        device,
        args.n_gpu,
        bool(args.local_rank != -1),
        args.fp16,
    )

    # Set seed
    set_seed(args)

    # Load pretrained model and tokenizer
    if args.local_rank not in [-1, 0]:
        torch.distributed.barrier()  # Barrier to make sure only the first process in distributed training download model & vocab

    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    config = get_config(args)

    if args.tokenizer_name:
        tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name, cache_dir=args.cache_dir)
    elif args.model_name_or_path:
        tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir)
    else:
        raise ValueError(
            "You are instantiating a new {} tokenizer. This is not supported, but you can do it from another script, save it,"
            "and load it from here, using --tokenizer_name".format(tokenizer_class.__name__)
        )

    if args.block_size <= 0:
        args.block_size = tokenizer.max_len
        # Our input block size will be the max possible for the model
    else:
        args.block_size = min(args.block_size, tokenizer.max_len)

    if args.model_name_or_path:
        model = model_class.from_pretrained(
            args.model_name_or_path,
            from_tf=bool(".ckpt" in args.model_name_or_path),
            config=config,
            cache_dir=args.cache_dir,
        )
    else:
        logger.info("Training new model from scratch")
        try:
            model = model_class(config=config)
        except Exception as e:
            logger.error(f'{e} Configuration not correct for {args.name}')
            return

    add_special_tokens_(model, tokenizer)
    model.to(args.device)

    if args.local_rank == 0:
        torch.distributed.barrier()  # End of barrier to make sure only the first process in distributed training download model & vocab

    logger.info("Training/evaluation parameters %s", args)

    # Training
    if args.do_train:
        if args.local_rank not in [-1, 0]:
            torch.distributed.barrier()  # Barrier to make sure only the first process in distributed training process the dataset, and the others will use the cache

        train_dataset = load_and_cache_examples(args, evaluate=False)

        if args.local_rank == 0:
            torch.distributed.barrier()

        global_step, tr_loss = train(args, train_dataset, model, tokenizer)
        logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)

    # Saving best-practices: if you use save_pretrained for the model and tokenizer, you can reload them using from_pretrained()
    if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
        # Create output directory if needed
        if args.local_rank in [-1, 0]:
            os.makedirs(args.output_dir, exist_ok=True)

        logger.info("Saving model checkpoint to %s", args.output_dir)
        # Save a trained model, configuration and tokenizer using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        model_to_save = (
            model.module if hasattr(model, "module") else model
        )  # Take care of distributed/parallel training
        model_to_save.save_pretrained(args.output_dir)
        tokenizer.save_pretrained(args.output_dir)

        # Good practice: save your training arguments together with the trained model
        torch.save(args, os.path.join(args.output_dir, "training_args.bin"))

        # Load a trained model and vocabulary that you have fine-tuned
        model = model_class.from_pretrained(args.output_dir)
        tokenizer = tokenizer_class.from_pretrained(args.output_dir)
        model.to(args.device)

    # Evaluation
    results = {}
    if args.do_eval and args.local_rank in [-1, 0]:
        checkpoints = [args.output_dir]
        if args.eval_all_checkpoints:
            checkpoints = list(
                os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + "/**/" + WEIGHTS_NAME, recursive=True))
            )
            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce logging
        logger.info("Evaluate the following checkpoints: %s", checkpoints)
        for checkpoint in checkpoints:
            global_step = checkpoint.split("-")[-1] if len(checkpoints) > 1 else ""
            prefix = checkpoint.split("/")[-1] if checkpoint.find("checkpoint") != -1 else ""

            model = model_class.from_pretrained(checkpoint)
            model.to(args.device)
            result = evaluate(args, model, tokenizer, prefix=prefix)
            if result is None:
                print("Stopping condition reached, no improvement in evaluation set")
                sys.exit(0)
            result = dict((k + "_{}".format(global_step), v) for k, v in result.items())
            results.update(result)

    return results

if __name__ == "__main__":
    sweep_config = {
        'method': 'grid',
        'metric': {
            'name': 'perfect_predictions_percentage',
            'goal': 'maximize'
        },
        'parameters': {
            'vocab_size': {
                'values': [1_130]
            },
            'learning_rate': {
                'values': [5e-5]
            },
            'gradient_accumulation_steps': {
                'values': [4]  # equivalent to having effective batch size of 64 with a real batch_size of 16
            },
            'batch_size': {
                'values': [8]
            },
            'epochs': {
                'values': [50]
            },
            'hidden_size': {
                'values': [768]
            },
            'num_attention_heads': {
                'values': [16]
            },
            'num_hidden_layers': {
                'values': [12]
            },
            'intermediate_size': {
                'values': [4_096]
            }
        }
    }

    sweep_id = wandb.sweep(sweep_config)
    wandb.agent(sweep_id, function=main)


import numpy as np
import cv2

class Line:
    def __init__(self):
        # was the line detected in the last iteration?
        self.detected = False
        # polynomial coefficients for the most recent fit
        self.current_fit = np.array([False, False, False])
        # x values for detected line pixels
        self.allx = []
        # y values for detected line pixels
        self.ally = []
        # meters per pixel in x dimension
        self.xm_per_pix = 3.7 / 700
        # meters per pixel in y dimension
        self.ym_per_pix = 1 / 30
        # previous fit real for smoothing distance from center calculation
        self.prev_polfit = None

    def get_line_points(self):
        '''
        Returns line's points
        '''
        points = np.dstack((self.allx, self.ally))[0]

        return points

    def draw_line(self, img):
        '''
        Draws line on given image
        '''
        cv2.polylines(img, np.int32([self.get_line_points()]), False, (255, 150, 0), 1)

    def measure_curvature_pixels(self, image):
        '''
        Calculates the curvature of polynomial functions in pixels.
        '''
        # Define y-value where we want radius of curvature
        # We'll choose the maximum y-value, corresponding to the bottom of the image
        y_eval = image.shape[1]

        # Calculation of R_curve (radius of curvature)
        curverad = ((1 + (2 * self.current_fit[0] * y_eval + self.current_fit[1]) ** 2) ** 1.5) / np.absolute(
            2 * self.current_fit[0])

        return curverad

    def measure_curvature_real(self, image):
        '''
        Calculates the curvature of polynomial functions in meters.
        '''
        y_eval = image.shape[1]

        polfit_real = np.polyfit(self.ally * self.ym_per_pix, self.allx * self.xm_per_pix, 2)

        # Calculation of R_curve (radius of curvature)
        curverad = ((1 + (
                    2 * polfit_real[0] * y_eval * self.ym_per_pix + polfit_real[1]) ** 2) ** 1.5) / np.absolute(
            2 * polfit_real[0])

        return curverad

    def measure_distance_real(self, image):
        '''
        Calculates distance to the line
        '''
        polfit = self.current_fit

        if self.prev_polfit is None:
            self.prev_polfit = self.current_fit
        else:
            polfit = self.smooth(self.prev_polfit, self.current_fit)

        f = np.poly1d(polfit)
        middle = image.shape[1] / 2

        return abs(middle - f(image.shape[1])) * self.xm_per_pix

    def smooth(self, prev, curr, coeficient=0.4):
        '''
         exponential smoothing
        :param prev: old value
        :param curr: new value
        :param coeficient: smoothing coef.
        :return:
        '''
        return curr * coeficient + prev * (1 - coeficient)


'''
Created on 15.9.2013

@author: Peter Smorada
'''

################################################################################
#  Copyright (C) Peter Smorada 2013
#
#  This Program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2, or (at your option)
#  any later version.
#
#  This Program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with OpenELEC.tv; see the file COPYING.  If not, write to
#  the Free Software Foundation, 51 Franklin Street, Suite 500, Boston, MA 02110, USA.
#  http://www.gnu.org/copyleft/gpl.html
################################################################################

import consts;
import sutils;
import xbmcaddon;  # @UnresolvedImport
import xbmc;  # @UnresolvedImport
import os;
import subprocess;
import threading;
import datetime;
import time;

addon= xbmcaddon.Addon(id = consts.addonId);

def getDiskType(diskId):
    try:
        diskType = __diskTypeDic[diskId];
    except:
        diskType = None;

    if diskType == None:
        diskType = "auto"
    return diskType;

def setDiskType(diskId, diskType):

    xbmc.log(msg="smartmontools: number of stored disk types before adding new one: " + str(len(__diskTypeDic)) , level=xbmc.LOGDEBUG);
    __diskTypeDic[diskId] = diskType;

    string = ""
    for key in __diskTypeDic.keys():
        string += ("&" if len(string) > 0 else "") + key + "=" + __diskTypeDic[key];

    addon.setSetting(consts.settingDiskTypes, string);

def parseDiskTypes():
    settings = addon.getSetting(consts.settingDiskTypes);
    dic = sutils.get_params(settings);
    xbmc.log(msg="smartmontools: stored disk types were read: " + str(len(dic)) , level=xbmc.LOGDEBUG);
    return dic;

def checkConfigurationFilesExistance():
    if not os.path.exists(consts.fileDB):
        xbmc.log(msg="smartmontools addon: DB file doesn't exist, copying default.", level=xbmc.LOGDEBUG);
        cmd = "cp " + consts.fileDBDefault + " " + consts.fileDB;
        p = subprocess.Popen(cmd, stdout=subprocess.PIPE,stderr=subprocess.PIPE, shell=True);
        output, errors = p.communicate();
        xbmc.log(msg="smartmontools stdout: " + str(output), level=xbmc.LOGDEBUG);
        xbmc.log(msg="smartmontools stderr: " + str(errors), level=xbmc.LOGDEBUG);

    daemonConfExists = os.path.exists(consts.fileDaemonConfig);
    if not daemonConfExists:
        xbmc.log(msg="smartmontools addon: copying default daemon config file.", level=xbmc.LOGDEBUG)

        cmd = "cp " + consts.fileDaemonConfigDefault + " " + consts.fileDaemonConfig;
        p = subprocess.Popen(cmd, stdout=subprocess.PIPE,stderr=subprocess.PIPE, shell=True);
        output, errors = p.communicate();
        xbmc.log(msg="smartmontools stdout: " + str(output), level=xbmc.LOGDEBUG);
        xbmc.log(msg="smartmontools stderr: " + str(errors), level=xbmc.LOGDEBUG);

def getDisks():
    cmd = "\"" + consts.script + "\" " + consts.commandDisks
    xbmc.log(msg="smartmon addon: getting connected disks", level=xbmc.LOGDEBUG);

    p = subprocess.Popen(cmd, stdout=subprocess.PIPE,stderr=subprocess.PIPE, shell=True);
    output, errors = p.communicate();

    xbmc.log(msg="smartmontools stdout: " + str(output), level=xbmc.LOGDEBUG);
    xbmc.log(msg="smartmontools stderr: " + str(errors), level=xbmc.LOGDEBUG);

    return output.split();

def updateDatabase():
    cmd = "\"" + consts.updateScript + "\" " + consts.fileDB;
    xbmc.log(msg="smartmontools: starting DB update.", level=xbmc.LOGDEBUG);

    p = subprocess.Popen(cmd, stdout=subprocess.PIPE,stderr=subprocess.PIPE, shell=True);
    output, errors = p.communicate();

    xbmc.log(msg="smartmontools stdout: " + str(output), level=xbmc.LOGDEBUG);
    xbmc.log(msg="smartmontools stderr: " + str(errors), level=xbmc.LOGDEBUG);

    # storing date when db update was performed
    today = datetime.date.today();
    todayString = str(today.year) + "-" + str(today.month) + "-" + str(today.day);
    xbmc.log(msg="smartmontools: updating date of last DB update check: " + todayString, level=xbmc.LOGDEBUG);
    addon.setSetting(consts.settingLastDBupdateChecked, todayString);
    # TODO: return true if updated

def shouldUpdateDB():
    allowed = True if addon.getSetting(consts.settingDBupdates) == "true" else False;
    #allowed = True;
    if allowed:
        lastUpdateDateString = addon.getSetting(consts.settingLastDBupdateChecked);
        #lastUpdateDateString = "2013-09-12"
        if not lastUpdateDateString:  # if check was never performed it should be done
            xbmc.log(msg="smartmontools: update of db was never performed.", level=xbmc.LOGDEBUG);
            return True;
        l = lastUpdateDateString.split("-");
        lastCheckDate = datetime.date(int(l[0]), int(l[1]), int(l[2]));
        today = datetime.date.today();
        td = today - lastCheckDate;
        difference = td.days;
        updateInterval = int(addon.getSetting(consts.settingDBupdateInterval));

        if difference >= updateInterval:
            return True
            xbmc.log(msg="smartmontools: updating db due to number of days past " + str(difference) +".", level=xbmc.LOGDEBUG);
        else:
            xbmc.log(msg="smartmontools: not updating db due to number of days past " + str(difference) +".", level=xbmc.LOGDEBUG);
            return False;
    else:
        xbmc.log(msg="smartmontools: updating db is not allowed.", level=xbmc.LOGDEBUG);
        return False;

class dbUpdater (threading.Thread):

    def __init__(self, shouldRun = True):
        self.shouldRun = shouldRun;
        threading.Thread.__init__(self);

    def run(self):
        xbmc.log(msg="smartmontools: update thread started", level=xbmc.LOGDEBUG);
        while self.shouldRun:
            xbmc.log(msg="smartmontools: update thread performing check if is time to execute update script.", level=xbmc.LOGDEBUG);
            update = shouldUpdateDB();
            if update:
                xbmc.log(msg="smartmontools: update thread executing update check.", level=xbmc.LOGDEBUG);
                updateDatabase();
            time.sleep(2 * 60 * 60);

# initiation of needed module variable
__diskTypeDic = parseDiskTypes();


from settings import *
from node import Node

class Grid():

    def __init__(self, rows, cols, start, goal):

        self.row= int(rows)
        self.col = int(cols)

        self.grid = [[Node(y, x) for x in range(self.col)] for y in range(self.row)]

        self.grid[start.x][start.y] = start
        self.grid[goal.x][goal.y] = goal

        self.generate_all_neighbors()

    def generate_all_neighbors(self):

        for row in range(len(self.grid)):
            for col in range(len(self.grid[row])):

                self.grid[row][col].find_neighbors(self.grid)


import logging

class Node:
  def __init__(self, name, value=0, children=None):
    self.logger = logging.getLogger(self.__class__.__qualname__)
    if children is None:
      children = {}
    self.name = name
    self.value = int(value)
    self.children = children

  def add(self, frames, value):
    self.value += int(value)

    if (frames and len(frames) > 0):
      head = frames[0]
      child = self.children[head] if (self.children and head in self.children) else None
      if (child is None):
        child = Node(head)
        self.children[head] = child
      frames[0:1] = []
      child.add(frames, value)

  def serialize(self):
    res = {'name': self.name, 'value': self.value}

    children = []

    for key in self.children:
      children.append(self.children[key].serialize())

    if (len(children) > 0):
      res['children'] = children

    return res