from __future__ import absolute_import
from __future__ import division

import tensorflow as tf
import tensorflow.contrib.layers as layers
import architectures
import numpy as np
from sklearn import metrics
import os

from flip_gradient import flip_gradient
import nn_utils

class PNet(object):

    def __init__(self,
                 config,
                 args,
                 word_vecs,
                 init = tf.random_uniform_initializer(minval=-0.01, maxval=0.01), # init = layers.xavier_initializer(),
                 name='PNet'):

        self.cfg       = config
        self.args      = args
        self.word_vecs =  word_vecs
        self.init      = init
        self.name      = name

        self.memory_size    = self.cfg.memory_size
        self.sent_size      = self.cfg.sent_size
        self.embed_size     = self.cfg.embed_size
        self.hidden_size    = self.cfg.hidden_size
        self.l2_reg_lambda  = self.cfg.l2_reg_lambda
        self.max_grad_norm  = self.cfg.max_grad_norm
        self.hops           = self.cfg.hops

        self.build_vars()
        self.build_eval_op()

    def build_vars(self):

        with tf.variable_scope(self.name):

            A = tf.convert_to_tensor(self.word_vecs)
            self.word2vec = tf.Variable(A, name="word2vec", trainable=True)
            self.nil_vars = set([self.word2vec.name])
            self.lr = tf.placeholder(tf.float32, [], name="learning_rate")
            self.adapt = tf.placeholder(tf.float32, [], name="adapt_rate")

            self.P_net  = architectures.HAN(self.cfg, self.word2vec, init=self.init, scope="P_net")

    def build_eval_op(self):

        self.reviews    = tf.placeholder(tf.int32, [None, self.memory_size, self.sent_size], name="reviews")
        self.labels     = tf.placeholder(tf.int32,   [None, 2], name="labels")

        with tf.variable_scope(self.name):

            self.p_reps,  self.p_word_attns,  self.p_sent_attns  = self.P_net(self.reviews, reuse=False)

            with tf.variable_scope('P_net'):
                self.sen_logits = nn_utils.fc_layer(self.p_reps, output_dim=2, scope="sentiment_classifier", reuse=False)
                self.sen_predictions = tf.argmax(self.sen_logits, 1, name="sen_predictions")
                sen_loss = tf.losses.softmax_cross_entropy(self.labels, self.sen_logits, weights=1.0, scope="sentiment_classifier")
                var_list = [var for var in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES) if "sentiment_classifier" in var.name]
                reg_loss = self.l2_reg_lambda * sum(tf.nn.l2_loss(var) for var in var_list)
                self.loss = sen_loss + reg_loss

                self.dom_logits = nn_utils.fc_layer(self.p_reps, output_dim=2, scope="domain_classifier", reuse=False)
                self.dom_predictions = tf.argmax(self.dom_logits, 1, name="dom_predictions")

    def build_graph(self, sen_reviews, sen_labels, dom_reviews, dom_labels):

        with tf.variable_scope(self.name):

            sen_reps, _, _  = self.P_net(sen_reviews, reuse=True)
            dom_reps, _, _  = self.P_net(dom_reviews, reuse=True)

            with tf.variable_scope('P_net'):
                self.sen_loss = self.add_clf_loss(sen_reps, sen_labels, scope='sentiment_classifier')
                self.dom_loss = self.add_dann_loss(dom_reps, dom_labels, self.adapt, scope='domain_classifier')
                self.pivot_loss = self.sen_loss + self.dom_loss

    def add_clf_loss(self, cls_reps, cls_labels, num_labels=2, weight=1.0, scope=None):

        cls_logits = nn_utils.fc_layer(cls_reps, output_dim=num_labels, scope=scope, reuse=True)
        clf_loss = tf.losses.softmax_cross_entropy(cls_labels, cls_logits, weights=weight, scope=scope)

        var_list = [var for var in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES) if scope in var.name]
        reg_loss = self.l2_reg_lambda * sum(tf.nn.l2_loss(var) for var in var_list)

        loss = clf_loss + reg_loss

        return loss

    def add_dann_loss(self, dom_reps, dom_labels, adapt_rate, weight=1.0, scope=None):

        dom_reps_grl = flip_gradient(dom_reps, adapt_rate)
        dom_logits = nn_utils.fc_layer(dom_reps_grl, output_dim=2, scope=scope, reuse=True)
        dom_loss = tf.losses.softmax_cross_entropy(dom_labels, dom_logits, weights=weight, scope=scope)

        var_list = [var for var in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES) if scope in var.name]
        reg_loss = self.l2_reg_lambda * sum(tf.nn.l2_loss(var) for var in var_list)

        loss = dom_loss + reg_loss

        return loss

    def create_train_op(self):

        with tf.name_scope('train'):

            self.opt = tf.train.MomentumOptimizer(self.lr, 0.9)
            self.train_op = nn_utils.train_network(self.opt, self.pivot_loss, self.nil_vars, self.max_grad_norm, "train_op")

    def eval_sen(self, sess, reviews, sen_labels, batch_size=None):

        preds = []
        batch_generator = nn_utils.batch_generator(reviews=reviews, batch_size=batch_size, shuffle=False)
        for i in range(batch_generator.batch_num):
            xb = batch_generator.next_batch()
            pred = self.predict_sen(sess, xb)
            preds += list(pred)
        preds = np.array(preds)
        acc = metrics.accuracy_score(preds, np.argmax(sen_labels, axis=1))

        return acc, preds

    def eval_dom(self, sess, src_reviews, tar_reviews, src_labels, tar_labels, steps=None, batch_size=None):

        dom_preds  = []
        dom_labels = []
        src_batch_generator = nn_utils.batch_generator(reviews=src_reviews, labels=src_labels, batch_size=batch_size, shuffle=False)
        tar_batch_generator = nn_utils.batch_generator(reviews=tar_reviews, labels=tar_labels, batch_size=batch_size, shuffle=False)
        for i in range(steps):
            xb_s, yb_s = src_batch_generator.next_batch()
            xb_t, yb_t = tar_batch_generator.next_batch()
            xb_d       = np.vstack((xb_s, xb_t))
            dom_label  = np.vstack((yb_s, yb_t))
            dom_pred   = self.predict_dom(sess, xb_d)
            dom_preds  += list(dom_pred)
            dom_labels += list(dom_label)

        dom_acc = metrics.accuracy_score(np.array(dom_preds), np.argmax(dom_labels, axis=1))

        return dom_acc

    def vis_attention(self, sess, reviews, batch_size=None):

        p_w_atttns,  p_s_attns  = [], []
        batch_generator = nn_utils.batch_generator(reviews, batch_size=batch_size, shuffle=False)
        for i in range(batch_generator.batch_num):
            xb  = batch_generator.next_batch()
            w1, s1 = self.get_attention(sess, xb)
            p_w_atttns += list(w1)
            p_s_attns += list(s1)

        return np.array(p_w_atttns), np.array(p_s_attns)

    def initialize_session(self, sess):
        sess.run(tf.global_variables_initializer())
        self.saver = tf.train.Saver()

        output_dir = "./work/models/"
        if not os.path.exists(output_dir):
            os.mkdir(output_dir)
        self.save_path="./work/models/" + self.args.source_domain + '_' + self.args.target_domain + "_PNet.ckpt"

    def save_model(self, sess):
        self.saver.save(sess, self.save_path)

    def load_model(self, sess):
        try:
            self.saver.restore(sess, self.save_path)
        except Exception as e:
            raise IOError("Failed to to load model " "from save path: %s" % self.save_path)
        self.saver.restore(sess, self.save_path)
        print("Successfully load model from save path: %s" % self.save_path)

    def predict_sen(self, sess, reviews):

        feed_dict = {
            self.reviews: reviews
        }
        sen_predictions = sess.run(self.sen_predictions, feed_dict=feed_dict)

        return sen_predictions

    def predict_dom(self, sess, reviews):
        feed_dict = {
            self.reviews: reviews,
        }
        dom_predictions = sess.run(self.dom_predictions, feed_dict=feed_dict)

        return dom_predictions

    def eval_loss(self, sess, reviews, labels):

        feed_dict = {
            self.reviews: reviews,
            self.labels: labels
        }
        loss = sess.run(self.loss, feed_dict=feed_dict)

        return loss

    def get_attention(self, sess, reviews):

        feed_dict = {
            self.reviews: reviews
        }
        p_word_attns, p_sent_attns = sess.run([self.p_word_attns, self.p_sent_attns], feed_dict=feed_dict)

        return p_word_attns, p_sent_attns


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# MIT License
#
# Copyright (c) 2019, 2020 Assured Information Security, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#

import numpy
from gnuradio import gr
import pmt

