import functools

from keras.layers import *
from keras.models import Model
from keras.wrappers.scikit_learn import KerasClassifier
from keras.regularizers import l2
from keras.optimizers import SGD
import numpy as np
import pandas as pd
import tensorflow as tf
import os
import pickle
import plac

from rectified_adam import RectifiedAdam

from sklearn.ensemble import AdaBoostClassifier
from sklearn.metrics import roc_auc_score

from cfg import *

OPTIMIZER = RectifiedAdam(lr=LR, decay=LR / (EPOCHS * BATCHES))

# OPTIMIZER = SGD(lr=LR, decay=LR/(EPOCHS*BATCHES), momentum=0.9)

def fix_tf():
    # Fix Tensorflow Memory Usage
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        try:
            # Currently, memory growth needs to be the same across GPUs
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
        except RuntimeError as e:
            # Memory growth must be set before GPUs have been initialized
            print(e)

def build_keras():
    ip = Input((INPUT_DIMS,))

    def f_lambda_smiles(x):
        x = x[:, SMILES_IDX_I:SMILES_IDX_F]
        x = tf.reshape(x, (-1, SMILES_LEN, SMILES_DIMS))
        return x

    def f_lambda_lines(x):
        return x[:, LINES_IDX_I:LINES_IDX_F]

    def f_lambda_maccs(x):
        return x[:, MACCS_IDX_I:MACCS_IDX_F]

    def f_lambda_target(x):
        return x[:, TARGET_IDX_I:TARGET_IDX_F]

    lambda_smiles = functools.partial(f_lambda_smiles)
    lambda_smiles.__name__ = 'lambda_smiles'
    lambda_lines = functools.partial(f_lambda_lines)
    lambda_lines.__name__ = 'lambda_lines'
    lambda_maccs = functools.partial(f_lambda_maccs)
    lambda_maccs.__name__ = 'lambda_maccs'
    lambda_target = functools.partial(f_lambda_target)
    lambda_target.__name__ = 'lambda_target'

    features = []

    # Cell Lines
    l_dropout = Dropout(LINES_DROPOUT)

    # forward
    lx = Lambda(lambda x: lambda_lines(x))(ip)
    if LINES_DROPOUT > 0:
        lx = l_dropout(lx)
    features.append(lx)

    if SMILES_ENABLE:
        s_conv1 = Conv1D(SMILES_KERNEL_DIMS, 11, padding='same',
                         kernel_regularizer=l2(SMILES_L2),
                         kernel_initializer=SMILES_KERNEL_INIT)
        s_bn1 = BatchNormalization()
        s_relu1 = ReLU()
        s_ap1 = AveragePooling1D(pool_size=5)
        s_conv2 = Conv1D(SMILES_KERNEL_DIMS, 11, padding='same',
                         kernel_regularizer=l2(SMILES_L2),
                         kernel_initializer=SMILES_KERNEL_INIT)
        s_bn2 = BatchNormalization()
        s_relu2 = ReLU()
        s_ap2 = AveragePooling1D(pool_size=5)
        s_gmp = GlobalMaxPooling1D()

        # forward
        sx = Lambda(lambda x: lambda_smiles(x))(ip)
        sx = s_conv1(sx)
        sx = s_bn1(sx)
        sx = s_relu1(sx)
        sx = s_ap1(sx)
        sx = s_conv2(sx)
        sx = s_bn2(sx)
        sx = s_relu2(sx)
        sx = s_ap2(sx)
        sx = s_gmp(sx)

        features.append(sx)

    # MACCS
    if MACCS_ENABLE:
        m_dropout = Dropout(LINES_DROPOUT)

        # forward
        mx = Lambda(lambda x: lambda_maccs(x))(ip)
        if MACCS_DROPOUT > 0:
            mx = m_dropout(mx)
        features.append(mx)

    if TARGET_ENABLE:
        t_dropout = Dropout(TARGET_DROPOUT)

        # forward
        tx = Lambda(lambda x: lambda_target(x))(ip)
        if MACCS_DROPOUT > 0:
            tx = t_dropout(tx)
        features.append(tx)

    j_concat = Concatenate()
    j_dropout = Dropout(HIDDEN_DROPOUT)
    j_hidden1 = Dense(HIDDEN_DIMS,
                      kernel_regularizer=l2(HIDDEN_L2),
                      kernel_initializer=HIDDEN_KERNEL_INIT)
    j_bn1 = BatchNormalization()
    j_act1 = ReLU()
    j_cls = Dense(1, activation='sigmoid')

    # Forward
    jx = j_concat(features)
    if HIDDEN_DIMS > 0:
        jx = j_hidden1(jx)
        jx = j_bn1(jx)
        jx = j_act1(jx)

    y = j_cls(jx)

    model = Model(inputs=ip, outputs=y)
    model.compile(optimizer=OPTIMIZER,
                  loss='binary_crossentropy',
                  metrics=['acc'],
                  loss_weights=[1e5])

    if SUMMARY:
        model.summary()

    return model

def build_sklearn():
    return KerasClassifier(build_fn=build_keras,
                           epochs=EPOCHS,
                           batch_size=BATCH_SIZE,
                           verbose=1)

def build_adaboost():
    sklearn_model = build_sklearn()
    model = AdaBoostClassifier(base_estimator=sklearn_model, n_estimators=N_ESTIMATORS)

    return model

def idgc20(y_true, y_pred, cname, dname, submission: bool = False):
    results = {
        'uid': [],
        'y_true': [],
        'y_pred': [],
        'cell': [],
        'drug': []
    }
    if submission:
        del results['y_true']

    for i in range(0, y_pred.shape[0]):
        results['uid'].append(i)
        if not submission:
            results['y_true'].append(y_true[i, 0])
        results['y_pred'].append(y_pred[i, 1])
        results['cell'].append(cname[i])
        results['drug'].append(dname[i])

    df = pd.DataFrame.from_dict(results)
    if submission:
        df['rank'] = 100

    cells_drugs = []

    all_idgc_20 = []
    for cell in df['cell'].unique():

        dfc = df[df['cell'] == cell]
        dfc.sort_values(by=['y_pred'], ascending=False, inplace=True)

        dgc_20 = 0
        for i in range(0, min(20, dfc.shape[0])):
            if submission:
                df.loc[df['uid'] == dfc.iloc[i]['uid'], 'rank'] = i
            if not submission:
                dgc_20 += dfc.iloc[i]['y_true'] / np.log2(i + 2)
        idgc_20 = 0
        for i in range(0, min(20, dfc.shape[0])):
            idgc_20 += 1 / np.log2(i + 2)

        idgc_20 = dgc_20 / idgc_20
        all_idgc_20.append(idgc_20)

    if submission:
        dfs = df[df['rank'] != 100]
        dfs.sort_values(by=['cell', 'rank'], ascending=True, inplace=True)
        dfs.drop(labels=['uid', 'y_pred'], axis=1, inplace=True)
        dfs.to_csv('submission.csv', index=False)

    return np.mean(all_idgc_20)

@plac.annotations(
    cmd=('train, val, rank', 'option', 'c', str)
)
def main(cmd: str = 'train'):
    fix_tf()

    X_train = np.load(os.path.join(TV_PATH, 'X_train.npy'), mmap_mode='r+')
    y_train = np.load(os.path.join(TV_PATH, 'y_train.npy'))
    cname_train = np.load(os.path.join(TV_PATH, 'cname_train.npy'))
    dname_train = np.load(os.path.join(TV_PATH, 'dname_train.npy'))

    X_val = np.load(os.path.join(TV_PATH, 'X_val.npy'), mmap_mode='r+')
    y_val = np.load(os.path.join(TV_PATH, 'y_val.npy'))

    cname_val = np.load(os.path.join(TV_PATH, 'cname_val.npy'))
    dname_val = np.load(os.path.join(TV_PATH, 'dname_val.npy'))

    if cmd == 'train':
        model = build_adaboost()

        model.fit(X_train, y_train)
        y_pred = np.array(model.predict_proba(X_val))

        idgc_20 = idgc20(y_val, y_pred, cname_val, dname_val)

        print("IDGC_20: %f" % idgc_20)

        print("Estimator Weights")
        print(model.estimator_weights_)

        for n, est in enumerate(model.estimators_):
            print("Saving estimator %d..." % n)
            est.model.save_weights('weights/%s_%d.h5' % (SESSION, n))

        model.estimators_ = []
        open('weights/%s.pkl' % SESSION, 'wb').write(pickle.dumps(model))

    elif cmd == 'val':
        model = pickle.loads(open('weights/%s.pkl' % SESSION, 'rb').read())

        for n in range(0, N_ESTIMATORS):
            m_n = build_sklearn()
            m_n.model = build_keras()
            m_n.model.load_weights('weights/%s_%d.h5' % (SESSION, n))
            model.estimators_.append(m_n)

        y_pred = np.array(model.predict_proba(X_val))

        idgc_20 = idgc20(y_val, y_pred, cname_val, dname_val)
        print("IDGC_20: %f" % idgc_20)

    elif cmd == 'rank':
        model = pickle.loads(open('weights/%s.pkl' % SESSION, 'rb').read())

        X_test = np.load(os.path.join(TV_PATH, 'X_test.npy'), mmap_mode='r+')
        cname_test = np.load(os.path.join(TV_PATH, 'cname_test.npy'))
        dname_test = np.load(os.path.join(TV_PATH, 'dname_test.npy'))

        for n in range(0, N_ESTIMATORS):
            m_n = build_sklearn()
            m_n.model = build_keras()
            m_n.model.load_weights('weights/%s_%d.h5' % (SESSION, n))
            model.estimators_.append(m_n)

        y_pred = np.array(model.predict_proba(X_test))

        idgc_20 = idgc20(None, y_pred, cname_test, dname_test, submission=True)

if __name__ == '__main__':
    plac.call(main)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Use text editor to edit the script and type in valid Instagram username/password

from InstagramAPI import InstagramAPI
import json

api = InstagramAPI("instamike59.officiel", "091h1scd")

if (api.login()):
	api.getSelfUserFeed()  # get self user feed
	data = api.LastJson
	print (type(data))
	json.dumps(data, indent=4)
	with open('data.json', 'w') as f:
    		f.write(json.dumps(data, indent=4))

#	print(data)
	print("Login succes!")
else:
	print("Can't login!")


from pyspark import SparkContext

sc = SparkContext("local[*]", "week9Assignment")
sc.setLogLevel("WARN")

def check_age(line):
    fields = line.split(",")
    if int(fields[1]) > 18:
        return fields[0], fields[1], fields[2], "Y"
    else:
        return fields[0], fields[1], fields[2], "N"

base_rdd = sc.textFile("C:/Users/chnis/Downloads/Big Data/week9 -Spark1/age.txt")

mapped_rdd = base_rdd.map(lambda line: check_age(line))

result = mapped_rdd.collect()
for x in result:
    print(x)


'''
@author: Dmitry
'''
import unittest
import numpy as np
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets('MNIST_data', reshape=False, one_hot=True, validation_size = 5000)

from SimpleCapsNet import *
from Trainer import *

class CapsTests(unittest.TestCase):
    def setUp(self):
        self.y = tf.Variable([1.0, 1.0], trainable = False)

    def addY(self, x):
        self.y = tf.subtract(self.y, tf.multiply(self.y, 0.01))
        return tf.add(x, self.y)

    @unittest.skip("skip basic")
    def test_basic_scenario(self):
        x = tf.Variable(tf.truncated_normal([2], 3.0, 1.0, dtype = tf.float32))

        z = self.addY(x)
        epochs = 250
        learning_rate = 0.1

        optimizer = tf.train.AdamOptimizer(learning_rate = learning_rate)
        sqr_minimizer = optimizer.minimize(tf.square(z))

        with tf.Session() as session:
            session.run(tf.global_variables_initializer())
            for i in range(epochs):
                #self.addY(x)
                _, val, val2 = session.run((x, x, self.addY(x)))
                print(val, val2)
        np.testing.assert_allclose(val, [0.0, 0.0], rtol = 0.0, atol = 0.0001)

    def test_1(self):
        placeholder = tf.placeholder(tf.float32, [None, 28, 28, 1])
        network = SimpleCapsNet()
        network.run(placeholder)
        saver = CustomSaver(['test', 'test/epochend'])
        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())
            saver.save_session(sess, True, (90))
        with tf.Session() as sess2:
            saver.restore_session(sess2, True)
            vals = sess2.run(network.caps_layer.W, {placeholder: mnist.train.images[0:2]})
        np.testing.assert_array_equal(list(vals.shape), [1152, 8, 160])

    def test_2(self):
        p = tf.placeholder(tf.float32, [None, 4])
        r = p[:, :3]
        with tf.Session() as sess:
            pr = sess.run(r, feed_dict = {p: [[1, 2, 3, 4], [5, 6, 7, 8]]})
        np.testing.assert_array_equal(pr, [[1, 2, 3], [5, 6, 7]])

    def test_3(self):
        x = tf.placeholder(tf.float32, [None, 2])
        y = tf.placeholder(tf.float32, [None, 2])
        z = tf.square(x)
        grads = tf.gradients(z, x, grad_ys = y)
        grads2 = tf.gradients(grads, x)

        with tf.Session() as sess:
            r = sess.run(grads2, feed_dict={x: [[2, 2], [4, 4]], y:[[3, 3], [5, 5]]})
        np.testing.assert_array_equal(r, [[[6, 6], [10, 10]]]) # 2 * y

    def test_4(self):
        x = tf.placeholder(tf.float32, [2, 3])
        with tf.Session() as sess:
            ext = tf.expand_dims(x, -1)
            y = sess.run(norm(ext - tf.transpose(x, [1, 0]), axis = -2), feed_dict={x: [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]})
        np.testing.assert_allclose(y, [[0.0, 5.19615221], [5.19615221, 0.0]], 1e-5)


