import os

import sys

import shutil

import itertools

from itertools import product

import json

import pandas as pd

import numpy as np

import random

from random import choices

from scipy.spatial import distance

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

import tensorflow as tf

#import tensorflow_probability as tfp

from scipy.stats.stats import pearsonr

from tensorflow.keras import backend as K

from sklearn.utils import class_weight

from sklearn.preprocessing import LabelEncoder

from sklearn.preprocessing import RobustScaler

from sklearn.cluster import OPTICS

from sklearn.cluster import KMeans

from sklearn import metrics

from Levenshtein import distance as lev_dist

from focal_loss import sparse_categorical_focal_loss

#from imblearn.under_sampling import RandomUnderSampler

#from imblearn.over_sampling import RandomOverSampler

#from imblearn.tensorflow import balanced_batch_generator

import matplotlib.pyplot as plt

#import torch

#from torch.utils.data import WeightedRandomSampler

from focal_loss import SparseCategoricalFocalLoss

tf.keras.metrics.SparseTopKCategoricalAccuracy(k=1)

import neural_network

PATH_KMER_F_DICT = "data/ncov_global/kmer_f_word_dictionaries.json"

PATH_KMER_R_DICT = "data/ncov_global/kmer_r_word_dictionaries.json"

#m_loss = neural_network.MaskedLoss()

#bce = tf.keras.losses.BinaryCrossentropy(from_logits=False)

#focal_loss_func = SparseCategoricalFocalLoss(gamma=2)

cross_entropy_loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False, reduction='auto')

#mae = tf.keras.losses.MeanAbsoluteError()

#mse = tf.keras.losses.MeanSquaredError()

test_tf_ratio = 0.0

max_norm = 1.0

enc_stddev = 1.0

dec_stddev = 0.0001

beta = 0.9999

amino_acid_codes = "QNKWFPYLMTEIARGHSDVC"

def decay_lr(lr, factor=0.95):

return factor * lr

def decayed_learning_rate(initial_learning_rate, step, decay_rate=0.95, decay_steps=1000000):

return initial_learning_rate * np.float_power(decay_rate, (step / decay_steps))

def loss_function(real, pred):

return loss

def make_kmers(seq, size):

return ["".join(seq[x:x+size]) for x in range(len(seq) - size + 1)]

def compute_Levenshtein_dist(seq_in, seq_out):

#return np.random.randint(1, 5)

return lev_dist(seq_in, seq_out)

def clean_up(list_folders):

continue

def add_padding_to_seq(seq, s_token):

return seq #"{},{}".format(str(s_token), seq)

def get_u_kmers(x_seq, y_seq, max_l_dist, len_aa_subseq, forward_dict, start_token, unrelated):

return fil_x, fil_y, kmer_f_dict, kmer_r_dict, l_distance, filtered_l_distance

def split_test_train(x, y, split_size):

return x_1, x_2, y_1, y_2

def generate_cross_product(x_df, y_df, in_clade, out_clade, max_l_dist, len_aa_subseq, forward_dict, rev_dict, start_token, cols=["X", "Y"], unrelated=False, unrelated_threshold=15, train_nation="USA", train_pairs=True, train_size=0.8):

return train_df, test_df, kmer_f_dict, kmer_r_dict

def transform_noise(noise):

return tf.convert_to_tensor(noise, dtype=tf.float32)

def get_all_possible_words(vocab, kmer_size=3):

return kmer_f_dict, kmer_r_dict

def convert_to_array(str_data):

return toarray

def pred_convert_to_array(str_data):

return tf.convert_to_tensor(toarray, dtype=tf.int32)

def format_POS_variations(var_pos):

for key in var_pos:

var_pos[key] = np.unique(var_pos[key].tolist())

print(var_pos[key])

def one_hot_encoding():

print(encoded_seq.shape)

def read_in_out(path):

return samples

def reconstruct_seq(kmers):

return "".join(reconstructed_seq)

def encode_sequences_kmers(f_dict, kmer_r_dict, x_seq, y_seq, s_kmer):

return in_seq, out_seq

def get_all_kmers(x_seq, y_seq, f_dict, s_kmer):

return list(set(all_kmers))

def read_wuhan_seq(wu_path, rev_dict):

return ",".join(enc_wu_seq)

def get_words_indices(word_list):

return forward_dictionary, reverse_dictionary

def save_as_json(filepath, data):

with open(filepath, 'w') as fp:

json.dump(data, fp)

def read_json(path):

with open(path, 'r') as fp:

f_content = json.loads(fp.readline())

return f_content

def format_clade_name(c_name):

return c_name.replace("/", "_")

def embedding_info(dict_json):

return len(dict_json)

def convert_to_string_list(l):

return l

def get_variation_loss(true_in, true_out, pred_logits):

return mse_y + (1.0 - mse_x)

def get_sequence_variation_percentage(true_in, pred_logits):

return x_percent_variation, pred_percent_variation, pair_percent_variation

def sample_unrelated_x_y(unrelated_X, unrelated_y, batch_size):

return convert_to_array(un_X), convert_to_array(un_y)

def scale_encodings(encoding):

encoding = encoding.numpy()

tf_enc = RobustScaler().fit_transform(encoding)

return tf.convert_to_tensor(tf_enc, dtype=tf.float32)

def clip_weights(tensor, clip_min=-1e-2, clip_max=1e-2):