class j2497_tagger(gr.sync_block):
    """
    docstring for block j2497_tagger
    """
    def __init__(self):
        gr.sync_block.__init__(self,
            name="j2497_tagger",
            in_sig=[numpy.complex64, numpy.float32],
            out_sig=[numpy.complex64])

        self.scanning = True
        self.threshold = 210
        self.peak_index = 0
        self.fill_start_buffer = False
        self.tag_offset = 0
        self.start_buffer = numpy.array([], dtype=numpy.float32)

    def work(self, input_items, output_items):
        in0 = input_items[0]
        in1 = input_items[1]

        out = output_items[0]

        # Tag First Peak by Detecting Four Peaks Spaced 100 Samples Apart
        if self.scanning is True and len(in1>0):
            for n in range(0,len(in1),5):

                # Detect Candidate Starts
                if in1[n] > self.threshold or self.fill_start_buffer is True:
                    self.fill_start_buffer = True

                    # Gather Small Buffer for Detecting Peaks
                    if len(self.start_buffer) < 500:
                        if len(self.start_buffer) + len(in1[n:]) >= 500:
                            self.start_buffer = numpy.append(self.start_buffer, in1[n:])
                        else:
                            self.start_buffer = numpy.append(self.start_buffer, in1[n:])
                            self.tag_offset = self.tag_offset-len(self.start_buffer)

                    # Start Buffer is Full
                    if len(self.start_buffer) >= 500:

                        # Find Peak in Start Buffer
                        get_max_index = 0
                        for m in range(0,100):

                            # Get Center of Peak
                            if self.start_buffer[m] > self.threshold:
                                get_max_index = numpy.argmax(self.start_buffer[m:m+20])
                                break

                        # Make Tag in Complex Stream
                        if (self.start_buffer[get_max_index+100] > self.threshold) and (self.start_buffer[get_max_index+200] > self.threshold) and (self.start_buffer[get_max_index+300] > self.threshold):
                            start_tag_estimate = self.nitems_written(0)+self.tag_offset+n+get_max_index
                            if start_tag_estimate-self.nitems_written(0) < 0:
                                start_tag_estimate = self.nitems_written(0)
                            key = pmt.intern("burst")
                            value = pmt.intern("start")
                            self.add_item_tag(0, start_tag_estimate, key, value)
                            self.peak_index = n+get_max_index
                            self.scanning = False
                            self.start_buffer = numpy.array([], dtype=numpy.float32)
                            self.fill_start_buffer = False
                            self.tag_offset = 0
                            break

                        # False Positive
                        else:
                            self.start_buffer = numpy.array([], dtype=numpy.float32)
                            self.fill_start_buffer = False
                            self.tag_offset = 0
                            n = n+50

        # Tag Last Peak
        if self.scanning is False:
            last_index = 0

            # Update for Windows Less than 100
            if self.peak_index > len(in1):
                self.peak_index = self.peak_index - len(in1)
            else:
                # Find the End of the Peaks
                for nn in range(self.peak_index,len(in1),100):
                    last_index = nn

                    # Make Tag in Complex Stream
                    if in1[nn] < self.threshold:
                        key = pmt.intern("burst")
                        value = pmt.intern("end")
                        self.add_item_tag(0, self.nitems_written(0)+nn, key, value)
                        self.scanning = True
                        break

                # Remember Index at the End of Window
                self.peak_index = 100-(len(in1)-last_index)

        out[:] = in0
        return len(output_items[0])


# Data https://www.physionet.org/content/eegmmidb/1.0.0/
# Follow the instruction to download the data, 109 subjects, 64 channels, 160 Hz
# Download data to data/PhysioNet/eegmmidb/1.0.0

import os
import glob
import numpy as np
import mne
import pickle
from sklearn.preprocessing import StandardScaler
from sklearn.utils import shuffle

PATH = '../../rcnn/data/PhysioNet/eegmmidb/1.0.0/'
SUBS = glob.glob(PATH + 'S[0-9]*')
FNAMES = sorted(x[-4:] for x in SUBS)

# 1. Baseline, eyes open
# 2. Baseline, eyes closed
# 3. 7. 11.  Task 1 (open and close left or right fist)
# 5. 9. 13.  Task 3 (open and close both fists or both feet)
# 4. 8. 12.  Task 2 (imagine opening and closing left or right fist)
# 6. 10. 14. Task 4 (imagine opening and closing both fists or both feet)
#
# T0 corresponds to rest,
# T1 corresponds to onset of motion (real or imagined) of
#     the left fist  (in runs 3, 7, 11 for ME, and 4, 8, 12 for MI)
#     both fists     (in runs 5, 9, 13 for ME, and 6, 10, 14 for MI)
# T2 corresponds to onset of motion (real or imagined) of
#     the right fist (in runs 3, 7, 11 for ME, and 4, 8, 12 for MI)
#     both feet      (in runs 5, 9, 13 for ME, and 6, 10, 14 for MI)
FNAMES.remove('S089')

def get_data(subj_num=FNAMES, l_freq=13, h_freq=55, tmin=0.0, tmax=4.0):
    """ Extract EEG time series and class labels from edf files
    All T1, T2 events have duration 4.1 or 4.2 s, use 0~4s for epoching
    :param subj_num:
    :param l_freq: low frequency for bandpass filter
    :param h_freq: high frequency for bandpass filter
    :param tmin:
    :param tmax:

    :return: Xs, ys: list (Ni, nc, T), (Ni, )
    """

    run_type_0 = ['02']
    run_type_1 = ['04', '08', '12']
    run_type_2 = ['06', '10', '14']
    run_types = run_type_0 + run_type_1 + run_type_2

    # fixed parameters
    nc = 64       # number of channels
    sfreq = 160
    n_type0 = 22  # each subject has 21~24 fist or feet events, set to use the same number of data points for baseline
    event_id = {'T1': 2, 'T2': 3}  # T0:0 rest event, not interested

    T = round(tmax * sfreq)

    print('Sub, (N,   nc,  T)')
    Xs, ys = [], []

    # Iterate over subjects: S001, S002, S003 ...
    for i, subj in enumerate(subj_num):
        freq_flag = False

        # Get file names for motor imagery data
        fnames = glob.glob(os.path.join(PATH, subj, subj+'R*.edf'))
        fnames = sorted([name for name in fnames if name[-6:-4] in run_types])

        subject_trials, subject_labels = [], []
        # Iterate over subject's experiment runs
        for i, fname in enumerate(fnames):

            run = fname[-6:-4]

            # Load data into MNE raw object
            raw = mne.io.read_raw_edf(fname, preload=True, verbose=False)

            if raw.info['sfreq'] != 160:
                freq_flag = True
                print('{} is sampled at {}Hz so will be excluded.'.format(subj, raw.info['sfreq']))
                break

            # By default picks all 64 channels
            picks = mne.pick_types(raw.info, eeg=True)
            assert len(picks) == nc

            # Apply notch filter and bandpass filter
            raw.notch_filter([60.])
            raw.filter(l_freq=l_freq, h_freq=h_freq, picks=picks, method='iir', verbose=False)

            if run == '02':
                # Baseline, eye closed
                # There is only the rest event, randomly slice n_type0=22 snippets of length T for this class
                data = raw.get_data() * 1e6
                X = np.zeros((n_type0, nc, T))
                y = np.zeros((n_type0, ), dtype=int)   # eye closed
                for i in range(n_type0):
                    offset = np.random.randint(0, data.shape[1] - T)
                    X[i] = data[:, offset: offset+T]

            else:
                # Epoching
                events, _ = mne.events_from_annotations(raw, verbose=False)
                epochs = mne.Epochs(raw, events=events, event_id=event_id, tmin=tmin, tmax=tmax, baseline=None,
                                    # reject={'eeg': reject_threshold*1e-6},
                                    picks=picks, preload=True, verbose=False, on_missing='warning')

                # Data
                X = epochs.get_data() * 1e6     # converts to uV, the L503 in epochs.py makes T=641
                X = X[:, :, :-1]                # make T = sfeq * (tmax-tmin)

                # specific treatment of classes/labels
                y = epochs.events[:, -1]
                if run in run_type_1:
                    y[y == 2] = 1   # left fist
                    y[y == 3] = 2   # right fist
                else:  # run_type_2
                    y[y == 3] = 4   # both feet
                    y[y == 2] = 3   # both fists

            subject_trials.append(X)
            subject_labels.append(y)

        if freq_flag:
            continue
        subject_trials = np.concatenate(subject_trials, axis=0)  # (None, nc=64, T=640)
        subject_labels = np.concatenate(subject_labels)
        # print({i: sum(subject_labels == i) for i in range(5)})
        print(subj, subject_trials.shape)

        Xs.append(subject_trials)
        ys.append(subject_labels)

    return Xs, ys

def prepare_data(X, y, split_ratio=[0.64, 0.16, 0.2], chunking=True, set_seed=42):
    """ Prepare data for DNN, CNN, RNN models

    Partition data by subjects
    (Optional) Partition each trial of 4 sec to 10 consecutive chunks of 0.4 sec

    :param X:  list of (None, nc, T), each for a subject
    :param y:  list (None, )
    :param split_ratio: [train, val, test] split ratio
    :param chunking:  default False, if True, chunk T to 10 consecutive chunks of T/10

    :return: X_train, y_train, X_val, y_val, X_test, y_test
    """

    print('There are {} subjects, time series length {}, {} channels'.format(len(X), X[0].shape[1], X[0].shape[2]))

    X = np.concatenate(X, axis=0)

    # Eliminate the influence of reference electrode
    # here minus the average across all subjects and trials
    X = np.transpose(X, (0, 2, 1))
    N, T, nc = X.shape
    X = X.reshape(N*T, nc)
    X -= X.mean(axis=0)
    X = X.reshape(N, T, nc)
    X = np.transpose(X, (0, 2, 1))   # (N, nc, T)

    # Z-score for every trial for each time step
    scaler = StandardScaler()
    X = np.stack([scaler.fit_transform(x) for x in X])

    # Reshape to list of subjects
    cutoffs = np.cumsum([len(y_) for y_ in y])[:-1]
    X = np.array_split(X, cutoffs)  # [(N1, nc, T), (N2, nc, T), ...]

    # Train/val/test split by subjects
    np.random.seed(set_seed)
    shuffle_indices = np.random.permutation(len(X))
    test_size = int(len(X) * split_ratio[2])
    train_size = int(len(X) * split_ratio[0])

    X_train = [X[i] for i in shuffle_indices[:train_size]]
    y_train = [y[i] for i in shuffle_indices[:train_size]]
    X_val = [X[i] for i in shuffle_indices[train_size:-test_size]]
    y_val = [y[i] for i in shuffle_indices[train_size:-test_size]]
    X_test = [X[i] for i in shuffle_indices[-test_size:]]
    y_test = [y[i] for i in shuffle_indices[-test_size:]]

    def chunk(X, y):
        # Partition each trial of 4 sec to 10 consecutive chunks of 0.4 sec
        X_new, y_new = [], []
        for x, y_ in zip(X, y):
            x = np.array_split(x, 10, axis=-1)  # 10 x (N, nc, T/10)
            x = np.concatenate(x, axis=0)       # (10N, nc, T/10)
            X_new.append(x)

            y_ = y_.reshape(1, -1).repeat(10, axis=0).ravel()
            y_new.append(y_)

        return X_new, y_new

    if chunking:
        X_train, y_train = chunk(X_train, y_train)
        X_val, y_val = chunk(X_val, y_val)
        X_test, y_test = chunk(X_test, y_test)

    X_train = np.concatenate(X_train, axis=0)
    y_train = np.concatenate(y_train)
    X_train, y_train = shuffle(X_train, y_train)

    X_val = np.concatenate(X_val, axis=0)
    y_val = np.concatenate(y_val)
    X_val, y_val = shuffle(X_val, y_val)

    X_test = np.concatenate(X_test, axis=0)
    y_test = np.concatenate(y_test)

    return X_train, y_train, X_val, y_val, X_test, y_test

