# TensorFlow version of NIPS2016 soundnet
# Required package: librosa: A python package for music and audio analysis.
# $ pip install librosa

from ops import batch_norm, conv2d, relu, maxpool
from util import preprocess, load_from_list, load_audio
from model import Model
from glob import glob

import tensorflow as tf
import numpy as np
import argparse
import time
import sys
import os

# Make xrange compatible in both Python 2, 3
try:
    xrange
except NameError:
    xrange = range

local_config = {
            'batch_size': 1,
            'train_size': np.inf,
            'epoch': 200,
            'eps': 1e-5,
            'learning_rate': 1e-3,
            'beta1': 0.9,
            'load_size': 22050*4,
            'sample_rate': 22050,
            'name_scope': 'SoundNet',
            'phase': 'train',
            'dataset_name': 'ESC50',
            'subname': 'mp3',
            'checkpoint_dir': 'checkpoint',
            'dump_dir': 'output',
            'model_dir': None,
            'param_g_dir': './models/sound8.npy',
            }

class Model():
    def __init__(self, session, config=local_config, param_G=None):
        self.sess           = session
        self.config         = config
        self.param_G        = param_G
        self.g_step         = tf.Variable(0, trainable=False)
        self.counter        = 0
        self.model()

    def model(self):
        # Placeholder
        self.sound_input_placeholder = tf.placeholder(tf.float32,
                shape=[self.config['batch_size'], None, 1, 1]) # batch x h x w x channel
        self.object_dist = tf.placeholder(tf.float32,
                shape=[self.config['batch_size'], None, 1000]) # batch x h x w x channel
        self.scene_dist = tf.placeholder(tf.float32,
                shape=[self.config['batch_size'], None, 401]) # batch x h x w x channel

        # Generator
        self.add_generator(name_scope=self.config['name_scope'])

        # KL Divergence
        self.object_loss = self.KL_divergence(self.layers[25], self.object_dist, name_scope='KL_Div_object')
        self.scene_loss = self.KL_divergence(self.layers[26], self.scene_dist, name_scope='KL_Div_scene')
        self.loss = self.object_loss + self.scene_loss

        # Summary
        self.loss_sum = tf.summary.scalar("g_loss", self.loss)
        self.g_sum = tf.summary.merge([self.loss_sum])
        self.writer = tf.summary.FileWriter("./logs", self.sess.graph)

        # variable collection
        self.g_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                                    scope=self.config['name_scope'])

        self.saver = tf.train.Saver(keep_checkpoint_every_n_hours=12,
                                    max_to_keep=5,
                                    restore_sequentially=True)

        # Optimizer and summary
        self.g_optim = tf.train.AdamOptimizer(self.config['learning_rate'], beta1=self.config['beta1']) \
                          .minimize(self.loss, var_list=(self.g_vars), global_step=self.g_step)

        # Initialize
        init_op = tf.global_variables_initializer()
        self.sess.run(init_op)

        # Load checkpoint
        if self.load(self.config['checkpoint_dir']):
            print(" [*] Load SUCCESS")
        else:
            print(" [!] Load failed...")

    def add_generator(self, name_scope='SoundNet'):
        with tf.variable_scope(name_scope) as scope:
            self.layers = {}

            # Stream one: conv1 ~ conv7
            self.layers[1] = conv2d(self.sound_input_placeholder, 1, 16, k_h=64, d_h=2, p_h=32, name_scope='conv1')
            self.layers[2] = batch_norm(self.layers[1], 16, self.config['eps'], name_scope='conv1')
            self.layers[3] = relu(self.layers[2], name_scope='conv1')
            self.layers[4] = maxpool(self.layers[3], k_h=8, d_h=8, name_scope='conv1')

            self.layers[5] = conv2d(self.layers[4], 16, 32, k_h=32, d_h=2, p_h=16, name_scope='conv2')
            self.layers[6] = batch_norm(self.layers[5], 32, self.config['eps'], name_scope='conv2')
            self.layers[7] = relu(self.layers[6], name_scope='conv2')
            self.layers[8] = maxpool(self.layers[7], k_h=8, d_h=8, name_scope='conv2')

            self.layers[9] = conv2d(self.layers[8], 32, 64, k_h=16, d_h=2, p_h=8, name_scope='conv3')
            self.layers[10] = batch_norm(self.layers[9], 64, self.config['eps'], name_scope='conv3')
            self.layers[11] = relu(self.layers[10], name_scope='conv3')

            self.layers[12] = conv2d(self.layers[11], 64, 128, k_h=8, d_h=2, p_h=4, name_scope='conv4')
            self.layers[13] = batch_norm(self.layers[12], 128, self.config['eps'], name_scope='conv4')
            self.layers[14] = relu(self.layers[13], name_scope='conv4')

            self.layers[15] = conv2d(self.layers[14], 128, 256, k_h=4, d_h=2, p_h=2, name_scope='conv5')
            self.layers[16] = batch_norm(self.layers[15], 256, self.config['eps'], name_scope='conv5')
            self.layers[17] = relu(self.layers[16], name_scope='conv5')
            self.layers[18] = maxpool(self.layers[17], k_h=4, d_h=4, name_scope='conv5')

            self.layers[19] = conv2d(self.layers[18], 256, 512, k_h=4, d_h=2, p_h=2, name_scope='conv6')
            self.layers[20] = batch_norm(self.layers[19], 512, self.config['eps'], name_scope='conv6')
            self.layers[21] = relu(self.layers[20], name_scope='conv6')

            self.layers[22] = conv2d(self.layers[21], 512, 1024, k_h=4, d_h=2, p_h=2, name_scope='conv7')
            self.layers[23] = batch_norm(self.layers[22], 1024, self.config['eps'], name_scope='conv7')
            self.layers[24] = relu(self.layers[23], name_scope='conv7')

            # Split one: conv8, conv8_2
            # NOTE: here we use a padding of 2 to skip an unknown error
            # https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/common_shape_fns.cc#L45
            self.layers[25] = conv2d(self.layers[24], 1024, 1000, k_h=8, d_h=2, p_h=2, name_scope='conv8')
            self.layers[26] = conv2d(self.layers[24], 1024, 401, k_h=8, d_h=2, p_h=2, name_scope='conv8_2')

    def train(self):
        """Train SoundNet"""

        start_time = time.time()

        # Data info
        data = glob('./data/ESC-50-master/audio/*.{}'.format(self.config['subname']))
        batch_idxs = min(len(data), self.config['train_size']) // self.config['batch_size']
        for epoch in xrange(self.counter//batch_idxs, self.config['epoch']):

            for idx in xrange(self.counter%batch_idxs, batch_idxs):

                # By default, librosa will resample the signal to 22050Hz. And range in (-1., 1.)
                sound_sample = load_from_list(data[idx*self.config['batch_size']:(idx+1)*self.config['batch_size']], self.config)

                # Update G network
                # NOTE: Here we still use dummy random distribution for scene and objects
                _, summary_str, l_scn, l_obj = self.sess.run([self.g_optim, self.g_sum, self.scene_loss, self.object_loss],
                    feed_dict={self.sound_input_placeholder: sound_sample, \
                            self.scene_dist: np.random.randint(2, size=(1, 1, 401)), \
                            self.object_dist: np.random.randint(2, size=(1, 1, 1000))})
                self.writer.add_summary(summary_str, self.counter)

                print ("[Epoch {}] {}/{} | Time: {} | scene_loss: {} | obj_loss: {}".format(epoch, idx, batch_idxs, time.time() - start_time, l_scn, l_obj))

                if np.mod(self.counter, 1000) == 1000 - 1:
                    self.save(self.config['checkpoint_dir'], self.counter)

                self.counter += 1

    #########################
    #          Loss         #
    #########################
    # Adapt the answer here: http://stackoverflow.com/questions/41863814/kl-divergence-in-tensorflow
    def KL_divergence(self, dist_a, dist_b, name_scope='KL_Div'):
        return tf.reduce_mean(-tf.nn.softmax_cross_entropy_with_logits(logits=dist_a, labels=dist_b))

    #########################
    #       Save/Load       #
    #########################
    @property
    def get_model_dir(self):
        if self.config['model_dir'] is None:
            return "{}_{}".format(
                self.config['dataset_name'], self.config['batch_size'])
        else:
            return self.config['model_dir']

    def load(self, ckpt_dir='checkpoint'):
        return self.load_from_ckpt(ckpt_dir) if self.param_G is None \
        else self.load_from_npy()

    def save(self, checkpoint_dir, step):
        """ Checkpoint saver """
        model_name = "SoundNet.model"
        checkpoint_dir = os.path.join(checkpoint_dir, self.get_model_dir)

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

        self.saver.save(self.sess,
                        os.path.join(checkpoint_dir, model_name),
                        global_step=step)

    def load_from_ckpt(self, checkpoint_dir='checkpoint'):
        """ Checkpoint loader """
        print(" [*] Reading checkpoints...")

        checkpoint_dir = os.path.join(checkpoint_dir, self.get_model_dir)

        ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
        if ckpt and ckpt.model_checkpoint_path:
            ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
            self.saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
            print(" [*] Success to read {}".format(ckpt_name))
            self.counter = int(ckpt_name.rsplit('-', 1)[-1])
            print(" [*] Start counter from {}".format(self.counter))
            return True
        else:
            print(" [*] Failed to find a checkpoint under {}".format(checkpoint_dir))
            return False

    def load_from_npy(self):
        if self.param_G is None: return False
        data_dict = self.param_G
        for key in data_dict:
            with tf.variable_scope(self.config['name_scope'] + '/'+ key, reuse=True):
                for subkey in data_dict[key]:
                    try:
                        var = tf.get_variable(subkey)
                        self.sess.run(var.assign(data_dict[key][subkey]))
                        print('Assign pretrain model {} to {}'.format(subkey, key))
                    except:
                        print('Ignore {}'.format(key))

        self.param_G.clear()
        return True

def main():

    args = parse_args()
    local_config['phase'] = args.phase

    # Setup visible device
    os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda_device

    # Make path
    if not os.path.exists(args.outpath):
        os.mkdir(args.outpath)

    # Load pre-trained model
    param_G = np.load(local_config['param_g_dir'], encoding='latin1').item() \
            if args.phase in ['finetune', 'extract'] \
            else None

    # Init. Session
    sess_config = tf.ConfigProto()
    sess_config.allow_soft_placement=True
    sess_config.gpu_options.allow_growth = True

    with tf.Session(config=sess_config) as session:
        # Build model
        model = Model(session, config=local_config, param_G=param_G)

        if args.phase in ['train', 'finetune']:
            # Training phase
            model.train()
        elif args.phase == 'extract':
            # import when we need
            from extract_feat import extract_feat

            # Feature extractor
            #sound_sample = np.reshape(np.load('./data/demo.npy', encoding='latin1'), [local_config['batch_size'], -1, 1, 1])

            import librosa
            audio_path = './data/demo.mp3'
            sound_sample, _ = load_audio(audio_path)
            sound_sample = preprocess(sound_sample, config=local_config)

            output = extract_feat(model, sound_sample, args)

def parse_args():
    """ Parse input arguments """
    parser = argparse.ArgumentParser(description='SoundNet')

    parser.add_argument('-o', '--outpath', dest='outpath', help='output feature path. e.g., [output]', default='output')

    parser.add_argument('-p', '--phase', dest='phase', help='demo or extract feature. e.g., [train, finetune, extract]', default='finetune')

    parser.add_argument('-m', '--layer', dest='layer_min', help='start from which feature layer. e.g., [1]', type=int, default=1)

    parser.add_argument('-x', dest='layer_max', help='end at which feature layer. e.g., [24]', type=int, default=None)

    parser.add_argument('-c', '--cuda', dest='cuda_device', help='which cuda device to use. e.g., [0]', default='0')

    feature_parser = parser.add_mutually_exclusive_group(required=False)
    feature_parser.add_argument('-s', '--save', dest='is_save', help='Turn on save mode. [False(default), True]', action='store_true')
    parser.set_defaults(is_save=False)

    args = parser.parse_args()

    return args

if __name__ == '__main__':
    main()


#!/usr/bin/env python
import numpy as np
np.random.seed(7)

import collections
import pandas as pd
import matplotlib.pyplot as plt
import time
import click

from gensim.models.keyedvectors import FastTextKeyedVectors
from sklearn.preprocessing import OneHotEncoder
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, roc_curve, auc, f1_score, precision_score, recall_score

from keras import optimizers, metrics
from keras.layers import LSTM, Dense, Dropout, Input, Activation, LeakyReLU, Dot, concatenate
from keras.models import Model
from keras.callbacks import History, ReduceLROnPlateau, EarlyStopping
import keras.backend as K
from keras.engine.topology import Layer

import warnings
warnings.simplefilter('ignore')

properties_names = {"":"None",
"http://www.wikidata.org/prop/direct/P1001":"juridiction concernée",
"http://www.wikidata.org/prop/direct/P101":"domaine d'activité",
"http://www.wikidata.org/prop/direct/P1018":"institution de normalisation",
"http://www.wikidata.org/prop/direct/P102":"parti politique",
"http://www.wikidata.org/prop/direct/P103":"langue maternelle",
"http://www.wikidata.org/prop/direct/P1038":"parentèle",
"http://www.wikidata.org/prop/direct/P1040":"monteur",
"http://www.wikidata.org/prop/direct/P105":"rang taxinomique",
"http://www.wikidata.org/prop/direct/P106":"occupation",
"http://www.wikidata.org/prop/direct/P1066":"élève de",
"http://www.wikidata.org/prop/direct/P1071":"lieu de fabrication",
"http://www.wikidata.org/prop/direct/P108":"employeur",
"http://www.wikidata.org/prop/direct/P110":"illustrateur",
"http://www.wikidata.org/prop/direct/P112":"fondateur",
"http://www.wikidata.org/prop/direct/P1142":"idéologie politique",
"http://www.wikidata.org/prop/direct/P118":"ligue",
"http://www.wikidata.org/prop/direct/P119":"lieu de sépulture",
"http://www.wikidata.org/prop/direct/P1269":"aspect de",
"http://www.wikidata.org/prop/direct/P127":"propriétaire",
"http://www.wikidata.org/prop/direct/P1302":"principales localités desservies",
"http://www.wikidata.org/prop/direct/P1303":"instrument de musique pratiqué",
"http://www.wikidata.org/prop/direct/P131":"localisation administrative",
"http://www.wikidata.org/prop/direct/P1321":"lieu d'origine (Suisse)",
"http://www.wikidata.org/prop/direct/P1327":"partenaire professionnel ou sportif",
"http://www.wikidata.org/prop/direct/P1336":"territoire revendiqué par",
"http://www.wikidata.org/prop/direct/P1344":"participant à",
"http://www.wikidata.org/prop/direct/P1346":"vainqueur",
"http://www.wikidata.org/prop/direct/P135":"mouvement",
"http://www.wikidata.org/prop/direct/P136":"genre",
"http://www.wikidata.org/prop/direct/P1365":"remplace",
"http://www.wikidata.org/prop/direct/P1366":"remplacé par",
"http://www.wikidata.org/prop/direct/P137":"opérateur",
"http://www.wikidata.org/prop/direct/P1376":"capitale de",
"http://www.wikidata.org/prop/direct/P138":"nommé en référence à",
"http://www.wikidata.org/prop/direct/P1382":"coïncide avec",
"http://www.wikidata.org/prop/direct/P1389":"certificat de produit",
"http://www.wikidata.org/prop/direct/P140":"religion",
"http://www.wikidata.org/prop/direct/P1408":"ville de licence",
"http://www.wikidata.org/prop/direct/P1412":"langues parlées, écrites ou signées",
"http://www.wikidata.org/prop/direct/P1431":"producteur délégué",
"http://www.wikidata.org/prop/direct/P1435":"statut patrimonial",
"http://www.wikidata.org/prop/direct/P144":"basé sur",
"http://www.wikidata.org/prop/direct/P1441":"présent dans l'œuvre",
"http://www.wikidata.org/prop/direct/P150":"contient les subdivisions territoriales administratives",
"http://www.wikidata.org/prop/direct/P1532":"nationalité sportive",
"http://www.wikidata.org/prop/direct/P1542":"a pour effet",
"http://www.wikidata.org/prop/direct/P155":"précédé par",
"http://www.wikidata.org/prop/direct/P1552":"caractérisé par",
"http://www.wikidata.org/prop/direct/P156":"suivi par",
"http://www.wikidata.org/prop/direct/P1589":"point le plus bas",
"http://www.wikidata.org/prop/direct/P159":"siège social",
"http://www.wikidata.org/prop/direct/P161":"distribution",
"http://www.wikidata.org/prop/direct/P162":"producteur",
"http://www.wikidata.org/prop/direct/P166":"distinction reçue",
"http://www.wikidata.org/prop/direct/P169":"directeur général",
"http://www.wikidata.org/prop/direct/P17":"pays",
"http://www.wikidata.org/prop/direct/P170":"créateur",
"http://www.wikidata.org/prop/direct/P171":"taxon supérieur",
"http://www.wikidata.org/prop/direct/P172":"groupe ethnique",
"http://www.wikidata.org/prop/direct/P175":"interprète",
"http://www.wikidata.org/prop/direct/P176":"fabricant",
"http://www.wikidata.org/prop/direct/P178":"développeur",
"http://www.wikidata.org/prop/direct/P179":"dans la série",
"http://www.wikidata.org/prop/direct/P180":"dépeint",
"http://www.wikidata.org/prop/direct/P1830":"propriétaire de",
"http://www.wikidata.org/prop/direct/P184":"directeur de thèse",
"http://www.wikidata.org/prop/direct/P185":"doctorant",
"http://www.wikidata.org/prop/direct/P186":"matériau",
"http://www.wikidata.org/prop/direct/P1877":"d'après une œuvre de",
"http://www.wikidata.org/prop/direct/P1889":"à ne pas confondre avec",
"http://www.wikidata.org/prop/direct/P19":"lieu de naissance",
"http://www.wikidata.org/prop/direct/P190":"jumelage ou partenariat",
"http://www.wikidata.org/prop/direct/P197":"gare voisine",
"http://www.wikidata.org/prop/direct/P1995":"spécialité médicale",
"http://www.wikidata.org/prop/direct/P20":"lieu de mort",
"http://www.wikidata.org/prop/direct/P205":"pays souverain sur le bassin versant",
"http://www.wikidata.org/prop/direct/P206":"baigné par",
"http://www.wikidata.org/prop/direct/P22":"père",
"http://www.wikidata.org/prop/direct/P2341":"autochtone de",
"http://www.wikidata.org/prop/direct/P2389":"organisation dont la fonction est à la tête",
"http://www.wikidata.org/prop/direct/P2416":"discipline sportive pratiquée en compétition",
"http://www.wikidata.org/prop/direct/P25":"mère",
"http://www.wikidata.org/prop/direct/P2541":"zone opérationnelle",
"http://www.wikidata.org/prop/direct/P26":"conjoint",
"http://www.wikidata.org/prop/direct/P2632":"lieu de détention",
"http://www.wikidata.org/prop/direct/P264":"label discographique",
"http://www.wikidata.org/prop/direct/P2650":"intéressé par",
"http://www.wikidata.org/prop/direct/P27":"pays de nationalité",
"http://www.wikidata.org/prop/direct/P272":"société de production",
"http://www.wikidata.org/prop/direct/P276":"lieu",
"http://www.wikidata.org/prop/direct/P2789":"se connecte avec",
"http://www.wikidata.org/prop/direct/P279":"sous-classe de",
"http://www.wikidata.org/prop/direct/P282":"système d'écriture",
"http://www.wikidata.org/prop/direct/P2925":"domaine du saint ou de la divinité",
"http://www.wikidata.org/prop/direct/P2936":"langue utilisée",
"http://www.wikidata.org/prop/direct/P30":"continent",
"http://www.wikidata.org/prop/direct/P3075":"religion officielle",
"http://www.wikidata.org/prop/direct/P3094":"se développe de",
"http://www.wikidata.org/prop/direct/P3095":"pratiqué par",
"http://www.wikidata.org/prop/direct/P31":"nature de l'élément",
"http://www.wikidata.org/prop/direct/P3373":"frère ou sœur",
"http://www.wikidata.org/prop/direct/P344":"directeur de la photographie",
"http://www.wikidata.org/prop/direct/P35":"chef d'État",
"http://www.wikidata.org/prop/direct/P355":"organisation filiale",
"http://www.wikidata.org/prop/direct/P36":"capitale",
"http://www.wikidata.org/prop/direct/P361":"partie de",
"http://www.wikidata.org/prop/direct/P366":"usage",
"http://www.wikidata.org/prop/direct/P37":"langue officielle",
"http://www.wikidata.org/prop/direct/P371":"présentateur",
"http://www.wikidata.org/prop/direct/P3729":"rang immédiatement inférieur",
"http://www.wikidata.org/prop/direct/P3730":"rang immédiatement supérieur",
"http://www.wikidata.org/prop/direct/P375":"lanceur",
"http://www.wikidata.org/prop/direct/P376":"sur le corps astronomique",
"http://www.wikidata.org/prop/direct/P38":"monnaie",
"http://www.wikidata.org/prop/direct/P3828":"élément de parure",
"http://www.wikidata.org/prop/direct/P39":"fonction",
"http://www.wikidata.org/prop/direct/P40":"enfant",
"http://www.wikidata.org/prop/direct/P403":"se jette dans",
"http://www.wikidata.org/prop/direct/P407":"langue de l'œuvre, du nom ou du terme",
"http://www.wikidata.org/prop/direct/P410":"grade militaire",
"http://www.wikidata.org/prop/direct/P412":"tessiture",
"http://www.wikidata.org/prop/direct/P413":"position de jeu/spécialité",
"http://www.wikidata.org/prop/direct/P449":"diffuseur original",
"http://www.wikidata.org/prop/direct/P451":"partenaire",
"http://www.wikidata.org/prop/direct/P457":"texte fondateur",
"http://www.wikidata.org/prop/direct/P460":"réputé identique à",
"http://www.wikidata.org/prop/direct/P461":"contraire",
"http://www.wikidata.org/prop/direct/P463":"membre de",
"http://www.wikidata.org/prop/direct/P495":"pays d’origine",
"http://www.wikidata.org/prop/direct/P4969":"œuvre dérivée",
"http://www.wikidata.org/prop/direct/P50":"auteur",
"http://www.wikidata.org/prop/direct/P501":"enclavé dans",
"http://www.wikidata.org/prop/direct/P520":"armement",
"http://www.wikidata.org/prop/direct/P527":"comprend",
"http://www.wikidata.org/prop/direct/P53":"famille",
"http://www.wikidata.org/prop/direct/P533":"cible",
"http://www.wikidata.org/prop/direct/P54":"membre de l'équipe de sport",
"http://www.wikidata.org/prop/direct/P551":"résidence",
"http://www.wikidata.org/prop/direct/P559":"extrémité",
"http://www.wikidata.org/prop/direct/P562":"banque émettrice",
"http://www.wikidata.org/prop/direct/P57":"réalisateur ou metteur en scène",
"http://www.wikidata.org/prop/direct/P58":"scénariste",
"http://www.wikidata.org/prop/direct/P6":"chef de l'exécutif",
"http://www.wikidata.org/prop/direct/P607":"conflit",
"http://www.wikidata.org/prop/direct/P61":"découvreur ou inventeur",
"http://www.wikidata.org/prop/direct/P610":"point culminant",
"http://www.wikidata.org/prop/direct/P611":"ordre religieux",
"http://www.wikidata.org/prop/direct/P641":"sport",
"http://www.wikidata.org/prop/direct/P647":"repêchage par",
"http://www.wikidata.org/prop/direct/P664":"organisateur",
"http://www.wikidata.org/prop/direct/P674":"personnages",
"http://www.wikidata.org/prop/direct/P676":"parolier",
"http://www.wikidata.org/prop/direct/P681":"composant cellulaire",
"http://www.wikidata.org/prop/direct/P6885":"région historique",
"http://www.wikidata.org/prop/direct/P69":"scolarité",
"http://www.wikidata.org/prop/direct/P706":"localisation géographique",
"http://www.wikidata.org/prop/direct/P710":"participant",
"http://www.wikidata.org/prop/direct/P737":"influencé par",
"http://www.wikidata.org/prop/direct/P740":"lieu de fondation",
"http://www.wikidata.org/prop/direct/P749":"organisation mère",
"http://www.wikidata.org/prop/direct/P750":"distributeur",
"http://www.wikidata.org/prop/direct/P793":"événement clé",
"http://www.wikidata.org/prop/direct/P800":"œuvre notable",
"http://www.wikidata.org/prop/direct/P802":"élève",
"http://www.wikidata.org/prop/direct/P807":"séparé de",
"http://www.wikidata.org/prop/direct/P828":"a pour cause",
"http://www.wikidata.org/prop/direct/P837":"jour de l'année",
"http://www.wikidata.org/prop/direct/P84":"architecte",
"http://www.wikidata.org/prop/direct/P840":"lieu de l'action",
"http://www.wikidata.org/prop/direct/P859":"sponsor",
"http://www.wikidata.org/prop/direct/P86":"compositeur",
"http://www.wikidata.org/prop/direct/P87":"librettiste",
"http://www.wikidata.org/prop/direct/P88":"commanditaire",
"http://www.wikidata.org/prop/direct/P921":"sujet ou thème principal",
"http://www.wikidata.org/prop/direct/P931":"localité desservie par cette infrastructure de transport",
"http://www.wikidata.org/prop/direct/P937":"lieu de travail",
"http://www.wikidata.org/prop/direct/P945":"allégeance",
"http://www.wikidata.org/prop/direct/P974":"affluent",
"http://www.wikidata.org/prop/direct/P98":"éditeur scientifique"}

###############################################################################################################
#Format data to the right input

def flatten(coll):
    for i in coll:
        if isinstance(i, collections.Iterable) and not isinstance(i, str):
            for subc in flatten(i):
                yield subc
        else:
            yield i

def get_lookup_tables(df):

    #for relation number
    relations = set(df.relation.tolist())
    relations = list(relations)
    global num_classes
    num_classes = len(relations)

        ##relation to number :
    label_encoder = LabelEncoder()
    integer_encoded = label_encoder.fit_transform(relations)

    global relation_code_to_name
    relation_code_to_name = dict(zip(integer_encoded, relations))

        ##number to one-hot-encoded vector
    integer_encoded = integer_encoded.reshape(len(integer_encoded), 1)
    onehot_encoder = OneHotEncoder(sparse=False)
    onehot_encoded = onehot_encoder.fit_transform(integer_encoded)

        ##relation name to encoding
    lookup_table_relations = dict(zip(relations, onehot_encoded))

    #for dependency tags
    deptags = ['', 'parataxis', 'ROOT', 'nsubj',
    'nsubjpass', 'nmod', 'dep', 'cop',
    'ccomp', 'acl', 'name', 'det', 'nummod', 'conj', 'expl',
     'aux', 'amod', 'compound', 'auxpass', 'advmod', 'cc', 'mwe', 'xcomp',
    'advcl', 'mark', 'punct',  'acl:relcl', 'case', 'csubj',
     'iobj', 'neg', 'nmod:poss', 'dobj', 'appos',  'discourse']
    deptags_codes = list(range(1, len(deptags)+1))
    lookup_table_deptags = dict(zip(deptags, deptags_codes))

    #for part-of-speech tags
    postags = ['', 'SCONJ', 'NUM', 'CONJ', 'NOUN', 'VERB', 'PRON', 'ADP', 'DET', 'ADV', 'PUNCT',
               'INTJ', 'SYM', 'PROPN', 'PART', 'AUX', 'X', 'ADJ']
    postags_codes = list(range(1, len(postags)+1))
    lookup_table_postags = dict(zip(postags, postags_codes))

    #for the words, they are 1-hot encoded too. This time taken straigth from the features, no list
    vocab = set()
    col1 = df.word_path1.apply(lambda x : x.split("/"))
    batchvocab = set(list(flatten(col1)))
    vocab = batchvocab | vocab

    col2 = df.word_path2.apply(lambda x : x.split("/"))
    batchvocab = set(list(flatten(col2)))
    vocab = batchvocab | vocab

    vocab = list(vocab)+["UKNWN"]
    vocab_codes = list(range(1, len(vocab)+1))

    global vocab_size
    vocab_size = len(vocab)

    lookup_table_vocab = dict(zip(vocab, vocab_codes))

    return(lookup_table_relations,lookup_table_deptags, lookup_table_postags, lookup_table_vocab)

def data_reshape(df, lookup_table_relations, lookup_table_deptags, lookup_table_postags, lookup_table_vocab, maxlen=30):
    y = []
    features_words1 = []
    features_words2 = []
    features_pos1 = []
    features_pos2 = []
    features_dep1 = []
    features_dep2 = []

    for row in df.itertuples(): #ie for each sample
        y.append(row.relation!='')
        seq_word1 =  [lookup_table_vocab[word] if ((type(word)==str) & (word in lookup_table_vocab)) else lookup_table_vocab["UKNWN"] for word in row.word_path1.split("/")]
        seq_word2 =  [lookup_table_vocab[word] if ((type(word)==str) & (word in lookup_table_vocab)) else lookup_table_vocab["UKNWN"] for word in row.word_path2.split("/")]
        seq_dep1 =  [lookup_table_deptags[dep.split(':')[0]] if "conj" in dep else lookup_table_deptags[dep] for dep in row.dep_path1.split("/")]
        seq_dep2 =  [lookup_table_deptags[dep.split(':')[0]] if "conj" in dep else lookup_table_deptags[dep] for dep in row.dep_path2.split("/")]
        seq_pos1 =  [lookup_table_postags[pos_tag] for pos_tag in row.pos_path1.split("/")]
        seq_pos2 =  [lookup_table_postags[pos_tag] for pos_tag in row.pos_path2.split("/")]
        features_words1.append(seq_word1)
        features_words2.append(seq_word2)
        features_dep1.append(seq_dep1)
        features_dep2.append(seq_dep2)
        features_pos1.append(seq_pos1)
        features_pos2.append(seq_pos2)
    padded_features_words1 = pad_sequences(features_words1, padding='post', maxlen=maxlen)
    padded_features_words2 = pad_sequences(features_words2, padding='post', maxlen=maxlen)
    padded_features_dep1 = pad_sequences(features_dep1, padding='post', maxlen=maxlen)
    padded_features_dep2 = pad_sequences(features_dep2, padding='post', maxlen=maxlen)
    padded_features_pos1 = pad_sequences(features_pos1, padding='post', maxlen=maxlen)
    padded_features_pos2 = pad_sequences(features_pos2, padding='post', maxlen=maxlen)
    X = [padded_features_words1, padded_features_words2, padded_features_dep1, padded_features_dep2,
         padded_features_pos1, padded_features_pos2]
    return(X, y)

def sets_test_train(feats_file = "/home/cyrielle/Codes/clean_code/Features/data/features_ratio1_definitve_XuBiLSTM.csv"):
    df = pd.read_csv(feats_file)
    df.fillna(value="", inplace=True)
    df.set_index(["entity1","entity2", "original_sentence"], inplace=True)

    with open('indexes_test.json') as json_file:
        ind_test = json.load(json_file)
        df_test = df[df.id_entities_pair.isin(ind_test)]

    with open('indexes_train.json') as json_file:
        ind_train = json.load(json_file)
        df_train = df[df.id_entities_pair.isin(ind_train)]

    #train
    lookup_table_relations2, lookup_table_deptags2, lookup_table_postags2, lookup_table_vocab2 = get_lookup_tables(df_train)

    dataX, datay = data_reshape(df_train, lookup_table_relations2, lookup_table_deptags2, lookup_table_postags2, lookup_table_vocab2)
    index_train, index_val = train_test_split(range(0, len(datay)), test_size=0.25, random_state=7)
    X_train = [[dataXj[i] for i in index_train] for dataXj in dataX]
    X_val = [[dataXj[i] for i in index_val] for dataXj in dataX]#[dataXi[index_val] for dataXi in dataX]
    y_train = [datay[i] for i in index_train]
    y_train = np.array(y_train)
    y_val = [datay[i] for i in index_val]
    y_val = np.array(y_val)

    #test
    X_test, y_test = data_reshape(df_test, lookup_table_relations2, lookup_table_deptags2, lookup_table_postags2, lookup_table_vocab2)

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

###############################################################################################################
#Define and train models

class ZeroMaskedEntries(Layer):
    """
    This layer is called after an Embedding layer.
    It zeros out all of the masked-out embeddings.
    It also swallows the mask without passing it on.
    You can change this to default pass-on behavior as follows:

    def compute_mask(self, x, mask=None):
        if not self.mask_zero:
            return None
        else:
            return K.not_equal(x, 0)

    author : sergeyf
    """

    def __init__(self, **kwargs):
        self.support_mask = True
        super(ZeroMaskedEntries, self).__init__(**kwargs)

    def build(self, input_shape):
        self.output_dim = input_shape[1]
        self.repeat_dim = input_shape[2]

    def call(self, x, mask=None):
        mask = K.cast(mask, 'float32')
        mask = K.repeat(mask, self.repeat_dim)
        mask = K.permute_dimensions(mask, (0, 2, 1))
        return x * mask

    def compute_mask(self, input_shape, input_mask=None):
        return None

def build_model_xu(X_train):
    word_embedding_dim = 200
    pos_embedding_dim = 50
    dep_embedding_dim = 50

    word_vocab_size = vocab_size+1
    pos_vocab_size = 18+1
    dep_vocab_size = 35+1

    hidden_layer_size =100
    word_state_size = 200
    other_state_size =50

    #inputs
    input_word1 = Input(shape = (None,))
    input_word2 = Input(shape = (None,))
    input_dep1 = Input(shape =  (None,))
    input_dep2 = Input(shape =  (None,))
    input_pos1 = Input(shape = (None,))
    input_pos2 = Input(shape = (None,))

    #embeddings
    embed_words = Embedding(word_vocab_size, word_embedding_dim, mask_zero=True)
    embed_pos = Embedding(pos_vocab_size, pos_embedding_dim, mask_zero=True)
    embed_dep = Embedding(dep_vocab_size, dep_embedding_dim, mask_zero=True)

    #embed the inputs of both sides with the same embedding layers
    embedded_words1 = embed_words(input_word1)
    embedded_words2 = embed_words(input_word2)
    embedded_pos1 = embed_pos(input_pos1)
    embedded_pos2 = embed_pos(input_pos2)
    embedded_dep1 = embed_dep(input_dep1)
    embedded_dep2 = embed_dep(input_dep2)

    #lstms for each channel
    lstm_word1 = LSTM(units = word_state_size, return_sequences=True)(embedded_words1)
    lstm_word1_nomask = ZeroMaskedEntries()(lstm_word1)
    pool_word1 = GlobalMaxPooling1D()(lstm_word1_nomask) #shape = num samples, num_features

    lstm_pos1 = LSTM(units = other_state_size, return_sequences=True)(embedded_pos1)
    lstm_pos1_nomask = ZeroMaskedEntries()(lstm_pos1)
    pool_pos1 = GlobalMaxPooling1D()(lstm_pos1_nomask)

    lstm_dep1 = LSTM(units = other_state_size, return_sequences=True)(embedded_dep1)
    lstm_dep1_nomask = ZeroMaskedEntries()(lstm_dep1)
    pool_dep1 = GlobalMaxPooling1D()(lstm_dep1_nomask)

    lstm_word2 = LSTM(units = word_state_size, return_sequences=True)(embedded_words2)
    lstm_word2_nomask = ZeroMaskedEntries()(lstm_word2)
    pool_word2 = GlobalMaxPooling1D()(lstm_word2_nomask)

    lstm_pos2 = LSTM(units = other_state_size, return_sequences=True)(embedded_pos2)
    lstm_pos2_nomask = ZeroMaskedEntries()(lstm_pos2)
    pool_pos2 = GlobalMaxPooling1D()(lstm_pos2_nomask)

    lstm_dep2 = LSTM(units = other_state_size, return_sequences=True)(embedded_dep2)
    lstm_dep2_nomask = ZeroMaskedEntries()(lstm_dep2)
    pool_dep2 = GlobalMaxPooling1D()(lstm_dep2_nomask)

    hidden_word_pre = Dot(axes =-1, normalize = True)([pool_word1, pool_word2]) #dot product of the two tensors
    hidden_word = Dense(word_state_size, name = "Dense_word")(hidden_word_pre)

    hidden_pos_pre = Dot(axes =-1, normalize = True)([pool_pos1, pool_pos2])
    hidden_pos = Dense(other_state_size, name= "Dense_pos")(hidden_pos_pre)

    hidden_dep_pre = Dot(axes =-1, normalize = True)([pool_dep1, pool_dep2])
    hidden_dep = Dense(other_state_size, name = "Dense_dep")(hidden_dep_pre)

    main_hidden_pre = concatenate([hidden_word, hidden_pos, hidden_dep])
    main_hidden = Dense(hidden_layer_size, name = "Dense_main")(main_hidden_pre)

    output = Dense(num_classes, activation='softmax')(main_hidden)

    #define model
    model = Model(inputs=[input_word1, input_word2, input_dep1, input_dep2, input_pos1,input_pos2], outputs=output)
    adam = optimizers.Adam(lr=0.001)
    model.compile(loss= "categorical_crossentropy",metrics = [metrics.mean_squared_error], optimizer = adam)
    return(model)

def train_model_xu(model, X_train, X_val, y_train, y_val, result_folder):
    #callbacks
    history = History()
    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5, min_lr=0.001)
    callback_list = [history, reduce_lr]

    #fit model
    model.fit(x= [X_train[0],X_train[1],X_train[2],X_train[3],X_train[4], X_train[5]],
          y= y_train,
          epochs=20,
          batch_size=1,
          validation_data=([X_val[0],X_val[1],X_val[2],X_val[3],X_val[4], X_val[5]], y_val),
          callbacks=callback_list)

    #save or not the model
    if result_folder is not None :
        checkpoint_filepath= result_folder +"model_{}.model".format(time.strftime("%Y-%m-%d-%H-%M-%S"))
        model.save(checkpoint_filepath)
    return(model)