return tf.clip_by_value(tensor, clip_value_min=clip_min, clip_value_max=clip_max)

#return tf.clip_by_norm(tensor, clip_norm=2.0)

def create_mut_balanced_dataset(X, Y, kmer_f_dict, seq_len, batch_size):

return mut_pattern, mut_pattern_dist, mut_pattern_dist_freq, mut_buckets

def find_cluster_indices(output_seqs, batch_size, datatype="train_y"):

return cluster_labels, cluster_indices_dict, scatter_df

def calculate_sample_weights(X_train, y_train, pos_variations_count):

return train_data_generator

def calculate_input_sample_weights(input_seq, x_pos_variations_count):

return weights_dict

def create_dirs(folder_path):

if not os.path.isdir(folder_path):

os.mkdir(folder_path)

def stateful_encoding(size_stateful, inputs, enc, training=False):

return enc_out, enc_state, enc

def loop_encode_decode_stateful(seq_len, batch_size, vocab_size, input_tokens, output_tokens, gen_encoder, gen_decoder, enc_units, tf_ratio, train_test, s_stateful, mut_freq, pos_variations_count, batch_step):

return global_logits, gen_encoder, gen_decoder, loss

def loop_encode_decode_predict_stateful(seq_len, batch_size, vocab_size, input_tokens, output_tokens, gen_encoder, gen_decoder, enc_units, tf_ratio, train_test, s_stateful, mut_freq):

return gen_logits, loss

def predict_sequence(tr_epoch, tr_batch, test_dataset_in, test_dataset_out, te_batch_size, n_te_batches, seq_len, vocab_size, enc_units, loaded_encoder, loaded_generator, s_stateful, train_type, save=False):

return np.mean(avg_test_loss), np.mean(avg_test_seq_var)

def generate_per_seq(tr_epoch, tr_batch, seq_len, te_batch_size, vocab_size, batch_x_test, batch_y_test, loaded_encoder, loaded_generator, enc_units):

true_predicted_multiple.to_csv(df_path, index=None)

def save_predicted_test_data(test_data_in, test_data_out, te_batch_size, enc_units, vocab_size, seq_len, size_stateful, epoch_type_name, enc_model_path, dec_model_path):

pred_dataframe.to_csv(df_filename, index=None)

def generator_step(seq_len, batch_size, vocab_size, gen_decoder, dec_state_h, dec_state_c, real_i, real_o, train_gen):

return pred_logits, gen_decoder, step_loss

def generated_output_seqs(seq_len, batch_size, vocab_size, gen_decoder, dec_state_h, dec_state_c, real_i, real_o, train_gen):

return pred_logits, gen_decoder, step_loss

def save_batch(batch_x, batch_y, batch_mut_distribution):

return batch_mut_distribution

def get_mutation_tr_indices(seq, kmer_f_dict, kmer_r_dict, f_dict, r_dict, type_seq="x"):

return parent_child_mut_indices, parent_child_pos_vars, parent_child_pos_vars_count

'''def get_mutation_tr_indices(train_in, train_out, kmer_f_dict, kmer_r_dict, f_dict, r_dict, parent_child_pos_vars=dict(), parent_child_pos_vars_count=dict()):

parent_child_mut_indices = dict()

for index, (x, y) in enumerate(zip(train_in, train_out)):

true_x = x.split(",")

true_y = y.split(",")

re_true_x = true_x

re_true_y = true_y

for i in range(len(true_x)):

first = re_true_x[i:i+1][0]

sec = re_true_y[i:i+1][0]

if first != sec:

key = "{}>{}>{}".format(first, (i+1), sec)

if key not in parent_child_mut_indices:

parent_child_mut_indices[key] = list()

parent_child_mut_indices[key].append(index)

key_pos_var = "{}".format(str(i))

if key_pos_var not in parent_child_pos_vars:

parent_child_pos_vars[key_pos_var] = list()

parent_child_pos_vars_count[key_pos_var] = dict()

if int(sec) not in parent_child_pos_vars[key_pos_var]:

parent_child_pos_vars[key_pos_var].append(int(sec))

if int(sec) not in parent_child_pos_vars_count[key_pos_var]:

parent_child_pos_vars_count[key_pos_var][int(sec)] = 0

parent_child_pos_vars_count[key_pos_var][int(sec)] += 1

save_as_json("data/generated_files/parent_child_pos_vars_{}.txt".format(str(np.random.randint(0, 1e10, 1)[0])), parent_child_pos_vars)

save_as_json("data/generated_files/parent_child_pos_vars_count_{}.txt".format(str(np.random.randint(0, 1e10, 1)[0])), parent_child_pos_vars_count)

return parent_child_mut_indices, parent_child_pos_vars, parent_child_pos_vars_count

def pairwise_dist(A, B):

# https://gist.github.com/mbsariyildiz/34cdc26afb630e8cae079048eef91865

# squared norms of each row in A and B

na = tf.reduce_sum(tf.square(A), 1)

nb = tf.reduce_sum(tf.square(B), 1)

# na as a row and nb as a co"lumn vectors

na = tf.reshape(na, [-1, 1])

nb = tf.reshape(nb, [1, -1])

# return pairwise euclidead difference matrix

D = na - 2*tf.matmul(A, B, False, True) + nb

D_min = tf.math.reduce_min(D)

D_max = tf.math.reduce_max(D)