if __name__ == '__main__':

    X, y = get_data(subj_num=FNAMES, tmin=0.0, tmax=4.0)

    X_train, y_train, X_val, y_val, X_test, y_test = prepare_data(
        X, y, split_ratio=[0.64, 0.16, 0.2], chunking=False)
    print('X_train', X_train.shape)
    print('y_train', y_train.shape)
    print('X_val', X_val.shape)
    print('X_test', X_test.shape)

    data = {'X_train': X_train, 'X_val': X_val, 'X_test': X_test,
            'y_train': y_train, 'y_val': y_val, 'y_test': y_test}

    with open('../data/data.pkl', 'wb') as f:
        pickle.dump(data, f)


import os
import random
import datetime
import argparse

import numpy as np
import tensorflow as tf
from tensorflow.keras import Model
from tensorboard.plugins.hparams import api as hp

import CSR_Net
import util

# confirm tf is using GPU
# print("Num GPUs Available: ", len(tf.config.experimental.list_physical_device$
# input('Press enter of gpu settings are good')

# gundersena@75.86.178.105:~/Desktop/crimata-super-res/train/logs/weights ~/Desktop
# scp rm -r gundersena.75.86.178.105:~/Desktop/crimata-super-res/main
# scp -r ~/Desktop/crimata-super-res/main gundersena@75.86.178.105:~/Desktop/crimata-super-res

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

def make_parser():
    """creates argument parser from train and eval"""
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(title='Commands')

    # train
    train_parser = subparsers.add_parser('train')
    train_parser.set_defaults(func=train)

    train_parser.add_argument('-i','--model-id')
    train_parser.add_argument('-c','--from_ckpt')
    train_parser.add_argument('-k','--new-data')
    train_parser.add_argument('-d','--dim-size',type=int)
    train_parser.add_argument('-x','--num-files',type=int)
    # train_parser.add_argument('-t','--train-file')
    # train_parser.add_argument('-v','--val-file')
    train_parser.add_argument('-e','--epochs',type=int)
    train_parser.add_argument('-b','--batch-size',type=int)
    train_parser.add_argument('-o','--cycle-length',type=int)
    train_parser.add_argument('-m','--max-lr',type=float)
    train_parser.add_argument('-n','--min-lr',type=float)

    # eval
    eval_parser = subparsers.add_parser('eval')
    eval_parser.set_defaults(func=eval)

    eval_parser.add_argument('-i','--model-id')
    eval_parser.add_argument('-n','--num-examples',type=int)
    eval_parser.add_argument('-w','--wavfile-list')
    eval_parser.add_argument('-r','--scale',type=int)
    eval_parser.add_argument('-s','--sample-rate',type=int)
    eval_parser.add_argument('-a','--make-audio')
    eval_parser.add_argument('-c','--from-ckpt', default='True')

    return parser

def train(args):
    """High-level method for training a model"""
    # load data
    x_train, y_train, n_sam = util.load_data(args, type='train', num_files=args.num_files, full_data=True)
    x_val, y_val = util.load_data(args, type='val', num_files=int(np.floor(args.num_files*0.3)))

    # callbacks
    checkpointer = tf.keras.callbacks.ModelCheckpoint(filepath=f'logs/weights/weights.{args.model_id}.tf',
        monitor='val_loss', save_best_only=True, save_weights_only=True, mode='auto')

    # smart_learn = CSR_Net.util.SGDRScheduler(min_lr=args.min_lr, max_lr=args.max_lr,
    #     steps_per_epoch=np.ceil(n_sam/args.batch_size), cycle_length=args.cycle_length)

    # lr_finder = CSR_Net.util.LRFinder(min_lr=1e-7, max_lr=3e-2,
    #     steps_per_epoch=np.ceil(n_sam/args.batch_size), epochs=args.epochs)

    # logdir = f'logs/fit/{args.model_id}-{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}'
    # hparams = {'max_lr':args.max_lr, 'min_lr':args.min_lr, 'cycle_length':args.cycle_length}
    # param_logger = hp.KerasCallback(logdir, hparams)
    #
    # tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logdir, histogram_freq=1,
    #     write_graph=True, update_freq='epoch')

    # make model
    model = make_model(args)

    # compile model
    optimizer = tf.keras.optimizers.Adam(learning_rate=args.max_lr)
    # optimizer = tf.keras.optimizers.SGD(learning_rate=args.max_lr, momentum=0.8, nesterov=False)
    model.compile(optimizer=optimizer, loss='mean_squared_error')

    # final review
    util.review_model(args, x_train, y_train)

    # train model
    model.fit(x=x_train, y=y_train, batch_size=args.batch_size, epochs=args.epochs,
        callbacks=[checkpointer],
        validation_data=[x_val, y_val], shuffle=True)

    # plot loss and lr metrics
    # lr_finder.plot_lr()
    # lr_finder.plot_loss()

def eval(args):
    """test the model on real audio"""
    # make model
    model = make_model(args)

    # create list of file names
    file_list = []
    with open(args.wavfile_list) as f:
        for line in f:
            file_list.append(line) # this is gonna get pretty big for a real dataset...

    # eval on random sample of files
    file_list = random.sample(file_list, args.num_examples)
    for idx, line in enumerate(file_list):
        file = line.rstrip('\n')
        CSR_Net.util.eval_wav(file, args, model)

def make_model(args):
    """define a graph and compile model"""
    model = CSR_Net.MlRes()

    if args.from_ckpt == 'True':
        model.load_weights((f'logs/weights/weights.{args.model_id}.tf'))

    return model

def main():
    parser = make_parser()
    args = parser.parse_args()
    args.func(args)

if __name__ == '__main__':
    main()


__author_='8.Ball'

#USAGE
#python object_detection.py --yolo YOLOV3

from imutils.video import VideoStream
import numpy as np
import argparse
import time
import cv2
import os

#construct the argument parse and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-y", "--yolo", required=True, help="base path to the YOLO directory")
ap.add_argument("-c", "--confidence", type=float, default=0.5, help="minimum probability to filter weak detections")
ap.add_argument("-t", "--threshold", type=float, default=0.3, help="threshold when applying non-maxima suppression")
args = vars(ap.parse_args())

#load the COCO class labels the YOLO model was trained on
labelsPath = os.path.sep.join([args["yolo"], "coco.names"])
LABELS = open(labelsPath).read().strip().split("\n")

#initialize a list of colors to represent each possible class label
np.random.seed(42)
COLORS = np.random.randint(0, 255, size=(len(LABELS), 3), dtype="uint8")

#paths to the YOLO weights and model configuration
weightsPath = os.path.sep.join([args["yolo"], "yolov3.weights"])
configPath = os.path.sep.join([args["yolo"], "yolov3.cfg"])

#Load the YOLO object detector trained on COCO dataset (80 classes)
print("[INFO] loading YOLO from disk...")
net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)

#Initialize the video stream and warm up the camera sensor
print("[INFO] starting video stream...")
vs = VideoStream(src=0).start()
time.sleep(1.0)

