'''Differential Evolution class definition
'''

import numpy as np
import ConfigSpace

class DEBase():
    def __init__(self, cs=None, f=None, dimensions=None, pop_size=None, max_age=None,
                 mutation_factor=None, crossover_prob=None, strategy=None, budget=None, **kwargs):
        # Benchmark related variables
        self.cs = cs
        self.f = f
        if dimensions is None and self.cs is not None:
            self.dimensions = len(self.cs.get_hyperparameters())
        else:
            self.dimensions = dimensions

        # DE related variables
        self.pop_size = pop_size
        self.max_age = max_age
        self.mutation_factor = mutation_factor
        self.crossover_prob = crossover_prob
        self.strategy = strategy
        self.budget = budget

        # Miscellaneous
        self.output_path = kwargs['output_path'] if 'output_path' in kwargs else ''

        # Global trackers
        self.inc_score = np.inf
        self.inc_config = None
        self.population = None
        self.fitness = None
        self.age = None
        self.history = []

    def reset(self):
        '''Can be called to reuse the same object for a new DE run
        '''
        self.inc_score = np.inf
        self.inc_config = None
        self.population = None
        self.fitness = None
        self.age = None
        self.history = []

    def init_population(self, pop_size=10):
        population = np.random.uniform(low=0.0, high=1.0, size=(pop_size, self.dimensions))
        return population

    def sample_population(self, size=3):
        selection = np.random.choice(np.arange(len(self.population)), size, replace=False)
        return self.population[selection]

    def boundary_check(self, vector, fix_type='random'):
        '''
        Checks whether each of the dimensions of the input vector are within [0, 1].
        If not, values of those dimensions are replaced with the type of fix selected.

        Parameters
        ----------
        vector : array
            The vector describing the individual from the population
        fix_type : str, {'random', 'clip'}
            if 'random', the values are replaced with a random sampling from [0,1)
            if 'clip', the values are clipped to the closest limit from {0, 1}

        Returns
        -------
        array
        '''
        violations = np.where((vector > 1) | (vector < 0))[0]
        if len(violations) == 0:
            return vector
        if fix_type == 'random':
            vector[violations] = np.random.uniform(low=0.0, high=1.0, size=len(violations))
        else:
            vector[violations] = np.clip(vector[violations], a_min=0, a_max=1)
        return vector

    def vector_to_configspace(self, vector):
        '''Converts numpy array to ConfigSpace object

        Works when self.cs is a ConfigSpace object and each component of the
        input vector is in the domain [0, 1].
        '''
        new_config = self.cs.sample_configuration()
        for i, hyper in enumerate(self.cs.get_hyperparameters()):
            if type(hyper) == ConfigSpace.OrdinalHyperparameter:
                ranges = np.arange(start=0, stop=1, step=1/len(hyper.sequence))
                param_value = hyper.sequence[np.where((vector[i] < ranges) == False)[0][-1]]
            elif type(hyper) == ConfigSpace.CategoricalHyperparameter:
                ranges = np.arange(start=0, stop=1, step=1/len(hyper.choices))
                param_value = hyper.choices[np.where((vector[i] < ranges) == False)[0][-1]]
            else:  # handles UniformFloatHyperparameter & UniformIntegerHyperparameter
                # rescaling continuous values
                param_value = hyper.lower + (hyper.upper - hyper.lower) * vector[i]
                if type(hyper) == ConfigSpace.UniformIntegerHyperparameter:
                    param_value = np.round(param_value).astype(int)   # converting to discrete (int)
            new_config[hyper.name] = param_value
        return new_config

    def f_objective(self):
        raise NotImplementedError("The function needs to be defined in the sub class.")

    def mutation(self):
        raise NotImplementedError("The function needs to be defined in the sub class.")

    def crossover(self):
        raise NotImplementedError("The function needs to be defined in the sub class.")

    def evolve_generation(self):
        raise NotImplementedError("The function needs to be defined in the sub class.")

    def run(self):
        raise NotImplementedError("The function needs to be defined in the sub class.")