###############################################################################################################
#Test and results

def save_classif_reports(y_test, y_pred,report_file= None, print=True):
    if report_file is not None :
        with open(report_file, 'w') as f :
            f.write(classification_report(y_test, y_pred))
            f.write("F1_score macro: "+str(f1_score(y_test, y_pred, labels=None, pos_label=1, average='macro')))
            f.write("F1_score weigthed: "+str(f1_score(y_test, y_pred, labels=None, pos_label=1, average='weighted')))
            f.write("Precision score macro: "+str(precision_score(y_test, y_pred, average="macro")))
            f.write("Precision score weighted: "+str(precision_score(y_test, y_pred, average="weighted")))
            f.write("Recall score macro: "+str(recall_score(y_test, y_pred, average="macro")))
            f.write("Recall score weighted: "+str(recall_score(y_test, y_pred, average="weighted")))
    if print == True :
        print("F1_score macro: "+str(f1_score(y_test, y_pred, labels=None, pos_label=1, average='macro')))
        print("F1_score weigthed: "+str(f1_score(y_test, y_pred, labels=None, pos_label=1, average='weighted')))
        print("Precision score macro: "+str(precision_score(y_test, y_pred, average="macro")))
        print("Precision score weighted: "+str(precision_score(y_test, y_pred, average="weighted")))
        print("Recall score macro: "+str(recall_score(y_test, y_pred, average="macro")))
        print("Recall score weighted: "+str(recall_score(y_test, y_pred, average="weighted")))