#loop over the video frames
while 1:
    frame = vs.read()
    #getting the spatial dimensions
    (H, W) = frame.shape[:2]

    #determine only the *output* layer names that is needed from YOLO
    ln = net.getLayerNames()
    ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]

    #construct a blob from the input frame and then perform a forward pass of the YOLO object detector, giving the bounding boxes and associated probabilities
    blob = cv2.dnn.blobFromImage(frame, 1 / 255.0, (416, 416), swapRB=True, crop=False)
    net.setInput(blob)
    start = time.time()
    layerOutputs = net.forward(ln)
    end = time.time()

    #show timing information on YOLO
    print("[INFO] YOLO took {:.6f} seconds".format(end - start))

    #initialize the lists of detected bounding boxes, confidence, and class IDs, respectively
    boxes = []
    confidences = []
    classIDs = []

    #loop over each of the layer outputs
    for output in layerOutputs:
        #loop over each of the detections
        for detection in output:
            #extract the class ID and confidence (i.e. probability) of the current object detection
            scores = detection[5:]
            classID = np.argmax(scores)
            confidence = scores[classID]

            #filter out weak predictions by ensuring the detected probability is greater than the minimum probability
            if confidence > args["confidence"]:
                #scale the bounding box coordinates back relative to the size of the image, keeping in mind the YOLO actually returns the center (x, y)-coordinates of the bounding box followerd by the boxes' width and height
                box = detection[0:4] * np.array([W, H, W, H])
                (centerX, centerY, width, height) = box.astype("int")

                #use the center (x, y)-coordinates to derive the top and left corner of the bounding box
                x = int(centerX - (width / 2))
                y = int(centerY - (height / 2))

                #update the list of bounding box coordinates, confidences, classIDs
                boxes.append([x, y, int(width), int(height)])
                confidences.append(float(confidence))
                classIDs.append(classID)

    #apply non-maxima suppression to suppress weak, overlapping bounding boxes
    idxs = cv2.dnn.NMSBoxes(boxes, confidences, args["confidence"], args["threshold"])
    #YOLO does not apply non-maxima suppression, so it is explicitly applied
    #Applying non-maxima suppression suppresses significantly overlapping bounding boxes, keeping only the most confident ones
    #NMS also ensures that we do not have any redundant or exrtraneous bounding boxes

    #ensure at least one detection exists
    if len(idxs) > 0:
        #loop over the indexes we are keepinng
        for i in idxs.flatten():
            #extract the bounding box coordinates
            (x, y) = (boxes[i][0], boxes[i][1])
            (w, h) = (boxes[i][2], boxes[i][3])

            #draw a bouding box rectangle and label on the frame
            color = [int(c) for c in COLORS[classIDs[i]]]
            cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
            text = "{}: {:.4f}".format(LABELS[classIDs[i]], confidences[i])
            cv2.putText(frame, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    #show the output frames
    cv2.imshow("Object Detection", frame)
    if cv2.waitKey(1) ==  27:
        break

cv2.destroyAllWindows()


"""credit: https://github.com/keithito/tacotron/blob/master/datasets/blizzard.py"""

from concurrent.futures import ProcessPoolExecutor
from functools import partial
import numpy as np
import os
from hparams import HParams as hp
import audio

from tqdm import tqdm
from multiprocessing import cpu_count

_max_out_length = 700
_end_buffer = 0.05
_min_confidence = 90

# Note: "A Tramp Abroad" & "The Man That Corrupted Hadleyburg" are higher quality than the others.
books = [
    'ATrampAbroad',
    'TheManThatCorruptedHadleyburg',
    # 'LifeOnTheMississippi',
    # 'TheAdventuresOfTomSawyer',
]

def build_from_path(in_dir, out_dir, num_workers=1, tqdm=lambda x: x):
    executor = ProcessPoolExecutor(max_workers=num_workers)
    futures = []
    index = 1
    for book in books:
        with open(os.path.join(in_dir, book, 'sentence_index.txt')) as f:
            for line in f:
                parts = line.strip().split('\t')
                if line[0] is not '#' and len(parts) == 8 and float(parts[3]) > _min_confidence:
                    wav_path = os.path.join(in_dir, book, 'wav', '%s.wav' % parts[0])
                    labels_path = os.path.join(in_dir, book, 'lab', '%s.lab' % parts[0])
                    text = parts[5]
                    task = partial(_process_utterance, out_dir, index, wav_path, labels_path, text)
                    futures.append(executor.submit(task))
                    index += 1
    results = [future.result() for future in tqdm(futures)]
    return [r for r in results if r is not None]

def _process_utterance(out_dir, index, wav_path, labels_path, text):
    # Load the wav file and trim silence from the ends:
    wav = audio.load_wav(wav_path)
    start_offset, end_offset = _parse_labels(labels_path)
    start = int(start_offset * hp.sr)
    end = int(end_offset * hp.sr) if end_offset is not None else -1
    wav = wav[start:end]
    max_samples = _max_out_length * hp.frame_shift * hp.sr
    if len(wav) > max_samples:
        # print(wav_path + ": wav too long")
        return None
    spectrogram = audio.spectrogram(wav).astype(np.float32)
    n_frames = spectrogram.shape[1]
    mel_spectrogram = audio.melspectrogram(wav).astype(np.float32)
    spectrogram_filename = 'blizzard-spec-%05d.npy' % index
    mel_filename = 'blizzard-mel-%05d.npy' % index
    np.save(os.path.join(out_dir, spectrogram_filename), spectrogram.T, allow_pickle=False)
    np.save(os.path.join(out_dir, mel_filename), mel_spectrogram.T, allow_pickle=False)
    return (spectrogram_filename, mel_filename, n_frames, text)

def _parse_labels(path):
    labels = []
    with open(os.path.join(path)) as f:
        for line in f:
            parts = line.strip().split(' ')
            if len(parts) >= 3:
                labels.append((float(parts[0]), ' '.join(parts[2:])))
    start = 0
    end = None
    if labels[0][1] == 'sil':
        start = labels[0][0]
    if labels[-1][1] == 'sil':
        end = labels[-2][0] + _end_buffer
    return (start, end)

def write_metadata(metadata, out_dir):
    with open(os.path.join(out_dir, 'train.txt'), 'w', encoding='utf-8') as f:
        for m in metadata:
            f.write('|'.join([str(x) for x in m]) + '\n')
    frames = sum([m[2] for m in metadata])
    hours = frames * hp.frame_shift / 3600
    print('Wrote %d utterances, %d frames (%.2f hours)' % (len(metadata), frames, hours))
    print('Max input length:  %d' % max(len(m[3]) for m in metadata))
    print('Max output length: %d' % max(m[2] for m in metadata))

if __name__ == "__main__":
    in_dir = "datasets/blizzard-taco-preprop"
    out_dir = "datasets/blizzard-taco-processed"
    num_workers = cpu_count()
    os.makedirs(out_dir, exist_ok=True)
    metadata = build_from_path(in_dir, out_dir, num_workers, tqdm=tqdm)
    write_metadata(metadata, out_dir)


import numpy as np
from slider.beatmap import Circle as SliderCircle
from slider.beatmap import Slider as SliderSlider
from slider.beatmap import Spinner as SliderSpinner
from slider.mod import circle_radius

from circleguard.mod import Mod

# We define our own hitobjects as the slider library's hitobjects have too many
# attributes and methods we don't care about, and they also lock position
# behind an extra attribute access (``hitobj.position.x``` vs ``hitobj.x``)
# which I'm not a fan of.
# Another necessary change is our hitobjects are "replay/map aware", which means
# they know how large they are (or potentially how difficult they are to hit
# as a result of OD if that is necessary in the future) and where they are (for
# HR) because we know with what mods and on what map the hitobject was played
# with.

class Hitobject:
    """
    A Hitobject in osu! gameplay, with a time and a position.
    """
    def __init__(self, time, xy):
        # TODO remove ``t`` in core 6.0.0, ``time`` should be preferred
        self.t = time
        self.time = time
        self.xy = xy
        self.x = xy[0]
        self.y = xy[1]

    @classmethod
    def from_slider_hitobj(cls, hitobj, replay, beatmap, already_converted=False):
        """
        Instantiates a circleguard hitobject from a
        :class:`slider.beatmap.HitObject`, a
        :class:`circleguard.loadables.Replay`, that the hitobject was hit on,
        and a :class:`slider.beatmap.Beatmap` that the hitobject is found in.

        The `already_converted` parameter is only to work around
        https://github.com/llllllllll/slider/issues/80 and will be removed when
        it is fixed.
        """
        easy = Mod.EZ in replay.mods
        hard_rock = Mod.HR in replay.mods
        CS = beatmap.cs(easy=easy, hard_rock=hard_rock)

        # Convert to ms.
        t = hitobj.time.total_seconds() * 1000
        # Due to floating point errors, ``t`` could actually be something
        # like ``129824.99999999999`` or ``128705.00000000001``, so round to the
        # nearest int.
        t = int(round(t))

        if hard_rock and not already_converted:
            hitobj = hitobj.hard_rock

        xy = [hitobj.position.x, hitobj.position.y]
        xy = np.array(xy)

        radius = circle_radius(CS)

        if isinstance(hitobj, SliderCircle):
            return Circle(t, xy, radius)
        if isinstance(hitobj, SliderSlider):
            return Slider(t, xy, radius)
        if isinstance(hitobj, SliderSpinner):
            return Spinner(t, xy)

    def __eq__(self, other):
        return self.time == other.time and self.xy == other.xy

    def __hash__(self):
        return hash((self.time, self.xy))

class Circle(Hitobject):
    """
    A circle in osu! gameplay, with a time, position, and radius.
    """
    def __init__(self, time, xy, radius):
        super().__init__(time, xy)
        self.radius = radius

    def __eq__(self, other):
        return (self.time == other.time and self.xy == other.xy and
            self.radius == other.radius)

    def __hash__(self):
        return hash((self.time, self.xy, self.radius))

class Slider(Hitobject):
    """
    A slider in osu! gameplay, with a time, position, and radius.
    """
    def __init__(self, time, xy, radius):
        super().__init__(time, xy)
        self.radius = radius

    def __eq__(self, other):
        return (self.time == other.time and self.xy == other.xy and
            self.radius == other.radius)

    def __hash__(self):
        return hash((self.time, self.xy, self.radius))

class Spinner(Hitobject):
    """
    A spinner in osu! gameplay, with a time and position.
    """
    def __init__(self, time, xy):
        super().__init__(time, xy)


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

"""
Test nematusLL for consistency with nematus
"""

import os
import unittest
import sys
import numpy as np
import logging
import Pyro4

nem_path = os.path.abspath(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../'))
sys.path.insert(1, nem_path)
from nematus.pyro_utils import setup_remotes, get_random_key, get_unused_port
from nematus.util import load_dict

from unit_test_utils import initialize

GPU_ID = 0
VOCAB_SIZE = 90000
SRC = 'ro'
TGT = 'en'
DATA_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'test_data')

model_options = dict(factors=1,  # input factors
                     dim_per_factor=None,
                     # list of word vector dimensionalities (one per factor): [250,200,50] for dimensionality of 500
                     encoder='gru',
                     decoder='gru_cond',
                     patience=10,  # early stopping patience
                     max_epochs=5000,
                     finish_after=10000000,  # finish after this many updates
                     map_decay_c=0.,  # L2 regularization penalty towards original weights
                     alpha_c=0.,  # alignment regularization
                     shuffle_each_epoch=True,
                     finetune=False,
                     finetune_only_last=False,
                     sort_by_length=True,
                     use_domain_interpolation=False,
                     domain_interpolation_min=0.1,
                     domain_interpolation_inc=0.1,
                     domain_interpolation_indomain_datasets=['indomain.en', 'indomain.fr'],
                     maxibatch_size=20,  # How many minibatches to load at one time
                     model_version=0.1,  # store version used for training for compatibility
                     pyro_key=None,  # pyro hmac key
                     pyro_port=None,  # pyro nameserver port
                     pyro_name=None,  # if None, will import instead of assuming a server is running
                     saveto='model.npz',
                     reload_=True,
                     dim_word=500,
                     dim=1024,
                     n_words=VOCAB_SIZE,
                     n_words_src=VOCAB_SIZE,
                     decay_c=0.,
                     clip_c=1.,
                     lrate=0.0001,
                     optimizer='adadelta',
                     maxlen=50,
                     batch_size=80,
                     valid_batch_size=80,
                     datasets=[DATA_DIR + '/corpus.bpe.' + SRC, DATA_DIR + '/corpus.bpe.' + TGT],
                     valid_datasets=[DATA_DIR + '/newsdev2016.bpe.' + SRC, DATA_DIR + '/newsdev2016.bpe.' + TGT],
                     dictionaries=[DATA_DIR + '/corpus.bpe.' + SRC + '.json',
                                   DATA_DIR + '/corpus.bpe.' + TGT + '.json'],
                     validFreq=10000,
                     dispFreq=10,
                     saveFreq=30000,
                     sampleFreq=10000,
                     use_dropout=False,
                     dropout_embedding=0.2,  # dropout for input embeddings (0: no dropout)
                     dropout_hidden=0.2,  # dropout for hidden layers (0: no dropout)
                     dropout_source=0.1,  # dropout source words (0: no dropout)
                     dropout_target=0.1,  # dropout target words (0: no dropout)
                     overwrite=False,
                     external_validation_script='./validate.sh',
                     )

x0 = np.array([[[3602],
                [8307],
                [7244],
                [7],
                [58],
                [9],
                [5893],
                [62048],
                [11372],
                [4029],
                [25],
                [34],
                [2278],
                [5],
                [4266],
                [11],
                [2852],
                [3],
                [2298],
                [2],
                [23912],
                [6],
                [16358],
                [3],
                [730],
                [2328],
                [5],
                [28],
                [353],
                [4],
                [0], ]])  # 0 = EOS

xx0 = np.tile(x0, [1, 1, 2])

x1 = np.array([[[3602],
                [8307],
                [7244],
                [7],
                [58],
                [9],
                [4],
                [0], ]])  # 0 = EOS

# False Truth for testing Gradients
y1 = np.array([[[1],
                [1],
                [1],
                [1],
                [1],
                [1],
                [1],
                [0], ]])  # 0 = EOS

x_0 = np.array([[[2], [0], [0], [0]]])

x_0_mask = np.array([[1], [1], [1], [0]], dtype=np.float32)

y_0 = np.array([[17], [9], [41], [120], [7], [117], [1087], [476], [1715], [62], [2], [0]])

y_0_mask = np.array([[1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1], [1]], dtype=np.float32)

per_sample_weight_1 = np.array([1, 1], dtype=np.float32)

xx_0 = np.array([[[4, 2], [4, 0], [3, 0], [2, 0]]])

xx_0_mask = np.array([[2, 1], [0, 1], [0, 1], [0, 0]], dtype=np.float32)

yy_0 = np.array([[22, 17], [24, 9], [31, 41], [420, 120], [37, 7], [127, 117], [1387, 1087], [446, 476], [1515, 1715],
                 [22, 62], [12, 2], [0, 0]])

yy_0_mask = np.array([[1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1], [1, 1]],
                     dtype=np.float32)

per_sample_weight_2 = np.array([0.5, 0.5, 0.5, 0.5], dtype=np.float32)

xx_1 = np.tile(xx_0, [2])

yy_1 = np.tile(yy_0, [2])

xx_1_mask = np.tile(xx_0_mask, [2])

yy_1_mask = np.tile(yy_0_mask, [2])

per_sample_weight = np.array([0.25, 0.25, 0.25, 0.25], dtype=np.float32)

x_1 = np.array([[[2, 2, 2, 2],
                 [0, 0, 0, 0],
                 [0, 0, 0, 0],
                 [0, 0, 0, 0]]])
x_1_mask = np.array([[1., 1., 1., 1.],
                     [1., 1., 1., 1.],
                     [1., 1., 1., 1.],
                     [0., 0., 0., 0.]], dtype=np.float32)
y_1 = np.array([[17, 17, 17, 17],
                [9, 9, 9, 9],
                [41, 41, 41, 41],
                [120, 120, 120, 120],
                [7, 7, 7, 7],
                [117, 117, 117, 117],
                [1087, 1087, 1087, 1087],
                [476, 476, 476, 476],
                [1715, 1715, 1715, 1715],
                [62, 62, 62, 62],
                [2, 2, 2, 2],
                [0, 0, 0, 0]])
y_1_mask = np.array([[1., 1., 1., 1.],
                     [1., 1., 1., 1.],
                     [1., 1., 1., 1.],
                     [1., 1., 1., 1.],
                     [1., 1., 1., 1.],
                     [1., 1., 1., 1.],
                     [1., 1., 1., 1.],
                     [1., 1., 1., 1.],
                     [1., 1., 1., 1.],
                     [1., 1., 1., 1.],
                     [1., 1., 1., 1.],
                     [1., 1., 1., 1.]], dtype=np.float32)

class PyroNematusTests(unittest.TestCase):
    # Forces the class to have these fields.
    logger = None
    context_mgr = None

    @classmethod
    def setUpClass(cls):
        cls.logger = logging.getLogger(__file__)
        cls.logger.info("========================================================================================")
        cls.logger.info("Setting up the pyro remote server as well as the nematus instance to test consistency.")
        cls.logger.info("========================================================================================")
        current_script_dir = os.path.dirname(os.path.abspath(__file__))
        remote_script = os.path.join(current_script_dir, '../nematus/nmt_remote.py')
        pyro_name = 'remote'

        pyro_key = get_random_key()
        pyro_port = get_unused_port()

        cls.context_mgr = setup_remotes(remote_metadata_list=[dict(script=remote_script,
                                                                   name=pyro_name, gpu_id=GPU_ID)],
                                        pyro_port=pyro_port,
                                        pyro_key=pyro_key)
        cls.context_mgr.__enter__()
        cls.remote_interface = initialize(model_options=model_options,
                                          pyro_port=pyro_port,
                                          pyro_name=pyro_name,
                                          pyro_key=pyro_key)

    @classmethod
    def tearDownClass(cls):
        cls.logger.info("========================================================================================")
        cls.logger.info("Tearing down the pyro remote server as well as the nematus instance")
        cls.logger.info("========================================================================================")
        cls.context_mgr.__exit__(None, None, None)

    def assert_params_same(self, params1, params2):
        k1 = params1.keys()
        k2 = params2.keys()
        k1.sort()
        k2.sort()
        self.assertTrue(k1 == k2)
        for k in k1:
            self.assertTrue(np.allclose(params1[k], params2[k]), 'value for %s should match' % k)

    def assert_params_different(self, params1, params2):
        k1 = params1.keys()
        k2 = params2.keys()
        k1.sort()
        k2.sort()
        self.assertTrue(k1 == k2)
        diff_vec = [not np.allclose(params1[k], params2[k]) for k in k1]
        self.logger.info("Difference vector is: " + str(diff_vec))
        self.assertTrue(any(diff_vec))

    def test_f_init_dims(self):
        """
        Best I can tell, f_init is only ever given one sentence, but it appears to be
        written to process multiple sentences.
        """
        self.logger.info("========================================================================================")
        self.logger.info("Starting the f_init_dims test to determine that x_f_init acts as expected.")
        self.logger.info("========================================================================================")
        x0_state0, x0_ctx0 = self.remote_interface.x_f_init(x0)  # (1, 1024) (31, 1, 2048)

        # If tile input, state/context should be tiled too
        xx0_state0, xx0_ctx0 = self.remote_interface.x_f_init(xx0)  # (2, 1024) (31, 2, 2048)
        self.assertTrue(np.allclose(np.tile(x0_state0, [2, 1]), xx0_state0))
        self.assertTrue(np.allclose(np.tile(x0_ctx0, [1, 2, 1]), xx0_ctx0))

        # Different inputs should create different state
        x1_state0, x1_ctx0 = self.remote_interface.x_f_init(x1)
        self.assertFalse(np.allclose(x0_state0, x1_state0))

        # Different inputs (of same length) should create different state and context
        x1_2_state0, x1_2_ctx0 = self.remote_interface.x_f_init(x1 * 2)
        self.assertFalse(np.allclose(x1_state0, x1_2_state0))
        self.assertFalse(np.allclose(x1_ctx0, x1_2_ctx0))

    def test_f_next_dims(self):
        self.logger.info("========================================================================================")
        self.logger.info("Starting the f_next_dims test to determine that x_f_next acts as expected.")
        self.logger.info("========================================================================================")
        self.remote_interface.set_noise_val(0)
        x0_state0, x0_ctx0 = self.remote_interface.x_f_init(x0)
        x0_prob1, x0_word1, x0_state1 = self.remote_interface.x_f_next(np.array([2893, ]), x0_ctx0, x0_state0)
        x0_prob2, x0_word2, x0_state2 = self.remote_interface.x_f_next(np.array([9023, ]), x0_ctx0, x0_state1)
        self.assertFalse(np.allclose(x0_state0, x0_state1), 'state should be changing')
        self.assertFalse(np.allclose(x0_prob1, x0_prob2), 'probability should be changing')
        # word might not change...

        self.logger.info('x0 prob shape, ' + str(x0_prob1.shape))
        self.logger.info('x0 word shape, ' + str(x0_word1.shape))
        self.logger.info('x0 state shape, ' + str(x0_state2.shape))

        xx0_state0, xx0_ctx0 = self.remote_interface.x_f_init(xx0)
        xx0_prob1, xx0_word1, xx0_state1 = self.remote_interface.x_f_next(np.array([2893, 2893]), xx0_ctx0, xx0_state0)
        xx0_prob2, xx0_word2, xx0_state2 = self.remote_interface.x_f_next(np.array([9023, 9023]), xx0_ctx0, xx0_state1)

        self.logger.info('xx0 prob shape, ' + str(xx0_prob1.shape))
        self.logger.info('xx0 word shape, ' + str(xx0_word1.shape))
        self.logger.info('xx0 state shape, ' + str(xx0_state2.shape))

        self.assertTrue(np.allclose(np.tile(x0_prob1, [2, 1]), xx0_prob1))
        self.assertTrue(np.allclose(np.tile(x0_prob2, [2, 1]), xx0_prob2))
        # jitter??
        # print 'same??', np.tile(x0_word1, [2]), xx0_word1
        # self.assertTrue(np.allclose(np.tile(x0_word1, [2]), xx0_word1))
        # self.assertTrue(np.allclose(np.tile(x0_word2, [2]), xx0_word2))
        self.assertTrue(np.allclose(np.tile(x0_state1, [2, 1]), xx0_state1))
        self.assertTrue(np.allclose(np.tile(x0_state2, [2, 1]), xx0_state2))

    def test_f_init_side_effects(self):
        self.logger.info("========================================================================================")
        self.logger.info("Starting a side effect test to determine that f_init has no side effects. ")
        self.logger.info("========================================================================================")
        self.remote_interface.x_f_update(lrate=0.1)  # should zero grads
        params0 = self.remote_interface.get_params_from_theano()
        _, _ = self.remote_interface.x_f_init(x0)
        _, _ = self.remote_interface.x_f_init(x1)
        _, _ = self.remote_interface.x_f_init(x0)
        params2 = self.remote_interface.get_params_from_theano()
        self.remote_interface.x_f_update(lrate=0.1)  # grads should still be zero, so this should not change params
        params1 = self.remote_interface.get_params_from_theano()
        self.assert_params_same(params0, params2)
        self.assert_params_different(params0, params1)

    def test_get_params_consistency(self):
        self.logger.info("========================================================================================")
        self.logger.info("Starting the param_consistency test to determine that params do not change unexpectedly.")
        self.logger.info("========================================================================================")
        params0 = self.remote_interface.get_params_from_theano()
        params1 = self.remote_interface.get_params_from_theano()
        params2 = self.remote_interface.get_params_from_theano()
        self.assert_params_same(params0, params1)
        self.assert_params_same(params0, params2)

        self.remote_interface.x_f_update(lrate=0.1)  # should zero grads
        params3 = self.remote_interface.get_params_from_theano()
        self.remote_interface.x_f_update(lrate=0.1)  # grads=0, so this should not change params
        params4 = self.remote_interface.get_params_from_theano()
        self.assert_params_different(params3, params4)

    def test_f_grad_shared(self):
        self.logger.info("========================================================================================")
        self.logger.info("Starting the f_next_dims test to determine that x_f_next acts as expected.")
        self.logger.info("========================================================================================")
        self.remote_interface.set_noise_val(0)
        params = self.remote_interface.get_params_from_theano()
        grad = self.remote_interface.x_f_grad_shared(x_0, x_0_mask, y_0, y_0_mask, per_sent_weight=[1])
        self.remote_interface.send_params_to_theano(params)
        grad2 = self.remote_interface.x_f_grad_shared(x_1, x_1_mask, y_1, y_1_mask, per_sent_weight=per_sample_weight)
        self.assertAlmostEqual(grad, grad2, places=1)
        self.remote_interface.send_params_to_theano(params)
        grad3 = self.remote_interface.x_f_grad_shared(xx_0, xx_0_mask, yy_0, yy_0_mask,
                                                      per_sent_weight=per_sample_weight_1)
        self.remote_interface.send_params_to_theano(params)
        grad4 = self.remote_interface.x_f_grad_shared(xx_1, xx_1_mask, yy_1, yy_1_mask,
                                                      per_sent_weight=per_sample_weight_2)
        self.assertAlmostEqual(grad3, grad4, places=1)

if __name__ == '__main__':
    unittest.main()


import numpy
__author__ = 'kitch'

def fftdot(a, b, *, out=None):
    #return numpy.ifft(numpy.fft(a))
    return


import mxnet as mx
import asyncio
import json
from PIL import Image,  ImageFont
from inference.base_inference_engine import AbstractInferenceEngine
from inference.exceptions import InvalidModelConfiguration, InvalidInputData, ApplicationError
import os
import numpy as np
from mxnet import image
import numpy
from gluoncv.model_zoo.segbase import *
from gluoncv.utils.parallel import *
from gluoncv.data.transforms.presets.segmentation import test_transform
from gluoncv.data.pascal_voc.segmentation import VOCSegmentation
from mxnet import ndarray as nd
import cv2
from io import BytesIO
import time

class InferenceEngine(AbstractInferenceEngine):

    def __init__(self, model_path):
        self.nclasses=0
        self.model="none"
        self.backbone="none"
        self.palette=[]
        self.font = ImageFont.truetype("/main/fonts/DejaVuSans.ttf", 20)
        super().__init__(model_path)

    def free(self):

        pass
    """
    Function that loads models for inference
    """
    def load(self):

        with open(os.path.join(self.model_path, 'configuration.json')) as f:
            data = json.load(f)
        try:

            self.set_model_configuration(data)
            self.validate_json_configuration(data)
        except ApplicationError as e:
            raise e
        self.nclasses=data["num_classes"]
        self.model=data["network"]
        self.backbone=data["backbone"]

        gpu_count=mx.util.get_gpu_count()
        if gpu_count>0:
            self.ctx1 = [mx.gpu()]
        else:
            self.ctx1= [mx.cpu()]
        VOCSegmentation.NUM_CLASS=self.nclasses
        self.net=get_segmentation_model(model=self.model,backbone=self.backbone,norm_layer=mx.gluon.nn.BatchNorm,norm_kwargs={},crop_size=480,aux=False,pretrained_base=False,ctx=self.ctx1)
        #self.net=get_model("fcn_resnet101_voc",pretrained=True,aux=False,crop_size=480,pretrained_base=True)
        self.net.initialize(force_reinit=True)
        net = DataParallelModel(self.net, self.ctx1, {})

        self.net.load_parameters(os.path.join(self.model_path,"model_best.params"),ctx=self.ctx1,ignore_extra=True)
        infile = open(os.path.join(self.model_path,'palette.txt'),'r')
        for line in infile:
            self.palette.append(int(line.strip('\n')))
        self.labels=data["classes_names"]
        self.background=None
        if "background" in self.labels:
            self.background=list(self.labels).index("background")

    """
    Function that decides if segments or a json response should be returned
    """

    async def infer(self, input_data, draw, predict_batch):

        response=None
        if draw:
            try:

                await self.inference(input_data.file.read())
            except ApplicationError as e:
                raise e
        else:
           response= await self.processing(input_data.file.read())

           return response

    """
    Function that returns segments
    """
    async def inference2(self,input_data):
        start_time=time.time()
        model=self.net

        bytes_as_np_array = np.frombuffer(input_data.file.read(), dtype=np.uint8)
        img=image.imdecode(bytes_as_np_array,cv2.IMREAD_COLOR)
        img = test_transform(img,self.ctx1[0])
        mx.nd.waitall()

        output = model.predict(img)
        if self.model=='fcn':
           predict = mx.nd.squeeze(mx.nd.argmax(nd.softmax(output[0],axis=1), 1)).asnumpy()
        else:
            argmax=mx.nd.argmax(output, 1)
            predict = mx.nd.squeeze(argmax)
            predict= predict.asnumpy()
        mx.nd.waitall()
        mask=predict
        mask=Image.fromarray(mask.astype('uint8'))
        mask.putpalette(self.palette)
        mask.save("/main/result.jpg","png")
        for ctx2 in self.ctx1:
            ctx2.empty_cache()
        print(time.time()-start_time)
    """
    Function that is called when infering multiple pictures
    """
    async def run_batch(self, input_data, draw, predict_batch):
        result_list = []
        for image in input_data:
            post_process = await self.infer(image, draw, predict_batch)
            if post_process is not None:
                result_list.append(post_process)
        return result_list

    """
    Function that returns segments pasted on top the original image
    """

    async def inference(self,input_data):
        model=self.net
        bytes_as_np_array = np.frombuffer(input_data, dtype=np.uint8)
        img=image.imdecode(bytes_as_np_array,cv2.IMREAD_COLOR)
        img = test_transform(img,self.ctx1[0])
        output = model.predict(img)
        if self.model=='fcn':
           predict = mx.nd.squeeze(mx.nd.argmax(nd.softmax(output[0],axis=1), 1)).asnumpy()
        else:
           predict = mx.nd.squeeze(mx.nd.argmax(nd.softmax(output,axis=1), 1)).asnumpy()

        mask=predict

        mask=Image.fromarray(mask.astype('uint8'))
        mask.putpalette(self.palette)
        img=Image.open(BytesIO(input_data))
        img=img.convert(mode='RGBA')

        mask=mask.convert(mode="RGBA")
        mask=np.array(mask)

        mask[:,:,3]=200
        if self.background != None :
            red, green, blue=mask[:,:,0], mask[:,:,1], mask[:,:,2]
            transfer=(red == self.palette[self.background*3]) & (green == self.palette[self.background*3+1]) & (blue ==self.palette[self.background*3+2])
            mask[:,:,:4][transfer]=[0,0,0,0]
        mask=Image.fromarray(mask.astype('uint8'))

        img.paste(mask,mask)
        img.save('/main/result.jpg','PNG')
        for ctx2 in self.ctx1:
            ctx2.empty_cache()

    """
    Function that returns a json response after infernece
    """

    async def processing(self, input_data):
        await asyncio.sleep(0.00001)
        model=self.net
        bytes_as_np_array = np.frombuffer(input_data, dtype=np.uint8)
        img=image.imdecode(bytes_as_np_array,cv2.IMREAD_COLOR)
        img = test_transform(img,self.ctx1[0])
        output = model.predict(img)
        if self.model=='fcn':
           predict = mx.nd.squeeze(mx.nd.argmax(nd.softmax(output[0],axis=1), 1)).asnumpy()
           confidences=nd.softmax(output[0],axis=1).asnumpy()
        else:
           predict = mx.nd.squeeze(mx.nd.argmax(nd.softmax(output,axis=1), 1)).asnumpy()
           confidences=nd.softmax(output,axis=1).asnumpy()

        mask=np.copy(predict)
        d={}
        d=numpy.unique(mask)

        boxes = []
        mask=Image.fromarray(mask.astype('uint8'))
        mask.putpalette(self.palette)
        mask=mask.convert(mode="RGB")
        image1=cv2.cvtColor(np.array(mask),cv2.COLOR_RGB2BGR)
        response1=[]
        for key in d :
            if key != self.background:

                lower=np.array([self.palette[(int(key)*3)+2],self.palette[(int(key)*3)+1],self.palette[(int(key)*3)]])
                upper=np.array([self.palette[(int(key)*3)+2]+1,self.palette[(int(key)*3)+1]+1,self.palette[(int(key)*3)]+1])

                mask1=cv2.inRange(image1,lower,upper)
                # Finding Contours
                # Use a copy of the image e.g. edged.copy()
                # since findContours alters the image
                contours, hierarchy = cv2.findContours(mask1,
                cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
                segments=[]
                for c in contours:
                    (x, y, w, h) = cv2.boundingRect(c)
                    if(len(np.array(c).flatten()) > 3):
                        boxes.append([x,y, x+w,y+h])
                        segments.append(np.array(c).flatten())
                confidence=[]
                confidence1=[]
                counting=0
                for b in boxes:

                    ind=int(key)
                    response=dict([("class_name",str(self.labels[ind])),("ObjectClassId",ind),("bbox",b),("segment",list(segments[counting])),("confidence","0.998")])
                    response=json.loads(str(response).replace("\'", "\""))
                    response1.append(response)
                    response1=json.loads(str(response1).replace("\'", "\""))
                    confidence1=[]
                    counting=counting+1
                boxes=[]

        for ctx2 in self.ctx1:
            ctx2.empty_cache()

        return response1

    def free(self):
        pass

    def validate_configuration(self):
        # check if weights and palette file exists
        if not os.path.exists(os.path.join(self.model_path, 'model_best.params')):
            raise InvalidModelConfiguration('model_best.params')
        if not os.path.exists(os.path.join(self.model_path, 'palette.txt')):
            raise InvalidModelConfiguration('palette.txt')
        return True

    #Load model configuration
    def set_model_configuration(self, data):
        self.configuration['type']=data['type']
        try:
            self.configuration['nclasses'] = data['num_classes']
            if not isinstance(self.configuration['nclasses'], int):
                raise InvalidModelConfiguration('nclasses field should be of type int')
        except ApplicationError as e:
            raise e
        except Exception as e:
            raise InvalidModelConfiguration('missing nclasses field in config.json')
        try:
            self.configuration['backbone'] = data['backbone']
            if not isinstance(self.configuration['backbone'], str):
                raise InvalidModelConfiguration('backbone field should be of type string')
        except ApplicationError as e:
            raise e
        except Exception as e:
            raise InvalidModelConfiguration('missing backbone field in config.json')
        try:
            self.configuration['model'] = data['network']
            if not isinstance(self.configuration['model'], str):
                raise InvalidModelConfiguration('model field should be of type string')
        except ApplicationError as e:
            raise e
        except Exception as e:
            raise InvalidModelConfiguration('missing model field in config.json')
        try:
            self.configuration['classesname'] = data['classes_names']
            if not isinstance(self.configuration['classesname'], list):
                raise InvalidModelConfiguration('classes field should be of type list')
        except ApplicationError as e:
            raise e
        except Exception as e:
            raise InvalidModelConfiguration('missing classesname field in config.json')

    #Validate the json configuration
    def validate_json_configuration(self, data):
        try:
            if not isinstance(data['num_classes'], int):
                raise InvalidModelConfiguration('classes field should be of type int')
        except ApplicationError as e:
            raise e
        except Exception as e:
            raise InvalidModelConfiguration('missing confidence field in config.json')
        try:
            if not isinstance(data['classes_names'], list):
                raise InvalidModelConfiguration('classesname field should be of type int')
        except ApplicationError as e:
            raise e
        except Exception as e:
            raise InvalidModelConfiguration('missing classesname field in config.json')


# Convert an arelequin .arp file to eigenstrat format
# Generate .arp file using fastsimcoal2 -s0 option
# Only DNA sequence suppoerted (with -s0), but
# multiple populations/chromosomes are allowed.
# Ascertained data automatically detected.
# usage: python arelequin2eigenstrat.py -a arp_file.arp(.gz) -o out_root [other options]
# will generate out_root.[snp,ind,geno].
# By default will convert to genotype data, so use the -p option if you
# want to keep it unphased.

from __future__ import division
import gzip, sys, getopt, gdc, pdb
import numpy as np

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

def parse_options():
    """
    Options are described by the help() function
    """
    options ={ "arp":None, "out":"out", "phased":False  }

    try:
        opts, args = getopt.getopt(sys.argv[1:], "a:o:p", ["arp", "out", "phased"])
        print opts, args
    except Exception as err:
        print str(err)
        sys.exit()

    for o, a in opts:
        print o,a
        if o in ["-a","--arp"]:         options["arp"] = a
        if o in ["-o","--out"]:         options["out"] = a
        elif o in ["-p","--phased"]:    options["phased"] = True

    print "found options:"
    print options
    return options

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

def load_from_arp(arp):
    """
    Parse an arp file into an internal format, returning site data, which is a
    list of lists, with the sites on each chromosome, and gt data which is a
    dict with infomation on the samples and genotypes.
    I'm really just guessing what the structure of the file is.
    """
    arp_file=gdc.open2(arp)

    N_chr=-1
    N_sites=-1
    site_data=[]
    gt_data={}
    ascertained=False
    for line in arp_file:
        if line.startswith("#Number of independent chromosomes"):
            N_chr=int(line.split()[-1])
        if line.startswith("#Total number of polymorphic sites") and not ascertained:
            N_sites=int(line.split()[-1])
        if line.startswith("#ASCERTAINED DATA"):
            ascertained=True
            N_sites=-1
        if line.startswith("#Number of polym. sites meeting ascertainment criterion:"):
            N_sites=int(line.split()[-1])
        if "polymorphic positions on chromosome" in line and not ascertained:
            npos=int(line.split()[1])
            chrom=int(line.split()[-1])
            line=next(arp_file)
            site_data.append(None)
            site_data[chrom-1]=[int(x.replace(",", "")) for x in line[1:].split()]
        if line.startswith("#Ascertained polymorphic positions on chromosome") and ascertained:
            chrom=int(line.split()[-1])
            line=next(arp_file)
            npos=len(line.split())
            site_data.append(None)
            site_data[chrom-1]=[int(x.replace(",", "")) for x in line[1:].split()]
        if "SampleName=" in line:
            sname=line.split("\"")[-2]
            line=arp_file.next()
            ssize=int(line.split("=")[-1])
            gt=np.zeros((N_sites, ssize), dtype='int')
            arp_file.next()               # "SampleData= {" line
            for i in range(ssize):
                line=arp_file.next()
                gt_string=line.split()[2]
                if len(gt_string)!=N_sites:
                    raise Exception("Wrong number of sites")
                gt[:,i]=[int(int(x)>0) for x in gt_string]   # converting {1,2,3}->1
            gt_data[sname]={"size":ssize, "gt":gt}

    if not N_chr==len(site_data):
        raise Exception( "Number of chromosomes does not match site data" )
    if not N_sites==sum([len(x) for x in site_data]):
        raise Exception( "Total number of sites does not match" )

    return site_data, gt_data

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

def unphase(gt_data):
    """
    Halve the number of samples and unphase.
    """
    for sample, data in gt_data.iteritems():
        oldrows, oldcols=data["gt"].shape
        if oldcols % 2:
            raise Exception("Can't unphase odd number of chromosomes")
        new_gt=np.zeros((oldrows, oldcols//2))
        for i in range(oldcols//2):
            new_gt[:,i]=data["gt"][:,2*i]+data["gt"][:,(2*i+1)]
        data["size"]=data["size"]//2
        data["gt"]=new_gt

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

def write_snp(site_data, options):
    """
    Write the .snp file - we only write 4 columns here,
    because we only use snp output (fastsimcoal -s0 option).
    Actually fastsimcoal can generate ACGT data, if we really
    wanted to.
    """
    snp_file=open(options["out"]+".snp", "w")
    for i, poss in enumerate(site_data):
        chrom=str(i+1)
        for pos in poss:
            snp_file.write( "\t"+chrom+"_"+str(pos)+"\t"+chrom+"\t0.0\t"+str(pos)+"\n")
    snp_file.close()

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

def write_ind(gt_data, options):
    """
    Write the .ind file.
    """
    ind_file=open(options["out"]+".ind", "w")

    for sample in sorted(gt_data):
        data=gt_data[sample]
        nsample=data["size"]
        sample_string=sample.replace(" ", "_")
        for i in range(1,nsample+1):
            ind_file.write("\t"+sample_string+"_"+str(i)+"\tU\t"+sample_string+"\n")
    ind_file.close()

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

def write_geno(gt_data, options):
    """
    Write the geno file.
    """
    geno_file=open(options["out"]+".geno", "w")

    pops=sorted(gt_data)

    joined_data=np.concatenate([gt_data[pop]["gt"] for pop in pops], axis=1)
    if options["phased"]:
        np.savetxt(geno_file, 1-joined_data, fmt="%d", delimiter="")
    else:
        np.savetxt(geno_file, 2-joined_data, fmt="%d", delimiter="")

    geno_file.close()

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

def main(options):
    """
    """
    site_data, gt_data=load_from_arp(options["arp"])

    if not options["phased"]:
       unphase(gt_data)

    write_snp(site_data, options)
    write_ind(gt_data, options)
    write_geno(gt_data, options)

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

if __name__=="__main__":
	options=parse_options()
	main(options)


#!/usr/bin/env python
'''
This Python module, originally forked from https://github.com/kiritbasu/Fake-Apache-Log-Generator,
aims to generate massive log events to Kafka. More concretelly, we will be
writing into two different topics.
'''
import time
import datetime
import pytz
import numpy
import random
import gzip
import zipfile
import sys
import argparse
from random import randrange
import os, ssl

# External imports:
# pip install faker
# pip install tzlocal
from faker import Faker
from tzlocal import get_localzone
local = get_localzone()

import threading, logging, time
# import multiprocessing

# Kafka Producer imports:
# pip install kafka
from kafka import KafkaProducer
from kafka.errors import KafkaError

import json

'''
To create the 'header' and 'ping' topics in Kafka:
./<kafka_path>/bin/kafka-server-start.sh config/server.propertieskafka-server-start.sh config/server.properties
./<kafka_path>/bin/kafka-topics.sh --create --zookeeper localhost:2181 --replication-factor 1 --partitions 1 --topic header
./<kafka_path>/bin/kafka-topics.sh --create --zookeeper localhost:2181 --replication-factor 1 --partitions 1 --topic ping

To check that the messages are reaching Kafka:
./<kafka_path>/bin/kafka-console-consumer.sh --zookeeper localhost:2181 --topic header

Note: We assume that you've already run the zookeeper server at localhost:2181. By default:
./<kafka_path>/bin/zookeeper-server-start.sh config/zookeeper.properties
'''

class Producer(threading.Thread):
    daemon = True

    def run(self):
        producer = KafkaProducer(bootstrap_servers='localhost:9092',
                value_serializer=lambda v: json.dumps(v).encode('utf-8'))

        while True:
            faker = Faker()

            timestr = time.strftime("%Y%m%d-%H%M%S")
            otime = datetime.datetime.now()

            response=["200","404","500","301"]
            verb=["GET","POST","DELETE","PUT"]
            resources=["/list","/wp-content","/wp-admin","/explore","/search/tag/list","/app/main/posts","/posts/posts/explore","/apps/cart.jsp?appID="]
            ualist=[faker.firefox, faker.chrome, faker.safari, faker.internet_explorer, faker.opera]

            # This field will be used for choosing between topics and, later,
            # for Spark to use them
            headers=["HEADER","PING"]

            ip = str(faker.ipv4())
            dt = str(otime.strftime('%d/%b/%Y:%H:%M:%S'))
            tz = str(datetime.datetime.now(local).strftime('%z'))
            vrb = str(numpy.random.choice(verb,p=[0.6,0.1,0.1,0.2]))

            #header_agg = random.choice(headers)
            header_agg = str(numpy.random.choice(headers,p=[0.3,0.7]))

            # uri = random.choice(resources)
            # if uri.find("apps")>0:
                # uri += `random.randint(1000,10000)`

        	# resp = numpy.random.choice(response,p=[0.9,0.04,0.02,0.04])
            resp = "200"
            byt = str(int(random.gauss(5000,50)))
            referer = str(faker.uri())
            useragent = str(numpy.random.choice(ualist,p=[0.5,0.3,0.1,0.05,0.05] )())

            # message = '%s - %s - - [%s %s] "%s %s HTTP/1.0" %s %s "%s" "%s"\n' % (header_agg,ip,dt,tz,vrb,uri,resp,byt,str(referer),str(useragent))
            message = '%s - %s - - [%s %s] "%s %s HTTP/1.0" %s "%s" "%s"\n' % (header_agg,ip,dt,tz,vrb,resp,byt,str(referer),str(useragent))

            message_json = {'header':header_agg, 'ip':ip, 'dt':dt, 'tz':tz, 'vrb':vrb, 'resp':resp, 'byt':byt, 'referer':referer, 'useragent':useragent}
            print(json.dumps(message_json))

            # testing generated JSON message (validated in https://jsonformatter.curiousconcept.com/#)
            # f = open('message_json', 'w')
            # f.write(json.dumps(message_json))

            if (header_agg=="HEADER"):
                producer.send('header', message_json)
            else:
                producer.send('ping', message_json)
            # time between each send action

            time.sleep(0.5)

def main():
    tasks = [
        Producer()
    ]

    for t in tasks:
        t.start()

    # how many times
    time.sleep(10)

if __name__ == "__main__":
    logging.basicConfig(
        format='%(asctime)s.%(msecs)s:%(name)s:%(thread)d:%(levelname)s:%(process)d:%(message)s',
        level=logging.INFO
        )
    main()


#!/usr/bin/python2.7

import numpy as np
from .grammar import PathGrammar
from .length_model import PoissonModel
import glob
import re

# Viterbi decoding
class Viterbi(object):

    ### helper structure ###
    class TracebackNode(object):
        def __init__(self, label, predecessor, boundary = False):
            self.label = label
            self.predecessor = predecessor
            self.boundary = boundary

    ### helper structure ###
    class HypDict(dict):
        class Hypothesis(object):
            def __init__(self, score, traceback):
                self.score = score
                self.traceback = traceback
        def update(self, key, score, traceback):
            if (not key in self) or (self[key].score <= score):
                self[key] = self.Hypothesis(score, traceback)

    # @grammar: the grammar to use, must inherit from class Grammar
    # @length_model: the length model to use, must inherit from class LengthModel
    # @frame_sampling: generate hypotheses every frame_sampling frames
    # @pruning_factor: value between 0 (no pruning) and 1 (out of all new segment start hypotheses prune all but the best hypothesis)
    # @max_segment_start_hyp: maximal number of new segment start hypotheses per time frame
    def __init__(self, grammar, length_model, frame_sampling = 1, pruning_factor = 0.0, max_segment_start_hyp = np.inf):
        self.grammar = grammar
        self.length_model = length_model
        self.frame_sampling = frame_sampling
        self.pruning_factor = pruning_factor
        self.max_segment_start_hyp = max_segment_start_hyp

    # Viterbi decoding of a sequence
    # @log_frame_probs: logarithmized frame probabilities
    #                   (usually log(network_output) - log(prior) - max_val, where max_val ensures negativity of all log scores)
    # @return: the score of the best sequence,
    #          the corresponding framewise labels (len(labels) = len(sequence))
    #          and the inferred segments in the form (label, length)
    def decode(self, scorer):
        # create initial hypotheses
        hyps = self.init_decoding()
        # decode each following time step
        for t in range(self.frame_sampling, scorer.length(), self.frame_sampling):
            hyps = self.decode_frame(t, hyps, scorer)
            scorer.clear()
        # transition to end symbol
        final_hyp = self.finalize_decoding(hyps, scorer)
        scorer.clear()
        labels, segments = self.traceback(final_hyp, scorer.length(), scorer)
        return final_hyp.score, labels, segments

    def prune(self, hyps, new_segment_keys):
        if self.pruning_factor > 0 and len(new_segment_keys) > 0:
            scores = sorted( [ hyps[key].score for key in new_segment_keys ] )
            reference_idx = min(int(self.pruning_factor * len(scores)), len(scores) - 1)
            reference_idx = max(reference_idx, len(scores) - self.max_segment_start_hyp)
            for key in new_segment_keys:
                if hyps[key].score < scores[reference_idx]:
                    del hyps[key]

    def init_decoding(self):
        hyps = self.HypDict()
        context = self.grammar.update_context((), self.grammar.start_symbol())
        for label in self.grammar.possible_successors(context):
            key = context + (label, self.frame_sampling)
            score = self.grammar.score(context, label)
            hyps.update(key, score, self.TracebackNode(label, None, boundary = True))
        return hyps

    def decode_frame(self, t, old_hyp, scorer):
        new_hyp = self.HypDict()
        new_segment_keys = set()
        for key, hyp in old_hyp.items():
            context, label, length = key[0:-2], key[-2], key[-1]
            # stay in the same label...
            if length + self.frame_sampling <= self.length_model.max_length():
                new_key = context + (label, length + self.frame_sampling)
                new_hyp.update(new_key, hyp.score, self.TracebackNode(label, hyp.traceback, boundary = False))
            # ... or go to the next label
            context = self.grammar.update_context(context, label)
            for new_label in self.grammar.possible_successors(context):
                if new_label == self.grammar.end_symbol():
                    continue
                new_key = context + (new_label, self.frame_sampling)
                score = hyp.score + scorer.get(t-length, t-1, label) + self.length_model.score(length, label) + self.grammar.score(context, new_label)
                new_hyp.update(new_key, score, self.TracebackNode(new_label, hyp.traceback, boundary = True))
                new_segment_keys.add(new_key)
        # prune among new hypothesized segments
        self.prune(new_hyp, new_segment_keys)
        # return new hypotheses
        return new_hyp

    def finalize_decoding(self, old_hyp, scorer):
        final_hyp = self.HypDict.Hypothesis(-np.inf, None)
        T = ((scorer.length() - 1) // self.frame_sampling) * self.frame_sampling + 1
        for key, hyp in old_hyp.items():
            context, label, length = key[0:-2], key[-2], key[-1]
            context = self.grammar.update_context(context, label)
            score = hyp.score + scorer.get(T-length, T-1, label) + self.length_model.score(length, label) + self.grammar.score(context, self.grammar.end_symbol())
            if score >= final_hyp.score:
                final_hyp.score, final_hyp.traceback = score, hyp.traceback
        # return final hypothesis
        return final_hyp

    def traceback(self, hyp, n_frames, scorer):
        class Segment(object):
            def __init__(self, label):
                self.label, self.length, self.score = label, 0, -np.inf
        traceback = hyp.traceback
        labels = []
        segments = [Segment(traceback.label)]
        while not traceback == None:
            segments[-1].length += self.frame_sampling
            labels += [traceback.label] * self.frame_sampling
            if traceback.boundary and not traceback.predecessor == None:
                segments.append( Segment(traceback.predecessor.label) )
            traceback = traceback.predecessor
        # the traceback node for the last frame contributes only length 1, not length self.frame_sampling!
        segments[0].length -= self.frame_sampling - 1
        labels = labels[:-self.frame_sampling + 1]
        # bring into forward order (traceback is a backward pass)
        labels, segments = list(reversed(labels)), list(reversed(segments))
        # score the segments
        offset = 0
        for s in segments:
            s.score = scorer.get(offset, offset + s.length - 1, s.label)
            offset += s.length
        # append labels/length for non processed frames (due to self.frame_sampling there might be a small unprocessed offset at the end)
        segments[-1].length += n_frames - len(labels)
        labels += [hyp.traceback.label] * (n_frames - len(labels))
        return labels, segments