'''Operations on finite state automata '''

import argparse
from typing import Optional, Tuple
from finitestateautomata.libfsa import Automaton
from finitestateautomata.libregex import RegEx,RegExTerm
from finitestateautomata.libltl import LTLFormula
import sys
from finitestateautomata.utils.utils import printStates, printSetOfStates
from finitestateautomata.utils.operations import Operations, AutomataOperations, LTLOperations, OperationDescriptions, RegExOperations, OP_FSA_ACCEPTS, OP_FSA_ALPHABET, OP_FSA_AS_DFA, OP_FSA_AS_DFA_NAME, OP_FSA_AS_REGEX, OP_FSA_AS_REGEX_NAME, OP_FSA_COMPLEMENT, OP_FSA_COMPLEMENT_NAME, OP_FSA_COMPLETE, OP_FSA_COMPLETE_NAME, OP_FSA_ELIMINATE_EPSILON, OP_FSA_ELIMINATE_EPSILON_NAME, OP_FSA_IS_DETERMINISTIC, OP_FSA_LANGUAGE_EMPTY, OP_FSA_LANGUAGE_EMPTY_BUCHI, OP_FSA_LANGUAGE_INCLUDED, OP_FSA_MINIMIZE, OP_FSA_MINIMIZE_BUCHI, OP_FSA_MINIMIZE_BUCHI_NAME, OP_FSA_MINIMIZE_NAME, OP_FSA_PRODUCT, OP_FSA_PRODUCT_BUCHI, OP_FSA_PRODUCT_BUCHI_NAME, OP_FSA_PRODUCT_NAME, OP_FSA_PRODUCT_STRICT_BUCHI_NAME, OP_FSA_REACHABLE_STATES, OP_FSA_RELABEL, OP_FSA_RELABEL_NAME, OP_FSA_STRICT_PRODUCT, OP_FSA_STRICT_PRODUCT_BUCHI, OP_FSA_STRICT_PRODUCT_NAME, OP_LTL_CONVERT_LTL, OP_LTL_CONVERT_LTL_NAME, OP_REGEX_CONVERT_FSA, OP_REGEX_CONVERT_OMEGA_REGEX

def main():

    # optional help flag explaining usage of each individual operation
    parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False)
    parser.add_argument('-oh', '--operationhelp', dest='opHelp', nargs="?", const=" ")
    options, remainder = parser.parse_known_args() # Only use options of parser above

    if options.opHelp: # Check if -oh has been called
        if options.opHelp not in Operations:
            print("Operation '{}' does not exist. List of operations:\n\t- {}".format(options.opHelp, "\n\t- ".join(Operations)))
        else:
            print("{}: {}".format(options.opHelp, OperationDescriptions[Operations.index(options.opHelp)]))
        exit(1)

    parser = argparse.ArgumentParser(
        description='Perform operations on finite state automata.\nhttps://computationalmodeling.info')
    parser.add_argument('automaton_or_regularexpression', help="the automaton or regular expression to analyze")
    parser.add_argument('-sa', '--secondaryautomaton', dest='secondaryAutomaton',
                        help="a secondary automaton for the operation")
    parser.add_argument('-op', '--operation', dest='operation',
                        help="the operation or analysis to perform, one of : {}.\nUse 'finitestateautomata -oh OPERATION' for information about the specific operation.".format("; \n".join(Operations)))
    parser.add_argument('-oa', '--outputautomaton', dest='outputAutomaton',
                        help="the outputfile to write output automata to")
    parser.add_argument('-re', '--regularexpression', dest='regularExpression',
                        help="the regular expression to use")
    parser.add_argument('-ltl', '--ltlformula', dest='ltlFormula',
                        help="the LTL formula to use")
    parser.add_argument('-w', '--word', dest='inputtrace',
                        help="an input word, a comma-separated list of symbols")

    args = parser.parse_args()

    if args.operation not in Operations:
        sys.stderr.write("Unknown operation: {}\n".format(args.operation))
        sys.stderr.write("Operation should be one of: {}.\n".format(", ".join(Operations)))
        exit(1)

    dsl = None
    if args.automaton_or_regularexpression:
        try:
            with open(args.automaton_or_regularexpression, 'r') as fsaFile:
                dsl = fsaFile.read()
        except FileNotFoundError as e:
            sys.stderr.write("File does not exist: {}.\n".format(args.automaton_or_regularexpression))
            exit(1)

    try:
        process(args, dsl)
    except Exception as e:
        sys.stderr.write("{}\n".format(e))
        # raise e
        exit(1)

    exit(0)

def requireSecondaryAutomaton(args)->Tuple[str,Automaton]:
    if args.secondaryAutomaton is None:
        raise Exception("A secondary automaton must be specified.")
    try:
        with open(args.secondaryAutomaton, 'r') as fsaFile:
            saDsl = fsaFile.read()
        nameB, B = Automaton.fromDSL(saDsl)
        return nameB, B
    except FileNotFoundError as e:
        raise Exception("Failed to read file: {}.\n".format(args.secondaryAutomaton))

def requireAutomaton(A:Optional[Automaton])->Automaton:
    if A is None:
        raise Exception("Automaton is not defined")
    return A

def requireRegEx(R: Optional[RegExTerm])->RegExTerm:
    if R is None:
        raise Exception("Regular expression is not defined")
    return R

def requireLTLFormula(F:Optional[LTLFormula])->LTLFormula:
    if F is None:
        raise Exception("LTL Formula is not defined")
    return F

def process(args, dsl):

    A = None
    R = None
    F = None
    name = None

    if args.operation in AutomataOperations:
        name, A = Automaton.fromDSL(dsl)

    if args.operation in RegExOperations:
        name, R = RegEx.fromDSL(dsl)

    if args.operation in LTLOperations:
        name, F = LTLFormula.fromDSL(dsl)

    if args.operation not in Operations:
        print("Unknown operation or no operation provided")
        print("Operation should be one of: {}.".format(", ".join(Operations)))
        exit(1)

    # accepts (requires input word)
    if args.operation == OP_FSA_ACCEPTS:
        res, path = requireAutomaton(A).acceptsWithPath(args.inputtrace)
        print(res)
        if res and path is not None:
            printStates(path)

    # isDeterministic
    if args.operation == OP_FSA_IS_DETERMINISTIC:
        print(requireAutomaton(A).isDeterministic())

    # asDFA
    if args.operation == OP_FSA_AS_DFA:
        res = requireAutomaton(A).asDFA()
        print(res.asDSL(OP_FSA_AS_DFA_NAME(name)))

    # eliminateEpsilon
    if args.operation == OP_FSA_ELIMINATE_EPSILON:
        res = requireAutomaton(A).eliminateEpsilonTransitions()
        print(res.asDSL(OP_FSA_ELIMINATE_EPSILON_NAME(name)))

    # alphabet
    if args.operation == OP_FSA_ALPHABET:
        print(requireAutomaton(A).alphabet())

    # complete
    if args.operation == OP_FSA_COMPLETE:
        res = requireAutomaton(A).complete()
        print(res.asDSL(OP_FSA_COMPLETE_NAME(name)))

    # complement
    if args.operation == OP_FSA_COMPLEMENT:
        res = requireAutomaton(A).complement()
        print(res.asDSL(OP_FSA_COMPLEMENT_NAME(name)))

    # product (requires secondary automaton)
    if args.operation == OP_FSA_PRODUCT:
        nameB, B = requireSecondaryAutomaton(args)
        res = requireAutomaton(A).product(B)
        print(res.asDSL(OP_FSA_PRODUCT_NAME(name, nameB)))

    # strictProduct (requires secondary automaton)
    if args.operation == OP_FSA_STRICT_PRODUCT:
        nameB, B = requireSecondaryAutomaton(args)
        res = requireAutomaton(A).strictProduct(B)
        print(res.asDSL(OP_FSA_STRICT_PRODUCT_NAME(name, nameB)))

    # productBuchi (requires secondary automaton)
    if args.operation == OP_FSA_PRODUCT_BUCHI:
        nameB, B = requireSecondaryAutomaton(args)
        res = requireAutomaton(A).asRegularBuchiAutomaton().productBuchi(B.asRegularBuchiAutomaton())
        print(res.asDSL(OP_FSA_PRODUCT_BUCHI_NAME(name, nameB)))

    # strictProductBuchi (requires secondary automaton)
    if args.operation == OP_FSA_STRICT_PRODUCT_BUCHI:
        nameB, B = requireSecondaryAutomaton(args)
        res = requireAutomaton(A).strictProductBuchi(B)
        print(res.asDSL(OP_FSA_PRODUCT_STRICT_BUCHI_NAME(name, nameB)))

    # languageEmpty
    if args.operation == OP_FSA_LANGUAGE_EMPTY:
        empty, word, path = requireAutomaton(A).languageEmpty()
        print(empty)
        if not empty:
            print(word)
            if path is not None:
                printStates(path)

    # languageEmptyBuchi
    if args.operation == OP_FSA_LANGUAGE_EMPTY_BUCHI:
        RA, stateMap = requireAutomaton(A).asRegularBuchiAutomatonWithStateMap()
        empty, wordprefix, wordrepeat, pathprefix, pathrepeat = RA.languageEmptyBuchi()
        print(empty)
        if not empty:
            print(wordprefix)
            print(wordrepeat)
            if pathprefix is not None and pathrepeat is not None:
                printStates(list(map(lambda s: stateMap[s], pathprefix)))
                printStates(list(map(lambda s: stateMap[s], pathrepeat)))

    # reachable states
    if args.operation == OP_FSA_REACHABLE_STATES:
        res = requireAutomaton(A).reachableStates()
        printSetOfStates(res)

    # languageIncluded
    if args.operation == OP_FSA_LANGUAGE_INCLUDED:
        _, B = requireSecondaryAutomaton(args)
        included, word = requireAutomaton(A).languageIncluded(B)
        print(included)
        if not included:
            print(word)

    # minimize
    if args.operation == OP_FSA_MINIMIZE:
        res = requireAutomaton(A).minimize()
        print(res.asDSL(OP_FSA_MINIMIZE_NAME(name)))

    # minimize Büchi
    if args.operation == OP_FSA_MINIMIZE_BUCHI:
        res = requireAutomaton(A).minimizeBuchi()
        print(res.asDSL(OP_FSA_MINIMIZE_BUCHI_NAME(name)))

    # relabel
    if args.operation == OP_FSA_RELABEL:
        res = requireAutomaton(A).relabelStates()
        print(res.asDSL(OP_FSA_RELABEL_NAME(name)))

    # convertRegEx (requires regular expression)
    if args.operation == OP_REGEX_CONVERT_FSA:
        res = requireRegEx(R).asFSA()
        print(res.asDSL(name))  # type: ignore

    # convertOmegaRegEx (requires omega-regular expression)
    if args.operation == OP_REGEX_CONVERT_OMEGA_REGEX:
        res = requireRegEx(R).asNBA()
        print(res.asDSL(name))  # type: ignore

    # convertLTL (requires ltl formula)
    if args.operation == OP_LTL_CONVERT_LTL:
        res = requireLTLFormula(F).asFSA()
        print(res.asDSL(OP_LTL_CONVERT_LTL_NAME(name)))

    # asRegEx
    if args.operation == OP_FSA_AS_REGEX:
        res = RegEx.fromFSA(A, OP_FSA_AS_REGEX_NAME(name))
        print(res.asDSL(OP_FSA_AS_REGEX_NAME(name)))

    # # test parser
    # if args.operation == 'testLTLParser':
    #     print(F.asDSL('name'))

if __name__ == "__main__":
    main()


#!/usr/bin/env python3
# coding: utf-8

import numpy as np
import pencil as pc
from scipy.optimize import curve_fit
from helperfuncs import *
import matplotlib as mpl
mpl.use('pgf')
import sys
from os import listdir, path
from mpl_toolkits.axes_grid1 import Grid
import argparse
from os import mkdir, path, listdir
from os.path import join, isdir, isfile

parser = argparse.ArgumentParser()
parser.add_argument('-v', '--verbose', action='store_true', default=False, help='add verbosity')
parser.add_argument('ddir', default='', type=str, help='datadir directory')
parser.add_argument('choice', default='', type=str, choices=('taylor', 'energy'), help='choice of plotting')

args = parser.parse_args()
figdir = 'figures'

dstart = args.ddir.split('_')[0]
dirs = [l for l in listdir('.') if l.startswith(dstart) and isdir(l)]
try:
    dirs = sorted(dirs, key=lambda a: float(a.split('_')[-1]))