def test_model(trained_model, X_test, y_test, result_folder):
    #predict the Y with model
    y_pred = trained_model.predict(X_test)
    y_pred_classes = y_pred.argmax(axis=-1)
    predictions = [properties_names[relation_code_to_name[i]] for i in y_pred_classes]

    #turn y_test back into variable
    y_true_classes = y_test.argmax(axis=-1)
    reel = [properties_names[relation_code_to_name[i]] for i in y_true_classes]
    if result_folder is not None :
        report_path = result_folder +"classif_report_{}".format(time.strftime("%Y-%m-%d-%H-%M-%S"))
        save_classif_reports(reel, predictions, report_file = report_path, print=True)
    else:
        save_classif_reports(reel, predictions, print=True)

###############################################################################################################
#Build, train and test
@click.command()
@click.argument(
    '--feats_file',
    type=click.File('rb'),
    help = 'Path where to find the .csv with the features for our model.')
@click.option(
    '--result_folder',
    type=click.File('wb'),
    help='Path of the file where to keep the .csv containing the features for our model.'
         'If not provided, no output of anything -not even the model- will be retained.',
)
def build_and_train_xu(feats_file, result_folder):
    X_train, X_val, X_test, y_train, y_val, y_test = sets_test_train(feats_file)
    model = build_model_xu(X_train)
    trained_model = train_model_xu(model, X_train, X_val,  y_train, y_val, result_folder)
    test_model(trained_model,X_test, y_test, result_folder)