class DE(DEBase):
    '''DE object

    Creates and initialises DE with the specified hyperparameters.
    Contains functions to run DE and perform DE operations.

    Parameters
    ----------
    cs : ConfigSpace
        ConfigSpace object defining the parameter space
    f : function
        The objective function that takes a configuration and (optional) budget as input
    dimensions : int
        The dimensionality of the parameter space
    pop_size : int
        Population size of DE for each generation
    max_age : float, optional
        Number of generations after which an individual will be replaced with random individual
    mutation_factor : float
        Scaling factor for mutation (F)
    crossover_prob : float
        Probability for crossover (Cr)
    strategy : str
        Strategies for mutation and crossover concatenated using '_'
        Default of 'rand1_bin' implies
            mutation strategy is 'rand1' and crossover strategy is 'bin' (binomial)
    budget : int, optional
        Can be used to define the budget to be used for evaluation of f()
    '''
    def __init__(self, cs=None, f=None, dimensions=None, pop_size=None, max_age=np.inf,
                 mutation_factor=None, crossover_prob=None, strategy='rand1_bin',
                 budget=None, **kwargs):
        super().__init__(cs=cs, f=f, dimensions=dimensions, pop_size=pop_size, max_age=max_age,
                         mutation_factor=mutation_factor, crossover_prob=crossover_prob,
                         strategy=strategy, budget=budget, **kwargs)
        if self.strategy is not None:
            self.mutation_strategy = self.strategy.split('_')[0]
            self.crossover_strategy = self.strategy.split('_')[1]
        else:
            self.mutation_strategy = self.crossover_strategy = None

    def f_objective(self, x, budget=None):
        '''Objective function that will evaluate x on budget
        '''
        if self.f is None:
            raise NotImplementedError("An objective function needs to be passed.")
        config = self.vector_to_configspace(x)
        if budget is not None:  # to be used when called by multi-fidelity based optimizers
            fitness, cost = self.f(config, budget=budget)
        else:
            fitness, cost = self.f(config)
        return fitness, cost

    def init_eval_pop(self, budget=None):
        '''Creates new population of 'pop_size' and evaluates individuals
        '''
        self.population = self.init_population(self.pop_size)
        self.fitness = np.array([np.inf for i in range(self.pop_size)])
        self.age = np.array([self.max_age] * self.pop_size)

        traj = []
        runtime = []
        history = []
        for i in range(self.pop_size):
            config = self.population[i]
            self.fitness[i], cost = self.f_objective(config, budget)
            if self.fitness[i] < self.inc_score:
                self.inc_score = self.fitness[i]
                self.inc_config = config
            traj.append(self.inc_score)
            runtime.append(cost)
            history.append((config.tolist(), float(self.fitness[i]), float(budget or 0)))

        return traj, runtime, history

    def mutation_rand1(self, r1, r2, r3):
        '''Performs the 'rand1' type of DE mutation
        '''
        diff = r2 - r3
        mutant = r1 + self.mutation_factor * diff
        return mutant

    def mutation_rand2(self, r1, r2, r3, r4, r5):
        '''Performs the 'rand2' type of DE mutation
        '''
        diff1 = r2 - r3
        diff2 = r4 - r5
        mutant = r1 + self.mutation_factor * diff1 + self.mutation_factor * diff2
        return mutant

    def mutation_currenttobest1(self, current, best, r1, r2):
        '''Performs the 'current-to-best' type of DE mutation
        '''
        diff1 = best - current
        diff2 = r1 - r2
        mutant = current + self.mutation_factor * diff1 + self.mutation_factor * diff2
        return mutant

    def mutation_rand2dir(self, r1, r2, r3):
        diff = r1 - r2 - r3
        mutant = r1 + self.mutation_factor * diff / 2
        return mutant

    def mutation(self, current=None, best=None):
        '''Performs DE mutation based on the strategy
        '''
        if self.mutation_strategy == 'rand1':
            r1, r2, r3 = self.sample_population(size=3)
            mutant = self.mutation_rand1(r1, r2, r3)

        elif self.mutation_strategy == 'rand2':
            r1, r2, r3, r4, r5 = self.sample_population(size=5)
            mutant = self.mutation_rand2(r1, r2, r3, r4, r5)

        elif self.mutation_strategy == 'rand2dir':
            r1, r2, r3 = self.sample_population(size=3)
            mutant = self.mutation_rand2dir(r1, r2, r3)

        elif self.mutation_strategy == 'best1':
            r1, r2 = self.sample_population(size=2)
            if best is None:
                best = self.population[np.argmin(self.fitness)]
            mutant = self.mutation_rand1(best, r1, r2)

        elif self.mutation_strategy == 'best2':
            r1, r2, r3, r4 = self.sample_population(size=4)
            if best is None:
                best = self.population[np.argmin(self.fitness)]
            mutant = self.mutation_rand2(best, r1, r2, r3, r4)

        elif self.mutation_strategy == 'currenttobest1':
            r1, r2 = self.sample_population(size=2)
            if best is None:
                best = self.population[np.argmin(self.fitness)]
            mutant = self.mutation_currenttobest1(current, best, r1, r2)

        elif self.mutation_strategy == 'randtobest1':
            r1, r2, r3 = self.sample_population(size=3)
            if best is None:
                best = self.population[np.argmin(self.fitness)]
            mutant = self.mutation_currenttobest1(r1, best, r2, r3)

        return mutant

    def crossover_bin(self, target, mutant):
        '''Performs the binomial crossover of DE
        '''
        cross_points = np.random.rand(self.dimensions) < self.crossover_prob
        if not np.any(cross_points):
            cross_points[np.random.randint(0, self.dimensions)] = True
        offspring = np.where(cross_points, mutant, target)
        return offspring

    def crossover_exp(self, target, mutant):
        '''Performs the exponential crossover of DE
        '''
        n = np.random.randint(0, self.dimensions)
        L = 0
        while ((np.random.rand() < self.crossover_prob) and L < self.dimensions):
            idx = (n+L) % self.dimensions
            target[idx] = mutant[idx]
            L = L + 1
        return target

    def crossover(self, target, mutant):
        '''Performs DE crossover based on the strategy
        '''
        if self.crossover_strategy == 'bin':
            offspring = self.crossover_bin(target, mutant)
        else:
            offspring = self.crossover_exp(target, mutant)
        return offspring

    def selection(self, trials, budget=None):
        '''Carries out a parent-offspring competition given a set of trial population
        '''
        assert len(self.population) == len(trials)
        traj = []
        runtime = []
        history = []
        for i in range(len(trials)):
            # evaluation of the newly created individuals
            fitness, cost = self.f_objective(trials[i], budget)
            # selection -- competition between parent[i] -- child[i]
            ## equality is important for landscape exploration
            if fitness <= self.fitness[i]:
                self.population[i] = trials[i]
                self.fitness[i] = fitness
                # resetting age since new individual in the population
                self.age[i] = self.max_age
            else:
                # decreasing age by 1 of parent who is better than offspring/trial
                self.age[i] -= 1
            # updation of global incumbent for trajectory
            if self.fitness[i] < self.inc_score:
                self.inc_score = self.fitness[i]
                self.inc_config = self.population[i]
            traj.append(self.inc_score)
            runtime.append(cost)
            history.append((trials[i].tolist(), float(fitness), float(budget or 0)))
        return traj, runtime, history

    def ranked_selection(self, trials, pop_size, budget=None, debug=False):
        '''Returns the fittest individuals from two sets of population
        '''
        # assert len(self.population) == len(trials)
        traj = []
        runtime = []
        history = []
        track = []
        trial_fitness = []
        # pop_size = len(trials)
        for i in range(len(trials)):
            fitness, cost = self.f_objective(trials[i], budget)
            trial_fitness.append(fitness)
            if fitness < self.inc_score:
                self.inc_score = fitness
                self.inc_config = trials[i]
            traj.append(self.inc_score)
            runtime.append(cost)
            history.append((trials[i].tolist(), float(fitness), float(budget or 0)))
        if debug:
            print("Ranking {} from {} vs. {}".format(pop_size, self.pop_size, len(trials)))
        tot_pop = np.vstack((self.population, trials))
        tot_fitness = np.hstack((self.fitness, trial_fitness))
        tot_age = np.hstack((self.age, [self.max_age] * len(trials)))
        rank = np.sort(np.argsort(tot_fitness)[:pop_size])
        self.population = tot_pop[rank]
        self.fitness = tot_fitness[rank]
        self.age = tot_age[rank]
        self.pop_size = pop_size
        return traj, runtime, history

    def kill_aged_pop(self, budget=None, debug=False):
        '''Replaces individuals with age older than max_age
        '''
        traj = []
        runtime = []
        history = []
        idxs = np.where(self.age <= 0)[0]
        if len(idxs) == 0:
            return traj, runtime, history
        if debug:
            print("Killing {} individual(s) for budget {}: {}".format(len(idxs), budget, self.age[idxs]))
        new_pop = self.init_population(pop_size=len(idxs))
        for i, index in enumerate(idxs):
            self.population[index] = new_pop[i]
            self.fitness[index], cost = self.f_objective(self.population[index], budget)
            self.age[index] = self.max_age
            if self.fitness[index] < self.inc_score:
                self.inc_score = self.fitness[index]
                self.inc_config = self.population[index]
            traj.append(self.inc_score)
            runtime.append(cost)
            history.append((self.population[index].tolist(), float(self.fitness[index]),
                            float(budget or 0)))
        return traj, runtime, history

    def evolve_generation(self, budget=None, best=None):
        '''Performs a complete DE evolution, mutation -> crossover -> selection
        '''
        trials = []
        for j in range(self.pop_size):
            target = self.population[j]
            donor = self.mutation(current=target, best=best)
            trial = self.crossover(target, donor)
            trial = self.boundary_check(trial)
            trials.append(trial)
        trials = np.array(trials)
        traj, runtime, history = self.selection(trials, budget)
        return traj, runtime, history

    def run(self, generations=1, verbose=False, budget=None):
        '''Main function that performs DE optimisation for the specified 'generations'

        Parameters
        ----------
        generations : int
            Number of generations to evolve, or iterations
        verbose : bool
            To print output during optimisation
        budget : int, optional
            Pass budget to override DE budget initialised with
            Allows reuse of the same DE object by calling reset()

        Returns
        -------
        traj : list
            List of fitness value of evaluated configurations
        runtime : list
            List of cost of evaluated configurations
        history : list (of tuples)
            Each tuple contains 3 further elements
                The vector/array form of the configuration (x)
                The fitness value at that iteration
                The budget at which the fitness was evaluated
        '''
        self.traj = []
        self.runtime = []
        self.history = []

        if budget is None:
            budget = self.budget

        if verbose:
            print("Initializing and evaluating new population...")
        self.traj, self.runtime, self.history = self.init_eval_pop(budget=budget)

        if verbose:
            print("Running evolutionary search...")
        for i in range(generations):
            if verbose:
                print("Generation {:<2}/{:<2} -- {:<0.7}".format(i+1, generations, self.inc_score))
            traj, runtime, history = self.evolve_generation(budget=budget)
            self.traj.extend(traj)
            self.runtime.extend(runtime)
            self.history.extend(history)

        if verbose:
            print("\nRun complete!")

        return np.array(self.traj), np.array(self.runtime), np.array(self.history)


from torch.utils.data import DataLoader
import torch
from tqdm import tqdm
import numpy as np
import pandas as pd
import os
import albumentations as A

import json

from dataset.dataset import *
from dataset.transforms import Transform
from model.models import *
from utils.setting import set_seed
from utils.collate import collate_fn
from utils.utils import get_result