import os
from datetime import datetime
import re
import logging
from os.path import join

def validate_date_prefix(date_prefix, formats, present):
    '''Returns True if the string date_prefix is a valid date, given the list
    of formats and if the date is smaller than the given present date.
    '''

    for format in formats:
        try:
            file_date = datetime.strptime(date_prefix, format)
            return file_date.date() < present.date()
        except Exception:
            pass
    logging.info(date_prefix + " is not a valid date to format "
                 + "for all formats or its in the future.")
    return False

def match_to_rename(reg, re_subs, re_date, file_name, formats, present):
    '''Verifies if the file_name matches to a valid date, that could not be
    greater than the present date given the list of formats of date.
    Returns the groups of matched regex substitution or None if do not match.
    '''

    if re.match(reg, file_name):
        date_prefix = re.sub(reg, re_date, file_name)
        if validate_date_prefix(date_prefix, formats, present):
            return re.sub(reg, re_subs, file_name)

def format_name(file_name, present):
    '''Returns formatted name if matches to a date or None.

    Parameters:
        file_name (str): The name of the file to be renamed
        present (datetime): The date on which the name cannot pass. Ex.: today

    Returns:
        formatted_name (str): A new str, formatted like '%Y-%m-%d_%H%M%S' if
        the initial name contains date and hour or '%Y-%m-%d' if it just
        contains date or None if does not match.
    '''

    # File name can contain prefix (ex: IMG, VID_)
    prefix = r"^(IMG|VID|Screenshot(\ from)?|InShot|PANO)?"
    # File name should have a valid date (%Y-%m-%d or %Y%m%d)
    date = r"(_|-|\ )?(\d{4})-?(\d{2})-?(\d{2})"
    # File name can contain suffix for date hours (ex: _293823)
    hour = r"(_|-|\ )?(\d{2})-?(\d{2})-?(\d{2})"
    # Suffix of file name and ext.
    suffix = r"(.*)?"

    date_hour_formats = ['%Y-%m-%d_%H%M%S',
                         '%Y-%m-%d_%H-%M-%S',
                         '%Y-%m-%d-%H-%M-%S',
                         '%Y-%m-%d %H-%M-%S']
    formatted_name = match_to_rename(reg=prefix+date+hour+suffix,
                                     re_subs=r"\4-\5-\6_\8\9\10\11",
                                     re_date=r"\4-\5-\6_\8\9\10",
                                     file_name=file_name,
                                     formats=date_hour_formats,
                                     present=present)

    # If the file name don't contain date and hour, it can contain only date
    if not formatted_name:
        formatted_name = match_to_rename(reg=prefix+date+suffix,
                                         re_subs=r"\4-\5-\6\7",
                                         re_date=r"\4-\5-\6",
                                         file_name=file_name,
                                         formats=['%Y-%m-%d'],
                                         present=present)

    # If not formatted both on date and hour or just date, it will be None
    return formatted_name

def preview_rename_file(file, present):
    old_name = os.path.split(os.path.abspath(file))[1]
    new_name = format_name(old_name, present)
    return [old_name, new_name]

def preview_rename_all(files):
    present = datetime.now()
    new_names = []
    preview_renamed_files = []
    will_not_be_renamed = []
    for name in sorted(files):
        names = preview_rename_file(name, present)
        if names:
            old_name = names[0]
            new_name = names[1]
            if new_name not in new_names:
                new_names.append(new_name)
                preview_renamed_files.append([old_name, new_name])
            else:
                will_not_be_renamed.append([old_name, new_name])
    return preview_renamed_files, will_not_be_renamed

def rename_file(file, final_path, old_name, new_name):
    '''Renames (and moves) the file to the final_path.
    Only renames if the file begins with prefix+date_hour or prefix+date
    and if a file with de same name already doesn't exists on destiny.
    '''
    # Only renames if the file begins with prefix+date_hour or prefix+date
    # and if a file with de same name already doesn't exists on destiny.
    try:
        formatted_path_file = os.path.join(final_path, new_name)
        if (not (os.path.exists(formatted_path_file))
           and file != formatted_path_file):
            os.rename(file, formatted_path_file)
            logging.info("| " + old_name + " | " + new_name + " |")
            return new_name
        else:
            # TODO Compare metadata of files
            logging.error("Renamed file " + old_name + " already exists"
                          + " in destiny directory. Nothing to do.")
    except Exception as e:
        logging.info("| " + old_name + " | ERROR: " + str(e) + " |")
        logging.error(e)

def rename_all(directory, final_path, files):
    '''Renames the files, removing the prefix like "IMG_", "VID-" and
    formatting the new date with hyphens. Ex. 2020-10-06_22-10-10.png.
    Returns the list of modified files new names.
    '''
    files = sorted(files)
    preview_renamed_files, not_renamed_files = preview_rename_all(files)

    print('These files will NOT be renamed:')
    print(not_renamed_files)

    print('These files will be renamed:')
    print(*preview_renamed_files, sep='\n')

    if len(preview_renamed_files) > 0:
        answer = input('Would you like to rename these '
                       + str(len(preview_renamed_files)) + ' files?'
                       + ' Type N for NO and any other for YES.\n')

        if answer.lower() != 'n':
            logging.info("Renaming all files...")
            logging.info("|     Name     |      New Name     |")
            renamed_files = []

            for old_name, new_name in preview_renamed_files:
                renamed = rename_file(join(directory, old_name),
                                      final_path, old_name, new_name)
                if renamed:
                    renamed_files.append(renamed)
            print(str(len(renamed_files)) + ' files renamed.')
            logging.info(str(len(renamed_files)) + " files renamed.")

            return renamed_files


# -*- coding: utf-8 -*-

import csv
import json
import re

# Make utility functions available
exec(compile(open("util.py", "r").read(), "util.py", 'exec'))

#open needed files and copy their content into lists
with open('../data/mp/stats/templates.txt', 'r') as file:
	reader = csv.reader(file, delimiter=',')
	templates = list(reader)

#open needed files and copy their content into lists
with open('../data/mp/stats/assignweapons.txt', 'r') as file:
	reader = csv.reader(file, delimiter=',')
	weapons = list(reader)

# Read file to populate names
with open("../data/mp/messages/strings/names.txt", 'r') as f:
	names = f.readlines()
	# Remove trailing \n
	names = [x.strip() for x in names]
	# Remove empty elements
	names = list(filter(None, names))

# Build dictionary for JSON
# Assignment of index to variable
# 0 -> id
# 1 -> unused/superfluous
# 2 -> body
# 3 -> brain
# 4 -> construct
# 5 -> ecm/unused?
# 6 -> unused/superfluous
# 7 -> propulsion
# 8 -> repair
# 9 -> type
# 10 -> sensor
# 11 -> nrWeapons/unused
# assignweapons.txt
# 0 -> id
# 1 -> weapon 0
# 2 -> weapon 1
# 3 -> weapon 2
# 4 -> unused/superfluous

obj = dict()
for line in templates:
	tempid = line[0]
	att = dict()
	att['id'] = line[0]
	att['body'] = line[2]
	if line[3] != 'ZNULLBRAIN':
		att['brain'] = line[3]
	if line[4] != 'ZNULLCONSTRUCT':
		att['construct'] = line[4]
	if line[5] != 'ZNULLECM':
		att['ecm'] = line[5]
	if line[7] != 'ZNULLPROP':
		att['propulsion'] = line[7]
	if line[8] != 'ZNULLREPAIR':
		att['repair'] = line[8]
	att['type'] = line[9]
	if (
		(line[10] != 'ZNULLSENSOR') and
		(line[10] != 'DefaultSensor1Mk1')
		):
		att['sensor'] = line[10]
	weaponList = []
	hasWeapon = False
	for weapon in weapons:
		if tempid == weapon[0]:
			hasWeapon = True
			if weapon[1] != 'NULL':
				weaponList.append(weapon[1])
			if weapon[2] != 'NULL':
				weaponList.append(weapon[2])
			if weapon[3] != 'NULL':
				weaponList.append(weapon[3])
	if hasWeapon:
		att['weapons'] = weaponList
	att['name'] = id2name(tempid, names)
	obj[tempid] = att

with open('../jsondata/mp/stats/templates.json', 'w') as f:
	f.write(json.dumps(obj, sort_keys=True, indent=4))


import sys

class Node:

    def __init__(self, mult):
        self.parent = None
        self.child_2 = None
        self.child_4 = None
        self.mult = mult
        self.cur_value = None

def function(width):

    return_list = []
    root = Node(1)
    root.cur_value = 1

    list = [root]
    leafs = []

    while len(list) > 0 :

        node = list[0]

        if node.cur_value == width:
            list.remove(node)
            leafs.append(node)

        elif node.cur_value < width:
            if node.child_2 != None and node.child_4 != None:
                list.remove(node)

            elif node.child_2 != None:
                if node.cur_value*4 == width:
                    new_node = Node(4)
                    new_node.parent = node
                    new_node.cur_value = node.cur_value*4
                    node.child_4 = new_node
                    leafs.append(new_node)
                    list.remove(node)

                elif node.cur_value*4 < width:
                    new_node = Node(4)
                    new_node.parent = node
                    new_node.cur_value = node.cur_value*4
                    node.child_4 = new_node
                    list.append(new_node)
                    list.remove(node)

                else:
                    list.remove(node)
            else:
                if node.cur_value*2 == width:
                    new_node = Node(2)
                    new_node.parent = node
                    new_node.cur_value = node.cur_value*2
                    node.child_2 = new_node
                    leafs.append(new_node)
                    list.remove(node)

                elif node.cur_value*2 < width:
                    new_node = Node(2)
                    new_node.parent = node
                    new_node.cur_value = node.cur_value*2
                    node.child_2 = new_node
                    list.append(new_node)

                else:
                    list.remove(node)

    for node in leafs:
        temp_list = []
        while node.parent != None:
            temp_list.append(node.mult)
            node = node.parent
        return_list.append(temp_list)

    return return_list

def print_lists(lists):
    for list in lists:
        print(list)

# TEST
#
# python <scritp_name> <width>

print_lists(function(int(sys.argv[1])))


import functools
from math import inf

_DEBUG = True