if __name__ == "__main__":
    build_and_train_xu()


#!/usr/bin/env python

import evdev
from evdev import ecodes
from scipy.interpolate import interp1d
import pygame
from pygame.locals import *
import os
from time import sleep

DIM = [480,320] # screen framebuffer dimensions
WHITE      = (255, 255, 255)
BLACK      = (  0,   0,   0)
BLUE       = (  0,   0, 255)
GREEN      = (  0, 255,   0)
RED        = (255,   0,   0)
ORANGE     = (255, 165,   0)
GREY       = (128, 128, 128)
YELLOW     = (255, 255,   0)
PINK       = (255, 192, 203)
LBLUE      = (191, 238, 244)

os.putenv('SDL_VIDEODRIVER', 'fbcon')
os.putenv('SDL_FBDEV', '/dev/fb0')
os.putenv('SDL_MOUSEDRV', 'TSLIB')
os.putenv('SDL_MOUSEDEV', '/dev/input/event0')
#os.putenv('SDL_MOUSEDEV', '/devices/virtual/input/input0')

pygame.init()
pygame.mouse.set_visible(False)
lcd = pygame.display.set_mode(DIM)
lcd.fill((0,0,0))
pygame.display.update()

font_big = pygame.font.Font(None, 50)

touch_buttons = {'1':(DIM[0]/3,DIM[1]/3), '2':(DIM[0]/3*2,DIM[1]/3), '3':(DIM[0]/3,DIM[1]/3*2), '4':(DIM[0]/3*2,DIM[1]/3*2)}

def draw_background():
    lcd.fill((0,0,0))
    pygame.display.update()
    for k,v in touch_buttons.items():
        text_surface = font_big.render('%s'%k, True, WHITE)
        rect = text_surface.get_rect(center=v)
        lcd.blit(text_surface, rect)
        pygame.display.update()

draw_background()

x = 0
y = 0

devices = [evdev.InputDevice(path) for path in evdev.list_devices()]
#print(devices)

for device in devices:
   print(device.path, device.name, device.phys)

device = evdev.InputDevice('/dev/input/event0')
print(device)
capabilities = device.capabilities(verbose=True, absinfo=True)
#print(capabilities[('EV_ABS', 3L)][0][1][1])

'''
{('EV_KEY', 1L): [('BTN_TOUCH', 330L)],
 ('EV_ABS', 3L): [(('ABS_X', 0L), AbsInfo(value=3401, min=110, max=3930, fuzz=0, flat=0, resolution=0)),
                  (('ABS_Y', 1L), AbsInfo(value=1894, min=170, max=3917, fuzz=0, flat=0, resolution=0)),
                  (('ABS_PRESSURE', 24L), AbsInfo(value=1, min=0, max=1, fuzz=0, flat=0, resolution=0))],
 ('EV_SYN', 0L): [('SYN_REPORT', 0L),
                  ('SYN_CONFIG', 1L),
                  ('SYN_DROPPED', 3L)]}
'''

fbX = 480
fbY = 320

#minX = 110
#minY = 170
#maxX = 3930
#maxY = 3917

minX = capabilities[('EV_ABS', 3L)][0][1][1]
minY = capabilities[('EV_ABS', 3L)][1][1][1]
maxX = capabilities[('EV_ABS', 3L)][0][1][2]
maxY = capabilities[('EV_ABS', 3L)][1][1][2]

print(minX,maxX,minY,maxY)