except ValueError:
    dirs = sorted(dirs, key=lambda a: float(a.split('_')[-1][1:]))
dirs.append('helical')
if args.verbose: print(dirs)

#if args.choice == 'energy':
#    print('energy plot not yet implemented, choose taylor instead')
#    exit(1)

def figsize(scale, ratio=None):
    #fig_width_pt = 418.25555
    fig_width_pt = 523.5307              # Get this from LaTeX using \the\textwidth
    inches_per_pt = 1.0/72.27                       # Convert pt to inch
    fig_width = fig_width_pt*inches_per_pt*scale    # width in inches
    if not ratio:
        golden_mean = (np.sqrt(5.0)-1.0)/2.0   #golden_mean
        ratio = golden_mean
    fig_height = fig_width*ratio             # height in inches
    fig_size = [fig_width,fig_height]
    return fig_size

# I make my own newfig and savefig functions
def newfig(width):
#    plt.clf()
    fig = plt.figure(figsize=figsize(width, ratio=0.75))
    ax = fig.add_subplot(111)
    return fig, ax

def savefig(fig, filename):
    fig.savefig('{}.pgf'.format(filename))
    fig.savefig('{}.pdf'.format(filename))

def make_tsplot(fig, ax):
    if args.choice == 'taylor':
        try:
            plotarr = 2*np.pi*tser.brms/tser.jrms
        except AttributeError:
            return
    elif args.choice == 'energy':
        plotarr = tser.brms**2
    ax.plot(tser.t, plotarr, label=dd.split('_')[-1], linewidth=1.5,
            color=plt.cm.Paired(next(clrindx)))
    ti = search_indx(tser.t, 1., eps=.05)
    ti2 = search_indx(tser.t, 0.8*tser.t[-1], eps=.05)
    if args.verbose: print('index from which fitting begins: ',ti)
    if ti is not None:
        po1, pco1 = curve_fit(lambda x, a, b: powerlaw(x, a, b, x0=0), tser.t[ti:], plotarr[ti:])
        #ax.plot(tser.t[ti:ti2], powerlaw(tser.t[ti:ti2],1.2*po1[0],po1[1]), linewidth=1.5,
        #        label=r'fit: $B\sim t^{%.1f}$' % po1[1], color=(0.72, 0.61, 0.46, 1.00))
        if args.verbose:
            print(dd, 'L ~ t^{1:.2f}'.format(*po1))

def setup_tsplot(ax):
    ax.set_yscale('log')
    ax.set_xscale('log')
    ax.set_xlabel('time $t$')
    ax.set_ylabel(r'$L_{\textrm{T}}$')

pgf_with_latex = {                      # setup matplotlib to use latex for output
    "pgf.texsystem": "pdflatex",        # change this if using xetex or lautex
    "text.usetex": True,                # use LaTeX to write all tex
    "font.family": "serif",
    "font.serif": [],                   # blank entries should cause plots to inherit fonts from the document
    "font.sans-serif": [],
    "font.monospace": [],
    "font.size": 10,
    "axes.labelsize": 10,               # LaTeX default is 10pt font.
    "legend.fontsize": 8,               # Make the legend/label fonts a little smaller
    "xtick.labelsize": 8,
    "ytick.labelsize": 8,
    "figure.figsize": figsize(0.95),     # default fig size of 0.9 textwidth
    "pgf.preamble": [
        r"\usepackage[utf8x]{inputenc}",    # use utf8 fonts becasue your computer can handle it :)
        r"\usepackage[T1]{fontenc}",        # plots will be generated using this preamble
        ]
    }
mpl.rcParams.update(pgf_with_latex)

import matplotlib.pyplot as plt

# Setup Figure and Axes

fig, ax  = newfig(0.45)
setup_tsplot(ax)#, xlim=(.1,tser.t[-1]))
tmax = 0.
clrindx = iter(np.linspace(0,1,len(dirs)))

for dd in dirs:
    if args.verbose: print(dd)
    tser = pc.read_ts(datadir=dd, quiet=not args.verbose)
    make_tsplot(fig, ax)
    if tser.t[-1] > tmax:
        ax.set_xlim(.1, tser.t[-1])
    ax.legend(loc='upper left', frameon=False)

filename = dstart + '_taylor'
fig.tight_layout(pad=0.3)
if args.verbose:
    print(filename)
savefig(fig, path.join(figdir, filename))
print('SUCCESS')


# -*- coding: utf-8 -*-
"""
Created on Wed Mar  2 12:30:05 2022

@author: yyh
"""
import numpy as np
import pandas as pd
import os
import h5py
region = np.load('open_region.npy')  #Import candidate region files (open chromatin regions or TSS regions)
name = pd.read_excel("sample_ID.xlsx",sheet_name='Sheet1',usecols=[0])  #Sample ID of the dataset
name_array = np.array(name)
name_list = name_array.tolist()
n = len(name_list)
End = np.zeros([561414,n])  #Create a matrix to store the results

for i in range(1,n+1):
    print(i)
    nam = name_list[i-1]
    nam2 = "".join(nam)
    path = os.path.join("dataset_path",nam2,'data_n')   #Full path to the splice sample file
    pos = 0

    for j in range(1,23):  #Treat each chromosome sequentially
        j_str = str(j)
        mat = ''.join(['chr',j_str,'.mat'])
        mat2 = os.path.join(path,mat)
        mat3 = h5py.File(mat2,'r')
        #print(mat3.keys())
        read = np.transpose(mat3['read'])   #Extract the reads of the sample
        num = read.shape[0]

        region_len = mat3['region_len']  #Total length of chromosome extracted
        size = int(region_len[0,0])
        chrom = np.zeros([size,1])

        reg = region[region[:,0]==j]
        num2 = reg.shape[0]

        for k in range(1,num+1):
            end = int(read[k-1,0]+read[k-1,1]-1) #Counting the position and number of fragment ends
            chrom[end-1,0] += 1

        for m in range (1,num2+1):  #Counting the total number of fragment ends in each region
            End[pos+m-1,i-1] = sum(chrom[(reg[m-1,1]-1):reg[m-1,2],0])

        pos += num2

np.save('end.npy',End)


from typing import Optional
import torch
import torch.nn as nn

class MaskedLanguageModeling(nn.Module):
    def __init__(self, bert: nn.Module, voc_size:int=30000):
        super(MaskedLanguageModeling, self).__init__()
        self.bert = bert
        d_model = bert.emb.tok_emb.weight.size(1)
        self.linear = nn.Linear(d_model, voc_size)

    def forward(self, input: torch.Tensor, seg: torch.Tensor) -> torch.Tensor:
        '''
        param:
            input: a batch of sequences of words
            seg: Segmentation embedding for input tokens
        dim:
            input:
                input: [B, S]
                seg: [B, S]
            output:
                result: [B, S, V]
        '''
        output = self.bert(input, seg) # [B, S, D_model]
        output = self.linear(output) # [B, S, voc_size]

        return output # [B, S, voc_size]

class NextSentencePrediction(nn.Module):
    def __init__(self, bert: nn.Module):
        super(NextSentencePrediction, self).__init__()
        self.bert = bert
        d_model = bert.emb.tok_emb.weight.size(1)
        self.linear = nn.Linear(d_model, 2)

    def forward(self, input: torch.Tensor, seg: torch.Tensor) -> torch.Tensor:
        '''
        param:
            input: a batch of sequences of words
            seg: Segmentation embedding for input tokens
        dim:
            input:
                input: [B, S]
                seg: [B, S]
            output:
                result: [B, S, V]
        '''
        output = self.bert(input, seg) # [B, S, D_model]
        output = self.linear(output) # [B, S, 2]

        return output[:, 0, :] # [B, 2]

class BertModel(nn.Module):
    def __init__(self, voc_size:int=30000, seq_len: int=512, d_model: int=768, d_ff:int=3072, pad_idx: int=1,
                num_encoder: int=12, num_heads: int=12, dropout: float=0.1):
        super(BertModel, self).__init__()
        self.pad_idx = pad_idx
        self.emb = BERTEmbedding(seq_len, voc_size, d_model, dropout)
        self.encoders = Encoders(seq_len, d_model, d_ff, num_encoder, num_heads, dropout)

    def forward(self, input: torch.Tensor, seg: torch.Tensor) -> torch.Tensor:
        '''
        param:
            input: a batch of sequences of words
        dim:
            input:
                input: [B, S]
            output:
                result: [B, S, V]
        '''
        pad_mask = get_attn_pad_mask(input, input, self.pad_idx)

        emb = self.emb(input, seg) # [B, S, D_model]
        output = self.encoders(emb, pad_mask) # [B, S, D_model]

        return output # [B, S, D_model]

class Encoders(nn.Module):
    def __init__(self, seq_len: int, d_model: int, d_ff:int, num_encoder: int, num_heads: int, dropout: float):
        super(Encoders, self).__init__()
        self.models = nn.ModuleList([Encoder(d_model, d_ff, seq_len, num_heads, dropout=dropout) for i in range(num_encoder)])

    def forward(self, tokens, enc_mask):
        # input: [B, S, D_model]
        for model in self.models:
            tokens = model(tokens, enc_mask)

        return tokens

class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, seq_len: int, dropout: int=0.1):
        super(PositionalEncoding, self).__init__()
        self.seq_len = seq_len
        self.dropout = nn.Dropout(dropout)
        self.emb = nn.Embedding(seq_len, d_model)

    def forward(self, x: torch.Tensor):
        # x: [B, S]. x is tokens
        pos = torch.arange(self.seq_len, dtype=torch.long, device=x.device) # [S]
        pos = pos.unsqueeze(0).expand(x.size()) # [1, S] -> [B, S]
        pos_emb = self.emb(pos)
        return self.dropout(pos_emb) # [B, S, D_model]

class BERTEmbedding(nn.Module):
    """
    Embeddings for BERT.
    It includes segmentation embedding, token embedding and positional embedding.
    I add dropout for every embedding layer just like the original transformer.
    """
    def __init__(self, seq_len: int=512, voc_size: int=30000, d_model: int=768, dropout: float=0.1) -> None:
        super(BERTEmbedding, self).__init__()
        self.tok_emb = nn.Embedding(num_embeddings=voc_size, embedding_dim=d_model)
        self.tok_dropout = nn.Dropout(dropout)
        self.seg_emb = nn.Embedding(2, d_model)
        self.seg_dropout = nn.Dropout(dropout)
        self.pos_emb = PositionalEncoding(d_model, seq_len, dropout)

    def forward(self, tokens: torch.Tensor, seg: torch.Tensor):
        """
        tokens: [B, S]
        seg: [B, S]. seg is binary tensor. 0 indicates that the corresponding token for its index belongs sentence A
        """
        tok_emb = self.tok_emb(tokens) # [B, S, d_model]
        seg_emb = self.seg_emb(seg) # [B, S, d_model]
        pos_emb = self.pos_emb(tokens) # [B, S, d_model]

        return self.tok_dropout(tok_emb) + self.seg_dropout(seg_emb) + pos_emb  # [B, S, d_model]

class Encoder(nn.Module):
    def __init__(self, d_model: int, d_ff:int, seq_len: int, num_enc_heads: int, dropout: float):
        super(Encoder, self).__init__()

        self.attn = MultiHeadAttn(d_model, num_enc_heads)
        self.layer_norm1 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.fc = PositionWiseFC(d_model, d_ff)
        self.layer_norm2 = nn.LayerNorm(d_model)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, pad_mask):
        # x+SubLayer(LayerNorm(x))
        layer_norm_x = self.layer_norm1(x)
        attn_score = self.attn(layer_norm_x, layer_norm_x, layer_norm_x, pad_mask) + x
        attn_score = self.dropout1(attn_score)

        # x+SubLayer(LayerNorm(x))
        attn_score = self.fc(self.layer_norm2(attn_score)) + attn_score
        attn_score = self.dropout2(attn_score)

        return attn_score # [B, S, D_model]