def inference(model, test_loader, device):
    size = 256
    transform = A.Compose([A.Resize(size, size)])
    print('Start prediction.')

    model.eval()

    file_name_list = []
    preds_array = np.empty((0, size*size), dtype=np.long)

    with torch.no_grad():
        for step, (imgs, image_infos) in enumerate(tqdm(test_loader)):

            # inference (512 x 512)
            outs = model(torch.stack(imgs).to(device))
            oms = torch.argmax(outs, dim=1).detach().cpu().numpy()

            # resize (256 x 256)
            temp_mask = []
            for img, mask in zip(np.stack(imgs), oms):
                transformed = transform(image=img, mask=mask)
                mask = transformed['mask']
                temp_mask.append(mask)

            oms = np.array(temp_mask)

            oms = oms.reshape([oms.shape[0], size*size]).astype(int)
            preds_array = np.vstack((preds_array, oms))

            file_name_list.append([i['file_name'] for i in image_infos])
    print("End prediction.")
    file_names = [y for x in file_name_list for y in x]

    return file_names, preds_array

if __name__ == "__main__":
    import warnings
    warnings.filterwarnings(action='ignore')

    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    cfg = json.load(open("cfg.json", "r"))

    batch_size = cfg["batch_size"]
    learning_rate = cfg["learning_rate"]
    num_epochs = cfg["epochs"]
    categories = cfg["categories"]
    seed = cfg["seed"]
    set_seed(seed)

    exp = cfg["exp"]

    transform = Transform()

    # Load test dataset
    test_path = os.path.join(cfg["data_dir"], cfg["ann_file"]["test"])
    test_dataset = CustomDataset(cfg["data_dir"], cfg["ann_file"]["test"], categories, mode='test', transform=transform.test)
    test_loader = DataLoader(dataset=test_dataset,
                                            batch_size=batch_size,
                                            num_workers=4,
                                            collate_fn=collate_fn)

    # 모델 호출
    model = torch.load(f"exp/{exp}/best.pt")

    # test set에 대한 prediction
    file_names, preds = inference(model, test_loader, device)

    get_result("output.csv", file_names, preds)


import sys
from PyQt5 import QtCore, QtWidgets
from PyQt5.QtWidgets import * #QMainWindow, QLabel, QVBoxLayout, QLineEdit, QWidget
from PyQt5.QtCore import QSize
import pyqtgraph as pg
from py_polynomial_solver.view.math_text_label import MathTextLabel
import numpy as np
from py_polynomial_solver.model.polynomial import Polynomial
from py_polynomial_solver.view.polynomial_editor_dialog import PolynomialEditorDialog

class PolynomialSolverApp(QMainWindow):
    def __init__(self):
        QMainWindow.__init__(self)

        self.setMinimumSize(QSize(640, 480))
        self.setWindowTitle("Hello world - pythonprogramminglanguage.com")

        self.polynomials = [Polynomial([10, 20, 30, 40])]

        self.editPolynomialButton = QPushButton("Edit")
        self.editPolynomialButton.clicked.connect(self.edit_button_clicked)

        self.graphWidget = pg.PlotWidget()

        # centralWidget = QWidget(self)
        # self.setCentralWidget(centralWidget)

        title = QLabel("Hello World from PyQt", self)
        title.setAlignment(QtCore.Qt.AlignCenter)

        #mathText=r'$X_k = \sum_{n=0}^{N-1} x_n . e^{\frac{-i2\pi kn}{N}}$'
        mathText=r'$x^2 - 10x + 10$'
        self.updateCharts()

        layout = QVBoxLayout(self)
        layout.addWidget(title)
        layout.addWidget(self.graphWidget)
        layout.addWidget(self.editPolynomialButton)
        layout.addWidget(MathTextLabel(mathText, self), alignment=QtCore.Qt.AlignHCenter)

        widget = QWidget()
        widget.setLayout(layout)
        self.setCentralWidget(widget)

    def updateCharts(self):
        x = np.arange(-10, 10, 1)
        y = np.power(self.polynomials[0].coefficients(), np.arange(0,4))

        self.graphWidget.plot(x, y)

    def edit_button_clicked(self):
        print([1,2,3,4,5])
        dlg = PolynomialEditorDialog(self.polynomials)
        dlg.exec()

        self.polynomials = dlg.polynomials
        self.updateCharts()

def main():
    app = QtWidgets.QApplication(sys.argv)
    mainWin = PolynomialSolverApp()
    mainWin.show()
    sys.exit(app.exec_())

if __name__ == "__main__":
    main()


# import the necessary packages
from scipy.spatial import distance as dist
from collections import OrderedDict

import numpy as np

class CentroidTracker():

    def __init__(self, maxDisappeared=50):
        self.nextObjectID = 0
        self.objects = OrderedDict()
        self.disappeared = OrderedDict()

        self.maxDisappeared = maxDisappeared

    def register(self, centroid):
        # when registering an object we use the next available object
        self.objects[self.nextObjectID] = centroid
        # ID to store the centroid
        self.disappeared[self.nextObjectID] = 0
        self.nextObjectID += 1

    def deregister(self, objectID):
        # to deregister an object ID we delete the object ID from
        # both of our respective dictionaries
        del self.objects[objectID]
        del self.disappeared[objectID]

    def update(self, rects):
        # check to see if the list of input bounding box rectangles
        # is empty
        if len(rects) == 0: #i.e there are no objects.
            # loop over any existing tracked objects and mark them
            # as disappeared
            for objectID in list(self.disappeared.keys()):
                self.disappeared[objectID] += 1
                if self.disappeared[objectID] > self.maxDisappeared:
                    self.deregister(objectID)

            return self.objects

        # initialize an array of input centroids for the current frame
        inputCentroids = np.zeros((len(rects), 2), dtype="int")
        # loop over the bounding box rectangles
        for (i, (startX, startY, endX, endY)) in enumerate(rects):
            # use the bounding box coordinates to derive the centroid
            cX = int((startX + endX) / 2.0)
            cY = int((startY + endY) / 2.0)
            inputCentroids[i] = (cX, cY)

        if len(self.objects) == 0: # if there are no objects.
            for i in range(len(inputCentroids)):
                self.register(inputCentroids[i])

        else:

            # grab the set of object IDs and corresponding centroids
            objectIDs = list(self.objects.keys())
            objectCentroids = list(self.objects.values())

            D = dist.cdist(np.array(objectCentroids), inputCentroids)

            rows = D.min(axis=1).argsort()

            cols = D.argmin(axis=1)[rows] # retuurns the indices of min value along axis=1

            usedRows = set()
            usedCols = set()

            for (row, col) in zip(rows, cols):

                if row in usedRows or col in usedCols:
                    continue

                objectID = objectIDs[row]
                self.objects[objectID] = inputCentroids[col]
                self.disappeared[objectID] = 0

                # indicate that we have examined each of the row and
                # column indexes, respectively
                usedRows.add(row)
                usedCols.add(col)

            # compute both the row and column index we have NOT yet
            # examined

            unusedRows = set(range(0, D.shape[0])).difference(usedRows)
            unusedCols = set(range(0, D.shape[1])).difference(usedCols)

            if D.shape[0] >= D.shape[1]:
                for row in unusedRows:
                    objectID = objectIDs[row]
                    self.disappeared[objectID] += 1
                if self.disappeared[objectID] > self.maxDisappeared:
                        self.deregister(objectID)

            else:
                for col in unusedCols:
                    self.register(inputCentroids[col])
        return self.objects


import logging
import telegram
import requests
import json
import get_covid_city_stats as gccs
import api_key
import call_openai
from datetime import datetime
from random import randrange
from telegram.ext import Updater, CommandHandler, MessageHandler, filters
from telegram import InputMediaPhoto

logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                    level=logging.INFO)

logger = logging.getLogger(__name__)
reply = False