X = interp1d([minX,maxX],[0,fbX])
Y = interp1d([minY,maxY],[0,fbY])

for event in device.read_loop():
    if event.type == ecodes.EV_ABS:
        #print(evdev.categorize(event))
        #print(dir(event))
        #print(event.type, ecodes.ABS_X, event.value, event.code, event.sec, event.timestamp, event.usec )
        if event.code == ecodes.ABS_X:
            try:
                x = int(X(event.value))
            except:
                pass
        if event.code == ecodes.ABS_Y:
            try:
                y = int(Y(event.value))
            except:
                pass
        #print(x,y)
        draw_background()
        pygame.draw.circle(lcd, RED, (x,y), 40)
        pygame.display.update()

'''
ABS_X
ABS_Y
BTN_TOUCH
ABS_PRESSURE
SYN_REPORT
'''


import numpy as np
from simtk.openmm.app import *
from simtk.openmm import *
import simtk.openmm as mm
from simtk.unit import *
import sys
from sys import stdout
import os
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('-i','--input-pdb',help='Start PDB structure path',required=True,dest='input_path')
parser.add_argument('-o','--out-path',help='Directory to output trajectories to. Set to directory of input pdb if not specified.',dest='out_path')
parser.add_argument('--no-restraints',help='Disables flat bottom harmonic restraints. Default includes restraints.',action='store_false',dest='restraints')
parser.set_defaults(restraints=True)
parser.add_argument('--iteration',help='Index for this trajectory (each must have a unique index to avoid overwritting)',type=int,required=True)
parser.add_argument('-l','--length',help='Set length of simulation in ns',type=int,default=100,dest='length')
parser.add_argument('--gpu',help='Specify indexes for GPUs to use',type=int,nargs='*',default=[0])
parser.add_argument('--nochk',help='Disables writing checkpoint files. Default writes a checkpoint every 5 ns.',action='store_false',dest='checkpoint')
parser.set_defaults(checkpoint=True)
parser.add_argument('-dr',help='Specify residues for strong distance restraints between corresponding CA atoms. Used to restrain ends around a gap in a structure. NOTE: Indexes must be numbered with N-terminus as residue 1 and no gaps in the numbering. Additionally, the protein must be split into two separate chains of the same model.',
            nargs=2,type=int,dest='dist_restraint')

args = parser.parse_args()
pdb_path = args.input_path
out_path = args.out_path
if out_path is None:
    out_path = os.path.dirname(os.path.abspath(pdb_path))
restraints_on = args.restraints
iteration = args.iteration
sim_length = args.length
gpu_idxs = args.gpu
gpu_string = ''
for i in gpu_idxs:
    gpu_string += str(i)+','
gpu_string = gpu_string[:-1]
print(f'Using GPUs: {gpu_string}')
write_checkpoint = args.checkpoint
if args.dist_restraint == None:
    has_gap = False
else:
    has_gap = True
    res1 = args.dist_restraint[0]-1
    res2 = args.dist_restraint[1]-1

def save_pdb(name):
    '''Writes out a pdb structure of the current state of the simulation'''
    position = simulation.context.getState(getPositions=True).getPositions()
    energy = simulation.context.getState(getEnergy=True).getPotentialEnergy()
    PDBFile.writeFile(simulation.topology, position, open(name+'.pdb','w'))
    print('Saved file: ' + name + '.pdb')
    print(f'Energy: {energy._value*KcalPerKJ:3.3f} kcal/mol')

#set up system for explicit water simulation
pdb = PDBFile(pdb_path)
forcefield = ForceField('amber14-all.xml', 'amber14/tip3pfb.xml')

#add solvent with physiological salt concentrations
modeller = Modeller(pdb.topology, pdb.positions)
modeller.addSolvent(forcefield, padding=0.9*nanometers, positiveIon='Na+', negativeIon='Cl-',
                ionicStrength=0.1*molar)

system = forcefield.createSystem(modeller.topology, nonbondedMethod=PME,
        nonbondedCutoff=1*nanometer, constraints=HBonds)

def add_restraints():
    #add flat bottom harmonic restraints to each CA atom
    restraint = CustomExternalForce("k0*(max(d-d0, 0.0))^2 ; d=periodicdistance(x, y, z, x0, y0, z0)")

    restraint.addGlobalParameter("d0", 5*angstroms)
    restraint.addGlobalParameter("k0", 0.025 * 12 * kilocalories_per_mole/angstroms**2) #mass weighted, 12 for Carbon!!
    restraint.addPerParticleParameter("x0")
    restraint.addPerParticleParameter("y0")
    restraint.addPerParticleParameter("z0")

    num_restraints = 0
    for res in modeller.topology.residues():
        for atom in res.atoms():
            if atom.name == 'CA':
                restraint.addParticle(atom.index,modeller.positions[atom.index])
                num_restraints += 1
                break

    system.addForce(restraint)
    print('Added flat-bottom harmonic restraints to each CA atom')

if restraints_on:
    #check # of forces in system before and after to ensure worked
    add_restraints()

def add_dist_restraint():
    print('Adding distance restraint.')
    force = HarmonicBondForce()
    atom1_found = False
    atom2_found = False
    for res in modeller.topology.residues():
        if res.index == res1:
            print(f"Res1: {res.name} on chain {res.chain}")
            for atom in res.atoms():
                if atom.name == 'CA':
                    caIdx1 = atom.index
                    atom2_found = True
                    break
        elif res.index == res2:
            print(f"Res2: {res.name} on chain {res.chain}")
            for atom in res.atoms():
                if atom.name == 'CA':
                    caIdx2= atom.index
                    atom2_found = True
                    break
        if atom1_found and atom2_found:
            break
    pos1 = np.array(modeller.positions[caIdx1])
    pos2 = np.array(modeller.positions[caIdx2])
    eq_dist = np.linalg.norm(pos1-pos2)
    #force constant, kJ/mol/nm^2
    k = 300000
    force.addBond(caIdx1,caIdx2,eq_dist,k)
    system.addForce(force)
    print(f'Added distance restraints between CA {caIdx1} and {caIdx2}')

if has_gap:
    add_dist_restraint()

dt = 0.002 #ps

print(f"dt set to {dt*10**3:3.3f} fs")

integrator = LangevinIntegrator(300*kelvin, 1/picosecond, dt*picoseconds)

platform = Platform.getPlatformByName('CUDA')
properties = {'DeviceIndex': gpu_string, 'Precision': 'mixed'}

simulation = Simulation(modeller.topology, system, integrator, platform, properties)
simulation.context.setPositions(modeller.positions)

print('Start energy minimization')
simulation.minimizeEnergy()
save_pdb(os.path.join(out_path,"minimized_"+str(iteration)))
print('End energy minimization')

print("Generating new simulation object for NPT equilibration force")
#get values from previous context
old_state = simulation.context.getState(getPositions=True,getVelocities=True,enforcePeriodicBox=True)
system = forcefield.createSystem(modeller.topology, nonbondedMethod=PME,
        nonbondedCutoff=1*nanometer, constraints=HBonds)
integrator = LangevinIntegrator(300*kelvin, 1/picosecond, dt*picoseconds)
platform = Platform.getPlatformByName('CUDA')
properties = {'DeviceIndex': gpu_string, 'Precision': 'mixed'}

#add restraint forces back in
if restraints_on:
    add_restraints()
if has_gap:
    add_dist_restraint()

# Set up NPT equilibration
pressure = 1 * atmosphere
temperature = 300 * kelvin
barostat_frequency = 1
barostat = MonteCarloBarostat(pressure, temperature, barostat_frequency)
system.addForce(barostat)

#generate new simulation object
simulation = Simulation(modeller.topology, system, integrator, platform, properties)
simulation.context.setState(old_state)

print("Start NPT Equilibration")
simulation.step(250*1000) #.5ns
save_pdb(os.path.join(out_path,"NPT_"+str(iteration)))
print("NPT Equilibration completed")

print('Generating new simulation object without NPT force')
#get values from previous context
old_state = simulation.context.getState(getPositions=True,getVelocities=True,enforcePeriodicBox=True)
system = forcefield.createSystem(modeller.topology, nonbondedMethod=PME,
        nonbondedCutoff=1*nanometer, constraints=HBonds)
integrator = LangevinIntegrator(300*kelvin, 1/picosecond, dt*picoseconds)
platform = Platform.getPlatformByName('CUDA')
properties = {'DeviceIndex': gpu_string, 'Precision': 'mixed'}
#add restraint forces back in
if restraints_on:
    add_restraints()
if has_gap:
    add_dist_restraint()

#generate new simulation object
simulation = Simulation(modeller.topology, system, integrator, platform, properties)
simulation.context.setState(old_state)

print("Start simulation with restraints", str(restraints_on))

#report the # of steps taken each ns
print(f"Divide step by {int(1/(dt*10**-3))} to get time in ns")
simulation.reporters.append(StateDataReporter(stdout, int(1/(dt*10**-3)), step=True))
if write_checkpoint:
    simulation.reporters.append(CheckpointReporter(os.path.join(out_path, f'chkpnt_{iteration}.chk'), int(5/(dt*10**-3)))) #write a checkpoint every 5 ns
    print("Will save a checkpoint every 5 ns")

#Write out trajectory frames every 20 ps in DCD format
simulation.reporters.append(DCDReporter(os.path.join(out_path,'refinement_trajectory_'+str(iteration)+'.dcd'), int(0.02/(dt*10**-3))))

simulation.step(sim_length/(dt*10**-3)) #run simulation for sim_length ns
print("Simulation done!")

save_pdb(os.path.join(out_path,"Final_MD_"+str(iteration)))


import copy
import glob
import os
import time

import gym
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler

from arguments import get_args
from baselines.common.vec_env.subproc_vec_env import SubprocVecEnv
from envs import make_env
from kfac import KFACOptimizer
from model import CNNPolicy,CNNPolicy2, MLPPolicy
from storage import RolloutStorage
from visualize import visdom_plot

args = get_args()

assert args.algo in ['a2c', 'ppo', 'acktr']
if args.algo == 'ppo':
    assert args.num_processes * args.num_steps % args.batch_size == 0

num_updates = int(args.num_frames) // args.num_steps // args.num_processes

torch.manual_seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)

try:
    os.makedirs(args.log_dir)
except OSError:
    files = glob.glob(os.path.join(args.log_dir, '*.monitor.json'))
    for f in files:
        os.remove(f)