class Processor:
    """Processor that can schedule jobs"""

    def __init__(self, schedule_cost=0, dispatch_cost=0, preemption_cost=0,
                 cache_warmup_time=None, warm_cache_rate=1):
        """
        :param schedule_cost: overhead to schedule a job
        :param dispatch_cost: overhead to dispatch a job
        :param preemption_cost: overhead to preempt/resume a job
        :param cache_warmup_time: time to completely warm up cache
        :param warm_cache_rate: rate of execution when cache is cold
        """
        self.schedule = Schedule()
        self.time = 0

        self.schedule_cost = schedule_cost
        self.dispatch_cost = dispatch_cost
        self.preemption_cost = preemption_cost

        self.cache_warmup_time = cache_warmup_time
        self.warm_cache_rate = warm_cache_rate
        self.execution_rate = warm_cache_rate

    def reset(self):
        self.schedule = Schedule()
        self.time = 0
        self.execution_rate = self.warm_cache_rate

    def last_job_scheduled(self):
        """Returns the last job scheduled or None if processor was idle"""
        if len(self.schedule) > 0:
            last_scheduled_job = self.schedule[-1]
            if last_scheduled_job.end_time == self.time:
                return last_scheduled_job.job
        return None  # idle

    def schedule_job(self, job):
        """Schedule a job for one time unit"""
        if job != self.last_job_scheduled():
            self.execution_rate = 1  # reset cache

            if not job.has_started():
                job.remaining_overhead += self.schedule_cost + self.dispatch_cost
            else:
                job.remaining_overhead += self.dispatch_cost + self.preemption_cost  # resume new job

            if self.last_job_scheduled() is not None:
                job.remaining_overhead += self.preemption_cost  # preempt last job

        self.schedule.add(job, self.time, self.time + 1)
        self.time += 1

        gain_cache_hit_ratio = not job.has_remaining_overhead()
        job.decrement_remaining_cost(self.execution_rate)

        if gain_cache_hit_ratio and self.cache_warmup_time is not None:
            # linearly increase execution rate to warm cache rate by the cache warmup time
            self.execution_rate += ((self.warm_cache_rate - 1) / self.cache_warmup_time)
            if self.execution_rate >= self.warm_cache_rate:
                self.execution_rate = self.warm_cache_rate

        if job.has_completed():
            self.schedule[-1].job_completed = True

    def idle_until(self, t):
        """Idle processor until specified time"""
        if _DEBUG:
            assert t >= self.time
        self.time = t

class Schedule:
    """Sequence of scheduled jobs"""

    def __init__(self):
        self.schedule = []

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

    def __iter__(self):
        return iter(self.schedule)

    def __getitem__(self, item):
        return self.schedule[item]

    def __eq__(self, other):
        if len(self) != len(other):
            return False
        if any(s1.start_time != s2.start_time for s1, s2 in zip(self, other)):
            return False
        if any(s1.end_time != s2.end_time for s1, s2 in zip(self, other)):
            return False
        return all(s1.job.task == s2.job.task for s1, s2 in zip(self, other))

    def __str__(self):
        return "\n".join([str(x) for x in self.schedule])

    def add(self, job, start_time, end_time):
        """Add a job to the schedule"""

        if _DEBUG:
            # With a fully general priority function, we can only schedule one time unit at a time
            assert end_time == start_time + 1

        if len(self.schedule) > 0 and job == self.schedule[-1].job:
            if _DEBUG:
                assert start_time == self.schedule[-1].end_time

            # extend the last scheduled job if this is a continued execution
            self.schedule[-1].end_time = end_time
        else:
            self.schedule.append(ScheduledJob(start_time, end_time, job))

class ScheduledJob:
    """A job scheduled during an interval of time"""

    def __init__(self, start_time, end_time, job):
        self.start_time = start_time
        self.end_time = end_time
        self.job = job
        self.job_completed = False

    def __str__(self):
        return f"{str(self.job)} executing in [{self.start_time}, {self.end_time}]"

class UniprocessorScheduler:
    """Entity that schedules on a single processor"""

    def __init__(self, priority_function, processor=None):
        """
        :param priority_function: job priority function to use
        :param processor: processor to schedule on. Defaults to a zero overhead CPU
        """
        self.priority_function = priority_function
        if processor is None:
            self.CPU = Processor()
        else:
            self.CPU = processor

    def generate_schedule(self, task_system, final_time=None):
        """Generate a schedule for a provided task system"""

        # If no final time is provided, compute the final time required to provably show the task system is schedulable
        if final_time is None:
            if all(task.phase == 0 for task in task_system.tasks) and \
                    all(task.relative_deadline <= task.period for task in task_system.tasks):
                # For synchronous task systems with relative deadlines not exceeding periods, a deadline miss must
                # occur by the hyperperiod
                final_time = task_system.hyperperiod
            else:
                # Result by Leung and Merrill: If a deadline is missed in a periodic task system with
                # utilization <= 1, then it will be missed by time 2*P + max(D_i) + max(s_i)
                final_time = 2 * task_system.hyperperiod + max(task.relative_deadline for task in task_system.tasks) + \
                             max(task.phase for task in task_system.tasks)

        CPU = self.CPU
        CPU.reset()
        released_jobs = []
        remaining_jobs = sorted([job for task in task_system.tasks
                                 for job in task.generate_jobs(final_time)], key=lambda job: -job.release)

        if task_system.utilization() > CPU.warm_cache_rate:
            return CPU.schedule, False  # not schedulable

        while CPU.time < final_time and len(remaining_jobs) + len(released_jobs) > 0:
            if len(released_jobs) != 0:
                job_to_schedule = CPU.last_job_scheduled()
                for job in released_jobs:
                    if job_to_schedule is None or job_to_schedule.has_completed():
                        job_to_schedule = job  # CPU was idle, so choose this job
                    elif self.priority_function(job, CPU.time) + 1e-10 < self.priority_function(job_to_schedule,
                                                                                CPU.time):
                        # strict inequality here favors continuing execution of previous job and addition of 1e-10
                        # allows for minor handling of floating point errors from the variable execution rate
                        job_to_schedule = job

                CPU.schedule_job(job_to_schedule)

                if job_to_schedule.has_completed():
                    released_jobs.remove(job_to_schedule)

                if CPU.time > job_to_schedule.deadline:
                    return CPU.schedule, False  # not schedulable
            elif len(remaining_jobs) > 0:
                # idle until next job release
                CPU.idle_until(remaining_jobs[-1].release)

            while len(remaining_jobs) > 0 and remaining_jobs[-1].release <= CPU.time:
                released_jobs.append(remaining_jobs.pop())

        if all(job.deadline > final_time for job in remaining_jobs) and \
                all(job.deadline > final_time for job in released_jobs):
            return CPU.schedule, True
        else:
            return CPU.schedule, False

class MultiprocessorScheduler:
    """Entity that schedules on a multiprocessor"""

    def __init__(self, priority_function, processors, restrict_migration=False):
        """
        :param priority_function: job priority function to use
        :param processor: processors to schedule on
        :param restrict_migration: whether job migration is restricted
        """
        self.priority_function = priority_function
        self.CPUs = processors
        self.num_processors = len(processors)
        self.restrict_migration = restrict_migration

    @staticmethod
    def has_idle_processors(CPUs, jobs_to_schedule):
        """Returns whether any of the CPUs are idle with the current set of jobs to schedule"""
        return any(jobs_to_schedule[CPU] is None for CPU in CPUs)

    @staticmethod
    def get_idle_processor(CPUs, jobs_to_schedule):
        """Returns an idle CPU with the current set of jobs to schedule (or None if no such CPU exists)"""
        for CPU in CPUs:
            if jobs_to_schedule[CPU] is None:
                return CPU
        return None

    def generate_schedule(self, task_system, final_time=None):
        """Generate a schedule for a provided task system"""

        if final_time is None:
            if all(task.phase == 0 for task in task_system.tasks) and \
                    all(task.relative_deadline <= task.period for task in task_system.tasks):
                # For synchronous task systems with relative deadlines not exceeding periods, a deadline miss must
                # occur by the hyperperiod
                final_time = task_system.hyperperiod
            else:
                # Result by Leung and Merrill: If a deadline is missed in a periodic task system with
                # utilization <= 1, then it will be missed by time 2*P + max(D_i) + max(s_i)
                final_time = 2 * task_system.hyperperiod + max(task.relative_deadline for task in task_system.tasks) + \
                             max(task.phase for task in task_system.tasks)

        CPUs = self.CPUs
        for CPU in CPUs:
            CPU.reset()
        released_jobs = []
        remaining_jobs = sorted([job for task in task_system.tasks
                                 for job in task.generate_jobs(final_time)], key=lambda job: -job.release)
        migration_restriction = {job: None for job in remaining_jobs}

        if task_system.utilization() > self.num_processors * max(CPU.warm_cache_rate for CPU in CPUs):
            return [CPU.schedule for CPU in CPUs], False  # not schedulable

        while CPUs[0].time < final_time and len(remaining_jobs) + len(released_jobs) > 0:
            if len(released_jobs) != 0:
                jobs_to_schedule = {CPU: CPU.last_job_scheduled() for CPU in CPUs}

                for CPU in CPUs:
                    if jobs_to_schedule[CPU] is None or jobs_to_schedule[CPU].has_completed():
                        jobs_to_schedule[CPU] = None

                if self.restrict_migration:
                    # Handle restricted migration jobs first
                    for job in released_jobs:
                        CPU_to_reschedule = migration_restriction[job]
                        if CPU_to_reschedule is not None:
                            current_job = jobs_to_schedule[CPU_to_reschedule]
                            if current_job is None or \
                                    self.priority_function(job, CPUs[0].time) + 1e-10 < \
                                    self.priority_function(current_job, CPUs[0].time):
                                # strict inequality here favors continuing execution of previous job and the 1e-10
                                # allows for minor handling of floating point errors from the variable execution rate
                                jobs_to_schedule[CPU_to_reschedule] = job

                # Handle all jobs whose migration is not (yet) restricted
                for job in released_jobs:
                    if job in jobs_to_schedule.values():
                        continue

                    if _DEBUG:
                        if not self.restrict_migration:
                            assert migration_restriction[job] is None

                    if migration_restriction[job] is None:
                        if self.has_idle_processors(CPUs, jobs_to_schedule):
                            # CPU was idle, so choose this job
                            jobs_to_schedule[self.get_idle_processor(CPUs, jobs_to_schedule)] = job
                        elif self.priority_function(job, CPUs[0].time) + 1e-10 < \
                                max(self.priority_function(job_to_schedule, CPUs[0].time)
                                    for job_to_schedule in jobs_to_schedule.values()):
                            # strict inequality here favors continuing execution of previous job and the 1e-10
                            # allows for minor handling of floating point errors from the variable execution rate
                            CPU_to_reschedule = max(jobs_to_schedule.items(),
                                                    key=lambda CPU_job:
                                                    self.priority_function(CPU_job[1], CPUs[0].time))[0]
                            jobs_to_schedule[CPU_to_reschedule] = job

                        if _DEBUG:
                            assert len(jobs_to_schedule) == self.num_processors

                last_time = CPUs[0].time

                for CPU, job in jobs_to_schedule.items():
                    if job is not None:
                        CPU.schedule_job(job)

                for CPU in CPUs:
                    CPU.idle_until(last_time + 1)

                for CPU in CPUs:
                    job_to_schedule = CPU.last_job_scheduled()

                    if self.restrict_migration and job_to_schedule is not None:
                        if _DEBUG:
                            if migration_restriction[job_to_schedule] is not None:
                                assert CPU == migration_restriction[job_to_schedule]

                        # set migration restriction since job has been scheduled
                        migration_restriction[job_to_schedule] = CPU

                        if _DEBUG:
                            if migration_restriction[job_to_schedule] is not None:
                                assert CPU == migration_restriction[job_to_schedule]

                    if job_to_schedule is not None and job_to_schedule.has_completed():
                        released_jobs.remove(job_to_schedule)

                if _DEBUG:
                    assert all(CPU.time == CPUs[0].time for CPU in CPUs)

                if any(CPU.last_job_scheduled() is not None and not CPU.last_job_scheduled().has_completed() and \
                       CPU.time > CPU.last_job_scheduled().deadline
                       for CPU in CPUs):
                    return [CPU.schedule for CPU in CPUs], False  # not schedulable
            elif len(remaining_jobs) > 0:
                for CPU in CPUs:
                    # idle until next job release
                    CPU.idle_until(remaining_jobs[-1].release)

            while len(remaining_jobs) > 0 and remaining_jobs[-1].release <= CPUs[0].time:
                released_jobs.append(remaining_jobs.pop())

        if all(job.deadline > final_time for job in remaining_jobs) and \
                all(job.deadline > final_time for job in released_jobs):
            return [CPU.schedule for CPU in CPUs], True
        else:
            return [CPU.schedule for CPU in CPUs], False


#! /usr/bin/env python
"""
Optimizer performance comparison
"""
import os
import sys
import subprocess
from multiprocessing import Pool

problems = {
    'xray': '../doc/examples/xray/model.py',
    'pemu': '../doc/examples/superlattice/PEMU-web.py',
    }