class MultiHeadAttn(nn.Module):
    def __init__(self, d_model: int, num_heads: int, attn_type: str='dot', dropout: float=0.1):
        super(MultiHeadAttn, self).__init__()

        self.attn_type = attn_type.lower()
        assert attn_type in ['dot', 'additive'], 'attn_type not supported: attn_type. Please select one of [dot, addictive]'

        assert d_model % num_heads == 0, 'the num_heads doesn\'t match'

        self.num_heads = num_heads
        self.d_k = int(d_model / num_heads)

        # Implements weights via `nn.Linear` instead `nn.Parameters()`
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.softmax = nn.Softmax(dim=-1)

        if self.attn_type == 'additive':
            self.fc = nn.Linear(2*self.d_k, self.d_k)
            self.v = nn.Parameter(torch.rand(1, self.d_k))

        self.W_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, k, v, q, mask):
        # k, v, q: [B, S, D_model]
        # mask: [B, S, S]
        attn_score = self.scaled_dot_product_attn(q, k, v, mask) # [B, S, D_model]
        attn_score = self.W_o(attn_score) # # [B, S, D_model]

        return attn_score

    def scaled_dot_product_attn(self, k, v, q, mask):
        # k, q, v: [B, S, D_model]
        # mask: [B, S, D_model]
        # FIXME: Attention에서도 사용할 수 있게 re-use 가능한가? ->
        batch_size, seq_len, _ = k.size()
        q = self.W_q(q).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) # [B, h, S, d_k]
        k = self.W_k(k).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) # [B, h, S, d_k]
        v = self.W_v(v).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) # [B, h, S, d_k]

        if self.attn_type == 'dot':
            qk = self.dot_attn(q, k)
        else:
            qk = self.additive_attn(q, k)

        # masking
        if mask != None:
            mask = mask.unsqueeze(1).repeat(1, self.num_heads, 1, 1) # [B, h, S, S]
            qk.masked_fill_(mask == True, -1e9)

        attn_score = self.softmax(qk) # [B, h, S, S]
        attn_score = attn_score.matmul(v) # [B, h, S, d_k]
        attn_score = attn_score.transpose(1, 2).reshape(batch_size, seq_len, self.d_k * self.num_heads) # [B, S, d_model]
        attn_score = self.dropout(attn_score)

        return attn_score

    def dot_attn(self, q, k):
        # $Attention(Q, K, V) = softmax(QK^T/\sqrt(d_k))V
        # q, k: # [B, h, S, d_k]
        qk = q.matmul(k.transpose(2, 3)) / torch.sqrt(torch.tensor(self.d_k, dtype=float)) # [B, h, S, S]
        return qk

    def additive_attn(self, q, k):
        # q, k: # [B, h, S, d_k]
        # Neural Machine Translation by Jointly Learning to Align and Translate
        # a(q, k) = w^T * tanh(W*[q:k])
        qk = torch.cat([q, k], dim=-1) # [B, h, S, 2*d_k]
        qk = self.fc(qk) # [B, h, S, d_k]
        qk = qk.permute(0, 1, 3, 2) # [B, d_k, S, d_k]
        v = self.v.reshape(1, 1, 1, -1) # [1, self.d_k] -> [1, 1, 1 self.d_k]
        v = v.repeat(q.size(0), -1, self.num_heads, 1).transpose(1, 2) # [B, h, S, self.d_k]
        qk = v.matul(qk) # [B, h, S, S]

        return qk

class PositionWiseFC(nn.Module):
    def __init__(self, d_model: int, d_ff: int, dropout: float=0.1):
        super(PositionWiseFC, self).__init__()

        self.fc1 = nn.Linear(d_model, d_ff)
        self.gelu = nn.GELU()
        self.fc2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # [B, S, D]
        x = self.fc1(x) # [B, S, D_ff]
        x = self.gelu(x) # [B, S, D_ff]
        x = self.fc2(x) # [B, S, D]
        x = self.dropout(x)

        return x

# reference: https://github.com/graykode/nlp-tutorial/blob/master/5-1.Transformer/Transformer(Greedy_decoder)_Torch.ipynb
def get_attn_pad_mask(seq_q, seq_k, pad_idx):
    batch_size, len_q = seq_q.size()
    batch_size, len_k = seq_k.size()
    pad_attn_mask = seq_k.data.eq(pad_idx).unsqueeze(1)  # [B, 1, S]
    return pad_attn_mask.expand(batch_size, len_q, len_k)  # [B, S_q, S_k]


# -*- coding:utf-8 -*-
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torchvision.transforms as transforms
from data.cifar import CIFAR10, CIFAR100
from data.mnist import MNIST
from data.datasets import input_dataset
from models import *
import argparse, sys
import numpy as np
import datetime
import shutil
from random import sample
from loss import loss_cross_entropy,loss_cores,f_beta

parser = argparse.ArgumentParser()
parser.add_argument('--lr', type = float, default = 0.05)
parser.add_argument('--lr_plan', type = str, help = 'base, cyclic', default = 'cyclic')
parser.add_argument('--loss', type = str, help = 'ce, cores', default = 'cores')
parser.add_argument('--result_dir', type = str, help = 'dir to save result txt files', default = 'results')
parser.add_argument('--noise_rate', type = float, help = 'corruption rate, should be less than 1', default = 0.2)
parser.add_argument('--noise_type', type = str, help='[pairflip, symmetric,instance]', default='pairflip')
parser.add_argument('--top_bn', action='store_true')
parser.add_argument('--ideal', action='store_true')
parser.add_argument('--dataset', type = str, help = 'mnist, cifar10, or cifar100', default = 'cifar10')
parser.add_argument('--model', type = str, help = 'cnn,resnet', default = 'cnn')
parser.add_argument('--n_epoch', type=int, default=100)
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--print_freq', type=int, default=50)
parser.add_argument('--num_workers', type=int, default=4, help='how many subprocesses to use for data loading')
parser.add_argument('--resume', '-r', action='store_true', help='resume from checkpoint')

def get_noise_pred(loss_div, args, epoch=-1, alpha=0.):
    #Get noise prediction
    print('DEBUG, loss_div', loss_div.shape)
    llast = loss_div[:, epoch]
    idx_last = np.where(llast>alpha)[0]
    print('last idx:', idx_last.shape)
    return idx_last

# Adjust learning rate and for SGD Optimizer
def adjust_learning_rate(optimizer, epoch,alpha_plan):
    for param_group in optimizer.param_groups:
        param_group['lr']=alpha_plan[epoch]/(1+f_beta(epoch))