def covid_city_stats(update, context):
    global reply
    #if(reply == True):
    try:
        aux = update.message.text
        city_stats_cmd = aux.split()
        user_input = city_stats_cmd[1]
        city_stats = gccs.get_covid_city_stats(user_input)

        for i in city_stats:
                judet = 'Judet: {} \n'.format(i['Judet'])
                localitate = 'Localitate: {} \n'.format(i['Localitate'])
                populatie = 'Populatie: {} \n'.format(i['Populatie'])
                cazuri = 'Cazuri: {} \n'.format(i['Cazuri'])
                incidenta = 'Incidenta: {} \n'.format(i['Incidenta'])
                last_updated = "Ultima actualizare: {} \n".format(i['last_updated'])
                stats = incidenta + localitate + judet + populatie + cazuri  + last_updated
                update.message.reply_text(stats)
    finally:
        log_action(
            update.message.chat.id,
            update.message.chat.title,
            update.message.chat.type,
            update.message.from_user.id,
            update.message.from_user.username,
            update.message.from_user.first_name,
            update.message.from_user.last_name, user_input
            )

def openai_jokes(update, context):
    global reply
    #if(reply == True):
    try:
        aux = update.message.text
        user_input = aux.split(' ',1)[1]
        completion = call_openai.req_openai(user_input)
        update.message.reply_text(completion)
    finally:
        log_action(
            update.message.chat.id,
            update.message.chat.title,
            update.message.chat.type,
            update.message.from_user.id,
            update.message.from_user.username,
            update.message.from_user.first_name,
            update.message.from_user.last_name,
            user_input
            )

def log_action(chat_id, title, chat_type, user_id , username, first_name, last_name, user_input):
    n0w = datetime.now()
    log = {"chat_id": None, "chat_name": None, "chat_type": None, "user_id": None, "username": None, "first_name": None, 'last_name': None, "user_input": None, "datetime": None}
    log['chat_id'] = chat_id
    log['chat_name'] = title
    log['chat_type'] = chat_type
    log['user_id'] = user_id
    log['username'] = username
    log['first_name'] = first_name
    log['last_name'] = last_name
    log['user_input'] = user_input.upper()
    log['datetime'] = str(n0w)
    logs = json.dumps(log)
    f = open("user_actions.log","a")
    f.write(logs + "\n")
    f.close()

def main():
    #TEST
    #updater = Updater(api_key.test(), use_context=True)

    #PROD
    updater = Updater(api_key.prod(), use_context=True)

    updater.dispatcher.add_handler(CommandHandler('city', covid_city_stats))
    updater.dispatcher.add_handler(CommandHandler('openai', openai_jokes))
    updater.start_polling()
    updater.idle()
if __name__ == '__main__':
    main()


import subprocess

from BCBio import GFF
from make_oligos import chromosomes_parsed, get_seqdict
from Bio.SeqFeature import FeatureLocation
from Bio import SeqIO
import pandas as pd
import numpy as np
import re
from ushuffle import shuffle, Shuffler
import random

def import_wiggles(wig_file_in_paths):
    wigfiles = []
    for w in wig_file_in_paths:
        df = pd.read_csv(w, sep="\t", names=["gene", "position", "pos2", "count"])
        df = df.iloc[:, [0, 1, 3]]
        wigfiles += [df]
    return wigfiles

# import files:
in_gff = "../data/reference_sequences/NC_000913.3.gff"
in_fasta = "../data/reference_sequences/NC_000913.3.fasta"
in_oligos_fasta = "../data/reference_sequences/oligos_cds_new.fasta"
wiggles_in = ["../data/wigglefiles/01_no_PNA_3primeend_new_mod.wig",
              "../data/wigglefiles/07_no_PNA_3primeend_new_mod.wig",
              "../data/wigglefiles/13_no_PNA_3primeend_new_mod.wig",
              "../data/wigglefiles/79_no_PNA_3primeend_new_mod.wig",
              "../data/wigglefiles/85_no_PNA_3primeend_new_mod.wig"]

# import the wiggle files:
wig_oligos = import_wiggles(wiggles_in)
all_oligos = set(wig_oligos[0]["gene"])
wig_invivo_F = "../data/wigglefiles/wiggles-papercomparison_2/GSM3455900_RET_BWK_U00096_3_F_new.wig"
wig_invivo_R = "../data/wigglefiles/wiggles-papercomparison_2/GSM3455900_RET_BWK_U00096_3_R_new.wig"
threshold = 5   # threshold above which a peak is called

df_invivo_F = pd.read_table(wig_invivo_F, skiprows=2, index_col=None, names=["position", "count"])
df_invivo_R = pd.read_table(wig_invivo_R, skiprows=2, index_col=None, names=["position", "count"])
df_invivo_R[1] = df_invivo_R["count"] * (-1)

# out_tsv = open("../analysis/check_ann_ss.tsv", "w")
# out_tsv.write("\t".join(["locus_tag", "gene", "start_position", "strand", "peak_R1", "peak_R2", "peak_R3", "peak_R4",
#                         "peak_R5", "peak_found_Hör", "peak_wmeydan" "peak_found_meydan"]) + "\n")

# output dataframeis created, capturing all peaks
out_df = pd.DataFrame(columns=["locus_tag", "gene", "start_position", "start_codon", "strand", "peak_R1", "peak_R2",
                               "peak_R3", "peak_R4", "peak_R5", "peak_found_Hör", "peak_meydan", "peak_found_meydan",
                               "pred_RBS_region", "p_value", "sequence_start", "upstream_seq", "MFE"])

gibbs_df = pd.DataFrame(columns=["locus_tag", "gene", "SD_seq", "pos_from_stc", "shuffled", "delta_G", "p_value"])

# add a table showing enriched start codons/sd motifs:
motif_df = pd.DataFrame(np.zeros((35, 8)), columns=["ATG", "GTG", "TTG", "CTG", "ATC", "ATT", "AGG", "GGA"])
re_motiv = r"(ATG)|(GTG)|(TTG)|(CTG)|(ATC)|(ATT)|(AGG)|(GGA)"