# Aim for evals; assume problem size is 10
root = 'out'
nrepeat, evals = 30, 1000000
#nrepeat, evals = 3, 10000
fitters = {
    'de'   : '--fit=de     --steps=%d --pop=10'%(evals//100),
    'ps'   : '--fit=ps     --steps=%d --pop=10'%(evals//100),
    'rl1'  : '--fit=rl     --steps=%d --pop=10 --starts=1'%(evals//100),
    'rln'  : '--fit=rl     --steps=2000 --starts=%d --pop=0.5'%(evals//10000),
    'pt'   : '--fit=pt     --steps=%d --pop=10 --burn=0'%(evals//100),
    'dream': '--fit=dream  --steps=100 --pop=10 --burn=%d'%(evals//100-100),
    # bfgs has an extra n evaluations for the derivatives and line search
    'bfgs' : '--fit=newton --steps=200 --starts=%d'%(evals//2000),
    # nelder-mead has contract option which sometimes shrinks the simplex
    'nm'   : '--fit=amoeba --steps=1000 --starts=%d'%(evals//1000),
    }

nthread = None # Use #cpus

def cmd(store,model,opts):
    pat = "nice refl1d --batch --stepmon --overwrite --random --store=%s '%s' %s"
    return pat%(store,model,opts)

def run(job):
    modelname,model,fitname,opts,k = job
    store = os.path.join(root, modelname, fitname, "T%d"%k)
    if not os.path.exists(store):
        os.makedirs(store)
    print("starting %d %s %s\n"%(k,model,opts))
    subprocess.call(cmd(store,model,opts), shell=True)
    print("done %d %s %s\n"%(k,model,opts))

def main(problems, fitters, nthread, root):
    if os.path.exists(root):
        raise RuntimeError("directory <%s> exists."%root)
    pool = Pool(processes=nthread)
    pool.map(run, [(i,pi,j,fj,k)
                   for k in range(nrepeat)
                   for i,pi in problems.items()
                   for j,fj in fitters.items()
                   ] )

if __name__ == "__main__":
    main(problems, fitters, nthread, root)


from __future__ import print_function, division
import pandas as pd
import numpy as np
# import matplotlib
# matplotlib.use('qt4agg')
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler, minmax_scale
def mae(y, prediction):
    return np.mean(np.abs(y - prediction))

# Read in full data set
X = pd.read_csv(
    'Bike-Sharing-Dataset/hour.csv', sep=',',
    dayfirst=False,
    # parse_dates={'datetime' : ['dteday', 'hr']},
    #names=['yr', 'mnth', 'holiday', 'workingday', 'weathersit', 'weekday', 'temp', 'atemp', 'hum', 'windspeed']
    )

# Add time series
X = X.sort_values(['dteday', 'hr'])
# Add ewma over same hour
cols = ['hr', 'weekday']
#cols = ['hr']
gb = X[['cnt'] + cols].groupby(cols)

for halflife in [1.5]:
    def ewma(X):
        X = X.shift()
        X = pd.DataFrame(pd.ewma(X['cnt'], halflife=halflife))
        return X
    x = gb.transform(ewma)
    x = x.fillna(0)
    print("MAE ewma", halflife, mae(X['cnt'], x['cnt']))

#X['cnt_ewma'] = x['cnt']
#from IPython import embed
#embed()

def train_test_split(X):
    n = len(X)
    index = np.random.permutation(np.arange(0, n))

    X_train = X.iloc[index[n // 10:]]
    X_test = X.iloc[index[:n // 10]]
    assert len(X) == len(X_train) + len(X_test)
    X_train = X_train.reset_index(drop=True)
    X_test = X_test.reset_index(drop=True)
    return X_train, X_test

df_train, df_test = train_test_split(X)

print("MAE trivial", mae(df_test['cnt'], np.mean(df_train['cnt'])))
print("MAE trivial median", mae(df_test['cnt'], np.median(df_train['cnt'])))

from sklearn.ensemble import RandomForestRegressor
from sklearn.ensemble import GradientBoostingRegressor
est = RandomForestRegressor(n_estimators=2000, n_jobs=4)
#est = GradientBoostingRegressor(n_estimators=2000)
def extract_target(X):
    y = np.asarray(X['cnt'])
    Xnp = X.copy()
    del Xnp['cnt']
    del Xnp['dteday']
    del Xnp['casual']
    del Xnp['registered']
    del Xnp['instant']
    columns = Xnp.columns
    Xnp = np.asarray(Xnp)
    return Xnp, y, columns

X_train, y_train, columns = extract_target(df_train)
#est.fit(X_train, y_train - df_train['cnt_ewma'])
est.fit(X_train, y_train)

X_test, y_test, _ = extract_target(df_test)

prediction = est.predict(X_test)
#prediction += df_test['cnt_ewma']
importances = pd.Series(est.feature_importances_, index=columns)
importances.plot(kind='bar')
plt.show()
print("MAE trivial random forest", mae(df_test['cnt'], prediction))

# fit
df_test['pred'] = prediction
gb = df_test.groupby('dteday').mean()
gb['cnt'] .plot()
gb['pred'].plot()
plt.show()
sm = gb
sm['cnt_day'] = sm['cnt']
sm = sm.reset_index()

# predict
Xm = X.merge(sm[['cnt_day', 'dteday']], on=['dteday'], how='left')


"""
# ==================================================================
#   Python module
#
#   This module provides methods to compute the lifetime and
#   branching ratio of HNLs given its mass and couplings as
#   input parameters.
#
#   Created: 30/11/2014 Elena Graverini (elena.graverini@cern.ch)
#
#   Updated: 07/10/2017 Kyrylo Bondarenko (bondarenko@lorentz.leidenuniv.nl)
#
#   Sample usage:
#     ipython -i hnl.py
#     In [1]: b = HNL(1.,[1.e-8, 2.e-8, 1.e-9],True)
#     HNLbranchings instance initialized with couplings:
#          U2e   = 1e-08
#          U2mu  = 2e-08
#          U2tau = 1e-09
#     and mass:
#          m = 1.0 GeV
#     In [2]: b.computeNLifetime()
#     Out[2]: 4.777721453160521e-05
#     In [3]: b.findBranchingRatio('N -> pi mu')
#     Out[3]: 0.11826749348890987
#
# ==================================================================
"""
from __future__ import print_function
from __future__ import division
import math
import ROOT, os
import shipunit as u

# Load PDG database
pdg = ROOT.TDatabasePDG.Instance()

def PDGname(particle):
    """
    Change particle name for use with the PDG database
    """
    if '1' in particle:
        particle = particle.replace('1',"'")
    if particle == 'nu': return 'nu_e' # simple workaround to handle 3nu channel
    return particle

def mass(particle):
    """
    Read particle mass from PDG database
    """
    particle = PDGname(particle)
    tPart = pdg.GetParticle(particle)
    return tPart.Mass()

def lifetime(particle):
    """
    Read particle lifetime from PDG database
    """
    particle = PDGname(particle)
    tPart = pdg.GetParticle(particle)
    return tPart.Lifetime()

class CKMmatrix():
    """
    CKM matrix, from http://pdg.lbl.gov/2017/reviews/rpp2016-rev-ckm-matrix.pdf
    """
    def __init__(self):
        self.Vud = 0.9743
        self.Vus = 0.2251
        self.Vub = 3.6e-03
        self.Vcd = 0.2249
        self.Vcs = 0.9735
        self.Vcb = 4.11e-02
        self.Vtd = 8.6e-03
        self.Vts = 4.0e-02
        self.Vtb = 0.999

class constants():
    """
    Store some constants useful for HNL physics
    """
    def __init__(self):
        self.decayConstant = {'pi+':0.130*u.GeV,
                            'pi0':0.130*u.GeV,
                            'eta':1.2*0.130*u.GeV,
                            'eta1':0.152*u.GeV,
                            'eta_c':0.335*u.GeV,
                            'rho+':0.209*u.GeV,
                            'rho0':0.209*u.GeV,
                            'omega':0.195*u.GeV,
                            'phi':0.229*u.GeV,
                            'D_s+':0.249*u.GeV,
                            'D*_s+':0.315*u.GeV} # decay constants f_h of pseudoscalar and vector mesons
        self.GF = 1.166379e-05/(u.GeV*u.GeV) # Fermi's constant (GeV^-2)
        self.s2thetaw = 0.23126 # square sine of the Weinberg angle
        self.heV = 6.58211928*pow(10.,-16) # no units or it messes up!!
        self.hGeV = self.heV * pow(10.,-9) # no units or it messes up!!

# Load some useful constants
c = constants()

class HNLbranchings():
    """
    Lifetime and total and partial decay widths of an HNL
    """
    def __init__(self, mass, couplings, debug=False):
        """
        Initialize with mass and couplings of the HNL

        Inputs:
        mass (GeV)
        couplings (list of [U2e, U2mu, U2tau])
        """
        self.U2 = couplings
        self.U = [math.sqrt(ui) for ui in self.U2]
        self.MN = mass*u.GeV
        self.CKM = CKMmatrix()
        self.CKMelemSq = {'pi+':self.CKM.Vud**2.,
                        'rho+':self.CKM.Vud**2.,
                        'D_s+':self.CKM.Vcs**2.,
                        'D*_s+':self.CKM.Vcs**2.,
                        'K+':self.CKM.Vus**2.,
                        # Quarks (up,down)
                        # up: 1=u, 2=c, 3=t
                        # down: 1=d, 2=s, 3=b
                        (1,1):self.CKM.Vud**2., (1,2):self.CKM.Vus**2., (1,3):self.CKM.Vub**2.,
                        (2,1):self.CKM.Vcd**2., (2,2):self.CKM.Vcs**2., (2,3):self.CKM.Vcb**2.,
                        (3,1):self.CKM.Vtd**2., (3,2):self.CKM.Vts**2., (3,3):self.CKM.Vtb**2.}
        self.decays = [ 'N -> nu nu nu',
                        'N -> e- e+ nu_e',
                        'N -> e- e+ nu_mu',
                        'N -> e- e+ nu_tau',
                        'N -> e- mu+ nu_mu',
                        'N -> mu- e+ nu_e',
                        'N -> pi0 nu_e',
                        'N -> pi0 nu_mu',
                        'N -> pi0 nu_tau',
                        'N -> pi+ e-',
                        'N -> mu- mu+ nu_e',
                        'N -> mu- mu+ nu_mu',
                        'N -> mu- mu+ nu_tau',
                        'N -> pi+ mu-',
                        'N -> eta nu_e',
                        'N -> eta nu_mu',
                        'N -> eta nu_tau',
                        'N -> rho0 nu_e',
                        'N -> rho0 nu_mu',
                        'N -> rho0 nu_tau',
                        'N -> rho+ e-',
                        'N -> omega nu_e',
                        'N -> omega nu_mu',
                        'N -> omega nu_tau',
                        'N -> rho+ mu-',
                        'N -> eta1 nu_e',
                        'N -> eta1 nu_mu',
                        'N -> eta1 nu_tau',
                        'N -> phi nu_e',
                        'N -> phi nu_mu',
                        'N -> phi nu_tau',
                        'N -> e- tau+ nu_tau',
                        'N -> tau- e+ nu_e',
                        'N -> mu- tau+ nu_tau',
                        'N -> tau- mu+ nu_mu',
                        'N -> D_s+ e-',
                        'N -> D_s+ mu-',
                        'N -> D*_s+ e-',
                        'N -> D*_s+ mu-',
                        'N -> eta_c nu_e',
                        'N -> eta_c nu_mu',
                        'N -> eta_c nu_tau' ]
        if self.MN>=1.:
            self.QCD_corr = self.QCD_correction()

        if debug:
            print("HNLbranchings instance initialized with couplings:")
            print("\tU2e   = %s"%self.U2[0])
            print("\tU2mu  = %s"%self.U2[1])
            print("\tU2tau = %s"%self.U2[2])
            print("and mass:")
            print("\tm = %s GeV"%(self.MN))

    def sqrt_lambda(self,a,b,c):
        """
        Useful function for decay kinematics. Returns 0 for kinematically forbidden region
        """
        l = a**2 + b**2 + c**2 - 2*a*b - 2*b*c - 2*a*c
        if l<0:
            return 0
        else:
            return math.sqrt(l)

    def QCD_correction(self):
        """
        Returns 3-loops QCD correction to HNL decay width into quarks
        """
        alpha_s = ROOT.TGraph( os.path.expandvars('$FAIRSHIP/python/alpha_s.dat') )
        a_s = alpha_s.Eval(self.MN)
        qcd_corr = a_s / math.pi
        qcd_corr += 5.2 * (a_s / math.pi)**2.
        qcd_corr += 26.4 * (a_s / math.pi)**3.
        return qcd_corr

    def Width_3nu(self):
        """
        Returns the HNL decay width into three neutrinos
        """
        width = (c.GF**2.)*(self.MN**5.)*sum(self.U2)/(192.*(u.pi**3.))
        width = 2.*width # Majorana case (charge conjugate channels)
        return width

    def Width_nu_f_fbar(self, alpha, beta):
        """
        Returns the HNL tree level decay width into the fermion-antifermion pair and a neutrino (decay through Z boson or Z and W)

        Inputs:
        - alpha is a flavour of neutrino (1,2 or 3), determine the mixing angle U_alpha
        - beta is a flavour of the fermions: (1, 2 or 3) for charge leptons or (4, 5, 6, 7, 8, 9) for quarks
        """
        if (alpha not in [1,2,3]) or (beta not in [1,2,3,4,5,6,7,8,9]):
            print('Width_nu_f_fbar ERROR: unknown channel alpha =',alpha,' beta =',beta)
            quit()
        l = [None,'e-','mu-','tau-','u','d','s','c','b','t']
        x = mass(l[beta]) / self.MN
        if x > 0.5: # the decay is kinematically forbidden
            return 0
        width = (c.GF**2.)*(self.MN**5.)*self.U2[alpha-1]/(192.*(math.pi**3.))
        L = 0.
        if x>0.01:
            logContent = (1. - 3.*x**2. - (1.-x**2.)*math.sqrt(1. - 4.*x**2.) ) / ( (x**2.)*(1 + math.sqrt(1. - 4.*x**2.)) )
        else:
            logContent = x**4 + 4.*x**6 + 14.*x**8
        if logContent > 0:
            L = math.log( logContent )
        if beta < 4: # lepton decay
            NZ = 1
            if alpha == beta: # interference case
                C1 = 0.25*(1. + 4.*c.s2thetaw + 8.*c.s2thetaw**2.)
                C2 = 0.5*c.s2thetaw*(2.*c.s2thetaw +1.)
            else: # no interference
                C1 = 0.25*(1. - 4.*c.s2thetaw + 8.*c.s2thetaw**2.)
                C2 = 0.5*c.s2thetaw*(2.*c.s2thetaw -1.)
        else: # decay into quarks
            NZ = 3
            if (beta in [4,7,9]): # up quarks
                C1 = 0.25*(1. - 8./3.*c.s2thetaw + 32./9.*c.s2thetaw**2.)
                C2 = 1./3.*c.s2thetaw*(4./3.*c.s2thetaw - 1.)
            if (beta in [5,6,8]): # down quarks
                C1 = 0.25*(1. - 4./3.*c.s2thetaw + 8./9.*c.s2thetaw**2.)
                C2 = 1./6.*c.s2thetaw*(2./3.*c.s2thetaw - 1.)
        width = width*NZ*( C1*( (1.-14.*x**2. -2.*x**4. -12.*x**6.)*math.sqrt(1-4.*x**2) +12.*x**4. *(-1.+x**4.)*L )
                        + 4.*C2*( x**2. *(2.+10.*x**2. -12.*x**4.) * math.sqrt(1.-4.*x**2) + 6.*x**4. *(1.-2.*x**2+2.*x**4)*L ) )
        width = 2.*width # Majorana case (charge conjugate channels)
        return width

    def Integrand(self,xx,xi):
        """
        Function to integrate needed for numerical integration using ROOT. Needed for 3-body decays trough W boson.
        First argument is a list of 1 number, argument. Second is the list of 3 numbers xi = mi/MN
        """
        x = xx[0]
        xl2 = xi[0]**2
        xu2 = xi[1]**2
        xd2 = xi[2]**2
        if x==0: # Workaround for division by zero
            return 0
        res = 1./x
        res *= (x - xl2 - xd2)
        res *= (1. + xu2 - x)
        res *= self.sqrt_lambda(x,xl2,xd2)
        res *= self.sqrt_lambda(1.,x,xu2)
        return res

    def I(self,x1,x2,x3):
        """
        Numerical integral needed for 3-body decays trough W boson.
        xi = mi/MN
        """
        theFunction = self.Integrand
        func = ROOT.TF1('func',theFunction,0,1,3) # Setting function
        func.SetParameters(x1,x2,x3)
        xmin = (x1 + x3)**2
        xmax = (1. - x2)**2
        wf1 = ROOT.Math.WrappedTF1(func) # Setting simple Gauss integration method
        ig = ROOT.Math.GaussIntegrator()
        ig.SetFunction(wf1)
        ig.SetRelTolerance(0.001)
        res = 12. * ig.Integral(xmin,xmax)
        return res

    def Width_l1_l2_nu2(self, alpha, beta):
        """
        Returns the HNL decay width into two different flavour leptons and a neutrino (decay through W boson)

        Inputs:
        - alpha is a flavour of the first lepton (1,2 or 3), determine the mixing angle U_alpha
        - beta is a flavour of the second lepton and neutrino (1, 2 or 3)
        """
        if (alpha not in [1,2,3]) or (beta not in [1,2,3]):
            print('Width_l1_l2_nu2 ERROR: unknown channel alpha =',alpha,' beta =',beta)
            quit()
        if alpha==beta: # The interference case is handled by Width_nu_f_fbar function, workaround for a total width calculation
            return 0
        l = [None,'e-','mu-','tau-']
        x1 = mass(l[alpha])/self.MN
        x2 = mass(l[beta])/self.MN
        if x1+x2>1: # The decay is kinematically forbidden
            return 0
        width = (c.GF**2.)*(self.MN**5.)*self.U2[alpha-1]/(192.*(math.pi**3.))
        width = width*self.I(x1,x2,0)
        width = 2.*width # Majorana case (charge conjugate channels)
        return width

    def Width_l_u_d(self, alpha, beta, gamma):
        """
        Returns the HNL tree level decay width into a charged lepton, up quark and down quark (decay through W boson)

        Inputs:
        - alpha is a flavour of the charged lepton (1,2 or 3), determine the mixing angle U_alpha
        - beta is the up quark generation (1, 2, 3 for u, c, t)
        - gamma is the down quark generation (1, 2, 3 for d, s, b)
        """
        if (alpha not in [1,2,3]) or (beta not in [1,2,3]) or (gamma not in [1,2,3]):
            print('Width_l_u_d ERROR: unknown channel alpha =',alpha,' beta =',beta,' gamma =',gamma)
            quit()
        l = [None,'e-','mu-','tau-']
        u = [None,'u','c','t']
        d = [None,'d','s','b']
        xl = mass(l[alpha])/self.MN
        xu = mass(u[beta])/self.MN
        xd = mass(d[gamma])/self.MN
        if xl+xu+xd>1: # The decay is kinematically forbidden
            return 0
        width = (c.GF**2.)*(self.MN**5.)*self.U2[alpha-1]/(192.*(math.pi**3.))
        width *= 3*self.CKMelemSq[(beta, gamma)]
        width *= self.I(xl,xu,xd)
        width = 2.*width # Majorana case (charge conjugate channels)
        return width

    def Width_H0_nu(self, H, alpha):
        """
        Returns the HNL decay width into neutral meson and neutrino

        Inputs:
        - H is a string (name of the meson)
        - alpha is a flavour of the nu (1, 2 or 3), determine the mixing angle U_alpha
        """
        if (H not in ['pi0','eta','rho0','omega','eta1','phi','eta_c']) or (alpha not in [1,2,3]):
            print('Width_H0_nu ERROR: unknown channel H0 =',H,' alpha =',alpha)
            quit()
        x = mass(H)/self.MN
        if x > 1: # The decay is kinematically forbidden
            return 0.
        width = (c.GF**2.)*(c.decayConstant[H]**2.)*(self.MN**3.)*self.U2[alpha-1]/(32.*math.pi)
        if H in ['pi0','eta','eta1','eta_c']: # pseudoscalar mesons
            width *= (1 - x**2.)**2.
        if H in ['rho0','omega','phi']: # vector mesons
            width *= (1. + 2*x**2.)*(1. - x**2.)**2.
            if H=='rho0':
                width *= (1. - 2.*c.s2thetaw)**2.
            if H=='omega':
                width *= (16.*c.s2thetaw**2.)/9.
            if H=='phi':
                width *= (1. - 4./3.*c.s2thetaw)**2.
        width = 2.*width # Majorana case (charge conjugate channels)
        return width

    def Width_H_l(self, H, alpha):
        """
        Returns the HNL decay width into charged meson and lepton

        Inputs:
        - H is a string (name of the positively charged meson)
        - alpha is the lepton flavour (1, 2 or 3), determine the mixing angle U_alpha
        """
        if (H not in ['pi+','rho+','D_s+','D*_s+']) or (alpha not in [1,2,3]):
            print('Width_H_l ERROR: unknown channel H =',H,' alpha =',alpha)
            quit()
        l = [None,'e-','mu-','tau-']
        xl = mass(l[alpha])/self.MN
        xh = mass(H)/self.MN
        if xl+xh > 1: # The decay is kinematically forbidden
            return 0.
        width = (c.GF**2.)*(c.decayConstant[H]**2.)*self.CKMelemSq[H]*(self.MN**3.)*self.U2[alpha-1]/(16.*math.pi)
        width *= self.sqrt_lambda(1., xl**2, xh**2)
        if H in ['pi+','D_s+']: # pseudoscalar mesons
            width *= ( (1. - xl**2.)**2. - xh**2. *(1. + xl**2.) )
        if H in ['rho+','D*_s+']: # vector mesons
            width *= ( (1. - xl**2.)**2. + xh**2.*(1. + xl**2.) - 2.*xh**4. )
        width = 2.*width # Majorana case (charge conjugate channels)
        return width

    def Width_charged_leptons(self):
        """
        Returns the total HNL leptonic decay width with charged leptons in final state
        """
        width = 0.
        for l1 in [1,2,3]:
            for l2 in [1,2,3]:
                width += self.Width_nu_f_fbar(l1,l2)
                width += self.Width_l1_l2_nu2(l1,l2)
        return width

    def Width_neutral_mesons(self):
        """
        Returns the total HNL decay width into a neutral meson and a neutrino
        """
        mesons = ['pi0','eta','rho0','omega','eta1','phi','eta_c']
        width = 0.
        for m in mesons:
            for l in [1,2,3]:
                width += self.Width_H0_nu(m,l)
        return width

    def Width_charged_mesons(self):
        """
        Returns the total HNL decay width into a charged meson and a charged lepton
        """
        mesons = ['pi+','rho+','D_s+','D*_s+']
        width = 0.
        for m in mesons:
            for l in [1,2,3]:
                width += self.Width_H_l(m,l)
        return width

    def Width_quarks_neutrino(self):
        """
        Returns the total HNL decay width with quarks and neutrino in final state. Uses 3-loops alpha_s correction
        """
        if self.MN<1.: # decay into quarks is not applicable at low HNL mass
            return 0
        width = 0.
        for l in [1,2,3]:
            for q in [4,5,6,7,8,9]:
                width += self.Width_nu_f_fbar(l,q)
        width *= (1. + self.QCD_corr)
        return width

    def Width_quarks_lepton(self):
        """
        Returns the total HNL decay width with quarks and charged lepton in final state. Uses 3-loops alpha_s correction
        """
        if self.MN<1.: # decay into quarks is not applicable at low HNL mass
            return 0
        width = 0.
        for l in [1,2,3]:
            for u in [1,2,3]:
                for d in [1,2,3]:
                    width += self.Width_l_u_d(l,u,d)
        width *= (1. + self.QCD_corr)
        return width

    def NDecayWidth(self):
        """
        Returns the total HNL decay width
        """
        leptonWidth = self.Width_3nu() + self.Width_charged_leptons()
        mesonWidth = self.Width_neutral_mesons() + self.Width_charged_mesons()
        quarkWidth = self.Width_quarks_neutrino() + self.Width_quarks_lepton()
        totalWidth = leptonWidth + max([mesonWidth, quarkWidth])
        return totalWidth

    def findBranchingRatio(self, decayString):
        """
        Returns the branching ratio of the selected HNL decay channel

        Inputs:
        - decayString is a string describing the decay, in the form 'N -> stuff1 ... stuffN'
        """
        br = 0.
        totalWidth = self.NDecayWidth()

        if (decayString not in self.decays) and (decayString not in ['N -> hadrons','N -> charged hadrons']):
            print('findBranchingRatio ERROR: unknown decay %s'%decayString)
            quit()

        if decayString == 'N -> nu nu nu' or decayString == 'N -> 3nu': br = self.Width_3nu() / totalWidth # inclusive
        if decayString == 'N -> e- e+ nu_e': br = self.Width_nu_f_fbar(1,1) / totalWidth
        if decayString == 'N -> e- e+ nu_mu': br = self.Width_nu_f_fbar(2,1) / totalWidth
        if decayString == 'N -> e- e+ nu_tau': br = self.Width_nu_f_fbar(3,1) / totalWidth
        if decayString == 'N -> e- mu+ nu_mu': br = self.Width_l1_l2_nu2(1,2) / totalWidth
        if decayString == 'N -> mu- e+ nu_e': br = self.Width_l1_l2_nu2(2,1) / totalWidth
        if decayString == 'N -> pi0 nu_e': br = self.Width_H0_nu('pi0',1) / totalWidth
        if decayString == 'N -> pi0 nu_mu': br = self.Width_H0_nu('pi0',2) / totalWidth
        if decayString == 'N -> pi0 nu_tau': br = self.Width_H0_nu('pi0',3) / totalWidth
        if decayString == 'N -> pi+ e-': br = self.Width_H_l('pi+',1) / totalWidth
        if decayString == 'N -> mu- mu+ nu_e': br = self.Width_nu_f_fbar(1,2) / totalWidth
        if decayString == 'N -> mu- mu+ nu_mu': br = self.Width_nu_f_fbar(2,2) / totalWidth
        if decayString == 'N -> mu- mu+ nu_tau': br = self.Width_nu_f_fbar(3,2) / totalWidth
        if decayString == 'N -> pi+ mu-': br = self.Width_H_l('pi+',2) / totalWidth
        if decayString == 'N -> eta nu_e': br = self.Width_H0_nu('eta',1) / totalWidth
        if decayString == 'N -> eta nu_mu': br = self.Width_H0_nu('eta',2) / totalWidth
        if decayString == 'N -> eta nu_tau': br = self.Width_H0_nu('eta',3) / totalWidth
        if decayString == 'N -> rho0 nu_e': br = self.Width_H0_nu('rho0',1) / totalWidth
        if decayString == 'N -> rho0 nu_mu': br = self.Width_H0_nu('rho0',2) / totalWidth
        if decayString == 'N -> rho0 nu_tau': br = self.Width_H0_nu('rho0',3) / totalWidth
        if decayString == 'N -> rho+ e-': br = self.Width_H_l('rho+',1) / totalWidth
        if decayString == 'N -> omega nu_e': br = self.Width_H0_nu('omega',1) / totalWidth
        if decayString == 'N -> omega nu_mu': br = self.Width_H0_nu('omega',2) / totalWidth
        if decayString == 'N -> omega nu_tau': br = self.Width_H0_nu('omega',3) / totalWidth
        if decayString == 'N -> rho+ mu-': br = self.Width_H_l('rho+',2) / totalWidth
        if decayString == 'N -> eta1 nu_e': br = self.Width_H0_nu('eta1',1) / totalWidth
        if decayString == 'N -> eta1 nu_mu': br = self.Width_H0_nu('eta1',2) / totalWidth
        if decayString == 'N -> eta1 nu_tau': br = self.Width_H0_nu('eta1',3) / totalWidth
        if decayString == 'N -> phi nu_e': br = self.Width_H0_nu('phi',1) / totalWidth
        if decayString == 'N -> phi nu_mu': br = self.Width_H0_nu('phi',2) / totalWidth
        if decayString == 'N -> phi nu_tau': br = self.Width_H0_nu('phi',3) / totalWidth
        if decayString == 'N -> e- tau+ nu_tau': br = self.Width_l1_l2_nu2(1,3) / totalWidth
        if decayString == 'N -> tau- e+ nu_e': br = self.Width_l1_l2_nu2(3,1) / totalWidth
        if decayString == 'N -> mu- tau+ nu_tau': br = self.Width_l1_l2_nu2(2,3) / totalWidth
        if decayString == 'N -> tau- mu+ nu_mu': br = self.Width_l1_l2_nu2(3,2) / totalWidth
        if decayString == 'N -> D_s+ e-': br = self.Width_H_l('D_s+',1) / totalWidth
        if decayString == 'N -> D_s+ mu-': br = self.Width_H_l('D_s+',2) / totalWidth
        if decayString == 'N -> D*_s+ e-': br = self.Width_H_l('D*_s+',1) / totalWidth
        if decayString == 'N -> D*_s+ mu-': br = self.Width_H_l('D*_s+',2) / totalWidth
        if decayString == 'N -> eta_c nu_e': br = self.Width_H0_nu('eta_c',1) / totalWidth
        if decayString == 'N -> eta_c nu_mu': br = self.Width_H0_nu('eta_c',2) / totalWidth
        if decayString == 'N -> eta_c nu_tau': br = self.Width_H0_nu('eta_c',3) / totalWidth

        if decayString == 'N -> hadrons':
            mesonWidth = self.Width_neutral_mesons() + self.Width_charged_mesons()
            quarkWidth = self.Width_quarks_neutrino() + self.Width_quarks_lepton()
            br = max([mesonWidth, quarkWidth]) / totalWidth
        if decayString == 'N -> charged hadrons':
            br = max([self.Width_charged_mesons(), self.Width_quarks_lepton()]) / totalWidth
        return br

    def allowedChannels(self):
        """
        Returns a dictionary of kinematically allowed decay channels

        Inputs:
        - decayString is a string describing the decay, in the form 'N -> stuff1 ... stuffN'
        """
        m = self.MN
        allowedDecays = {'N -> nu nu nu':'yes'}
        if m > 2.*mass('e-'):
            allowedDecays.update({'N -> e- e+ nu_e':'yes'})
            allowedDecays.update({'N -> e- e+ nu_mu':'yes'})
            allowedDecays.update({'N -> e- e+ nu_tau':'yes'})
        if m > mass('e-') + mass('mu-'):
            allowedDecays.update({'N -> e- mu+ nu_mu':'yes'})
            allowedDecays.update({'N -> mu- e+ nu_e':'yes'})
        if m > mass('pi0'):
            allowedDecays.update({'N -> pi0 nu_e':'yes'})
            allowedDecays.update({'N -> pi0 nu_mu':'yes'})
            allowedDecays.update({'N -> pi0 nu_tau':'yes'})
        if m > mass('pi+') + mass('e-'):
            allowedDecays.update({'N -> pi+ e-':'yes'})
        if m > 2.*mass('mu-'):
            allowedDecays.update({'N -> mu- mu+ nu_e':'yes'})
            allowedDecays.update({'N -> mu- mu+ nu_mu':'yes'})
            allowedDecays.update({'N -> mu- mu+ nu_tau':'yes'})
        if m > mass('pi+') + mass('mu-'):
            allowedDecays.update({'N -> pi+ mu-':'yes'})
        if m > mass('eta'):
            allowedDecays.update({'N -> eta nu_e':'yes'})
            allowedDecays.update({'N -> eta nu_mu':'yes'})
            allowedDecays.update({'N -> eta nu_tau':'yes'})
        if m > mass('rho0'):
            allowedDecays.update({'N -> rho0 nu_e':'yes'})
            allowedDecays.update({'N -> rho0 nu_mu':'yes'})
            allowedDecays.update({'N -> rho0 nu_tau':'yes'})
        if m > mass('rho+') + mass('e-'):
            allowedDecays.update({'N -> rho+ e-':'yes'})
        if m > mass('omega'):
            allowedDecays.update({'N -> omega nu_e':'yes'})
            allowedDecays.update({'N -> omega nu_mu':'yes'})
            allowedDecays.update({'N -> omega nu_tau':'yes'})
        if m > mass('rho+') + mass('mu-'):
            allowedDecays.update({'N -> rho+ mu-':'yes'})
        if m > mass('eta1'):
            allowedDecays.update({'N -> eta1 nu_e':'yes'})
            allowedDecays.update({'N -> eta1 nu_mu':'yes'})
            allowedDecays.update({'N -> eta1 nu_tau':'yes'})
        if m > mass('phi'):
            allowedDecays.update({'N -> phi nu_e':'yes'})
            allowedDecays.update({'N -> phi nu_mu':'yes'})
            allowedDecays.update({'N -> phi nu_tau':'yes'})
        if m > mass('e-') + mass('tau-'):
            allowedDecays.update({'N -> e- tau+ nu_tau':'yes'})
            allowedDecays.update({'N -> tau- e+ nu_e':'yes'})
        if m > mass('mu-') + mass('tau-'):
            allowedDecays.update({'N -> mu- tau+ nu_tau':'yes'})
            allowedDecays.update({'N -> tau- mu+ nu_mu':'yes'})
        if m > mass('D_s+') + mass('e-'):
            allowedDecays.update({'N -> D_s+ e-':'yes'})
        if m > mass('D_s+') + mass('mu-'):
            allowedDecays.update({'N -> D_s+ mu-':'yes'})
        if m > mass('D*_s+') + mass('e-'):
            allowedDecays.update({'N -> D*_s+ e-':'yes'})
        if m > mass('D*_s+') + mass('mu-'):
            allowedDecays.update({'N -> D*_s+ mu-':'yes'})
        if m > mass('eta_c'):
            allowedDecays.update({'N -> eta_c nu_e':'yes'})
            allowedDecays.update({'N -> eta_c nu_mu':'yes'})
            allowedDecays.update({'N -> eta_c nu_tau':'yes'})

        for decay in self.decays:
            if decay not in allowedDecays:
                allowedDecays.update({decay:'no'})
        return allowedDecays

class HNL(HNLbranchings):
    """
    HNL physics according to the nuMSM
    """
    def __init__(self, mass, couplings, debug=False):
        """
        Initialize with mass and couplings of the HNL

        Inputs:
        mass (GeV)
        couplings (list of [U2e, U2mu, U2tau])
        """
        HNLbranchings.__init__(self, mass, couplings, debug)
    def computeNLifetime(self,system="SI"):
        """
        Compute the HNL lifetime

        Inputs:
        - system: choose between default (i.e. SI, result in s) or FairShip (result in ns)
        """
        self.NLifetime = c.hGeV / self.NDecayWidth()
        if system == "FairShip": self.NLifetime *= 1.e9
        return self.NLifetime


import torch
import torch.nn as nn

from networks.layers import (
    RecurrentConvolutionalLayer,
    ConvLayer,
    Dense
)

class SharkRCNN(nn.Module):
    def __init__(self, num_steps):
        super(SharkRCNN, self).__init__()

        self.conv = ConvLayer(1, 32, (7,3), (2,1), (3,1))
        self.pool1 = nn.MaxPool2d((2,1),(2,1))

        self.dropout1 = nn.Dropout(.5)

        self.rcl1 = RecurrentConvolutionalLayer(32, 64, (3,2), (1,1), num_steps)
        self.dropout2 = nn.Dropout(.4)

        self.rcl2 = RecurrentConvolutionalLayer(64, 128, (5,2), (2,1), num_steps)
        self.dropout3 = nn.Dropout(.3)

        self.pool2 = nn.MaxPool2d((2,1),(2,1))

        self.rcl3 = RecurrentConvolutionalLayer(128, 256, (7,1), (3,0), num_steps)
        self.dropout4 = nn.Dropout(.2)

        self.rcl4 = RecurrentConvolutionalLayer(256, 512, (1,1), (0,0), num_steps)
        self.gap_pool = nn.AdaptiveAvgPool2d((1,1))

        self.fc = Dense([512, 256, 64, 4])

    def forward(self, x):
        x = self.conv(x)
        x = self.pool1(x)
        x = self.dropout1(x)

        x = self.rcl1(x)
        x = self.dropout2(x)

        x = self.rcl2(x)
        x = self.dropout3(x)

        x = self.pool2(x)

        x = self.rcl3(x)
        x = self.dropout4(x)

        x = self.rcl4(x)
        x = self.gap_pool(x)

        x = self.fc(x)

        return x


import pyspark
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, asc,desc

spark = SparkSession.builder.appName('SparkByExamples.com').getOrCreate()

simpleData = [("James","Sales","NY",90000,34,10000), \
    ("Michael","Sales","NY",86000,56,20000), \
    ("Robert","Sales","CA",81000,30,23000), \
    ("Maria","Finance","CA",90000,24,23000), \
    ("Raman","Finance","CA",99000,40,24000), \
    ("Scott","Finance","NY",83000,36,19000), \
    ("Jen","Finance","NY",79000,53,15000), \
    ("Jeff","Marketing","CA",80000,25,18000), \
    ("Kumar","Marketing","NY",91000,50,21000) \
  ]
columns= ["employee_name","department","state","salary","age","bonus"]

df = spark.createDataFrame(data = simpleData, schema = columns)

df.printSchema()
df.show(truncate=False)
pivotDF = df.groupBy("employee_name").pivot("department").sum("salary")
pivotDF.show()
pivotDF.show()


from django.urls import path, re_path, register_converter
from .views import *
from core.routers import PetuniRouter
from .converters import CRMIdConverter

app_name = 'crm'
register_converter(CRMIdConverter, 'crmid')
router = PetuniRouter()

urlpatterns = [
    path('contact/<crmid:crm_id>/pull/', CRMContactView.as_view(), name='contact'),
    path('account/<crmid:crm_id>/pull/', CRMAccountPullView.as_view(), name='account-pull'),
    path('account-schedule/<crmid:crm_id>/pull/', CRMAccountScheduleView.as_view(),
         name='account_schedule'),
    path('clinic-service-offer/<crmid:crm_id>/pull/', CRMClinicServiceOfferView.as_view(),
         name='clinic_service_offer'),
    path('user/<uuid:pk>/', CRMPetuniUserView.as_view(), name='user-detail'),
    path('user/<uuid:pk>/warn/', CRMPetuniUserWarnView.as_view(), name='user-warn'),
    re_path(r'shelter/(?P<pk>[0-9a-f-]+)/registration-request-(?P<action>(approve)|(reject))/$',
        CRMShelterApproveView.as_view(), name='shelter-registration-request'),
    re_path(r'(?P<post_type>(adoptionpost)|(shelterpost)|(clinicreview))/(?P<pk>[0-9a-f-]+)/$',
        CRMPostView.as_view(), name='post_delete'),
    re_path(r'(?P<comment_type>(comment))/(?P<pk>[0-9a-f-]+)/$',
        CRMCommentView.as_view(), name='comment-detail'),
    path('pet/<uuid:pk>/', CRMPetDeleteView.as_view(), name='pet-delete'),
]
urlpatterns += router.urls


def calc_gain(dataset, attr, target_attr):
    """
    Calculate Information Gain
    @param dataset: Dataset
    @param attr: super attribute
    @param target_attr: tarrget attribute
    """

    #super information entropy and sub-class entropy
    mclass_entropy = calc_entroy(dataset, attr)
    sclass_entropy = 0.0

    for i in dataset[target_attr].unique():
        # different attribute's different class's sub set
        sub_dataset = dataset[dataset[target_attr] == i]
        sclass_entropy += len(sub_dataset) / len(dataset) * calc_entroy(sub_dataset, attr)

    return mclass_entropy - sclass_entropy

A-Decision Tree/ppt/Calculate Information Gain.py


import db

class FormattedData():
    database = db.MyDatabase(db.SQLITE, dbname='taxi.db')
    def get_drivers(self,lname=None, fine=None, commendation=None):
        try:
            if lname:
                query = f"SELECT * FROM driver WHERE last_name = '{lname}';"
            elif fine:
                query = f"SELECT * FROM driver WHERE fine = 1;"
            elif commendation:
                query = f"SELECT * FROM driver  WHERE commendation = 1;"
            else:
                query = f"SELECT * FROM driver;"

            drivers = self.database.get_data(query=query)

            formatted_drivers = []

            for driver in drivers:
                auto_ids = self.database.get_data(query=f"SELECT auto_id FROM call WHERE driver_id = {driver['id']};")
                autos = ''
                for auto_id in auto_ids:
                    autos += f"| {self.get_auto(auto_id=auto_id[0])} |"

                formatted_drivers.append(f"{driver['id']} {driver['first_name']} {driver['last_name']} - " \
                                        f"Experience {driver['experience']} years - " \
                                        f"Fine {driver['fine']} - " \
                                        f"Commendation {driver['commendation']} - " \
                                        f"{autos}")
            return formatted_drivers
        except Exception as e:
            return f"Failed {e}"

    def get_by_area(self, area):
        try:
            query = f"SELECT driver_id FROM call WHERE call_area = '{area}';"
            drivers = self.database.get_data(query=query)
            formatted_drivers = []
            for driver in drivers:
                query = f"SELECT * FROM driver WHERE id = '{driver[0]}';"
                driver = self.database.get_data(query=query)
                formatted_drivers.append(f"{driver[0]['id']} {driver[0]['first_name']} {driver[0]['last_name']} - " \
                                        f"Experience {driver[0]['experience']} years - " \
                                        f"Fine {driver[0]['fine']} - " \
                                        f"Commendation {driver[0]['commendation']} - ")
            return formatted_drivers
        except Exception as e:
            return f"Failed {e}"

    def get_auto(self, number=None, auto_id=None):
        try:
            query = f"SELECT * FROM auto WHERE id = {auto_id};"
            auto = self.database.get_data(query=query)
            formatted_auto = f"{auto[0]['color']} - {auto[0]['brand']} - {auto[0]['number']}"
            return formatted_auto
        except Exception as e:
            return f"Failed -> {e}"

    def get_auto_by_num(self, number):
        try:
            query = f"SELECT * FROM auto WHERE number = '{number}';"
            auto = self.database.get_data(query=query)
            formatted_auto = f"{auto[0]['color']} - {auto[0]['brand']} - {auto[0]['number']}"
            return formatted_auto
        except Exception as e:
            return f"Failed -> {e}"


import warnings
from itertools import islice
from math import log2, pow, sqrt
from copy import copy

from PIL import Image, ImageChops
from numpy import array, zeros, append
from scipy.fftpack import dct, idct

from qtar.core.imageqt import ImageQT, ImageQTPM
from qtar.core.curvefitting import fit_cfregions, CFRegions, draw_cf_map
from qtar.core.quantizationmatrix import generate_flat_matrix
from qtar.core.adaptiveregions import adapt_regions
from qtar.core.permutation import reverse_permutation, fix_diff, get_diff_fix
from qtar.core.container import Container, Key
from qtar.core.matrixregion import MatrixRegions, divide_into_equal_regions, draw_borders_on
from qtar.core.zigzag import zigzag_embed_to_regions, zigzag_extract_from_regions, zigzag_embed_to_cfregions, zigzag_extract_from_cfregions

DEFAULT_PARAMS = {
    'homogeneity_threshold': 0.4,
    'min_block_size':        8,
    'max_block_size':        512,
    'quant_power':           1,
    'ch_scale':              4.37,
    'offset':                (0, 0),
    'pm_mode':               False,
    'cf_mode':               False,
    'cf_grid_size':          8,
    'wmdct_mode':            False,
    'wmdct_block_size':      32,
    'wmdct_scale':           0.1,
}

class QtarStego:
    def __init__(self,
                 homogeneity_threshold=DEFAULT_PARAMS['homogeneity_threshold'],
                 min_block_size=DEFAULT_PARAMS['min_block_size'],
                 max_block_size=DEFAULT_PARAMS['max_block_size'],
                 quant_power=DEFAULT_PARAMS['quant_power'],
                 ch_scale=DEFAULT_PARAMS['ch_scale'],
                 offset=DEFAULT_PARAMS['offset'],
                 pm_mode=DEFAULT_PARAMS['pm_mode'],
                 cf_mode=DEFAULT_PARAMS['cf_mode'],
                 cf_grid_size=DEFAULT_PARAMS['cf_grid_size'],
                 wmdct_mode=DEFAULT_PARAMS['wmdct_mode'],
                 wmdct_block_size=DEFAULT_PARAMS['wmdct_block_size'],
                 wmdct_scale=DEFAULT_PARAMS['wmdct_scale']):
        self.homogeneity_threshold = homogeneity_threshold
        self.min_block_size = min_block_size
        self.max_block_size = max_block_size
        self.quant_power = quant_power
        self.ch_scale = ch_scale
        self.offset = offset
        self.pm_mode = pm_mode
        self.cf_mode = cf_mode
        self.cf_grid_size = cf_grid_size
        self.wmdct_mode = wmdct_mode
        self.wmdct_block_size = wmdct_block_size
        self.wmdct_scale = wmdct_scale

    def embed(self, img_container, img_watermark, resize_to_fit=True, stages=False):
        size = int(pow(2, int(log2(min(img_container.width, img_container.height)))))

        max_b = min(self.max_block_size, *img_container.size)

        img_container = self.__prepare_image(img_container, shape=(size, size), offset=self.offset)
        chs_container = self.__convert_image_to_chs(img_container)
        container_image_mode = img_container.mode

        key = Key(ch_scale=self.ch_scale,
                  offset=self.offset,
                  pm_mode=self.pm_mode,
                  container_shape=(size, size) if self.pm_mode else None,
                  cf_mode=self.cf_mode,
                  cf_grid_size=self.cf_grid_size if self.cf_mode else None,
                  wmdct_mode=self.wmdct_mode,
                  wmdct_block_size=self.wmdct_block_size if self.wmdct_mode else None,
                  wmdct_scale=self.wmdct_scale if self.wmdct_mode else None)
        container = Container(key=key)

        chs_regions = []
        chs_pm_key = []
        for ch, ch_image in enumerate(chs_container):
            if self.pm_mode:
                regions = ImageQTPM(ch_image, self.min_block_size, max_b, self.homogeneity_threshold)
                chs_pm_key.append(regions.permutation)
            else:
                regions = ImageQT(ch_image, self.min_block_size, max_b, self.homogeneity_threshold)
            qt_key = regions.key

            regions_dct = self.__dct_regions(regions)

            if self.cf_mode:
                regions_embed = fit_cfregions(regions_dct, self.quant_power, self.cf_grid_size)
                key.chs_cf_key.append(regions_embed.curves)
            else:
                regions_embed, ar_indexes = adapt_regions(regions_dct, q_power=self.quant_power)
                key.chs_ar_key.append(ar_indexes)

            chs_regions.append(regions)
            container.chs_regions_dct.append(regions_dct)
            container.chs_regions_dct_embed.append(regions_embed)
            key.chs_qt_key.append(qt_key)

        available_space = container.available_space

        if available_space == 0:
            raise NoSpaceError("There is no space for embedding. Try other parameters.")

        wm_shape = img_watermark.size
        if resize_to_fit:
            wm_size = int(sqrt(available_space))
            wm_shape = (wm_size, wm_size)
        img_watermark = self.__prepare_image(img_watermark, shape=wm_shape, mode=img_container.mode)
        chs_watermark = self.__convert_image_to_chs(img_watermark)
        key.wm_shape = wm_shape

        chs_stego_img = []
        chs_embedded_dct_regions = []

        chs_count = len(chs_container)
        for ch in range(chs_count):
            regions_dct = container.chs_regions_dct[ch]
            embed_dct_regions = container.chs_regions_dct_embed[ch]
            wm_ch = chs_watermark[ch]

            if self.wmdct_mode:
                wm_regions = divide_into_equal_regions(wm_ch, self.wmdct_block_size)
                wm_dct_regions = self.__dct_regions(wm_regions)
                wm_ch = self.__quant_regions(wm_dct_regions, self.wmdct_scale)

            embedded_dct_regions = self.__embed_in_regions(embed_dct_regions, wm_ch)
            chs_embedded_dct_regions.append(embedded_dct_regions)

            idct_regions = MatrixRegions(regions_dct.rects, embedded_dct_regions.matrix)
            stego_img_mx = self.__idct_regions(idct_regions).matrix

            if self.pm_mode:
                pm_key = chs_pm_key[ch]
                qt_key = key.chs_qt_key[ch]

                stego_img_mx = reverse_permutation(stego_img_mx, pm_key)

                compressed_stego_img_mx = array(Image.fromarray(copy(stego_img_mx)).convert('L'))
                distorted_regions = ImageQTPM(compressed_stego_img_mx, key=qt_key)
                ch_pm_fix = get_diff_fix(pm_key, distorted_regions.permutation, distorted_regions.rects)

                key.chs_pm_fix_key.append(ch_pm_fix)

            chs_stego_img.append(stego_img_mx)

        img_stego = self.__chs_to_image(chs_stego_img, container_image_mode)
        img_stego = self.__prepare_image(img_stego, offset=(-self.offset[0], -self.offset[1]))
        img_container = self.__prepare_image(img_container, offset=(-self.offset[0], -self.offset[1]))
        stages_imgs = None
        if stages:
            stages_imgs = {
                "1-container": img_container,
                "2-quad_tree": self.__regions_to_image(chs_regions, container_image_mode,
                                                       borders=True, only_right_bottom=True),
                "3-adaptive_regions": self.__regions_to_image(container.chs_regions_dct_embed, container_image_mode,
                                                              borders=True, factor=20),
                "4-dct": self.__regions_to_image(chs_embedded_dct_regions, container_image_mode,
                                                 factor=20),
                "5-watermark": img_watermark,
                "6-stego_image": img_stego
            }
        return StegoEmbedResult(img_stego, key, container.fact_bpp, img_container, img_watermark, stages_imgs)

    @staticmethod
    def __prepare_image(image, shape=None, offset=None, mode=None):
        if shape is not None:
            image = image.resize(shape, Image.BILINEAR)
        if offset is not None:
            x, y = offset
            image = ImageChops.offset(image, -x, -y)
        if mode is not None:
            image = image.convert(mode)
        return image

    @staticmethod
    def __dct_regions(regions, inverse=False):
        size = regions.matrix.shape[0]
        mx_result = zeros((size, size))
        regions_dct = MatrixRegions(regions.rects, mx_result)
        for i, region in enumerate(regions):
            if inverse:
                region_dct = idct(idct(region.T, norm='ortho').T, norm='ortho')
            else:
                region_dct = dct(dct(region, norm='ortho').T, norm='ortho').T
            regions_dct[i] = region_dct
        return regions_dct

    @staticmethod
    def __quant_regions(regions, scale):
        quantized = MatrixRegions(regions.rects, copy(regions.matrix))
        for i, region in enumerate(quantized):
            height, width = region.shape
            q_mx = generate_flat_matrix(max((width, height)))
            quantized[i] = region / (q_mx[0:height, 0:width] / scale)
        return quantized

    @staticmethod
    def __dequant_regions(regions, scale):
        dequantized = MatrixRegions(regions.rects, copy(regions.matrix))
        for i, region in enumerate(dequantized):
            height, width = region.shape
            q_mx = generate_flat_matrix(max((width, height)))
            dequantized[i] = region * (q_mx[0:height, 0:width] / scale)
        return dequantized

    @classmethod
    def __idct_regions(cls, regions):
        return cls.__dct_regions(regions, True)

    def __embed_in_regions(self, regions, ch_watermark):
        if self.cf_mode:
            regions = CFRegions(regions.rects, copy(regions.matrix), regions.curves, self.cf_grid_size)
            if self.wmdct_mode:
                return zigzag_embed_to_cfregions(ch_watermark, regions)

        else:
            regions = MatrixRegions(regions.rects, copy(regions.matrix))
            if self.wmdct_mode:
                return zigzag_embed_to_regions(ch_watermark, regions)

        wm_iter = ch_watermark.flat

        for i in range(0, len(regions)):
            region = regions[i]
            size = region.size
            shape = region.shape
            region_stego = array(list(islice(wm_iter, size)))

            if region_stego.size == 0:
                if size == 0:
                    continue
                else:
                    return regions

            region_stego = region_stego / 255 * self.ch_scale
            region_stego = append(region_stego, region.flat[region_stego.size:size])
            region_stego = region_stego.reshape(shape)
            regions[i] = region_stego

        try:
            next(wm_iter)
        except StopIteration:
            pass
        else:
            warnings.warn("Container capacity is not enough for embedding given secret image."
                          "The extracted secret image will be incomplete.")
        return regions

    @classmethod
    def extract(cls, img_stego, key, stages=False):
        img_stego = cls.__prepare_image(img_stego, offset=key.offset)
        chs_stego = cls.__convert_image_to_chs(img_stego)
        stego_image_mode = img_stego.mode

        chs_regions = []
        chs_regions_extract = []
        chs_watermark = []
        for ch, ch_stego in enumerate(chs_stego):
            qt_key = key.chs_qt_key[ch]

            if key.pm_mode:
                pm_fix_key = key.chs_pm_fix_key[ch]
                distorted_regions = ImageQTPM(copy(ch_stego), key=qt_key)
                distorted_pm_mx_regions = MatrixRegions(distorted_regions.rects, distorted_regions.permutation)
                pm_key = fix_diff(distorted_pm_mx_regions, pm_fix_key).matrix
                regions = ImageQTPM(copy(ch_stego), key=qt_key, permutation=pm_key)
            else:
                regions = ImageQT(ch_stego, key=qt_key)

            regions_dct = cls.__dct_regions(regions)

            if key.cf_mode:
                cf_key = key.chs_cf_key[ch]
                regions_extract = CFRegions.from_regions(regions_dct, cf_key, key.cf_grid_size)
            else:
                ar_indexes = key.chs_ar_key[ch]
                regions_extract, ar_indexes = adapt_regions(regions_dct, ar_indexes=ar_indexes)

            if key.wmdct_mode:
                wm_quantized_regions = cls.__extract_from_regions(regions_extract, key)
                wm_dct_regions = cls.__dequant_regions(wm_quantized_regions, key.wmdct_scale)
                wm_regions = cls.__idct_regions(wm_dct_regions)
                ch_watermark = wm_regions.matrix
            else:
                ch_watermark = cls.__extract_from_regions(regions_extract, key)

            chs_regions.append(regions)
            chs_regions_extract.append(regions_extract)
            chs_watermark.append(ch_watermark)

        if stages:
            stages_imgs = {
                "7-quad_tree": cls.__regions_to_image(chs_regions, stego_image_mode,
                                                      borders=True, only_right_bottom=True),
                "8-adaptive_regions": cls.__regions_to_image(chs_regions_extract, stego_image_mode,
                                                             borders=True, factor=10),
                "9-extracted_watermark": cls.__chs_to_image(chs_watermark, stego_image_mode)
            }
            return stages_imgs
        else:
            return cls.__chs_to_image(chs_watermark, stego_image_mode)

    @classmethod
    def __extract_from_regions(cls, regions, key):
        if key.wmdct_mode:
            if key.cf_mode:
                return zigzag_extract_from_cfregions(regions, key.wm_shape, key.wmdct_block_size)
            else:
                return zigzag_extract_from_regions(regions, key.wm_shape, key.wmdct_block_size)

        region_stego = array([])
        wm_size = key.wm_shape[0] * key.wm_shape[1]
        for region in regions:
            region = region.flatten()
            region_stego = append(region_stego, region)
            if region_stego.size >= wm_size:
                break
        else:
            region_stego = append(region_stego, zeros(wm_size - region_stego.size))
        ch_watermark = region_stego[0:wm_size].reshape(key.wm_shape)
        return ch_watermark * 255 / key.ch_scale

    @staticmethod
    def __convert_image_to_chs(image):
        chs_image = [array(ch_image)
                     for ch_image in image.split()]
        return chs_image

    @staticmethod
    def __chs_to_image(chs_matrix, mode, offset=None, factor=1):
        chs_image = [Image.fromarray(ch_matrix * factor).convert('L')
                     for ch_matrix in chs_matrix]
        image = Image.merge(mode, chs_image)
        if offset is not None:
            x, y = offset
            image = ImageChops.offset(image, -x, -y)

        return image

    @classmethod
    def __draw_borders(cls, chs_regions, only_right_bottom=False):
        chs_matrix = [ch_regions.matrix for ch_regions in chs_regions]
        max_value = max([ch_matrix.max() for ch_matrix in chs_matrix])

        is_cf_regions = type(chs_regions[0]) is CFRegions

        for ch_id in range(0, len(chs_matrix)):
            for ch_regions in chs_regions:
                if is_cf_regions:
                    chs_matrix[ch_id] = draw_cf_map(
                        CFRegions(ch_regions.rects, chs_matrix[ch_id], ch_regions.curves, ch_regions.grid_size),
                        0, only_right_bottom)
                else:
                    chs_matrix[ch_id] = draw_borders_on(chs_matrix[ch_id], ch_regions.rects, 0, only_right_bottom)
            if is_cf_regions:
                chs_matrix[ch_id] = draw_cf_map(
                    CFRegions(chs_regions[ch_id].rects, chs_matrix[ch_id], chs_regions[ch_id].curves, chs_regions[ch_id].grid_size),
                    max_value, only_right_bottom)
            else:
                chs_matrix[ch_id] = draw_borders_on(chs_matrix[ch_id], chs_regions[ch_id].rects, max_value,
                                                    only_right_bottom)
        return chs_matrix

    @classmethod
    def __regions_to_image(cls, chs_regions, mode, offset=None, factor=1, borders=False, only_right_bottom=False):
        if borders:
            chs_matrix = cls.__draw_borders(chs_regions, only_right_bottom)
        else:
            chs_matrix = [ch_regions.matrix for ch_regions in chs_regions]

        return cls.__chs_to_image(chs_matrix, mode, offset, factor)

    @property
    def params(self):
        return {
            'homogeneity_threshold': self.homogeneity_threshold,
            'min_block_size': self.min_block_size,
            'max_block_size': self.max_block_size,
            'quant_power': self.quant_power,
            'ch_scale': self.ch_scale,
            'offset': self.offset,
            'pm_mode': self.pm_mode,
            'cf_mode': self.cf_mode,
            'cf_grid_size': self.cf_grid_size,
            'wmdct_mode': self.wmdct_mode,
            'wmdct_block_size': self.wmdct_block_size,
            'wmdct_scale': self.wmdct_scale
        }

class StegoEmbedResult:
    def __init__(self, img_stego, key, bpp, img_container=None, img_watermark=None, stages_imgs=None):
        self.img_stego = img_stego
        self.img_container = img_container
        self.img_watermark = img_watermark
        self.key = key
        self.bpp = bpp
        self.stages_imgs = stages_imgs

class NoSpaceError(Exception): pass


import socket
import select
import mtbot.protocol as p
from mtbot.botpackage import *

class MTServer:
    """Some kind of proxy."""
    psock = {}
    saddr = {}
    pdst = {}

    def __init__(self, adr, dst):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.bind(adr)
        self.destination = dst

    def create_socket(self, addr):
        nsocket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.psock[str(addr[1])] = nsocket
        self.saddr[str(nsocket.fileno())] = addr
        print("Created socket")

    def receive(self):
        sockets = [self.sock]
        for value in self.psock.values():
            sockets.append(value)

        read_sockets, write_sockets, error_sockets = select.select(sockets, [], [])
        for sock in read_sockets:
            try:
                data, addr = sock.recvfrom(1000)
                if data:
                    if sock == self.sock:
                        self.toserverhandler(sock, data, addr)
                    else:
                        self.toclienthandler(sock, data, addr)

            except Exception as e:
                print("Error: " + str(e))
                return

    def disconnect(self, sock):
        addr = self.saddr[str(sock.fileno())]
        del self.psock[str(addr[1])]
        del self.saddr[str(sock.fileno())]
        sock.close()
        print("disconnected socket")

    def close(self):
        self.sock.close()
        for sock in self.psock:
            sock.close()

    def toserverhandler(self, sock, data, addr):
        #print("ToserverData : %s" % data)
        send = True
        peer_id, channel, type, reliable, seqnum, dat = readPacket(data)
        #TODO find better index!!!
        ad = str(addr[1])
        #preprogressing
        print("t" + str(len(data)) + ":" + type + str(reliable) + str(seqnum) + "dat:" + str(dat))

        if not ad in self.psock:
            self.create_socket(addr)
            self.pdst[ad] = self.destination

        if send:
            self.psock[ad].sendto(data, self.pdst[ad])

        #postprogressing
        if type == "control" and dat:
             ctype = dat[0]
             if ctype == "disco":
                 self.disconnect(self.psock[ad])

    def toclienthandler(self, sock, data, addr):
        #print("ToclientData : %s" % data)
        send = True
        peer_id, channel, type, reliable, seqnum, dat = readPacket(data)
        #preprogressing
        print("f" + str(len(data)) + ":" + type + str(reliable) + str(seqnum) + "dat:" + str(dat))
        if type == "original":
            if dat[0] == p.tc["chat_message"]:
                d = dat[1]
                #len = d[:2]
                msg = d[5:]
                print("msg: " + str(msg))

        if send:
            self.sock.sendto(data, self.saddr[str(sock.fileno())])

        #postprogressing
        if type == "control" and dat:
             ctype = dat[0]
             if ctype == "disco":
                 self.disconnect(sock)

mt = MTServer(("0.0.0.0", 30001), ("185.234.72.95", 30001))

while (1):
    mt.receive()


N = int(input())

days = [[-1, -1]] + [[] for _ in range(N)]
for i in range(1, N+1):
    days[i] = list(map(int, input().split()))

maxV = 0
def dfs(x, res):
    global maxV
    for i in range(x+1, N+1):
        if i + days[i][0] - 1 < N+1:
            dfs(i+days[i][0]-1, res+days[i][1])
    maxV = max(res, maxV)
    return

for i in range(1, N+1):
    if i + days[i][0]-1 < N+1:
        dfs(i+days[i][0]-1, days[i][1])
print(maxV)

백준/삼성SW역량테스트기출문제/14501. 퇴사.py