def main():
    print("#######")
    print(
        "WARNING: All rewards are clipped so you need to use a monitor (see envs.py) or visdom plot to get true rewards")
    print("#######")

    os.environ['OMP_NUM_THREADS'] = '1'

    if args.vis:
        from visdom import Visdom
        viz = Visdom()
        win = None

    envs = SubprocVecEnv([
        make_env(args.env_name, args.seed, i, args.log_dir)
        for i in range(args.num_processes)
    ])

    obs_shape = envs.observation_space.shape
    print (obs_shape)
    obs_shape = (obs_shape[0] * args.num_stack, *obs_shape[1:])

    if len(envs.observation_space.shape) == 3:
        if args.custom == 1:
            print('## Using Custom CNN Policy....'+str(obs_shape[0]))
            actor_critic = CNNPolicy2(obs_shape[0], envs.action_space)
        else:
            print('## Using CNN Policy....'+str(obs_shape[0]))
            actor_critic = CNNPolicy(obs_shape[0], envs.action_space)
    else:
        print('## Using MLP Policy....')
        actor_critic = MLPPolicy(obs_shape[0], envs.action_space)

    if envs.action_space.__class__.__name__ == "Discrete":
        action_shape = 1
    else:
        action_shape = envs.action_space.shape[0]

    if args.cuda:
        actor_critic.cuda()

    if args.algo == 'a2c':
        optimizer = optim.RMSprop(actor_critic.parameters(), args.lr, eps=args.eps, alpha=args.alpha)
    elif args.algo == 'ppo':
        optimizer = optim.Adam(actor_critic.parameters(), args.lr, eps=args.eps)
    elif args.algo == 'acktr':
        optimizer = KFACOptimizer(actor_critic)

    rollouts = RolloutStorage(args.num_steps, args.num_processes, obs_shape, envs.action_space)
    current_state = torch.zeros(args.num_processes, *obs_shape)

    def update_current_state(state):
        shape_dim0 = envs.observation_space.shape[0]
        state = torch.from_numpy(state).float()
        if args.num_stack > 1:
            current_state[:, :-shape_dim0] = current_state[:, shape_dim0:]
        current_state[:, -shape_dim0:] = state

    state = envs.reset()
    update_current_state(state)

    rollouts.states[0].copy_(current_state)

    # These variables are used to compute average rewards for all processes.
    episode_rewards = torch.zeros([args.num_processes, 1])
    final_rewards = torch.zeros([args.num_processes, 1])
    if args.cuda:
        current_state = current_state.cuda()
        rollouts.cuda()

    if args.algo == 'ppo':
        old_model = copy.deepcopy(actor_critic)

    for j in range(num_updates):
        for step in range(args.num_steps):
            # Sample actions
            value, action = actor_critic.act(Variable(rollouts.states[step], volatile=True))
            cpu_actions = action.data.squeeze(1).cpu().numpy()

            # Obser reward and next state
            state, reward, done, info = envs.step(cpu_actions)
            reward = torch.from_numpy(np.expand_dims(np.stack(reward), 1)).float()
            episode_rewards += reward
            if done:
                start_time = time.process_time()
            # If done then clean the history of observations.
            masks = torch.FloatTensor([[0.0] if done_ else [1.0] for done_ in done])
            final_rewards *= masks
            final_rewards += (1 - masks) * episode_rewards
            episode_rewards *= masks

            if args.cuda:
                masks = masks.cuda()

            if current_state.dim() == 4:
                current_state *= masks.unsqueeze(2).unsqueeze(2)
            else:
                current_state *= masks

            update_current_state(state)
            rollouts.insert(step, current_state, action.data, value.data, reward, masks)

        next_value = actor_critic(Variable(rollouts.states[-1], volatile=True))[0].data

        if hasattr(actor_critic, 'obs_filter'):
            actor_critic.obs_filter.update(rollouts.states[:-1].view(-1, *obs_shape))

        rollouts.compute_returns(next_value, args.use_gae, args.gamma, args.tau)

        if args.algo in ['a2c', 'acktr']:
            values, action_log_probs, dist_entropy = actor_critic.evaluate_actions(
                Variable(rollouts.states[:-1].view(-1, *obs_shape)), Variable(rollouts.actions.view(-1, action_shape)))

            values = values.view(args.num_steps, args.num_processes, 1)
            action_log_probs = action_log_probs.view(args.num_steps, args.num_processes, 1)

            advantages = Variable(rollouts.returns[:-1]) - values
            value_loss = advantages.pow(2).mean()

            action_loss = -(Variable(advantages.data) * action_log_probs).mean()

            if args.algo == 'acktr' and optimizer.steps % optimizer.Ts == 0:
                # Sampled fisher, see Martens 2014
                actor_critic.zero_grad()
                pg_fisher_loss = -action_log_probs.mean()

                value_noise = Variable(torch.randn(values.size()))
                if args.cuda:
                    value_noise = value_noise.cuda()

                sample_values = values + value_noise
                vf_fisher_loss = -(values - Variable(sample_values.data)).pow(2).mean()

                fisher_loss = pg_fisher_loss + vf_fisher_loss
                optimizer.acc_stats = True
                fisher_loss.backward(retain_graph=True)
                optimizer.acc_stats = False

            optimizer.zero_grad()
            (value_loss * args.value_loss_coef + action_loss - dist_entropy * args.entropy_coef).backward()

            if args.algo == 'a2c':
                nn.utils.clip_grad_norm(actor_critic.parameters(), args.max_grad_norm)

            optimizer.step()
        elif args.algo == 'ppo':
            advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
            advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-5)

            old_model.load_state_dict(actor_critic.state_dict())
            if hasattr(actor_critic, 'obs_filter'):
                old_model.obs_filter = actor_critic.obs_filter

            for _ in range(args.ppo_epoch):
                sampler = BatchSampler(SubsetRandomSampler(range(args.num_processes * args.num_steps)),
                                       args.batch_size * args.num_processes, drop_last=False)
                for indices in sampler:
                    indices = torch.LongTensor(indices)
                    if args.cuda:
                        indices = indices.cuda()
                    states_batch = rollouts.states[:-1].view(-1, *obs_shape)[indices]
                    actions_batch = rollouts.actions.view(-1, action_shape)[indices]
                    return_batch = rollouts.returns[:-1].view(-1, 1)[indices]

                    # Reshape to do in a single forward pass for all steps
                    values, action_log_probs, dist_entropy = actor_critic.evaluate_actions(Variable(states_batch),
                                                                                Variable(actions_batch))

                    _, old_action_log_probs, _ = old_model.evaluate_actions(Variable(states_batch, volatile=True),
                                                                            Variable(actions_batch, volatile=True))

                    ratio = torch.exp(action_log_probs - Variable(old_action_log_probs.data))
                    adv_targ = Variable(advantages.view(-1, 1)[indices])
                    surr1 = ratio * adv_targ
                    surr2 = torch.clamp(ratio, 1.0 - args.clip_param, 1.0 + args.clip_param) * adv_targ
                    action_loss = -torch.min(surr1, surr2).mean()  # PPO's pessimistic surrogate (L^CLIP)

                    value_loss = (Variable(return_batch) - values).pow(2).mean()

                    optimizer.zero_grad()
                    (value_loss + action_loss - dist_entropy * args.entropy_coef).backward()
                    optimizer.step()

        rollouts.states[0].copy_(rollouts.states[-1])

        if j % args.save_interval == 0 and args.save_dir != "":
            save_path = os.path.join(args.save_dir, args.algo)
            try:
                os.makedirs(save_path)
            except OSError:
                pass

            # A really ugly way to save a model to CPU
            save_model = actor_critic
            if args.cuda:
                save_model = copy.deepcopy(actor_critic).cpu()
            torch.save(save_model, os.path.join(save_path, args.env_name + ".pt"))

        if j % args.log_interval == 0:
            end_time = time.process_time()
            print(
                "Updates {}, num frames {}, mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}, entropy {:.5f}, value loss {:.5f}, policy loss {:.5f}, time_spent {:.5f}".
                format(j, (j + 1) * args.num_processes * args.num_steps,
                       final_rewards.mean(),
                       final_rewards.median(),
                       final_rewards.min(),
                       final_rewards.max(), -dist_entropy.data[0],
                       value_loss.data[0], action_loss.data[0],
                       end_time - start_time))

        if j % args.vis_interval == 0:
            win = visdom_plot(viz, win, args.log_dir, args.env_name, args.algo)

if __name__ == "__main__":
    main()


from __future__ import (absolute_import, unicode_literals, division,
                        print_function)

import numpy.random as ra
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.gridspec import GridSpec
from scipy.ndimage.interpolation import shift
import re
channel_re = re.compile(r'.*CH([0-9]+)$')

class FeedAligner():
    '''Align images from different feeds, interactively
    Inputs:
        imgs:       a dictionary containing images to align,
                    like this: {'Ch0': img0, 'Ch1': img1, ...}
        ax:         a pyplot.axis instance where the image will be plotted
    '''
    def __init__(self, imgs, ax, xextent=None, yextent=None, feeds=None,
                 radius=None, angle=None):

        self.imgs = imgs
        self.ax = ax
        if xextent is None:
            xextent = [-1, 1]
        if yextent is None:
            yextent = [-1, 1]
        self.xextent = xextent
        self.yextent = yextent
        keys = imgs.keys()
        if feeds is None:
            idxs = [int(channel_re.match(k).group(1)) for k in keys]
            feeds = dict(zip(keys, idxs))
        self.feeds = feeds
        self.xoffsets = dict(zip(keys, np.zeros(len(keys))))
        self.yoffsets = dict(zip(keys, np.zeros(len(keys))))
        if radius is None:
            radius = 0
        self.radius = radius
        if angle is None:
            angle = 0
        self.angle = 0
        self.mode = 'ROUGH'
        self.variable_to_change = 'A'
        self.radius_step = (xextent[1] - xextent[0]) / 50
        self.angle_step = 5.
        self.step_factor = 1.
        self.ax.figure.canvas.mpl_connect('key_press_event', self.on_key)
        self.plotted = None
        self.plot_imgs()

    def on_key(self, event):
        x, y = event.xdata, event.ydata
        key = event.key

        if key == 'm':
            self.switch_mode()
        if key == 'r':
            self.variable_to_change = 'R'
        if key == 'a':
            self.variable_to_change = 'A'
        if key == 'down':
            self.makestep(-1)
        if key == 'up':
            self.makestep(1)

        self.print_variables()
        return x, y, key

    def makestep(self, sign):
        if self.variable_to_change == 'R':
            self.radius += sign * self.step_factor * self.radius_step
        if self.variable_to_change == 'A':
            self.angle += sign * self.step_factor * self.angle_step
        self.recalculate_offsets()

    def print_variables(self):
        print('Radius: {}; Angle: {}'.format(self.radius, self.angle))

    def recalculate_offsets(self):
        xoffsets, yoffsets = \
            standard_offsets(self.radius, -self.angle,
                             len(self.imgs.keys()))
        for k in self.imgs.keys():
            feed = self.feeds[k]

            self.xoffsets[k] = xoffsets[feed]
            self.yoffsets[k] = yoffsets[feed]

        self.plot_imgs()

    def switch_mode(self):
        if self.mode == 'ROUGH':
            self.mode = 'FINE'
            self.step_factor = 0.05
        else:
            self.mode = 'ROUGH'
            self.step_factor = 1

    def plot_imgs(self):
        self.ax.cla()
        total_image = 0
        keys = list(self.imgs.keys())
#        self.fig = plt.figure('Increase')
#        gs = GridSpec(4, len(keys) // 2)

#        half = len(keys) // 2

        for ik, k in enumerate(keys):
#            col = ik % half
#            row = ik // half
            img = self.imgs[k]
            xextent = self.xextent
            yextent = self.yextent
            xoff = \
                self.xoffsets[k] / (xextent[1] - xextent[0]) * len(img[0, :])
            yoff = \
                self.yoffsets[k] / (yextent[1] - yextent[0]) * len(img[:, 0])

            shifted = shift(img, [-xoff, -yoff])
            total_image += shifted