# go through all oligos and check for annotated TIS
for oligo in SeqIO.parse(in_oligos_fasta, "fasta"):
    if oligo.id.startswith("b"):
        start_codon = oligo.seq[58:61]  # extract start codon
        id = oligo.id
        strand = id.split("_")[3]
        locus_tag = id.split("_")[0]
        gene = id.split("_")[1]
        start_pos = int(id.split("_")[2])

        newline = [locus_tag, gene, start_pos, str(start_codon), strand, 0, 0, 0, 0, 0, False, 0, False, False, False,
                   False, False, False]
        dfs = [data_frame[data_frame["gene"] == id] for data_frame in wig_oligos]

        # calculate sd-binding regions:
        region_sd = oligo.seq[58 - 25: 58]

        print(region_sd)
        # calculate sd-binding regions:
        fasta_pathname = "../data/sd_search/freescan_files_annot_ss/seq" + "_" + locus_tag
        with open(fasta_pathname + ".fasta", "w") as fasta:
            fasta.write(">Seq_SD\n")
            fasta.write(region_sd.__str__())
        sd_bits = str.encode(region_sd.__str__())

        # freescan determine delta gibbs etc:
        ASD = "auuccuccacuag"
        output_freescan = subprocess.run(["free_scan", "-e", ASD, fasta_pathname + ".fasta"], capture_output=True).stdout
        # get list with all free energies:
        list_values = output_freescan.decode("utf-8").split("\n")
        list_values = list_values[:-2]
        list_values = [float(i) for i in list_values]
        lowest_val = min(list_values)
        ind_lowest_val = list_values.index(lowest_val)
        best_binding_region = region_sd[ind_lowest_val:ind_lowest_val + len(ASD)]
        print(locus_tag, best_binding_region, lowest_val, ind_lowest_val, "original")
        with open(fasta_pathname + ".fasta", "w") as fasta:
            fasta.write(">Seq_SD\n")
            fasta.write(best_binding_region.__str__())

        shuffled_gibbs = np.array([])
        for i in range(5):
            bits_best_binding = str.encode(best_binding_region.__str__())
            shuffled = shuffle(bits_best_binding, 2).decode("utf-8")
            print(gene, shuffled, "shuffled", i)
            with open(fasta_pathname + ".fasta", "w") as fasta:
                fasta.write(">Seq_SD\n")
                fasta.write(shuffled + "\n")
            output_freescan = subprocess.run(["free_scan", "-e", "auuccuccacuag", fasta_pathname + ".fasta"], capture_output=True).stdout  # freescan
            print(output_freescan)
            # get list with all free energies:
            val = float(output_freescan.decode("utf-8")[:-2])
            newline_gibbs = [locus_tag, gene, shuffled, ind_lowest_val - 13, True, val]
            # gibbs_df = pd.concat([gibbs_df, pd.DataFrame(columns=gibbs_df.columns, data=[newline_gibbs])])
            shuffled_gibbs = np.append(shuffled_gibbs, val)
        # end freescan determine delta gibbs etc:

        less_eq = np.less_equal(shuffled_gibbs, lowest_val)
        print(lowest_val)
        pval = sum(less_eq)/len(shuffled_gibbs)
        print("pvalue for gene", gene, pval)

        newline_gibbs = [locus_tag, gene, best_binding_region.__str__(), ind_lowest_val - 13, False, lowest_val, pval]
        gibbs_df = pd.concat([gibbs_df, pd.DataFrame(columns=gibbs_df.columns, data=[newline_gibbs])])
        newline[13] = best_binding_region.__str__()
        newline[14] = pval
        newline[15] = region_sd.__str__()

        # check for known start sites (at position 15+-3):
        df_kss = [df[(df["count"] > threshold) & df.position.isin(list(range(70, 77)))] for df in dfs]
        rec_peaks = 0

        # loop through 5 dfs and check whether theres a peak per replicate:
        for df in range(len(df_kss)):
            if len(df_kss[df]) > 0:
                df_kss[df]["position"] = df_kss[df]["position"] - 58
                peak_int = df_kss[df].loc[df_kss[df]["count"].idxmax()]["count"]
                peak_pos = df_kss[df].loc[df_kss[df]["count"].idxmax()]["position"]
                newline[df+5] = peak_int
                rec_peaks += 1
        # check whether there are 2 or more peaks -> validate TIS:
        if rec_peaks > 1:
            newline[10] = True
            print(peak_int, gene, peak_pos, locus_tag)
        # check in vivo dataframes to get whether it was also detected in their dataset:
        if strand == "1":
            start_seq = start_pos - 58
            end_seq = start_pos + 184
            d = df_invivo_F[df_invivo_F.iloc[:, 0].isin(list(range(start_seq, end_seq)))]
            newcol = d.iloc[:, 0] - start_pos
            d.iloc[:, 0] = newcol
        else:
            end_seq = start_pos + 58
            start_seq = start_pos - 184
            d = df_invivo_R[df_invivo_R.iloc[:, 0].isin(list(range(start_seq, end_seq)))]
            newcol = (d.iloc[:, 0] - start_pos) * (-1)
            d.iloc[:, 0] = newcol
        # check for peaks above threshold:
        df_invivo_kss = d[(d["count"] > threshold) & d.position.isin(list(range(12, 19)))]
        if len(df_invivo_kss) > 0:
            peak_int = df_invivo_kss.loc[df_invivo_kss["count"].idxmax()]["count"]
            peak_pos = df_invivo_kss.loc[df_invivo_kss["count"].idxmax()]["position"]
            newline[11] = peak_int
            newline[12] = True

        if newline[10]:
            upstream_seq = oligo.seq[int(peak_pos) + 58 - 37:int(peak_pos) + 58].__str__()

            for i in range(len(upstream_seq)-2):
                codon = upstream_seq[i:i+3]
                found_motif = re.search(re_motiv, codon)
                if found_motif is not None:
                    fm = found_motif[0]
                    motif_df.loc[i, fm] = motif_df.loc[i, fm] + 1
        else:
            upstream_seq = oligo.seq[58 + 15 - 37: 58+15].__str__()
        newline[16] = upstream_seq
        # calculate MFE of upstream_seq:
        mfe_seq = oligo.seq[58 - 30: 58 + 15].__str__()
        cmd = " ".join(["echo", "'" + mfe_seq + "'", "|", "RNAfold", "|", "grep", "-Ev", "'A|U|G|C'", "|",
                        "sed", "-E", "'s/.* \\(([^\\)]+).*$/\\1/'"])
        ps = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        MFE = float(ps.communicate()[0].decode()[:-2].strip())
        newline[17] = MFE

        out_df = out_df.append(pd.DataFrame(columns=out_df.columns, data=[newline]))
        #out_tsv.write("\t".join(newline)+"\n")

motif_df_normalized = motif_df / motif_df.sum().sum()
motif_df.to_csv("../data/motif_df.tsv", sep="\t")
motif_df_normalized.to_csv("../data/motif_df_normalizd.tsv", sep="\t")

out_df.to_csv("../analysis/annotated_sites.csv")

print("done with first comp.")

# now we want to identify alt. TIS:
tot_length = len(out_df)
found_in_our_ds = sum(out_df["peak_found_Hör"])
found_in_other_ds = sum(out_df["peak_found_meydan"])
found_both = sum(out_df["peak_found_meydan"] & out_df["peak_found_Hör"])
start_codons = ["ATG", "GTG", "TTG", "CTG", "ATC", "ATT"]
stop_codons = ["TGA", "TAA", "TAG"]
sc_re = r"(ATG)|(GTG)|(TTG)|(CTG)|(ATC)|(ATT)"

# initiate DF for identification of novel TIS:
out_df_ss = pd.DataFrame(columns=["locus_tag", "gene", "start_position", "start_codon", "pos_from_annot_start",
                                  "in_frame", "strand", "peak", "nr_peaks_Hör", "peak_found_Hör", "peak_meydan",
                                  "peak_found_meydan", "pred_RBS_region", "p_value", "sequence_start", "upstream_seq",
                                  "MFE"])
# initiate DF for SD search:
sd_df = pd.DataFrame(columns=["location", "gene", "Gibbs"])

# add a table showing enriched start codons/sd motifs in alternative TIS:
motif_df_alt = pd.DataFrame(np.zeros((35, 8)), columns=["ATG", "GTG", "TTG", "CTG", "ATC", "ATT", "AGG", "GGA"])