def accuracy(logit, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    output = F.softmax(logit, dim=1)
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res

# Train the Model
def train(epoch, num_classes, train_loader,model, optimizer,loss_all,loss_div_all,loss_type, noise_prior = None):
    train_total=0
    train_correct=0
    print(f'current beta is {f_beta(epoch)}')
    v_list = np.zeros(num_training_samples)
    idx_each_class_noisy = [[] for i in range(num_classes)]
    if not isinstance(noise_prior, torch.Tensor):
        noise_prior = torch.tensor(noise_prior.astype('float32')).cuda().unsqueeze(0)
    for i, (images, labels, indexes) in enumerate(train_loader):
        ind=indexes.cpu().numpy().transpose()
        batch_size = len(ind)
        class_list = range(num_classes)

        images = Variable(images).cuda()
        labels = Variable(labels).cuda()

        # Forward + Backward + Optimize
        logits = model(images)
        prec, _ = accuracy(logits, labels, topk=(1, 5))
        train_total+=1
        train_correct+=prec
        if loss_type=='ce':
            loss = loss_cross_entropy(epoch,logits, labels,class_list,ind, noise_or_not, loss_all, loss_div_all)
        elif loss_type=='cores':
            loss, loss_v = loss_cores(epoch,logits, labels,class_list,ind, noise_or_not, loss_all, loss_div_all, noise_prior = noise_prior)
            v_list[ind] = loss_v
            for i in range(batch_size):
                if loss_v[i] == 0:
                    idx_each_class_noisy[labels[i]].append(ind[i])
        else:
            print('loss type not supported')
            raise SystemExit
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if (i+1) % args.print_freq == 0:
            print ('Epoch [%d/%d], Iter [%d/%d] Training Accuracy: %.4F, Loss: %.4f'
                  %(epoch+1, args.n_epoch, i+1, len(train_dataset)//batch_size, prec, loss.data))

    class_size_noisy = [len(idx_each_class_noisy[i]) for i in range(num_classes)]
    noise_prior_delta = np.array(class_size_noisy)
    print(noise_prior_delta)

    train_acc=float(train_correct)/float(train_total)
    return train_acc, noise_prior_delta

# Evaluate the Model
def evaluate(test_loader,model,save=False,epoch=0,best_acc_=0,args=None):
    model.eval()    # Change model to 'eval' mode.
    print('previous_best', best_acc_)
    correct = 0
    total = 0
    for images, labels, _ in test_loader:
        images = Variable(images).cuda()
        logits = model(images)
        outputs = F.softmax(logits, dim=1)
        _, pred = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (pred.cpu() == labels).sum()
    acc = 100*float(correct)/float(total)

    if save:
        if acc > best_acc_:
            state = {'state_dict': model.state_dict(),
                     'epoch':epoch,
                     'acc':acc,
            }
            torch.save(state,os.path.join(save_dir,args.loss + args.noise_type + str(args.noise_rate)+'best.pth.tar'))
            #np.save(save_dir + '/' + args.loss + args.noise_type + str(args.noise_rate)+'loss_div_all_best.npy',loss_div_all)
            #np.save(save_dir + '/' + args.loss + args.noise_type + str(args.noise_rate)+'loss_all_best.npy',loss_all)
            best_acc_ = acc
        if epoch == args.n_epoch -1:
            state = {'state_dict': model.state_dict(),
                     'epoch':epoch,
                     'acc':acc,
            }
            torch.save(state,os.path.join(save_dir,args.loss + args.noise_type + str(args.noise_rate)+'last.pth.tar'))
            #np.save(save_dir + '/' + args.loss + args.noise_type + str(args.noise_rate)+'loss_div_all_last.npy',loss_div_all)
            #np.save(save_dir + '/' + args.loss + args.noise_type + str(args.noise_rate)+'loss_all_best.npy',loss_all)
    return acc, best_acc_

#####################################main code ################################################
args = parser.parse_args()
# Seed
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)

# Hyper Parameters
batch_size = 64
learning_rate = args.lr

# load dataset
train_dataset,test_dataset,num_classes,num_training_samples = input_dataset(args.dataset,args.noise_type,args.noise_rate)
noise_prior = train_dataset.noise_prior
noise_or_not = train_dataset.noise_or_not
print('train_labels:', len(train_dataset.train_labels), train_dataset.train_labels[:10])
# load model
print('building model...')
if args.model == 'cnn':
    model = CNN(input_channel=3, n_outputs=num_classes)
else:
    model = ResNet34(num_classes)
print('building model done')
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
# Creat loss and loss_div for each sample at each epoch
loss_all = np.zeros((num_training_samples,args.n_epoch))
loss_div_all = np.zeros((num_training_samples,args.n_epoch))
### save result and model checkpoint #######
save_dir = args.result_dir +'/' +args.dataset + '/' + args.model
if not os.path.exists(save_dir):
    os.system('mkdir -p %s' % save_dir)
train_loader = torch.utils.data.DataLoader(dataset = train_dataset,
                                   batch_size = batch_size,
                                   num_workers=args.num_workers,
                                   shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                  batch_size = 64,
                                  num_workers=args.num_workers,
                                  shuffle=False)
alpha_plan = [0.1] * 50 + [0.01] * 50
#alpha_plan = []
#for ii in range(args.n_epoch):
#    alpha_plan.append(learning_rate*pow(0.95,ii))
model.cuda()
txtfile=save_dir + '/' +  args.loss + args.noise_type + str(args.noise_rate) + '.txt'
if os.path.exists(txtfile):
    os.system('rm %s' % txtfile)
with open(txtfile, "a") as myfile:
    myfile.write('epoch: train_acc test_acc \n')

epoch=0
train_acc = 0
best_acc_ = 0.0
#print(best_acc_)
# training
noise_prior_cur = noise_prior
for epoch in range(args.n_epoch):
# train models
    adjust_learning_rate(optimizer, epoch, alpha_plan)
    model.train()
    train_acc, noise_prior_delta = train(epoch,num_classes,train_loader, model, optimizer,loss_all,loss_div_all,args.loss,noise_prior = noise_prior_cur)
    noise_prior_cur = noise_prior*num_training_samples - noise_prior_delta
    noise_prior_cur = noise_prior_cur/sum(noise_prior_cur)
# evaluate models
    test_acc, best_acc_ = evaluate(test_loader=test_loader, save=True, model=model,epoch=epoch,best_acc_=best_acc_,args=args)
# save results
    #det_by_loss  = det_acc(save_dir,best_ratio,args,loss_all,noise_or_not,epoch,sum_epoch = False)
    #det_by_loss_div = det_acc(save_dir,best_ratio,args,loss_div_all,noise_or_not,epoch,sum_epoch = False)
    print('train acc on train images is ', train_acc)
    print('test acc on test images is ', test_acc)
    #print('precision of labels by loss is', det_by_loss)
    #print('precision of labels by loss div is', det_by_loss_div)
    with open(txtfile, "a") as myfile:
        myfile.write(str(int(epoch)) + ': '  + str(train_acc) +' ' + str(test_acc) + "\n")
    np.save(save_dir + '/' + args.loss + args.noise_type + str(args.noise_rate)+'loss_all.npy',loss_all)
    np.save(save_dir + '/' + args.loss + args.noise_type + str(args.noise_rate)+'loss_div_all.npy',loss_div_all)
    np.save(save_dir + '/' + args.loss + args.noise_type + str(args.noise_rate)+'noise_or_not.npy',noise_or_not)
    np.save(save_dir + '/' + args.loss + args.noise_type + str(args.noise_rate)+'train_noisy_labels.npy',train_dataset.train_noisy_labels)
    if epoch ==40:
        idx_last = get_noise_pred(loss_div_all, args, epoch=epoch)
        np.save(save_dir + '/' + args.loss + args.noise_type + str(args.noise_rate)+'_noise_pred.npy',idx_last)


from flaskr import config
from flask import jsonify, render_template, g, request
from flaskr.home import blueprint

def parse_post_data(device, form):
    def update_name(name):
        if name != device.data.name or name != device.data.capabilities['name']:
            device.data.name = name
            device.data.capabilities['name'] = name
            device.data.last_properties['name'] = name
            device.data.set_name(name)
            return True
        return False

    def update_effect(effect):
        if effect != device.data.effect:
            device.data.effect = effect
            return True
        return False

    def update_duration(duration):
        if duration != device.data.duration:
            device.data.duration = duration
            return True
        return False

    def update_color_flow(color_flow):
        if color_flow != device.data.color_flow_mode:
            device.data.color_flow_mode = color_flow
            return True
        return False

    data = dict(
        name=update_name,
        effect=update_effect,
        duration=update_duration,
        color_flow=update_color_flow
    )

    should_update = False
    for arg, func in data:
        should_update = arg in form and func(form[arg])

    if should_update:
        config.set_device(device)

@blueprint.route('/device/<string:device>/', methods=['GET', 'POST'])
@blueprint.route('/device/<string:device>', methods=['GET', 'POST'])
def device_update(device, **post_data):
    if request.method == 'POST':
        parse_post_data(config.devices[device], **post_data)

@blueprint.route('/device/<string:dev_id>/toggle')
def toggle(dev_id):
    device = config.devices[dev_id]
    state = device.capabilities['power']
    if device.toggle() == 'ok':
        device.capabilities['power'] = 'off' if state == 'on' else 'on'
    return jsonify(original_state=state, new_state=device.capabilities['power'])


import os
import multiprocessing as mp
import sys
import time
import subprocess
import io
import logging
import argparse
from pathlib import Path
import csv
import pandas as pd
import gzip
from functools import partial
import json

from matplotlib.backends.backend_pdf import PdfPages
from matplotlib import (pyplot as plt)
from scipy.stats import entropy
import matplotlib
matplotlib.use('Agg')

def define_parser():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '-i', '--input',
        dest='input',
        type=Path,
        default='None',
        help="Input path"
        )

    parser.add_argument(
        '-o', '--output',
        dest='output',
        type=Path,
        default='./',
        help="Output path"
        )

    parser.add_argument(
        '-t', '--threads',
        dest='threads',
        type=int,
        default=1,
        help="Number of threads"
        )

    parser.add_argument(
        '-c', '--ref_clean',
        dest='ref_clean',
        type=str,
        default='',
        help="references_to_clean"
        )

    parser.add_argument(
        '-r', '--viral_ref',
        dest='viral_ref',
        type=str,
        default='',
        help="viral reference"
        )

    return parser

def verify_path(args, arg_name):
    if arg_name not in args:
        raise ValueError(f"{arg_name} not provided.")
    _path = args.__dict__[arg_name]
    if not _path.exists():
        raise FileNotFoundError(f"{arg_name}: {_path} doesn't exist")
    return _path

#%% GLOBAL

class MP():

    def __init__(self, num_threads):
        self.num_threads = num_threads

    def mltiprocess_command(self, func,list_paths,arguments):
        sema = mp.Semaphore(self.num_threads)
        pss = []
        for path in list_paths:
            p = mp.Process(target=func, args=(arguments,path,sema))
            p.start()
            pss.append(p)

        for p in pss:
            p.join()

    def set_threads(self,t):
        self.num_threads = t

#%% FASTQ FILTERING

FILTER_FASTQ = 'bwa mem -v 1%(p_param)s -a -t 20 %(genome)s %(fastq)s %(fastq2)s | samtools view -b -F 2 - | samtools fastq - > %(path)s/%(name)s.fastq'

def filter_fastq_command(arguments,fastq,sema):
    former_key,path,genome = arguments
    sema.acquire()
    name = os.path.basename(fastq)
    if former_key:
        name = name.replace('.fastq','')
        fastq2 = ''
        p_param = ' -p'
    else:
        name = name.replace('_R1_001.fastq.gz','')
        fastq2 = fastq.replace('_R1','_R2')
        p_param = ''

    print(name)
    print('bwa mem -v 1%(p_param)s -a -t 20 %(genome)s %(fastq)s %(fastq2)s | samtools view -b -F 2 - | samtools fastq - > %(path)s/%(name)s.fastq' % dict(name=name,p_param=p_param,path=path,genome=genome,fastq=fastq,fastq2=fastq2))
    subprocess.call(FILTER_FASTQ % dict(name=name,p_param=p_param,path=path,genome=genome,fastq=fastq,fastq2=fastq2), shell=True)
    sema.release()

def filter_fastq(args,MP_obj, offset=30):
    path_list = [os.path.join(r, f) for r, d, fs in os.walk(args.input) for f in fs if f.endswith('_R1_001.fastq.gz') and 'singelton' not in f]
    path_new_dir = None

    references_to_clean = json.loads(args.ref_clean)
    former_key = None
    for key in list(references_to_clean.keys()):
        dir_name = 'fastq/no' + key.split('_')[0]
        path_new_dir = os.path.join(args.output,dir_name)
        os.makedirs(path_new_dir,exist_ok=True)

        if former_key:
            path_list = [os.path.join(former_key, file) for file in os.listdir(former_key)]

        if 'bacteria' in references_to_clean.get(key):
            MP_obj.set_threads(1)
        MP_obj.mltiprocess_command(filter_fastq_command,path_list,[former_key,path_new_dir,references_to_clean.get(key)])
        if MP_obj.num_threads != args.threads:
            MP_obj.set_threads(args.threads)

        former_key = path_new_dir
    return(path_new_dir)

#%% FASTQ METASPADES SCAFFOLDS

SCAFFOLD = 'metaspades.py --12 %(fastq)s --threads 10 -o %(assembly_dir)s/%(name)s'

def scaffold_command(assembly_dir,fastq,sema):
    sema.acquire()
    name = os.path.basename(fastq)
    name = name.replace('.fastq','')
    print(name)
    subprocess.call(SCAFFOLD % dict(name=name,fastq=fastq,assembly_dir=assembly_dir), shell=True)
    sema.release()

def scaffold(filtered_fastq,out,MP_obj):
    filtered_fastq = json.loads(filtered_fastq)
    filtered_fastq_dir = list(filtered_fastq)[-1]
    dir_name = 'fastq/no' + filtered_fastq_dir.split('_')[0]
    filtered_fastq_dir = os.path.join(out,dir_name)
    path_list = [os.path.join(r, f) for r, d, fs in os.walk(filtered_fastq_dir) for f in fs]
    assembly_path = os.path.join(out,'assembly')
    os.makedirs(assembly_path, exist_ok=True)
    MP_obj.mltiprocess_command(scaffold_command,path_list,assembly_path)
    return (assembly_path)

#%% MAPPING AND REPORT the contigs

# MAP_AND_REPORT = '''bwa mem -v 1 -t 10 -a %(genome)s %(assembly_dirs)s/scaffolds.fasta \
#  | samtools view -b - | tee %(contig_path)s/%(name)s.bam | samtools view -b -F 4 - \
#      | tee %(contig_path)s/%(name)s.mapped.bam | samtools view - | cut -f3 | sort | uniq -c \
#          | awk -v i=$i '{print i$1"\t"$2}' > %(report_file)s '''
MAP_AND_REPORT = '''bwa mem -v 1 -t 10 -a %(genome)s %(assembly_dirs)s/scaffolds.fasta \
  | samtools view -b - | tee %(contig_path)s/%(name)s.bam | samtools view -b -F 4 - \
      | tee %(contig_path)s/%(name)s.mapped.bam | samtools view -q 1 -F 260 - | cut -f3 | sort | uniq -c \
          | awk -v i=$i '{print i$1"\t"$2}' > %(report_file)s'''

def map_and_report_command(arguments,assembly_dirs,sema):
    sema.acquire()
    report_path,contig_path,genome = arguments
    name = os.path.basename(assembly_dirs)
    # IN REPORT FILE: 1. how many alignments does the BAM file contain. 2. sequence of virus
    report_file = name + '.txt'
    report_file = open(os.path.join(report_path,report_file),"w+")
    report_file = report_path + '/' + name + '.txt'
    subprocess.call(MAP_AND_REPORT % dict(name=name,genome=genome,report_file=report_file,contig_path=contig_path,assembly_dirs=assembly_dirs), shell=True)
    sema.release()

def map_and_report(out,viral_ref,MP_obj):
    assembly_dir = os.path.join(out,'assembly')
    path_list = [os.path.join(assembly_dir,dir) for dir in os.listdir(assembly_dir)]
    print(path_list)
    contig_path = os.path.join(out,'BAM/virome_contigs')
    os.makedirs(contig_path, exist_ok=True)
    report_path = os.path.join(out,'reports')
    os.makedirs(report_path, exist_ok=True)
    MP_obj.mltiprocess_command(map_and_report_command,path_list,[report_path,contig_path,viral_ref])

#%% SORT BAM from contigs

SORT = 'samtools sort %(bam)s -o %(output_bam)s'

def sort_bam_command(noparam,bam,sema):
    sema.acquire()
    output_bam = bam.replace('.bam','.sorted.bam')
    print(output_bam)
    subprocess.call(SORT % dict(bam=bam,output_bam=output_bam), shell=True)
    sema.release()

def sort_bam(out,MP_obj):
    BAM_DIR = os.path.join(out,'BAM/virome_contigs')
    path_list = [os.path.join(r, f) for r, d, fs in os.walk(BAM_DIR) for f in fs if f.endswith('.mapped.bam')]
    MP_obj.mltiprocess_command(sort_bam_command,path_list,'')

#%% DEPTH

DEPTH = 'cut -f2 %(reports)s > %(refs_dir)s/%(sample)s.txt; \
    samtools faidx %(genome)s -r %(refs_dir)s/%(sample)s.txt -o %(refs_dir)s/%(sample)s.fasta; \
        bwa index %(refs_dir)s/%(sample)s.fasta; \
            bwa mem -a -t 10 -v 1 -a %(refs_dir)s/%(sample)s.fasta %(fastq)s %(fastq2)s \
                | samtools view -b - | samtools sort - | tee %(bam_dir)s/%(sample)s.mapped.sorted.bam \
                    | samtools depth -a - > %(depth_dir)s/%(sample)s.depth;'# \
                        #python /mnt/data3/code/bard/automation/depth_allviruses.py %(depth_dir)s/%(sample)s.depth %(depth_dir)s/%(sample)s.pdf'
"""
depth graph of all viruses refseq (all viruses identified from contigs)
vs. original raw sample
"""
def depth_allviruses(depth_dir,sample,genome):
    genome = os.path.dirname(genome)
    genome = os.path.join(genome,'sequences.csv')
    annotations = pd.read_csv(genome)
    depth_file = depth_dir + '/' + sample + '.depth'
    pdf_file = depth_file.replace('.depth', '.pdf')
    file = pd.read_csv(depth_file, header=None, delimiter='\t', names=['acc', 'pos', 'depth'])
    file = file.groupby('acc')
    viruses = [file.get_group(x) for x in file.groups]
    all_accs = sorted(viruses, key=lambda x: entropy(x['depth']), reverse=True)
    with PdfPages(pdf_file) as pdf:
        for df in all_accs:
            depth_col = df['depth']
            enrpy = entropy(depth_col)

            size = len(depth_col)
            coverage = float(sum([x for x in depth_col if x > 0]))/size
            # add annotations! Host, species!

            species = annotations.loc[annotations['Accession'] == df['acc'].iloc[0].split('.')[0], 'Species'].item()
            host = annotations.loc[annotations['Accession'] == df['acc'].iloc[0].split('.')[0], 'Host'].item()

            plt.plot(range(size), depth_col, 'ok', markersize=0.5, rasterized=True)
            plt.xlabel('position')
            plt.ylabel('depth')
            plt.title(f"Species:{species}   Host:{host}")
            plt.suptitle(f"Entropy = {enrpy}    Acc:{df['acc'].iloc[0]}")
            pdf.savefig()
            plt.close()

def depth_from_source_command(arguments,reports,sema):
    sema.acquire()
    fastq_dir,refs_dir,depth_dir,bam_dir,genome = arguments
    sample = os.path.basename(reports)
    sample = sample.replace('.txt','')
    fastq = sample + '_R1_001.fastq.gz'
    fastq = os.path.join(fastq_dir,fastq)
    fastq2 = fastq.replace('_R1','_R2')
    #genome = references_to_clean.get('viral_ref')
    print(sample)
    subprocess.call(DEPTH % dict(sample=sample,genome=genome,refs_dir=refs_dir,depth_dir=depth_dir,fastq=fastq,fastq2=fastq2,reports=reports,bam_dir=bam_dir), shell=True)
    depth_allviruses(depth_dir,sample,genome)
    sema.release()

def depth_from_source(args,MP_obj):
    refs_path = os.path.join(args.output,'refs/scaffoldsVirome')
    os.makedirs(refs_path, exist_ok=True)
    depth_path = os.path.join(args.output,'depth')
    os.makedirs(depth_path, exist_ok=True)
    bam_path = os.path.join(args.output,'BAM/virome_fastq')
    os.makedirs(bam_path, exist_ok=True)

    report_dir = os.path.join(args.output,'reports')
    path_list = [os.path.join(r, f) for r, d, fs in os.walk(report_dir) for f in fs if f.endswith('.txt')]
    MP_obj.mltiprocess_command(depth_from_source_command,path_list,[args.input,refs_path,depth_path,bam_path,args.viral_ref])

#%% INDEX BAM OF FASTQ (and contig)

INDEX = '''samtools index %(bam)s'''
def index_bam_command(no,bam_fastq,sema):
    sema.acquire()
    print(bam_fastq)
    bam_contig = bam_fastq.replace('virome_fastq','virome_contigs')
    subprocess.call(INDEX % dict(bam=bam_fastq), shell=True)
    subprocess.call(INDEX % dict(bam=bam_contig), shell=True)
    sema.release()

def index_bam(out,MP_obj):

    BAM_DIR = os.path.join(out,'BAM/virome_fastq')
    path_list = [os.path.join(r, f) for r, d, fs in os.walk(BAM_DIR) for f in fs if f.endswith('.mapped.sorted.bam')]
    MP_obj.mltiprocess_command(index_bam_command,path_list,'')

#%% REPORT FROM FASTQ BAM

COVERAGE_STATS = '''samtools coverage -H -r %(virus)s %(out)s/BAM/virome_fastq/%(sample)s.mapped.sorted.bam;'''
CONTIG_NUM = '''samtools view -q 1 -c -F 260 %(out)s/BAM/virome_contigs/%(sample)s.mapped.sorted.bam %(virus)s'''
LARGEST_CONTIG = '''samtools stats %(out)s/BAM/virome_contigs/%(sample)s.mapped.sorted.bam %(virus)s | grep ^RL | cut -f 2-'''

def depthEntropy(out_file):
    with open(out_file, 'r') as depth_file:
        depth_dist = [int(x.split()[2]) for x in depth_file]
        entrpy = entropy(depth_dist)
        return entrpy

def report_command(arguments,depth,sema):
    sema.acquire()
    out,breadth_dir,genome = arguments
    sample = os.path.basename(depth)
    sample = sample.replace('.depth','')

    df = pd.DataFrame(columns = ["virusAcc","startpos","length","numreadsmapped","covbases","coverage","meandepth","meanbaseq","meanmapq","numcontigsmapped","longestcontig","entropy"])

    depth_table = pd.read_table(depth,header=None)

    viruses = depth_table.iloc[:,0]
    viruses = sorted(set(viruses))
    for virus in viruses:
        out_file = os.path.join(breadth_dir,sample+'tmp.depth')
        depth_table[depth_table[0] == virus].to_csv(out_file,sep='\t', index=False, header=False)
        entropy = depthEntropy(out_file)
        coverage_stats = subprocess.check_output(COVERAGE_STATS % dict(virus=virus,out=out,sample=sample), shell=True).strip().decode("utf-8")
        contig_num = subprocess.check_output(CONTIG_NUM % dict(virus=virus,out=out,sample=sample), shell=True).strip().decode("utf-8")
        largest_contig = subprocess.check_output(LARGEST_CONTIG % dict(virus=virus,out=out,sample=sample), shell=True).strip().decode("utf-8")
        largest_contig = largest_contig.split('\n')[-1].split('\t')[0]
        new_row = coverage_stats + '\t' + str(contig_num) + '\t' + str(largest_contig) + '\t' + str(entropy)
        new_row = new_row.split('\t')
        df.loc[len(df)] = new_row

    df['Accession'] = df['virusAcc'].str.split('.').str.get(0)
    genome = os.path.dirname(genome)
    genome = os.path.join(genome,'sequences.csv')
    annotations = pd.read_csv(genome)
    annotations = annotations[['Accession','Species','Genus','Family','Host']]
    df = pd.merge(df, annotations, how="left", on='Accession')
    df = df.drop(['startpos','Accession'],axis=1)
    os.remove(os.path.join(breadth_dir,sample+'tmp.depth'))
    out_file = os.path.join(breadth_dir,sample+'.csv')
    df.to_csv(out_file,index=False)
    sema.release()

def report(out,viral_ref,MP_obj):
    breadth_path = os.path.join(out,'breadths')
    os.makedirs(breadth_path, exist_ok=True)

    depth_DIR = os.path.join(out,'depth')
    path_list = [os.path.join(r, f) for r, d, fs in os.walk(depth_DIR) for f in fs if f.endswith('.depth')]
    MP_obj.mltiprocess_command(report_command,path_list,[out,breadth_path,viral_ref])
    #report_command([out,breadth_path,viral_ref],path_list[0],1)

#%% COUNT READS IN FILTERED FASTQ

def count_reads_command(fastq,out,references_to_clean):

    #fastq,out = arguments
    print(out)
    fastq2 = fastq.replace('_R1','_R2')
    name = os.path.basename(fastq).replace('_R1_001.fastq.gz','')
    col_df = ['sample','totalReads']
    # TOTAL NUM OF READS
    tot_reads1 = sum(1 for line in gzip.open(fastq)) / 4
    tot_reads2 = sum(1 for line in gzip.open(fastq2)) / 4
    former_reads = int(tot_reads1)+int(tot_reads2)
    tot_reads = [name,former_reads]
    for key in list(references_to_clean.keys()):
        dir_name = 'no' + key.replace('_ref','')
        fastq_path = os.path.join(out,'fastq',dir_name)
        fastq = [os.path.join(r, f) for r, d, fs in os.walk(fastq_path) for f in fs if f.startswith(name)]
        no_reads = sum(1 for line in open(fastq[0])) / 4
        tot_reads.append(former_reads - no_reads)
        former_reads = no_reads
        col_df.append(key.replace('_ref','')+'Reads')

    df = pd.DataFrame(columns = col_df)

    df.loc[len(df)] = tot_reads

    return(df)

def count_reads(args):
    path_list = [os.path.join(r, f) for r, d, fs in os.walk(args.input) for f in fs if f.endswith('_R1_001.fastq.gz') and 'singelton' not in f]

    references_to_clean = json.loads(args.ref_clean)
    with mp.Pool(args.threads) as pool:
        dfs = pool.map(partial(count_reads_command, out=args.output, references_to_clean=references_to_clean), path_list)
    #dfs = map(partial(count_reads_command, out=args.output), path_list)

    merged_samples = pd.concat(dfs, axis=0) \
        .reset_index()
    merged_samples = merged_samples.drop(columns='index')
    out_path = args.output / ('read_count.csv')
    merged_samples.to_csv(out_path,index=False)

#%% Main
if __name__ == "__main__":
    DEBUG = False
    parser = define_parser()
    if DEBUG:
        inline = ['-i','/home/orz/virom_new_ref/raw_fastq/', '-o', '/home/orz/virom_new_ref/', '-c', '{"HG":"/mnt/data3/code/bard/references/Human_Hg38/hg38.fa","Bacteria":"/mnt/data3/code/bard/references/bacteria/bacteria_fullGenomes.fna.gz"}', '-r',"/mnt/data3/code/bard/references/virome/refseq_new_bwa/viral.genomic.all.fasta" ]
        args = parser.parse_args(inline)

    else:
        args = parser.parse_args()

    MP_obj = MP(args.threads)
    filter_fastq(args,MP_obj)
    scaffold(args.ref_clean,args.output,MP_obj)
    map_and_report(args.output,args.viral_ref,MP_obj)
    sort_bam(args.output,MP_obj)
    depth_from_source(args,MP_obj)
    index_bam(args.output,MP_obj)
    report(args.output,args.viral_ref,MP_obj)
    count_reads(args)


import sys
import itertools

# reference: https://cseweb.ucsd.edu/classes/wi05/cse206b/notes/L3.pdf
# https://github.com/erexhepa/Rosalind/blob/master/042_REAR.py
#
# Problem
# A reversal of a permutation creates a new permutation by inverting some interval of the permutation; (5,2,3,1,4), (5,3,4,1,2), and (4,1,2,3,5) are all reversals of (5,3,2,1,4). The reversal distance between two permutations π and σ, written drev(π,σ), is the minimum number of reversals required to transform π into σ (this assumes that π and σ have the same length).
#
# Given: A collection of at most 5 pairs of permutations, all of which have length 10.
#
# Return: The reversal distance between each permutation pair.
#
# Sample Dataset
# 1 2 3 4 5 6 7 8 9 10
# 3 1 5 2 7 4 9 6 10 8
#
# 3 10 8 2 5 4 7 1 6 9
# 5 2 3 1 7 4 10 8 6 9
#
# 8 6 7 9 4 1 3 10 2 5
# 8 2 7 6 9 1 5 3 10 4
#
# 3 9 10 4 1 8 6 7 5 2
# 2 9 8 5 1 7 3 4 6 10
#
# 1 2 3 4 5 6 7 8 9 10
# 1 2 3 4 5 6 7 8 9 10
# Sample Output
# 9 4 5 7 0
#
# def read_file(filename):
#     with open(filename) as file:
#         content = file.readlines()
#         content = [numbers.replace("\n","").split(" ") for numbers in content]
#         pairs = zip(content[::3], content[1::3])
#     return pairs
#
# def arrange(pair):
#     counter = 0
#     new_string = pair[0]
#     for i in range(len(pair[1])):
#         target = new_string.index(pair[1][i])
#         if pair[1][i] != new_string[i]:
#             counter += 1
#             new_string = new_string[0:i] + new_string[i:target+1][::-1] + new_string[target+1:]
#     return counter
#
# def SimpleReversalSort(filename):
#     ans = []
#     tup_lis = read_file(filename)
#     for pairs in tup_lis:
#         print(pairs)
#         ans.append(arrange(pairs))
#     return ans

def file_read(filename):
    """return a list of pair_in_list, [[[1,2],[3,4]] , [[5,6],[7,8]]]"""
    # different from the reference source which generate map object. (Map object can only iterate over once)
    with open(filename) as file:
        content = [pair.split("\n") for pair in file.read().strip().split("\n\n")]
        for index,pair in enumerate(content):
            content[index] = [list(map(int, single.split())) for single in pair]
    return content

def breakpoint_count(number_list):
    """Adding 0 and n+1 to the number list.
    Sum of breakpoint, adjecent numbers abs(n - n+1)
    """
    number_list = [0] + list(number_list) + [len(number_list)+1]
    return sum(map(lambda x, y: abs(x - y) != 1, number_list[1:], number_list[:-1]))

def breakpoint_index(number_list):
    """Return index of breakpoints in number_list"""
    number_list = [0] + list(number_list) + [len(number_list) + 1]
    ans = itertools.compress(range(len(number_list)-1), map(lambda x, y: abs(x - y) != 1, number_list[1:], number_list[:-1]))
    return list(ans)

def better_greedy_breakpoint(perm1, perm2):
    to_identity = {value: i + 1 for i, value in enumerate(perm2)}
    normalized_perm = [to_identity[value] for value in perm1]

    # Quick lambda function to reverse a region in the permutation.
    rev_perm = lambda perm, i, j: perm[:i] + perm[i:j + 1][::-1] + perm[j + 1:]

    # Initialize Variables
    normalized_perm = tuple(normalized_perm)
    current_perms = [normalized_perm]
    min_breaks = breakpoint_count(normalized_perm)
    dist = 0

    # Run the greedy BFS breakpoint reduction sorting.
    while min_breaks>0:
        new_perms = []
        dist += 1
        # Iterate over all combinations of breakpoint indices for all  current minimal permutations.
        for perm in current_perms:
            for rev_ind in itertools.product(breakpoint_index(perm), repeat=2):
                # Store some temporary variables for the given iteration.
                temp_perm = tuple(rev_perm(perm, rev_ind[0], rev_ind[1] - 1))
                temp_breaks = breakpoint_count(temp_perm)

                # Done we have no breakpoints.
                if temp_breaks == 0:
                    return dist

                # Create a new dictionary and update the minimum number of breakpoints if we've found a reduction.
                elif temp_breaks < min_breaks:
                    min_breaks = temp_breaks
                    new_perms = [temp_perm]

                # Add to the dictionary if the current breakpoints match the minimum number.
                elif temp_breaks == min_breaks:
                    new_perms.append(temp_perm)

        current_perms = new_perms

    else:
        return 0

content = file_read(sys.argv[1])

min_dists = [str(min(better_greedy_breakpoint(p1, p2),better_greedy_breakpoint(p2,p1))) for p1, p2 in content]
print(*min_dists)


#!/usr/bin/env python

# set the default Django settings module
import os
import django
os.environ.setdefault('DJANGO_SETTINGS_MODULE', "bibThings.settings")
# Setup django project
django.setup()

import codecs
import re

from django.db import models
from django.db.models import Q

from django.conf import settings
from bibBot.models import Verse

bookName=[]
chapterNo, verseNoBegin, verseNoEnd = 0, 0, 0

def get_bookstat (book_name):
    for book_info in settings.BOOK_STAT:
        if book_info[0] == book_name:
            return book_info
    return False

# Obsolete definition
"""
def Overse_exists (book_name, chapter_number, verse_number):
    book_info = get_bookstat(book_name)

    if (chapter_number <= book_info[1]) and (verse_number <= book_info[chapter_number + 1]):
        return True
    else:
        print("Check Chapter & Verse number [{} {}:{}]".format(book_name, str(chapter_number), str(verse_number)))
        return False
"""

def verses_exist (book_name, chapter_number, verse_number1, verse_number2):
    book_info = get_bookstat(book_name)

    total_verses = book_info[chapter_number + 1]
    if (chapter_number <= book_info[1]) and (verse_number1 < verse_number2) and \
            (verse_number1 <= total_verses) and (verse_number2 <= total_verses):
        return True
    else:
        print("Check Chapter & Verse number [{} {}:{}-{}]".format(book_name, str(chapter_number), str(verse_number1),
                                                                  str(verse_number2)))
        return False

# Obsolete definition
"""
def Ochapter_exists (book_name, chapter_number):
    book_info = get_bookstat(book_name)
    if chapter_number > book_info[1]:
        print("Check Chapter number: " + str(chapter_number))
        return False
    else:
        return True
"""

# Matching a single verse
def match_verse (match):

    initial_bookName = match.group('bookName').lower()
    #print(initial_bookName)
    #print(Verse.book_exists("", settings.VERSION, initial_bookName))

    book_found = False
    for book_aliases in settings.BN_ALIASES:
        if initial_bookName in book_aliases:
            EN_bookName = book_aliases[0]
            KR_bookName = book_aliases[2]
            book_found = True
            break

    if book_found:

        chapterNo = int(match.group('chapterNo'))
        verseNoBegin = int(match.group('verseNoBegin'))

        if Verse.verse_exists("", settings.VERSION, EN_bookName, chapterNo, verseNoBegin):
            verse = Verse.get_verse("", settings.VERSION, EN_bookName, chapterNo, verseNoBegin)
            print("[{} {}:{}] {}".format(KR_bookName, str(chapterNo), str(verseNoBegin), verse[0].verseText))

        else:
            print("Wrong chapter/verse number:{} {}:{}".format(KR_bookName, str(chapterNo), str(verseNoBegin)))

    else:
        print("Wrong Book name: " + initial_bookName)

# Matching multiple contiguous verses
def match_verses (match):

    initial_bookName = match.group('bookName').lower()

    book_found = False
    for book_aliases in settings.BN_ALIASES:
        if initial_bookName in book_aliases:
            EN_bookName = book_aliases[0]
            KR_bookName = book_aliases[2]
            book_found = True
            break

    if book_found:
        chapterNo = int(match.group('chapterNo'))
        verseNoBegin = int(match.group('verseNoBegin'))
        verseNoEnd = int(match.group('verseNoEnd'))

        #print(KR_bookName + " " + str(chapterNo) + ":" + str(verseNoBegin) + "-" + str(verseNoEnd))
        if verses_exist(EN_bookName, chapterNo, verseNoBegin, verseNoEnd):
            verses = Verse.get_verses("", settings.VERSION, EN_bookName, chapterNo, verseNoBegin, verseNoEnd)
            if verses:
                i = verseNoBegin
                print("[{} {}:{}-{}]".format(KR_bookName, str(chapterNo), str(verseNoBegin), str(verseNoEnd)))
                for verse in verses:
                    print("[{}:{}] {}".format(str(chapterNo), str(i), verse.verseText))
                    i += 1
    else:
        print("Wrong Book name: " + initial_bookName)

# Matching q chapter
def match_chapter (match):

    initial_bookName = match.group('bookName').lower()

    book_found = False
    for book_aliases in settings.BN_ALIASES:
        if initial_bookName in book_aliases:
            EN_bookName = book_aliases[0]
            KR_bookName = book_aliases[2]
            book_found = True
            break

    if book_found:
        chapterNo = int(match.group('chapterNo'))

        if Verse.chapter_exists("", settings.VERSION, EN_bookName, chapterNo):

            verses = Verse.get_verses_bc("", settings.VERSION, EN_bookName, chapterNo)  # EN_bookName  "204 Jn"
            if verses:
                i = 1
                print("[{}{}장]".format(KR_bookName, str(chapterNo)))
                for verse in verses:
                    print("[{}:{}] {}".format(str(chapterNo), str(i), verse.verseText))
                    i += 1
            else:
                print("verses of the chapter not found")
        else:
            print("Wrong chapter number: " + str(chapterNo))
    else:
        print("Wrong Book name: " + initial_bookName)

# Matching a query text (allowing space-insensitive)

def match_words(input_q_word):
    bookName=[]
    chapterNo, verseNoBegin, verseNoEnd = 0, 0, 0

    f_path = '/Users/hosoolee/Desktop/bibleBot/BibleText/bible.txt'

    q_word = input_q_word.replace(" ", "")
    re_q_word = re.compile(q_word)

    with codecs.open(f_path, 'r', encoding='euc-kr') as f:
        line = f.readline()
        # print("0" + line)
        visit_count = 0
        count = 0

        while line:
            visit_count += 1
            match = re_q_word.search(line)  # .groupdict()
            if match:
                count += 1

                matched_verse = line.split()
                KR_bookName = matched_verse[0]
                EN_bookName = settings.KR_BN_TO_EN_BN_DICT[KR_bookName].lower()  # 나중에 소문자로 바꿀 것 !
                chapterNo = matched_verse[1]
                verseNo = matched_verse[2]

                verse = Verse.get_verse("", settings.VERSION, EN_bookName, chapterNo, verseNo)

                if verse:
                    print("[{} {}:{}] {}".format(settings.KRBN_ABBR2FULL_DICT[KR_bookName], str(chapterNo), str(verseNo), verse[0].verseText))
                else:
                    #print(verse[0] + str(chapterNo) + ":" + str(verseNo) + "Irregular verses")
                    pass

            line = f.readline()

        print("No of Matches: " + str(count))


import os
import logging
import json
import boto3
from botocore import exceptions

# Initiate logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)
logging.basicConfig(
    format="%(levelname)s %(threadName)s [%(filename)s:%(lineno)d] %(message)s",
    datefmt="%Y-%m-%d:%H:%M:%S",
    level=logging.INFO,
)
logger.info("Lambda initialization completed")

# Obtaining relevant Lambda environment variables
CROSS_ACCOUNT_ROLE_NAME = os.environ["CROSS_ACCOUNT_ROLE_NAME"]

def get_client(service, account_id):
    """
    Gets boto3 client after assuming the Config role in the target AWS account

    Parameters:
        service (str): the identifier of the AWS service to create a client for
        account_id (str): the AWS account ID to assume role in

    Returns:
        boto3 client: a boto3 client object
    """

    # Constructing role ARN from account ID and role name
    role_arn = "arn:aws:iam::" + account_id + ":role/" + CROSS_ACCOUNT_ROLE_NAME

    # Obtaining credentials by assuming role
    sts_client = boto3.client("sts")
    credentials = sts_client.assume_role(RoleArn=role_arn, RoleSessionName="configLambdaExecution")["Credentials"]

    return boto3.client(
        service,
        aws_access_key_id=credentials["AccessKeyId"],
        aws_secret_access_key=credentials["SecretAccessKey"],
        aws_session_token=credentials["SessionToken"],
    )

def get_ou_path(organizations_client, account_id):
    """
    Constructs the full OU path of the target account

    Parameters:
        organizations_client: boto3 client for AWS Organizations
        account_id (str): the target AWS account ID

    Returns:
        ou_path (str): the full OU path of the target account
    """

    ou_path = ""
    type = ""

    # Looping through all levels of the Organization until complete path os obtained
    while type != "ROOT":
        response = organizations_client.list_parents(ChildId=account_id)
        ou_id = response["Parents"][0]["Id"]
        type = response["Parents"][0]["Type"]

        if type != "ROOT":
            ou_name = organizations_client.describe_organizational_unit(OrganizationalUnitId=ou_id)[
                "OrganizationalUnit"
            ]["Name"]
            account_id = ou_id
            ou_path = "/" + ou_name + ou_path
            response = organizations_client.list_parents(ChildId=account_id)
            type = response["Parents"][0]["Type"]

    # Constructing full OU path using root OU name
    root_name = organizations_client.list_roots()["Roots"][0]["Name"]
    ou_path = "/" + root_name + ou_path

    logger.info("Starting Network Firewall remediation for account: " + account_id + " in OU: " + ou_path)

    return ou_path

def get_rule_group_mapping(ssm_client, spoke_ou_path, account_id):
    """
    Obtains the rule group mapping containing each rule group to apply to OU

    Parameters:
        ssm_client: boto3 client for AWS Systems Manager in the central security account
        spoke_ou_path (str): the full OU path of the target account
        account_id (str): the target AWS account ID

    Returns:
        mapping (dict): raw mapping of OU to rule group
    """

    # Reading SSM parameter containing rule groups to apply to the OU
    try:
        response = ssm_client.get_parameter(Name=spoke_ou_path)
    except ssm_client.exceptions.ParameterNotFound:
        mapping = "None"
        logger.info("No mapping found for OU path: " + spoke_ou_path + ". Evaluation skipped.")
    else:
        ou_baseline_mapping = json.loads(response["Parameter"]["Value"])

        # Iterating through SSM parameter and obtaining mapping for target account
        for account, rules in ou_baseline_mapping.items():
            if (account == "all_accounts" and account_id not in ou_baseline_mapping.keys()) or account == account_id:
                mapping = rules
                return mapping

    return

def get_rule_group_arns(raw_mapping):
    """
    Obtains the ARNs of each rule group to apply to each policy

    Parameters:
        raw_mapping: raw mapping of OU to rule group

    Returns:
        mapping (dict): mapping of required rule group types to rule group ARNs
    """

    processed_mapping = {}

    for rule_group_type, rule_group_list in raw_mapping.items():
        if rule_group_type not in processed_mapping:
            processed_mapping[rule_group_type.upper()] = []

        for rule_group_arn in rule_group_list:
            if rule_group_type.upper() == "STATEFUL" and type(rule_group_list) is list:
                processed_mapping[rule_group_type.upper()].append({"ResourceArn": rule_group_arn})
            else:
                processed_mapping[rule_group_type.upper()].append(
                    {
                        "ResourceArn": rule_group_arn,
                        "Priority": rule_group_list[rule_group_arn],
                    }
                )

    logger.info("Obtained the following firewall rule group mapping: " + json.dumps(processed_mapping))

    return processed_mapping

def remediate_rule_order(
    strict_order,
    policy,
    evaluation,
    result_token,
    spoke_anf_client,
    spoke_config_client,
):
    """
    Remediates non-compliant policies in the account if strict ordering is required.
    This function creates a new firewall policy with the correct rule order, and associates it with any firewalls in the account.

    Parameters:
        policy: description of AWS Network Firewall policy to be modified
        spoke_anf_client: boto3 client for AWS Network Firewall in the target account
        strict_order (str): whether or not to apply strict ordering to policy
    """

    policy_name = policy["Name"]
    spoke_policy = spoke_anf_client.describe_firewall_policy(FirewallPolicyName=policy_name)
    spoke_policy_details = spoke_policy["FirewallPolicy"]
    resource_arn = spoke_policy["FirewallPolicyResponse"]["FirewallPolicyArn"]
    rule_order = spoke_policy_details["StatefulEngineOptions"]["RuleOrder"]
    remediated_policy_name = None

    # Checking if policy has the correct stateful rule order, and modifying if not
    if strict_order == "true" and rule_order == "DEFAULT_ACTION_ORDER":
        # Getting details of existing policy for creation of new policy
        policy_arn = spoke_policy["FirewallPolicyResponse"]["FirewallPolicyArn"]
        description = (
            spoke_policy["FirewallPolicyResponse"]["Description"]
            if "Description" in spoke_policy["FirewallPolicyResponse"]
            else "Policy created by central firewall management"
        )
        tags = (
            spoke_policy["FirewallPolicyResponse"]["Tags"]
            if spoke_policy["FirewallPolicyResponse"]["Tags"] != []
            else [{"Key": "Name", "Value": policy_name}]
        )
        encryption_config = spoke_policy["FirewallPolicyResponse"]["EncryptionConfiguration"]
        remediated_policy_name = policy_name + "-baselined"

        # Setting rule order to strict, and removing all stateful rule groups from policy (these are not compatible with strict ordering)
        spoke_policy_details["StatefulEngineOptions"]["RuleOrder"] = "STRICT_ORDER"
        spoke_policy_details["StatefulEngineOptions"]["StreamExceptionPolicy"] = "DROP"
        spoke_policy_details["StatefulRuleGroupReferences"] = []

        # Creating new policy with strict rule ordering, as rule ordering on an existing policy cannot be modified
        try:
            response = spoke_anf_client.create_firewall_policy(
                FirewallPolicyName=remediated_policy_name,
                FirewallPolicy=spoke_policy_details,
                Description=description,
                Tags=tags,
                EncryptionConfiguration=encryption_config,
            )
        except exceptions.ClientError as e:
            logger.error("Could not create a new policy with strict rule ordering: " + remediated_policy_name)
            logger.error(e)
        else:
            remediated_policy_arn = response["FirewallPolicyResponse"]["FirewallPolicyArn"]
            logger.info("Created new policy with strict rule ordering: " + remediated_policy_name)

        # Updating firewalls in the account to associate the remediated firewall policy with
        firewall_list = spoke_anf_client.list_firewalls()["Firewalls"]
        for firewall in firewall_list:
            firewall_arn = firewall["FirewallArn"]
            spoke_firewall = spoke_anf_client.describe_firewall(FirewallArn=firewall_arn)
            update_token = spoke_firewall["UpdateToken"]

            if spoke_firewall["Firewall"]["FirewallPolicyArn"] == policy_arn:
                try:
                    response = spoke_anf_client.associate_firewall_policy(
                        UpdateToken=update_token,
                        FirewallArn=firewall_arn,
                        FirewallPolicyArn=remediated_policy_arn,
                    )
                except exceptions.ClientError as e:
                    logger.error(
                        "Could not associate remediated policy with strict rule ordering on following firewall: "
                        + firewall_arn
                    )
                    logger.error(e)
                else:
                    logger.info(
                        "Successfully associated remediated policy "
                        + remediated_policy_name
                        + " with firewall: "
                        + firewall_arn
                    )

        # Deleting non-compliant firewall policy
        try:
            response = spoke_anf_client.delete_firewall_policy(FirewallPolicyName=policy_name)
        except exceptions.ClientError as e:
            logger.error("Could not delete non-compliant firewall policy: " + policy_name)
            logger.error(e)
        else:
            logger.info("Deleted non-compliant firewall policy: " + policy_name)

            # If remediation is successful, reporting to AWS Config
            put_new_evaluation(
                spoke_config_client,
                evaluation,
                resource_arn,
                result_token,
            )

    return remediated_policy_name

def reprioritise_rule_groups(strict_order, rule_group_mapping, rule_group_list):
    """
    Modifies priority of existing rule groups to ensure baseline rule groups are evaluated first

    Parameters:
        rule_group_mapping (dict): mapping of baseline rule groups to apply to policy
        rule_group_list: list of current rule groups attached to policy
        strict_order (str): whether or not to apply strict ordering to policy

    Returns:
        rule_group_list (list): list of reprioritised rule groups
    """

    # Obtaining the maximum priority of baseline rule groups
    if strict_order == "true":
        priority_list = []
        for rule_group in rule_group_mapping:
            priority_list.append(int(rule_group["Priority"]))

            # Checking that priority of baseline rule groups is correct
            rule_group_list[:] = [x for x in rule_group_list if x["ResourceArn"] != rule_group["ResourceArn"]]

        max_baseline_priority = max(priority_list)
        min_priority = min(rule_group_list, key=lambda x: x["Priority"]) if len(rule_group_list) > 0 else 0

        # Increasing the priority of each existing rule group by the max baseline rule group, thus avoiding priority conflicts
        if min_priority != 0 and min_priority["Priority"] <= max_baseline_priority:
            delta = max_baseline_priority - min_priority["Priority"] + 1
            for i in range(len(rule_group_list)):
                rule_group_list[i]["Priority"] += delta

    return rule_group_list

def remediate_baseline_rule_groups(
    event,
    policy_list,
    strict_order,
    rule_group_mapping,
    spoke_config_client,
    spoke_anf_client,
):
    """
    Remediates non-compliant policies in the account if baseline rule groups are not applied

    Parameters:
        event: invoking event
        spoke_anf_client: boto3 client for AWS Network Firewall in the target account
        spoke_config_client: boto3 client for AWS Config in the target account
        rule_group_mapping (dict): mapping of baseline rule groups to apply to policy
        strict_order (str): whether or not to apply strict ordering to policy

    Returns:
        modified_resources (list): list of remediated resources
    """

    evaluations_list = event["Evaluations"]
    result_token = event["ResultToken"]
    modified_resources = []

    # Looping through non-compliant firewall policies to apply the correct baseline rule groups
    for evaluation in evaluations_list:
        resource_name = evaluation["ComplianceResourceId"]
        resource_type = evaluation["ComplianceResourceType"]
        compliance = evaluation["ComplianceType"]

        if compliance == "NON_COMPLIANT":
            for policy in policy_list:

                # Modifying rule order if required
                remediated_policy_name = remediate_rule_order(
                    strict_order,
                    policy,
                    evaluation,
                    result_token,
                    spoke_anf_client,
                    spoke_config_client,
                )
                if remediated_policy_name is None:
                    policy_name = policy["Name"]
                else:
                    policy_name = remediated_policy_name

                resource_name = "policy/" + policy_name
                spoke_policy = spoke_anf_client.describe_firewall_policy(FirewallPolicyName=policy_name)
                update_token = spoke_policy["UpdateToken"]
                spoke_policy_details = spoke_policy["FirewallPolicy"]

                if "STATEFUL" in rule_group_mapping:
                    # Reprioritising rule groups to enforce baseline rule groups first
                    stateful_rule_group_list = reprioritise_rule_groups(
                        strict_order,
                        rule_group_mapping["STATEFUL"],
                        spoke_policy_details["StatefulRuleGroupReferences"],
                    )

                    stateful_rule_group_list.extend(rule_group_mapping["STATEFUL"])

                    # Removing duplicate rule groups
                    unique_stateful_rules = [dict(t) for t in {tuple(d.items()) for d in stateful_rule_group_list}]

                    spoke_policy_details["StatefulRuleGroupReferences"] = unique_stateful_rules

                if "STATELESS" in rule_group_mapping:
                    # Reprioritising rule groups to enforce baseline rule groups first
                    stateless_rule_group_list = reprioritise_rule_groups(
                        strict_order,
                        rule_group_mapping["STATELESS"],
                        spoke_policy_details["StatelessRuleGroupReferences"],
                    )

                    stateless_rule_group_list.extend(rule_group_mapping["STATELESS"])

                    # Removing duplicate rule groups
                    unique_stateless_rules = [
                        dict(t)
                        for t in {tuple(d.items()) for d in spoke_policy_details["StatelessRuleGroupReferences"]}
                    ]

                    spoke_policy_details["StatelessRuleGroupReferences"] = unique_stateless_rules
                try:
                    response = spoke_anf_client.update_firewall_policy(
                        FirewallPolicyName=policy_name,
                        FirewallPolicy=spoke_policy_details,
                        UpdateToken=update_token,
                    )
                except exceptions.ClientError as e:
                    logger.error("Could not add rule groups to non-compliant policy: " + resource_name)
                    logger.error(e)
                else:
                    resource_arn = response["FirewallPolicyResponse"]["FirewallPolicyArn"]

                    # If remediation is successful, reporting to AWS Config
                    put_new_evaluation(
                        spoke_config_client,
                        evaluation,
                        resource_arn,
                        result_token,
                    )

                    # Creating list of modified firewall policies
                    modified_resources.append({"ResourceArn": resource_arn, "ResourceType": resource_type})

    return modified_resources

def remediate_all_rule_groups(
    event,
    policy_list,
    strict_order,
    rule_group_mapping,
    spoke_config_client,
    spoke_anf_client,
):
    """
    Remediates non-compliant policies in the account if excess rule groups are attached to a policy that are not centrally defined

    Parameters:
        event: invoking event
        spoke_anf_client: boto3 client for AWS Network Firewall in the target account
        spoke_config_client: boto3 client for AWS Config in the target account
        rule_group_mapping (dict): mapping of rule groups to apply to policy
        strict_order (str): whether or not to apply strict ordering to policy

    Returns:
        modified_resources (list): list of remediated resources
    """

    evaluations_list = event["Evaluations"]
    result_token = event["ResultToken"]
    modified_resources = []

    # Looping through non-compliant firewall policies to apply the correct baseline rule groups
    for evaluation in evaluations_list:
        resource_name = evaluation["ComplianceResourceId"]
        resource_type = evaluation["ComplianceResourceType"]
        compliance = evaluation["ComplianceType"]

        if compliance == "NON_COMPLIANT":
            for policy in policy_list:
                # Modifying rule order if required
                remediated_policy_name = remediate_rule_order(
                    strict_order,
                    policy,
                    evaluation,
                    result_token,
                    spoke_anf_client,
                    spoke_config_client,
                )

                if remediated_policy_name is None:
                    policy_name = policy["Name"]
                else:
                    policy_name = remediated_policy_name

                policy_name = policy["Name"]
                resource_name = "policy/" + policy_name
                spoke_policy = spoke_anf_client.describe_firewall_policy(FirewallPolicyName=policy_name)
                update_token = spoke_policy["UpdateToken"]
                spoke_policy_details = spoke_policy["FirewallPolicy"]

                if "STATEFUL" in rule_group_mapping:
                    spoke_policy_details["StatefulRuleGroupReferences"] = rule_group_mapping["STATEFUL"]

                if "STATELESS" in rule_group_mapping:
                    spoke_policy_details["StatelessRuleGroupReferences"] = rule_group_mapping["STATELESS"]

                try:
                    response = spoke_anf_client.update_firewall_policy(
                        FirewallPolicyName=policy_name,
                        FirewallPolicy=spoke_policy_details,
                        UpdateToken=update_token,
                    )
                except exceptions.ClientError as e:
                    logger.error("Could not enforce rule groups on non-compliant policy: " + resource_name)
                    logger.error(e)
                else:
                    resource_arn = response["FirewallPolicyResponse"]["FirewallPolicyArn"]

                    # If remediation is successful, reporting to AWS Config
                    put_new_evaluation(
                        spoke_config_client,
                        evaluation,
                        resource_arn,
                        result_token,
                    )

                    # Creating list of modified firewall policies
                    modified_resources.append({"ResourceArn": resource_arn, "ResourceType": resource_type})

                    logger.info("Successfully enforced centrally defined rule groups on policy: " + resource_name)

    return modified_resources

def put_new_evaluation(spoke_config_client, evaluation, resource_id, result_token):
    """
    Reports new evaluations to AWS Config after non-compliant resources have been remediated

    Parameters:
        evaluation (dict): original resource evaluations to modify
        spoke_anf_client: boto3 client for AWS Config in the target account
        resource_id: ID of resource that was remediated
        result_token (str): AWS Config results token

    Returns:
        eval_cc (dict): resource evalations for AWS Config
    """

    evaluation["ComplianceType"] = "COMPLIANT"
    evaluation["Annotation"] = resource_id

    response = spoke_config_client.put_evaluations(ResultToken=result_token, Evaluations=[evaluation])

    return response

def handler(event, context):
    account_id = event["accountId"]

    # Obtaining required boto3 clients
    spoke_anf_client = get_client("network-firewall", account_id)
    spoke_config_client = get_client("config", account_id)
    ssm_client = boto3.client("ssm")
    organizations_client = boto3.client("organizations")
    spoke_ou_path = get_ou_path(organizations_client, account_id)

    # Reading SSM parameter to obtain raw rule group mapping
    raw_rule_group_mapping = get_rule_group_mapping(ssm_client, spoke_ou_path, account_id)

    # Reads whether or not to enforce baselines only, then removing from rule group mapping
    baseline_plus = raw_rule_group_mapping["baseline_plus"]
    del raw_rule_group_mapping["baseline_plus"]

    # Reads whether or not to enforce strict order, then removing from rule group mapping
    strict_order = raw_rule_group_mapping["strict_order"]
    del raw_rule_group_mapping["strict_order"]

    # Obtaining ARNs of firewall rule groups for mapping
    processed_rule_group_mapping = get_rule_group_arns(raw_rule_group_mapping)

    # Get list of all firewall policies in target account
    policy_list = spoke_anf_client.list_firewall_policies()["FirewallPolicies"]

    # Remediating all firewall policies
    if baseline_plus == "true":
        modified_resources = remediate_baseline_rule_groups(
            event,
            policy_list,
            strict_order,
            processed_rule_group_mapping,
            spoke_config_client,
            spoke_anf_client,
        )
    else:
        modified_resources = remediate_all_rule_groups(
            event,
            policy_list,
            strict_order,
            processed_rule_group_mapping,
            spoke_config_client,
            spoke_anf_client,
        )

    logger.info(
        "The following resources were modified in account " + account_id + ": " + json.dumps(modified_resources)
    )

    return


import torch
import torch.nn as nn
from torch.nn.parameter import Parameter
from operators.functions.unsorted_segment_sum import UnsortedSegmentSumFunction

class FDUnsortedSegmentSum(nn.Module):

  def __init__(self, num_segments):
    super(FDUnsortedSegmentSum, self).__init__()
    self.num_segments = num_segments

  def forward(self, data, v, segment_index, dim=1):

    assert (data.size() == v.size()), "{} vs {}".format(data.size(), v.size())

    if dim != 1:
      v = v.transpose(dim, 1).contiguous()
    size_prev = list(v.size())
    v = v.view(size_prev[0], size_prev[1], -1)
    v = UnsortedSegmentSumFunction.apply(v, segment_index, self.num_segments)
    size_prev[1] = self.num_segments
    v = v.view(size_prev)
    if dim != 1:
      v = v.transpose(dim, 1).contigous()
    return v

class UnsortedSegmentSum(nn.Module):

  def __init__(self, num_segments):
    super(UnsortedSegmentSum, self).__init__()
    self.num_segments = num_segments

  def forward(self, data, segment_index):
    return UnsortedSegmentSumFunction.apply(data, segment_index,
                                            self.num_segments)


import os
import random
from numpy.random import choice
import pandas as pd
import psycopg2
from tqdm import tqdm
import statsmodels.api as sm
from statsmodels.formula.api import ols
def N_col_sql(num_col, sorted_colset, order, f, f_sql, df, tables, dictalias, conn):
    df.columns = [c.lower() for c in df.columns]
    conn.autocommit = True
    cursor = conn.cursor()
    ops = ['=', '<', '>']
    joins = []
    col = sorted_colset[order - 1:num_col + order - 2]
    col.append('fps')

    for i in tqdm(range(10000)):
        questr = 'SELECT COUNT(*) FROM '
        questr = questr + ",".join(tables) + " WHERE "
        component = []
        result = []
        df_temp = df
        for k in range(len(col)):
            op = choice(ops)
            dist_val_list = list(set(df_temp[col[k]]))
            val = choice(dist_val_list)
            questr_temp = questr + dictalias['fps'][0] + '.' + str(col[k]) + op + str(val)
            if k < len(col) - 1:
                questr_temp = questr_temp.replace('COUNT(*)', col[k + 1])  # only select required columns, reduce I/O
            else:
                questr_temp = questr_temp.replace('COUNT(*)', col[k])
            df_temp = pd.read_sql(questr_temp, conn)
            # if len(df_temp) == 0:
            #    df_temp = df
            count = 0
            while (len(df_temp) == 0):
                # df_temp = df
                op = choice(ops)
                val = choice(dist_val_list)
                questr_temp = questr + dictalias['fps'][0] + '.' + str(col[k]) + op + str(val)
                if k < len(col) - 1:
                    questr_temp = questr_temp.replace('COUNT(*)',
                                                      col[k + 1])  # only select required columns, reduce I/O
                else:
                    questr_temp = questr_temp.replace('COUNT(*)', col[k])
                df_temp = pd.read_sql(questr_temp, conn)
                count = count + 1
                if count > 2:
                    if len(df_temp) == 0:
                        df_temp = df
                    break
            questr = questr + dictalias['fps'][0] + '.' + str(col[k]) + op + str(
                val) + ' AND '
            component.append(dictalias['fps'][0] + '.' + str(col[k]))
            component.append(op)
            component.append(val)
        questr = questr[:len(questr) - 5]
        questr += ';'
        # df_temp = pd.read_sql(questr, conn)
        # card = df['count'].values[0]
        cursor.execute(questr)
        card = cursor.fetchall()[0][0]
        questr += f',{card}\n'
        f.write(",".join(tables) + '#' + ','.join(joins) + '#' + ",".join(map(str, component)) + '#' + str(card) + '\n')
        f_sql.write(questr)
    f.close()
    f_sql.close()

df_fps = pd.read_csv('data/gpu_game_fps/fps_num_lower.csv', escapechar='\\', encoding='utf-8',
                     low_memory=False, quotechar='"',
                     usecols=['id', 'cpunumberofcores', 'cpunumberofthreads',
                              'cpubaseclock', 'cpucachel1', 'cpucachel2', 'cpucachel3', 'cpudiesize',
                              'cpufrequency', 'cpumultiplier', 'cpumultiplierunlocked',
                              'cpuprocesssize', 'cputdp', 'cpunumberoftransistors', 'cputurboclock',
                              'gpubandwidth', 'gpubaseclock', 'gpuboostclock', 'gpubusnterface',
                              'gpunumberofcomputeunits', 'gpudiesize', 'gpunumberofexecutionunits',
                              'gpufp32performance', 'gpumemorybus', 'gpumemorysize', 'gpumemorytype', 'gpupixelrate',
                              'gpuprocesssize', 'gpunumberofrops', 'gpushadermodel', 'gpunumberofshadingunits',
                              'gpunumberoftmus', 'gputexturerate', 'gpunumberoftransistors',
                              'gpuvulkan', 'gamename', 'gameresolution', 'gamesetting', 'fps'])

model = ols('fps ~ C(cpunumberofcores) + C(cpunumberofthreads) + C(cpubaseclock) + C(cpucachel1) + C(cpucachel2) + C('
            'cpucachel3) + C(cpudiesize) + C(cpufrequency) + C(cpumultiplier) + C(cpumultiplierunlocked) + C(cpuprocesssize) + '
            'C(cputdp) + C(cpunumberoftransistors) + C(cputurboclock) + C(gpubandwidth) + C(gpubaseclock) + C(gpuboostclock) + '
            'C(gpubusnterface) + C(gpunumberofcomputeunits) + C(gpudiesize) + C(gpunumberofexecutionunits) + C(gpufp32performance) '
            '+ C(gpumemorybus) + C(gpumemorysize) + C(gpumemorytype) + C(gpupixelrate) + C(gpuprocesssize) + C(gpunumberofrops) + '
            'C(gpushadermodel) + C(gpunumberofshadingunits) + C(gpunumberoftmus) + C(gputexturerate) + C(gpunumberoftransistors) + '
            'C(gpuvulkan) + C(gamename) + C(gameresolution) + C(gamesetting)', data=df_fps).fit()

anova = sm.stats.anova_lm(model, typ=2)
sorted_colset = list(anova.iloc[:, -1].sort_values().index)
sorted_colset = sorted_colset[:-1]

model = ols('fps ~ C(cpunumberofcores) + C(cpubaseclock) + C(cpufrequency) + C(cputdp)', data=df_fps).fit()
anova_cpu = sm.stats.anova_lm(model, typ=2)
model = ols(
    'fps ~ C(gpuboostclock) + C(gpumemorybus) + C(gpumemorytype) + C(gpushadermodel) +C(gputexturerate) + C(gamesetting)',
    data=df_fps).fit()
anova_gpu = sm.stats.anova_lm(model, typ=2)

anova = pd.concat([anova_cpu.iloc[:-1,:],anova_gpu.iloc[:-1,:]])
anova.iloc[:,-1].sort_values()

dictalias = {'fps': ['f']}
tables = ['fps f']

conn = psycopg2.connect(database="fps",
                        user='postgres', password='',
                        host='localhost', port='5432'
                        )

# 2 cols
i=5
os.makedirs('data/gpu_game_fps/2_col_sql', exist_ok=True)
while i <= len(sorted_colset):
    f = open(f"data/gpu_game_fps/2_col_sql/fps_2col_sql_{i}.csv", 'w')
    f_sql = open(f"data/gpu_game_fps/2_col_sql/fps_2col_sql_{i}.sql", 'w')
    N_col_sql(2, sorted_colset, i, f, f_sql, df_fps, tables, dictalias, conn)
    i = i+8

# 5~37 cols
colset = []
i=0
while i <= len(sorted_colset):
    colset.append(sorted_colset[i])
    i=i+4

i=5
while i <= len(sorted_colset):
    os.makedirs(f'data/gpu_game_fps/{i}_col_sql', exist_ok=True)
    f = open(f"data/gpu_game_fps/{i}_col_sql/fps_{i}col_sql_1.csv", 'w')
    f_sql = open(f"data/gpu_game_fps/{i}_col_sql/fps_{i}col_sql_1.sql", 'w')
    N_col_sql(i, sorted_colset, 1, f, f_sql, df_fps, tables, dictalias, conn)
    i = i+4

def gen_train_test(path_sql_file, path_train_file, path_test_file, num_test):
    with open(path_sql_file, "r") as input:
        lines = input.readlines()
        with open(path_train_file, "w") as output_train:
            with open(path_test_file, "w") as output_test:
                i = 0
                testlines = random.sample(range(len(lines)), num_test)
                for line in lines:
                    if i in testlines:
                        output_test.write(line)
                    else:
                        output_train.write(line)
                    i = i + 1
            output_test.close()
        output_train.close()
    input.close()

i = 1
while i<38:
    path_sql_file = f"data/gpu_game_fps/2_col_sql/fps_2col_sql_{i}.csv"
    path_train_file = f'data/gpu_game_fps/2_col_sql/fps_2col_{i}_train.csv'
    path_test_file = f'data/gpu_game_fps/2_col_sql/fps_2col_{i}_test.csv'
    num_test = 500
    gen_train_test(path_sql_file, path_train_file, path_test_file, num_test)
    i = i+4

i = 5
while i<38:
    os.makedirs(f'data/gpu_game_fps/{i}_col_sql', exist_ok=True)
    path_sql_file = f"data/gpu_game_fps/{i}_col_sql/fps_{i}col_sql_{i}.csv"
    path_train_file = f'data/gpu_game_fps/{i}_col_sql/fps_{i}col_{i}_train.csv'
    path_test_file = f'data/gpu_game_fps/{i}_col_sql/fps_{i}col_{i}_test.csv'
    num_test = 500
    gen_train_test(path_sql_file, path_train_file, path_test_file, num_test)
    i = i+4