#            inc_ax = plt.subplot(gs[row * 2, col])
#            inc_ax.imshow(total_image,
#                          extent=xextent + yextent,
#                          vmin=np.percentile(total_image, 40))
#            inc_ax = plt.subplot(gs[row * 2 + 1, col])
#            inc_ax.imshow(shifted,
#                          extent=xextent + yextent,
#                          vmin=np.percentile(self.imgs[keys[0]], 40))
        self.ax.imshow(total_image,
                       extent=xextent + yextent,
                       vmin=np.percentile(self.imgs[keys[0]], 40)*len(keys))
        plt.draw()

def standard_offsets(radius=3., angle=0, nfeeds=7):
    if nfeeds == 1:
        return [0], [0]

    # 0 for feed 0, radius for the other six
    radii = np.array([0] + [radius]*(nfeeds - 1))
    # Feeds 1--6 are at angles -60, -120, etc. Here I use angle 0 for
    # convenience for feed 0, but it has no effect since radii[0] is 0
    feed_angles = \
        -np.arange(0, nfeeds, 1) * np.pi * 2/(nfeeds - 1) + np.radians(angle)

    xoffsets = radii * np.cos(feed_angles)
    yoffsets = radii * np.sin(feed_angles)
    return xoffsets, yoffsets

def test_with_point_source():
        # make the colormaps
    colors = 'white,red,green,blue,magenta,cyan,yellow'.split(',')
    cmaps = []
    for i in range(7):
        cmap = mpl.colors.LinearSegmentedColormap.from_list('cmap{}'.format(i),
                                                            ['black', colors[i]],
                                                            256)
        cmaps.append(cmap)

    real_xoffsets, real_yoffsets = standard_offsets(2.4, 45)

    imgs = {}
    xbins = np.linspace(-30, 30, 101)
    ybins = np.linspace(-30, 30, 101)
    for i in range(7):
        xoff = real_xoffsets[i]
        yoff = real_yoffsets[i]

        mean = [-xoff - 3, -yoff - 1]
        mean2 = [-xoff + 4, -yoff + 5]
        cov = [[1, 0], [0, 1]]  # diagonal covariance, points lie on x or y-axis

        x, y = ra.multivariate_normal(mean, cov, 50000).T
        x2, y2 = ra.multivariate_normal(mean2, cov, 50000).T

        x3 = ra.uniform(-30, 30, 50000)
        y3 = ra.uniform(-30, 30, 50000)

        hist, _, _ = np.histogram2d(x, y, bins=[xbins, ybins])
        hist2, _, _ = np.histogram2d(x2, y2, bins=[xbins, ybins])
        hist3, _, _ = np.histogram2d(x3, y3, bins=[xbins, ybins])

        img = hist.T + hist2.T + hist3.T
        channel = 'CH{}'.format(i)
        imgs[channel] = img[::-1]
        plt.imshow(img[::-1], extent=[-30, 30, -30, 30], alpha=1/7,
                   cmap=cmaps[i])
        plt.scatter([-xoff], [-yoff])

    fig = plt.figure('Interactive')
    ax = fig.add_subplot(111)
    FeedAligner(imgs, ax, [-30, 30], [-30, 30])
    plt.show()

if __name__ == '__main__':
    import astropy.io.fits as pf
    hdulist = pf.open('try_2.fits')

    imgs = {}
    feeds = {}
    for h in hdulist[1:]:
        name = h.name
        img = h.data
        imgs[name] = img[::-1]
        chnum = int(channel_re.match(name).group(1))
        feeds[name] = chnum // 2

    nx = hdulist['IMGCH0'].header['NAXIS1']
    ny = hdulist['IMGCH0'].header['NAXIS2']

    dx = np.abs(hdulist['IMGCH0'].header['CDELT1'])
    dy = np.abs(hdulist['IMGCH0'].header['CDELT2'])

    cx = hdulist['IMGCH0'].header['CRPIX1']
    cy = hdulist['IMGCH0'].header['CRPIX2']

    xextent = [-cx * dx, (nx - cx) * dx]
    yextent = [-cy * dy, (ny - cy) * dy]

    print(xextent, yextent)
    fig = plt.figure('Interactive')
    ax = fig.add_subplot(111)
    FeedAligner(imgs, ax, xextent, yextent, feeds)
    plt.show()


"""
Module containing a useable interface for the data output of Budapest-Florida code
in Python.

# Current features

- History class providing support for fort.18 files
- Model class providing support for fort.95 files
- LimmitCycle class providing support for fort.19 files
- bpfDataRead function provides easy initialization in given directory

Created by: Gábor B. Kovács 2021
"""
from astropy.constants import cgs
from math import pi
import numpy as np
import astropy.constants as constants
from astropy import units
#import calcion
#from . import calcion

class BaseData:
    """
    Base class for the three most common Bupest-Florida Code output, such as
     - fort.95 which contains the static starting model profile
     - fort.19 which contains the profiles for two full periods
     - fort.18 which contains the photosphere data, for the entire run
    """
    def __init__(self, data=None, col_names=None):
        self.datablock=data
        self.column_names=col_names

    def data(self,key):
        if key in self.datablock:
            return self.datablock[key]
        else:
            raise KeyError('There is no'+str(key)+'data member.')

    def insertColumn(self,data,columnnames):
        if isinstance(data,dict):
            #print('dict')
            self.datablock.update(data)
        else:
            #print(type(data))
            self.datablock.update(dict(zip(columnnames,data)))
        self.column_names += columnnames
        return self

    def __getattr__(self,key):
        return self.data(key)

class History(BaseData):
    """
    This class provides standard interface for working with fort.18 files.
    Column names:
     - time:    Elapsed time from hydro start in days
     - vel_rad: Radial velocity in the photosphere
     - L_phot:  Photosphere luminosity
     - radius:  Radius of surface (One zone below nzn.)
     - L_surf:  Luminosity of the surface zone
    """
    columnnames=('time','vel_rad','L_phot','radius','L_surf')
    def __init__(self,path):
        data,colnames=self.read_data(path)
        super().__init__(data,colnames)

    def read_data(self,path):
            datafile=open(path)
            dataLines=datafile.readlines()
#This should be static!
            columns=[]
            for i in range(0,len(dataLines[0].split())):
                columns.append([float(x.split()[i]) for x in dataLines])
                columns[i]=np.array(columns[i])
            datafile.close()
            return dict(zip(self.columnnames,columns)), self.columnnames

class Model(BaseData):
    """
    Class providing interface for working with fort.95 files.
    """
    columnnames=('zone','radius','pressure','spec_vol','energy','temperature','turbulent_energy',
                    'dm','L_c','p_turb','L_r','sk','L_turb','s','cs','fcsl','st','cla','taup','taud','opacity',
                    'edt','cp')
    def __init__(self,path):
        data,colnames=self.read_data(path)
        super().__init__(data,colnames)

    def read_data(self,path):
        datafile=open(path)
        dataLines=datafile.readlines()
        columns=[]
        for i in range(0,len(dataLines[0].split())-1):
            columns.append([float(x.split()[i].replace('D','E')) for x in dataLines])
            columns[i]=np.array(columns[i])
        datafile.close()
        return dict(zip(self.columnnames,columns)), self.columnnames

class DataPoint(BaseData):
    """
    Helper class for fort.19 data access, holds one record of fort.19 data.
    Fields are:
    phase       - phase of pulsation
    zone        - zone number of mass cell
    velocity    - velocity of mass cell
    L_r         - radiative luminosity
    radius      - radius at mass cell
    temperature - temperature of the cell
    dm          - size of the cell
    opacity     - opacity in the cell
    energy      - internal energy
    e_t         - turbulent energy
    entropy     - specific entropy
    F_t         - turbulent luminosity
    F_c         - convective luminosity
    pressure    - gas pressure in cell
    p_t         - turbulent pressure
    p_eddy      - viscous pressure of the cell
    mass        - mass (Langrangian) coordinate of the cell.
    """
    columnnames=('phase','zone','velocity','L_r','radius','temperature','dm','opacity','energy','e_t','entropy',
                'F_t','F_c','pressure','p_t','p_eddy','mass')

    def __init__(self,Line : str):
        #TODO: d D e E csere!

        data=dict(zip(DataPoint.columnnames,[float(x.translate({"d": "e","D" : "e"})) for x in Line.split()]))
        #for x in Line.split(): print(float(x))
        super().__init__(data,DataPoint.columnnames)

    @classmethod
    def createDataPointList(cls, Template = None):
        return dict(zip(Template.column_names, [ list() for i in enumerate(Template.column_names)]))

    @classmethod
    def fillDataPointList(cls,datapoints : list):
        if not all(isinstance(x,cls) for x in datapoints):
            raise TypeError("datapoints should be a list of DataPoint")
        template = datapoints[0]
        datapoint_list=cls.createDataPointList(template)
        for point in datapoints:
            for key in template.column_names:
                data = point.data(key)
                datapoint_list[key].append(data)
        for key in datapoint_list:
                datapoint_list[key]=np.array(datapoint_list[key])
        return datapoint_list

class RawProfiles:
    """
    Container holding DataPoints extracted from fort.19
    """
    def __init__(self, filename):
        self.datablock=None
        self.num_time_series=0
        self.num_profiles=0
        self.read_file(filename)
        self.CalcSpecVol()

    def read_file(self,filename):
        infile=open(filename)
        lines = infile.readlines()
        self.datablock = [DataPoint(line) for line in lines]
        self.num_time_series = int(self.datablock[-1].zone)
        self.num_profiles = int(self.datablock[-1].phase)
        infile.close()

    def __getitem__(self,key) -> DataPoint:
        return self.datablock[key]

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

    def CalcSpecVol(self):
        for i in range(len(self.datablock)):
            #print(i)
            r = self.datablock[i].radius
            if self.datablock[i].zone == 1:
                volume = 4 * r ** 3 * pi / 3
                spec_vol = volume / self.datablock[i].mass
            else:
                r_nm1 = self.datablock[i-1].radius
                volume = 4* pi /3 * ( r**3 - r_nm1 ** 3)
                if self.datablock[i].dm != 0:
                    spec_vol = volume / self.datablock[i].dm
                else:
                    spec_vol = volume / self.datablock[i-1].dm
            self.datablock[i].insertColumn([spec_vol],tuple(['spec_vol']))

class TimeSeries(BaseData):
    """
    Holding one zones time series of the limit cycle
    """
    def __init__(self,rawprofile_obj : RawProfiles,zone_number):
        if not all([isinstance(rawprofile_obj,RawProfiles),isinstance(zone_number,(int,float))]):
            raise TypeError("Wrong arguments")
        super().__init__(self.createTimeSeries(rawprofile_obj,zone_number),DataPoint.columnnames)

    def createTimeSeries(self,rawprofile_obj : RawProfiles, zone_number) -> dict:
        datapoints=rawprofile_obj[zone_number::rawprofile_obj.num_time_series]
        #print(type(datapoints),zone_number,rawprofile_obj.num_time_series,len(rawprofile_obj.datablock))
        data=DataPoint.fillDataPointList(datapoints)
        return data