# screen for TIS
for oligo in SeqIO.parse(in_oligos_fasta, "fasta"):
    if oligo.id.startswith("b"):
        id = oligo.id
        strand = id.split("_")[3]
        locus_tag = id.split("_")[0]
        gene = id.split("_")[1]
        start_pos = int(id.split("_")[2])

        dfs = [data_frame[data_frame["gene"] == id] for data_frame in wig_oligos]
        # get total read counts for relative density:
        tot_counts = [sum(data_frame["count"]) for data_frame in dfs]
        for n in range(len(dfs)):
            dfs[n]["norm_count"] = dfs[n]["count"] / sum(dfs[n]["count"])

        # check for start sites which have norm_value (rel_density) more than 0.1 and >5 cpm etc.:
        df_ss = [df[(df["count"] > threshold) & (df["norm_count"] > 0.1) &
                    df.position.isin(list(range(30, 70)) + list(range(79, 210)))] for df in dfs]
        df_ss = [df.reset_index(drop=True) for df in df_ss]
        rec_peaks = 0
        # loop through replicates, get peaks
        for df in range(len(df_ss)):
            df_ss[df]["norm_counts"] = df_ss[df]["count"] / tot_counts[df]
            if len(df_ss[df]) > 0:
                nr = 1
                # just get the peaks ,if within 5 nt there's another peak, the highest is taken:
                while nr < len(df_ss[df]):
                    row = df_ss[df].iloc[nr]
                    prev_row = df_ss[df].iloc[nr - 1]
                    if row["position"] - prev_row["position"] < 5:
                        index_del = df_ss[df].iloc[nr - 1:nr + 1]["count"].idxmin()
                        df_ss[df] = df_ss[df].drop(index_del)
                        df_ss[df].reset_index(drop=True)
                    else:
                        nr += 1
                df_ss[df] = df_ss[df].reset_index(drop=True)

        for df in range(len(df_ss)):
            for peak in range(len(df_ss[df])):
                peakpos = df_ss[df].loc[peak]["position"]
                region_sc = oligo.seq[peakpos-18:peakpos-11]
                # remove if there's no start codon in region before:
                if not any(sc in region_sc for sc in start_codons):
                    df_ss[df] = df_ss[df].drop(peak)

            df_ss[df] = df_ss[df].reset_index(drop=True)
            df_ss[df]["position"] = df_ss[df]["position"] - 58

            # write to resulting df (if theres start codon)!
            for peak in range(len(df_ss[df])):
                peakpos = df_ss[df].loc[peak]["position"]
                region_sc = oligo.seq[peakpos + 58 - 18:peakpos + 58 - 11]

                start_c = re.search(sc_re, str(region_sc))[0]
                pos_from_peak = re.search(sc_re, str(region_sc)).start() + peakpos - 18

                startsite_pos = peakpos + 58 - 18 + re.search(sc_re, str(region_sc)).start()
                region_sd = oligo.seq[startsite_pos-25: startsite_pos]
                sequence_start = oligo.seq[startsite_pos-20: startsite_pos + 10]
                upstream_seq = oligo.seq[int(peakpos) + 58 - 37:int(peakpos) + 58].__str__()
                # calculate MFE of upstream_seq:
                mfe_seq = oligo.seq[int(peakpos) + 58 - 45:int(peakpos) + 58].__str__()
                cmd = " ".join(["echo", "'" + mfe_seq + "'", "|", "RNAfold", "|", "grep", "-Ev", "'A|U|G|C'", "|",
                                "sed", "-E", "'s/.* \\(([^\\)]+).*$/\\1/'"])
                ps = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
                MFE = float(ps.communicate()[0].decode()[:-2].strip())

                # calculate sd-binding regions:
                fasta_pathname = "../data/sd_search/freescan_files_alt_ss/seq" + "_" + locus_tag + "_" + str(peakpos)
                with open(fasta_pathname + ".fasta", "w") as fasta:
                    fasta.write(">Seq_SD\n")
                    fasta.write(region_sd.__str__())
                sd_bits = str.encode(region_sd.__str__())
      #          for i in range(10):
      #              with open("../data/sd_search/freescan_shuffled_alt/seq" + "_" + locus_tag + "_" + str(
      #                      i) + ".fasta",
       #                       "w") as fasta:
       #                 fasta.write(">Seq_SD\n")
       #                 fasta.write(shuffle(sd_bits, 2).decode("utf-8"))

                # freescan determine delta gibbs etc:
                ASD = "auuccuccacuag"
                output_freescan = subprocess.run(["free_scan", "-e", ASD, fasta_pathname + ".fasta"],
                                                 capture_output=True).stdout
                # get list with all free energies:
                list_values = output_freescan.decode("utf-8").split("\n")
                list_values = list_values[:-2]
                list_values = [float(i) for i in list_values]
                lowest_val = min(list_values)
                ind_lowest_val = list_values.index(lowest_val)
                best_binding_region = region_sd[ind_lowest_val:ind_lowest_val + len(ASD)]
                print(locus_tag)
                with open(fasta_pathname + ".fasta", "w") as fasta:
                    fasta.write(">Seq_SD\n")
                    fasta.write(best_binding_region.__str__())

                shuffled_gibbs = np.array([])
                for i in range(10):
                    bits_best_binding = str.encode(best_binding_region.__str__())
                    shuffled = shuffle(bits_best_binding, 2).decode("utf-8")
                    with open(fasta_pathname + ".fasta", "w") as fasta:
                        fasta.write(">Seq_SD\n")
                        fasta.write(shuffled + "\n")
                    output_freescan = subprocess.run(["free_scan", "-e", "auuccuccacuag", fasta_pathname + ".fasta"],
                                                     capture_output=True).stdout  # freescan
                    # get list with all free energies:
                    val = float(output_freescan.decode("utf-8")[:-2])
                    newline_gibbs = [locus_tag, gene, shuffled, ind_lowest_val - 13, True, val]
                    # gibbs_df = pd.concat([gibbs_df, pd.DataFrame(columns=gibbs_df.columns, data=[newline_gibbs])])
                    shuffled_gibbs = np.append(shuffled_gibbs, val)
                # end freescan determine delta gibbs etc:

                less_eq = np.less_equal(shuffled_gibbs, lowest_val)
                pval = sum(less_eq) / len(shuffled_gibbs)

                newline_gibbs = [id, gene, best_binding_region.__str__(), ind_lowest_val - 13, False, lowest_val,
                                 pval]
                gibbs_df = pd.concat([gibbs_df, pd.DataFrame(columns=gibbs_df.columns, data=[newline_gibbs])])

                newline = [locus_tag, gene, start_pos, start_c, 0, "in_frame", strand, 0, 1, False, 0, False,
                best_binding_region.__str__(), pval, sequence_start.__str__(), upstream_seq, MFE]
                if strand == "1":
                    pos_pg = df_ss[df].loc[peak].position + start_pos
                    iv_d = df_invivo_F[df_invivo_F["position"].isin(list(range(pos_pg - 5, pos_pg + 5)))]
                else:
                    pos_pg = df_ss[df].loc[peak].position*(-1) + start_pos
                    iv_d = df_invivo_F[df_invivo_F["position"].isin(list(range(pos_pg - 5, pos_pg + 5)))]
                if any(iv_d["count"] > 5):
                    newline[11] = True
                    newline[10] = iv_d.loc[iv_d["count"].idxmax()]["count"]

                peak_int = df_ss[df].iloc[peak]["count"]
                peak_pos = df_ss[df].iloc[peak]["position"]
                newline[4] = pos_from_peak
                newline[7] = peak_int
                rec_peaks += 1
                # check in frame:
                if any(sc in oligo.seq[peak_pos+58:210] for sc in stop_codons):
                    newline[5] = "out_of_frame"
                out_df_ss = out_df_ss.append(pd.DataFrame(columns=out_df_ss.columns, data=[newline]))

                for i in range(len(upstream_seq) - 2):
                    codon = upstream_seq[i:i + 3]
                    found_motif = re.search(re_motiv, codon)
                    if found_motif is not None:
                        fm = found_motif[0]
                        motif_df_alt.loc[i, fm] = motif_df_alt.loc[i, fm] + 1

motif_df_alt.to_csv("../data/motif_df_alt.tsv", sep="\t")

out_df_ss = out_df_ss.sort_values(["locus_tag", "start_position", "pos_from_annot_start"])
out_df_ss = out_df_ss.reset_index(drop=True)

reduced_df_ss = out_df_ss.copy()
# just get the peaks:
nr = 1
# just get the peaks:
while nr < len(reduced_df_ss):
    row = reduced_df_ss.iloc[nr]
    prev_row = reduced_df_ss.iloc[nr - 1]
    if abs(row["pos_from_annot_start"] - prev_row["pos_from_annot_start"]) < 5 and row["locus_tag"] == prev_row["locus_tag"]:
        # drop lower one:
        index_del = reduced_df_ss.iloc[nr - 1:nr + 1]["peak"].idxmin()
        index_max = reduced_df_ss.iloc[nr - 1:nr + 1]["peak"].idxmax()
        reduced_df_ss = reduced_df_ss.drop(index_del)
        reduced_df_ss.loc[index_max, "peak_found_Hör"] = True
        reduced_df_ss.loc[index_max, "nr_peaks_Hör"] = reduced_df_ss.loc[index_max, "nr_peaks_Hör"] + 1
        print(row["locus_tag"], prev_row["locus_tag"])
        reduced_df_ss.reset_index(drop=True)
    else:
        nr += 1

reduced_df_ss = reduced_df_ss.reset_index(drop=True)

reduced_df_ss.to_csv("../analysis/alt_tis_w_threshold.csv")
#sum(reduced_df_ss["peak_found_Hör"])

#meydan_genes = pd.read_csv("../data/wigglefiles/meydan_alt_tis.csv")

chrom = get_seqdict("../data/reference_sequences/NC_000913.3.fasta")
gffs = chromosomes_parsed("../data/reference_sequences/NC_000913.3.gff", chrom)

comparison_start_sites = open("../analysis/comparison_startsites.tab", "w")

with open(in_gff) as in_handle:
    chroms = chromosomes_parsed(in_handle, chrom)
    genes = 0
    for record in chroms:
        # loop through all features (genes/sRNAs/...) for each chromosome record (record):
        for feature in record.features:
            # now loop through all sub-features (cds) of the genes:
            if feature.type == "gene":
                for cds in feature.sub_features:
                    if cds.type == "CDS":

                        print(feature.qualifiers['locus_tag'])
                        print(cds.location.strand)
                        print(record.seq[cds.location.start:cds.location.start+3])
                        genes += 1
    print(genes)

comparison_start_sites.close()

# create random sequences:
alphabet = ["A", "U", "G", "C"]

with open("../analysis/revision_06_2022/random_seqs_MFE.csv", "w") as file:
    file.write("MFE,peak_type\n")
    with open("../data/shuffled_seqs_esl.txt") as f_shuff:
        for line in f_shuff:
            print(line)
            cmd = " ".join(["echo", "'" + line + "'", "|", "RNAfold", "|", "grep", "-Ev", "'A|U|G|C'", "|",
                            "sed", "-E", "'s/.* \\(([^\\)]+).*$/\\1/'"])
            ps = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
            MFE = float(ps.communicate()[0].decode()[:-2].strip())
            file.write(str(MFE) + ",random sequence" + "\n")

    for oligo in SeqIO.parse("../data/reference_sequences/NC_000913.3.fasta", "fasta"):
        for i in range(10000):
            s = oligo.seq.__str__()
            length = len(s)
            start_rn = random.randint(0, length-46)
            ranseq = s[start_rn:start_rn + 45]

            bits_ranseq = str.encode(ranseq.__str__())
            shuffled = shuffle(bits_ranseq, 2).decode("utf-8")
            print(ranseq, shuffled, "shuffled")
            cmd = " ".join(["echo", "'" + shuffled + "'", "|", "RNAfold", "|", "grep", "-Ev", "'A|U|G|C'", "|",
                            "sed", "-E", "'s/.* \\(([^\\)]+).*$/\\1/'"])
            ps = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
            MFE = float(ps.communicate()[0].decode()[:-2].strip())
            file.write(str(MFE) + ",random sequence" + "\n")

    #for i in range(10000):
    #    sequence = ''.join(random.choice(alphabet) for i in range(45))
    #    cmd = " ".join(["echo", "'" + sequence + "'", "|", "RNAfold", "|", "grep", "-Ev", "'A|U|G|C'", "|",
    #                    "sed", "-E", "'s/.* \\(([^\\)]+).*$/\\1/'"])
    #    ps = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    #    MFE = float(ps.communicate()[0].decode()[:-2].strip())
    #    file.write(str(MFE) + ",random sequence" + "\n")


import glob
import os
import json
import numpy as np
import nltk
from nltk.stem import WordNetLemmatizer
import math
from scipy import spatial

lemmatizer = WordNetLemmatizer()
dic={}
p_index={}
term_count_in_each_doc={}
wfreq={}
doclist=[]
tfidf={}
stwl = {}

def Readstopwords():
    global stwl
    fObj = open('Stopword-List.txt',mode='r',encoding='utf-8')
    stwl = fObj.read()
    stwl = stwl.split()

    files = glob.glob('ShortStories/*')  #to read file names for doc id
    for i in files:
        p=os.path.basename(i)
        p=p.split('.')[0]
        doclist.append(int(p))
    #Creating a list of documents
    doclist.sort()

def RemovePun(text):
    return text.translate({ord(i): "" for i in "!@#$%^&*()[]{};:,./<>?\"'|——`~-=_’+”“"})

def ReadFiles():

    for x in doclist:
        obj = open("ShortStories/"+str(x)+".txt",mode='r',encoding='utf-8').read()
        fullfile = RemovePun(obj.lower())
        fullfile = fullfile.split()
        # file path names to use as doc id
        temp = []
        for word in fullfile:
            if word not in stwl:
                temp.append(lemmatizer.lemmatize(word))
        dic[x]=temp

def calpindex():
    for x in dic.keys():
        for words in dic[x]:
            if words not in wfreq:
                wfreq[words] = {}
                wfreq[words]['df'] = 1
                wfreq[words][x] = dic[x].count(words)
            elif x not in wfreq[words]:
                wfreq[words][x] = dic[x].count(words)
                wfreq[words]['df'] += 1
    p_file = open("wfreq.txt",mode='w',encoding='utf-8')
    p_file.write(json.dumps(wfreq))

def DocVec():
    for words in wfreq.keys():
        idf = math.log10(wfreq[words]['df'])/len(doclist)
        for ids in doclist:
            if ids not in tfidf.keys():
                tfidf[ids] = []
            tfid = idf * wfreq[words].get(ids,0)
            tfidf[ids].append(round(tfid,6))

    t_file = open("TFIDF.txt",mode='w',encoding='utf-8')
    t_file.write(json.dumps(tfidf))

def Readtfidf():
    global wfreq , tfidf
    try:

        T_file = open("TFIDF.txt",mode='rb')
        F_file = open("wfreq.txt",mode='rb')
        T_index = json.loads(T_file.read())
        F_index = json.loads(F_file.read())
        if len(T_index) == 0 or len(F_index) == 0:
            ReadFiles()
            calpindex()
            DocVec()
        else:
            tfidf = T_index.split()
            wfreq = F_file.split()
    except:
        ReadFiles()
        calpindex()
        DocVec()

def CosineSimilarity(queryVec,c):
    cosineVector={}
    for i in tfidf.keys():
        result = round(1 - spatial.distance.cosine(tfidf[i], queryVec),3)
        cosineVector[i]=result

    counter=1
    for i in cosineVector:
        if(cosineVector[i] > c):
            print(counter,"----> Doc",i," ---->",cosineVector[i])
            counter=counter+1
    print("Length: ", counter -1)

def Query():

    search = input("Enter the query: ")
    alpha = float(input("Enter the cutoff value : "))
    search = RemovePun(search.lower()).split()
    query = []
    queryVec = []

    for words in search:
        if words not in stwl:
            query.append(lemmatizer.lemmatize(words))
    print(query)

    for words in wfreq.keys():
        widf = query.count(words) * math.log10(wfreq[words].get('df',0))/len(doclist)
        queryVec.append(round(widf,6))
    CosineSimilarity(queryVec,alpha)

def main():
    Readstopwords()
    ReadFiles()
    calpindex()
    DocVec()
    while 1:
        Query()

main()