class Profiles(BaseData):
    """
    Holding a limit cycle profile at given phase number.
    """
    def __init__(self,rawprofile_obj : RawProfiles,phase_number):
        super().__init__(self.createProfile(rawprofile_obj,phase_number),DataPoint.columnnames)

    def createProfile(self,rawprofile_obj : RawProfiles,phase_number):
        datapoints=[point for point in rawprofile_obj if point.phase == phase_number]
        data=DataPoint.fillDataPointList(datapoints)
        return data

class LimitCycle:
    """
    Class providing user interface for the fort.19 data, both timeseries and profiles extraction, using the same interface.
    For columnnames see DataPoint class.
    """
    columnames=DataPoint.columnnames
    def __init__(self,infile_name):
        if isinstance(infile_name,str):
            rawdata=RawProfiles(infile_name)
        elif isinstance(infile_name,RawProfiles):
            rawdata = infile_name
        else:
            raise TypeError("Parameter should be either string or RawProfile object.")
        self.timeSeries=[TimeSeries(rawdata,zone_number) for zone_number in range(1,rawdata.num_time_series+1)]
        self.profiles=[Profiles(rawdata,phase_number) for phase_number in range(1,rawdata.num_profiles+1)]
        self.num_time_series = len(self.timeSeries)
        self.num_profiles=len(self.profiles)

#testing utilities
if(__name__ == '__main__'):
    from matplotlib import pyplot as plt
    import math
    a_point = DataPoint("45\t136\t12\t50\t5.6\t2e5\t4.13e+29\t125.68\t1.25e40\t1.2e38\t12000\t1e29\t12e5\t1.5e4\t1e2\t0.98\n")
    fort19_path='/home/gabesz/SPHERLS_playground/bpf_konv/konv-b/fort.19'#'/home/gabesz/Dokumentumok/VS_Code_workspace/calcion/fort.19'
    fort19_data=RawProfiles(fort19_path)
    print(a_point.datablock)
    print(a_point.dm)
    print(fort19_data[15].spec_vol)
  #  for cell_obj in fort19_data:
            #print(cell_obj.zone)

    #Testing Timeseries:
    #surface:
    surface_element=fort19_data.num_time_series-2
    print("érték:",fort19_data.num_time_series)
    surface_series=TimeSeries(fort19_data,surface_element)
    #print(surface_series.radius)
    mylog=np.vectorize(math.log10)
    #plt.plot(surface_series.phase/77.,surface_series.F_r/constants.L_sun.cgs)

    aProfile=Profiles(fort19_data,1)

    fort19_handler=LimitCycle(fort19_path)

    for phase in range(0):#,len(fort19_handler.profiles)):
        #phase=40
        ts_xdata=fort19_handler.timeSeries[fort19_handler.num_time_series-2].phase
        ts_ydata=fort19_handler.timeSeries[fort19_handler.num_time_series-2].L_r
        xdata=fort19_handler.profiles[phase].zone
        ydata=fort19_handler.profiles[phase].F_c
        #print(xdata,ydata)
        plt.subplot(2,1,1)
        plt.plot(xdata,ydata/constants.L_sun.cgs)
        plt.plot(xdata,fort19_handler.profiles[phase].L_r/constants.L_sun.cgs)
        plt.plot(xdata,fort19_handler.profiles[phase].F_t/constants.L_sun.cgs)
        #plt.plot(aProfile.zone[1:146],mylog(aProfile.temperature[1:146]))
        plt.subplot(2,1,2)
        plt.plot(ts_xdata/77,ts_ydata/constants.L_sun.cgs)
        plt.scatter(ts_xdata[phase]/77,ts_ydata[phase]/constants.L_sun.cgs,color="red")
        plt.show()
        plt.pause(5)

    #Testing reader

    #Other tests
    print(len(lim.profiles), lim.num_profiles)

import numpy,import astropy,from astropy


#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# vim:fenc=utf-8

"""
About: YOLOv2 Detector
"""

import json
import os
import time
import timeit
import struct
import gc

import cv2
import psutil
import numpy as np

# Disable printing too much warning messages
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import tensorflow as tf

import preprocessor
from utils.config import class_names, anchors
from utils.imgutils import decode_result, postprocess, image_to_feature_maps

class Detector(object):
    """YOLOv2 object detector
    Code here really needs refactoring...
    """

    def __init__(self, mode="raw", model_path="./model/yolo.pb"):
        self._name = "yolo"
        self._model_path = model_path
        self._mode = mode

        if mode == "preprocessed":
            self._input_tensor_name = "Pad_5:0"
            self._output_tensor_name = "output:0"
        elif mode == "raw":
            self._input_tensor_name = "input:0"
            self._output_tensor_name = "output:0"
        else:
            raise RuntimeError("Invalid server mode!")

        self.dtype_header = np.float16
        self.dtype_payload = np.uint8
        self.shape_splice = (8, 16)

        self.sess = self._read_model(self._model_path, self._name, is_onecore=True)
        self._init_tensors(self._input_tensor_name, self._output_tensor_name)

    @staticmethod
    def _read_model(path, name, is_onecore=True):
        """
        path: the location of pb file path
        name: name of tf graph
        return: tf.Session()
        """
        sess = tf.Session()
        # use one cpu core
        if is_onecore:
            session_conf = tf.ConfigProto(
                intra_op_parallelism_threads=1, inter_op_parallelism_threads=1
            )
            sess = tf.Session(config=session_conf)

        with tf.gfile.FastGFile(path, "rb") as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            sess.graph.as_default()
            tf.import_graph_def(graph_def, name=name)
        return sess

    def _init_tensors(self, input_tensor_name, output_tensor_name):
        self.input1 = self.sess.graph.get_tensor_by_name(
            "{}/{}".format(self._name, input_tensor_name)
        )
        self.output1 = self.sess.graph.get_tensor_by_name(
            "{}/{}".format(self._name, output_tensor_name)
        )

    @staticmethod
    def read_img_jpeg_bytes(img_path):
        img = cv2.imread(img_path)
        img = cv2.resize(img, (608, 608))
        encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 100]
        _, encimg = cv2.imencode(".jpg", img, encode_param)
        encimg = encimg.tobytes()
        return encimg

    @staticmethod
    def preprocess_image(image, image_size=(608, 608)):
        image_cp = np.copy(image).astype(np.float32)
        # resize image
        image_rgb = cv2.cvtColor(image_cp, cv2.COLOR_BGR2RGB)
        image_resized = cv2.resize(image_rgb, image_size)
        # normalize
        image_normalized = image_resized.astype(np.float32) / 225.0
        # expand dimension
        image_expanded = np.expand_dims(image_normalized, axis=0)
        return image_expanded

    def inference_raw(self, data):
        """Inference the raw image data in bytes"""
        data = np.frombuffer(data, np.uint8)
        img = cv2.imdecode(data, 1)
        img_preprossed = self.preprocess_image(img)
        res = self.sess.run(self.output1, feed_dict={self.input1: img_preprossed})
        bboxes, obj_probs, class_probs = decode_result(
            model_output=res,
            output_sizes=(608 // 32, 608 // 32),
            num_class=len(class_names),
            anchors=anchors,
        )
        bboxes, scores, class_max_index = postprocess(
            bboxes, obj_probs, class_probs, image_shape=(432, 320)
        )

        return bboxes, scores, class_max_index, class_names

    def inference_preprocessed(self, data):
        l1 = struct.unpack("<H", data[2:4])[0]
        # lp = struct.unpack(">I", data[4:8])[0]
        header_tmp = np.frombuffer(data[8 : 8 + l1], self.dtype_header)
        payload_tmp = np.frombuffer(data[8 + l1 :], self.dtype_payload)
        # predict
        feature_maps = image_to_feature_maps(
            [(payload_tmp, header_tmp)], self.shape_splice
        )
        res = self.sess.run(self.output1, feed_dict={self.input1: feature_maps})
        bboxes, obj_probs, class_probs = decode_result(
            model_output=res,
            output_sizes=(608 // 32, 608 // 32),
            num_class=len(class_names),
            anchors=anchors,
        )
        # image_shape: original image size for displaying
        bboxes, scores, class_max_index = postprocess(
            bboxes, obj_probs, class_probs, image_shape=(432, 320)
        )

        return bboxes, scores, class_max_index, class_names

    def inference(self, data):
        if self._mode == "preprocessed":
            ret = self.inference_preprocessed(data)
            return ret
        elif self._mode == "raw":
            ret = self.inference_raw(data)
            return ret

    @staticmethod
    def get_detection_results(bboxes, scores, class_max_index, class_names, thr=0.3):
        """get_detection_results"""
        results = list()
        for i, box in enumerate(bboxes):
            if scores[i] < thr:
                continue
            cls_indx = class_max_index[i]
            r = {
                "object": class_names[cls_indx],
                "score": float(scores[i]),
                "position": (int(box[0]), int(box[1]), int(box[2]), int(box[3])),
            }
            results.append(r)
        results = json.dumps(results)
        return results

def test():
    raw_img_data = Detector.read_img_jpeg_bytes("./pedestrain.jpg")
    det = Detector(mode="raw")
    ret = det.inference(raw_img_data)
    resp = det.get_detection_results(*ret)
    print("*** Inference result of raw image!")
    print(resp)
    del det
    gc.collect(1)
    gc.collect(2)

    # Test detection of preprocessed image
    prep = preprocessor.Preprocessor()
    compressed_img_data = prep.inference(raw_img_data, 70)
    print(
        f"*** Raw image size: {len(raw_img_data)}B, preprocessed image size: {len(compressed_img_data)}B"
    )
    det = Detector(mode="preprocessed")
    ret = det.inference(compressed_img_data)
    resp_prep = det.get_detection_results(*ret)
    print("*** Inference result of preprocessed image!")
    print(resp_prep)

def benchmark():
    """benchmark"""
    print("* Run benchmark of detector performance")
    setup = """
from __main__ import Detector
data = Detector.read_img_jpeg_bytes("./pedestrain.jpg")
detector = Detector(mode="raw")
# Warm-up the session
ret = detector.inference(data)
    """
    stmt = """
ret = detector.inference(data)
    """
    number = 17
    ret = timeit.timeit(stmt=stmt, setup=setup, number=number)
    ret = ret / number
    print(f"Inference time for one raw image: {ret} s")

    gc.collect()

    setup = """
from __main__ import Detector, preprocessor
data = Detector.read_img_jpeg_bytes("./pedestrain.jpg")
prep = preprocessor.Preprocessor()
compressed_img_data = prep.inference(data, 70)
detector = Detector(mode="preprocessed")
# Warm-up the session
ret = detector.inference(compressed_img_data)
    """
    stmt = """
ret = detector.inference(compressed_img_data)
    """
    number = 17
    ret = timeit.timeit(stmt=stmt, setup=setup, number=number)
    ret = ret / number
    print(f"Inference time for one preprocessed image: {ret} s")

def main():
    """main"""
    # Use RSS to estimate the memory usage of detector
    # data = Detector.read_img_jpeg_bytes("./pedestrain.jpg")
    # detector = Detector(mode="raw")
    # _ = detector.inference(data)
    # mem = psutil.Process().memory_info().rss / (1024 ** 2)
    # print(f"Memory usage: {mem} MB")

    test()
    # benchmark()

if __name__ == "__main__":
    main()