# A script to run a folder of malware executables.
# Takes one argument, the folder name (default: samples).
# Any messages created by this script during runtime will be output to IsRunningOutput.txt,
# but any output from the subprocesses themselves get sent directly to stdout.
# Will terminate (if still running) each program and then move them to a subfolder based on their status after an amount of time (5s).
# Each of the possible subfolders:
# running - was running after end time.
# not_running - was not running after end time.
# not_in_tasklist - process status was still running, but filename could not be found in task manager list.
# could_not_run - trying to create the process caused an exception.
# Processes that cannot be killed cannot have their executables moved (this error will be logged), and
# executables that cause some kind of error / exception when launched, but launch succesfully do NOT go in the could_not_run folder.
# There currently exists an issue where programs that create error messages (e.g. dll not found)
# will be marked as running, because they will not terminate until the messagebox is clicked.
# Exceptions caught by windows itself are not subject to this problem.
# Any program that has an open GUI window of any kind when it is being killed will be marked
# in the output log, but this is not possible in cygwin or other environments where windows modules cannot be inported.

# print function for python2
from __future__ import print_function
import os, re, subprocess, sys, threading, time
# windows specific imports, if they are missing try pip install pypiwin32, they will NOT work with cygwin
try:
  import win32api, win32con, win32gui, win32process
  # prevent windows from creating dialog box errors in subprocess
  win32api.SetErrorMode(win32con.SEM_FAILCRITICALERRORS | win32con.SEM_NOOPENFILEERRORBOX | win32con.SEM_NOGPFAULTERRORBOX)
  win32_imported = True
except ImportError:
  win32_imported = False
  print("Could not import win32 modules.")

exe_path = "samples"
if len(sys.argv) == 2:
  exe_path = sys.argv[1]
elif len(sys.argv) > 2:
  print("Invalid arguments.")
  sys.exit(1)
if not os.path.exists(exe_path):
  print("Folder " + exe_path + " does not exist")
  sys.exit(1)
# in s
wait_time = 5
running_dir = os.path.join(exe_path, "running")
not_running_dir = os.path.join(exe_path, "not_running")
not_in_tasklist_dir = os.path.join(exe_path, "not_in_tasklist")
could_not_run_dir = os.path.join(exe_path, "could_not_run")
# this includes the main thread, the max number of spawned processes is max_threads - 1
max_threads = 256
# only files that contain any in this array will be run
exe_filter = ['.exe']
# outfile
print_out = "IsRunningOutput.txt"
# log file handle
log = None

def create_dir(dir_name):
  if not os.path.exists(dir_name):
    os.makedirs(dir_name)

def is_proc_in_tasklist(name):
  # windows only, name can be a pid or an exe name (including .exe extension)
  # may get false positives with pids from the session # column
  if re.search(r"\s" + re.escape(str(name)) + r"\s", os.popen("tasklist /NH").read()):
    return True
  return False

def main_thread_wait():
  # this won't necessarily get the oldest thread, but getting any thread is good enough
  threads = threading.enumerate()
  if len(threads) > 1:
    # wait until thread completes
    if threads[0] is not threading.current_thread():
      threads[0].join()
    else:
      threads[1].join()

def get_open_windows():
  # for each open window, return a tuple containing the starting process, title, and handle
  if not win32_imported:
    return []
  values = []
  def foreach_window(hwnd, param):
    try:
      win_text = win32gui.GetWindowText(hwnd)
      # window visible and can accept input
      if win32gui.IsWindowVisible(hwnd) and win32gui.IsWindowEnabled(hwnd):
        # get the process id (not thread id) and return a handle to it
        p = win32api.OpenProcess(win32con.PROCESS_VM_READ | win32con.PROCESS_QUERY_INFORMATION, False, win32process.GetWindowThreadProcessId(hwnd)[1])
        if p != 0:
          # get the exe that started the process
          values.append((str(os.path.basename(win32process.GetModuleFileNameEx(p, 0))), win_text, hwnd))
          win32api.CloseHandle(p)
    except win32api.error:
      # probably a permissions error
      pass
  win32gui.EnumWindows(foreach_window, 0)
  return values

def move_if_running(proc, file_name):
  # kills if still running and moves to running dir, else moves to not running dir
  # if proc still running, but not in task manager, pur in not in tasklist dir
  old_path = os.path.join(exe_path, file_name)
  # if proc.poll() is None, this means it is still running
  if proc.poll() is None:
    proc_in_tasklist = is_proc_in_tasklist(file_name)
    windows = [win for win in get_open_windows() if win[0] == file_name]
    if windows:
      print(str(windows), file=log, end='\r\n')
    try:
      proc.kill()
    except OSError:
      print("Warning: could not kill " + file_name + ", likely because the process completed.", file=log, end='\r\n')
    try:
      if proc_in_tasklist:
        os.rename(old_path, os.path.join(running_dir, file_name))
      else:
        os.rename(old_path, os.path.join(not_in_tasklist_dir, file_name))
    except OSError:
      print("Warning: could not move " + file_name + ". Proc was running, found in tasklist: " + str(proc_in_tasklist) + ".", file=log, end='\r\n')
  else:
    try:
      os.rename(old_path, os.path.join(not_running_dir, file_name))
    except OSError:
      print("Warning: could not move " + file_name + ". Proc was not running.", file=log, end='\r\n')

with open(print_out, 'w') as log:
  create_dir(running_dir)
  create_dir(not_running_dir)
  create_dir(not_in_tasklist_dir)
  create_dir(could_not_run_dir)
  # get all files, ignore directories
  files = [f for f in os.listdir(exe_path) if any(s in f for s in exe_filter) and
    os.path.isfile(os.path.join(exe_path, f))]
  for f in files:
    # handle when the max thread count has been reached
    while threading.active_count() >= max_threads:
      main_thread_wait()
    try:
      # spawn a new process
      file_path = os.path.join(exe_path, f)
      proc = subprocess.Popen([file_path, "ARG1", "ARG2", "ARG3", "ARG4"])
      # create the thread that checks on the process after a delay
      t = threading.Timer(wait_time, move_if_running, (proc, f))
      t.daemon = False
      t.start()
    except OSError:
      os.rename(file_path, os.path.join(could_not_run_dir, f))
  # wait until all threads complete
  while threading.active_count() > 1:
    main_thread_wait()


#!/usr/bin/python3
# -*- coding: utf-8 -*-
# @Time    : 2020/7/18 14:04
# @Author  : lirixiang
# @Email   : 565539277@qq.com
# @File    : 63-GetMedian.py
"""
如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有
数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数
值排序之后中间两个数的平均值。
我们使用Insert()方法读取数据流，使用GetMedian()方法获取当前读取数据的中位数。
"""
# -*- coding:utf-8 -*-
# -*- coding:utf-8 -*-
class Solution:
    def __init__(self):
        self.numList = []

    def Insert(self, num):
        # write code here

        self.numList.append(num)
        self.numList.sort()

    def GetMedian(self,numList):
        # write code here
        length = len(self.numList)
        if length & 1 == 1:
            mid = (length ) // 2
            midNum = self.numList[mid]
            return midNum
        else:
            mid1 = length // 2 - 1
            mid2 = length // 2
            midNum = (self.numList[mid1] + self.numList[mid2]) / 2.0
            return midNum


from pe._constants import Operator
from pe._definition import Definition

def disarm(defn: Definition) -> Definition:
    """Remove semantic effects like actions, binds, or captures."""
    op = defn.op
    if op in {Operator.RUL, Operator.BND, Operator.CAP}:
        return disarm(defn.args[0])
    elif op.type == 'Primary':
        return defn
    elif op == Operator.DBG:
        subdef = disarm(defn.args[0])
        if subdef.op == Operator.DBG:
            return subdef  # collapse debug nodes
        return Definition(op, (subdef,))
    elif op.is_unary():
        args = defn.args
        return Definition(op, (disarm(args[0]), *args[1:]))
    else:
        args = defn.args
        return Definition(
            op,
            ([disarm(arg) for arg in args[0]], *args[1:])
        )