import collections
import torch
import numpy as np
import os
import h5py
import time
import math
import platform
import typing
from torch import Tensor

def count_net_params(net):
    n_param = 0
    for name, param in net.named_parameters():
        layer_param = torch.numel(param)
        print("::: Layer Name: ", name, "Layer Size", layer_param)
        n_param += layer_param
    return n_param

def count_rnn_params(net):
    n_param = 0
    for name, param in net.named_parameters():
        if 'rnn' in name:
            sizes = 1
            for el in param.size():
                sizes = sizes * el
            n_param += sizes
    return n_param

def get_dataset_path(args):
    # Dictionary for datasets
    if args.dataset_name == 'gscdv1':
        print("Preparing Google Speech Command Dataset...")
        dataset_name = 'speech_commands_v0.01'
        testset_name = 'speech_commands_test_set_v0.01'
    elif args.dataset_name == 'gscdv2':
        print("Preparing Google Speech Command Dataset...")
        dataset_name = 'speech_commands_v0.02'
        testset_name = 'speech_commands_test_set_v0.02'
    elif args.dataset_name == 'timit':
        print("Preparing TIMIT Dataset...")
        dataset_name = 'TIMIT/TIMIT'
        testset_name = None
    elif args.dataset_name == 'tidigits':
        print("Preparing TIDIGITS Dataset...")
        dataset_name = 'tidigits'
        testset_name = None
    else:
        raise RuntimeError('Dataset not supported.')

    data_dir = args.data_dir
    dataset_path = os.path.join(data_dir, dataset_name)
    testset_path = os.path.join(data_dir, testset_name)

    return dataset_path, testset_path

def load_h5py_data(f):
    dict_data = {}
    for k in f.keys():
        dict_data[k] = np.array(f[k])
    return dict_data

def log_lut(x, qi_in, qf_in, qi_out, qf_out):
    lut_in = quantize_array(x, qi_in, qf_in, 1)
    lut_out = np.log10(lut_in + 1)
    lut_out = quantize_array(lut_out, qi_out, qf_out, 1)
    return lut_out

def gen_description_file(augment_noise, target_snr):
    if augment_noise:
        description_file = 'description_snr' + str(target_snr) + '.csv'
    else:
        description_file = 'description.csv'
    return description_file

def piecewise_linear_log(x,
                         k0, b0,
                         k1, b1,
                         k2, b2,
                         k3, b3,
                         k4, b4,
                         k5, b5,
                         k6, b6,
                         k7, b7,
                         k8, b8,
                         k9, b9
                         ):
    return np.piecewise(x=x,
                        condlist=[x >= 0,
                                  x >= 6,
                                  x >= 12,
                                  x >= 18,
                                  x >= 24,
                                  x >= 108,
                                  x >= 192,
                                  x >= 276,
                                  x >= 1549,
                                  x >= 2792],
                        funclist=[lambda x: k0 * x + b0,
                                  lambda x: k1 * x + b1,
                                  lambda x: k2 * x + b2,
                                  lambda x: k3 * x + b3,
                                  lambda x: k4 * x + b4,
                                  lambda x: k5 * x + b5,
                                  lambda x: k6 * x + b6,
                                  lambda x: k7 * x + b7,
                                  lambda x: k8 * x + b8,
                                  lambda x: k9 * x + b9]
                        )

def create_folder(folder_list):
    for folder in folder_list:
        try:
            os.makedirs(folder)
        except:
            pass

def load_model(proj, net, model_path):
    if proj.use_cuda and torch.cuda.is_available():
        model_location = 'cuda:' + str(proj.gpu_device)
    else:
        model_location = 'cpu'
    pretrained_model = torch.load(model_path, map_location=model_location)
    dict_pretrained_model = dict(pretrained_model.items())
    list_pretrained_model = list(dict_pretrained_model.items())
    dict_converted_model = {}
    net_state_dict = net.state_dict()
    list_trained_model = list(net_state_dict.items())
    num_pretrain_keys = len(list_pretrained_model)
    num_train_keys = len(list_trained_model)

    # Convert Model
    print("num_pretrain_keys: ", num_pretrain_keys)
    print("num_train_keys:    ", num_train_keys)
    if num_pretrain_keys != num_train_keys:
        for k_model, v_model in dict_pretrained_model.items():
            if 'ih' in k_model:
                if 'weight' in k_model:
                    param_ih = v_model
                    continue
                if 'bias' in k_model:
                    param_ih = v_model
                    continue
            elif 'hh' in k_model:
                if 'weight' in k_model:
                    param_hh = v_model
                    param_cat = torch.cat((param_ih, param_hh), dim=-1)
                if 'bias' in k_model:
                    param_hh = v_model
                    param_cat = param_ih + param_hh
                dict_converted_model[k_model] = param_cat
            else:
                dict_converted_model[k_model] = v_model
        list_converted_model_params = list(dict_converted_model.items())
    else:
        list_converted_model_params = list_pretrained_model
    # Copy param of the pretrained model to the training network
    new_state_dict = collections.OrderedDict()
    count = 0
    for key, value in net.state_dict().items():
        layer_name, weights = list_converted_model_params[count]
        net_state_dict[key] = weights
        print("Pre-trained Layer: %s - Loaded into new layer: %s" % (layer_name, key))
        count += 1

    net.load_state_dict(net_state_dict)
    return net

# def load_model(net, model_path, gpu_device):
#     # Get adapted DeltaRNN dimensions
#     # adapted_input_size = int(math.ceil(input_size / float(num_array)) * num_array)
#     # adapted_hidden_size = int(math.ceil(hidden_size / float(num_array)) * num_array)
#     # Load pretrained model
#     pretrained_net = torch.load(model_path, map_location='cuda:' + str(gpu_device))
#     pretrained_net = list(pretrained_net.items())
#
#     # Iterate throughput retrain net parameters to load pretrained net parameters correctly
#     new_state_dict = collections.OrderedDict()
#     count = 0
#     num_param_key = len(pretrained_net)
#     for key, value in net.state_dict().items():
#         if count >= num_param_key:
#             break
#         layer_name, weights = pretrained_net[count]
#         new_state_dict[key] = weights
#         print("Pre-trained Layer: %s - Loaded into new layer: %s" % (layer_name, key))
#         count += 1
#     net.load_state_dict(new_state_dict)
#     return net

def timeSince(since):
    now = time.time()
    s = now - since
    m = math.floor(s / 60)
    s -= m * 60
    return '%dm %ds' % (m, s)

def create_context_window(tensor, window_size=2):
    """
    Create a new tensor with context window
    :param tensor: 3D tensor (T, N, F)
    :param window_size: Integer
    :return:
    """
    tensor = tensor.transpose(0, 1)
    dim_T = tensor.shape[0]
    dim_N = tensor.shape[1]
    dim_F = tensor.shape[2]
    dim_FC = int(dim_F / 3)
    dim_C = int(window_size * 3)
    pad_len = window_size - 1
    pad = torch.zeros((pad_len, dim_N, dim_F)).cuda()
    padded_dim_T = dim_T + pad_len
    feature = torch.cat((tensor, pad), dim=0)
    feature = torch.reshape(feature, (padded_dim_T, dim_N, 3, dim_FC))
    output = []
    for tstep in range(0, dim_T):
        temp = feature[tstep:tstep + window_size, :, :, :]
        # print(temp.shape)
        temp = temp.transpose(0, 1)
        temp = torch.reshape(temp, (dim_N, dim_C, dim_FC))
        output.append(temp)
    output = torch.stack(output)
    return output, dim_C, dim_FC

def write_hdf5_normal(filename, X, y):
    # fun
    feature_lens = np.asarray([sample.shape[1] for sample in X]).astype(np.float32)
    label_lens = np.asarray([len(target) for target in y]).astype(np.float32)
    features = np.concatenate(X, axis=1).T.astype(np.float32)
    labels = np.concatenate(y).astype(np.float32)

    with h5py.File(filename, 'w') as f:
        f.create_dataset('features', data=features)
        f.create_dataset('feature_lens', data=feature_lens)
        f.create_dataset('labels', data=labels)
        f.create_dataset('label_lens', data=label_lens)

def print_size_of_model(model):
    torch.save(model.state_dict(), "temp.p")
    print('Size (MB):', os.path.getsize("temp.p") / 1e6)
    os.remove('temp.p')

def to_sparse(x):
    """ converts dense tensor x to sparse format """
    x_typename = torch.typename(x).split('.')[-1]
    sparse_tensortype = getattr(torch.sparse, x_typename)

    indices = torch.nonzero(x)
    if indices.nelement() == 0:  # if all elements are zeros
        print("1", indices)
        return sparse_tensortype(*x.shape)
    else:
        print("2", indices)
        indices = indices.torch()
        values = x[tuple(indices[i] for i in range(indices.shape[0]))]
        return sparse_tensortype(indices, values, x.size())

def quantize_tensor(x: Tensor, qi: int, qf: int, en: int = 0, floor=False):
    """
    :param x: input tensor
    :param qi: number of integer bits before the decimal point
    :param qf: number of fraction bits after the decimal point
    :param en: Type of Quantziation (0 - none, 1 - round, 2 - floor)
    :param floor: Whether use floor() instead of round()
    :return: tensor quantized to fixed-point precision
    """
    if en == 0:
        return x
    else:
        power = float(2. ** qf)
        clip_val = float(2. ** (qi + qf - 1))
        if floor:
            value = torch.floor(x * power)
        else:
            value = torch.round(x * power)
        value = torch.clamp(value, -clip_val, clip_val - 1)  # saturation arithmetic
        value = value / power
        return value

def quantize_array(x, qi, qf, enable, unsigned=0, floor=False):
    """
    :param x: input numpy array
    :param qi: number of integer bits before the decimal point
    :param qf: number of fraction bits after the decimal point
    :param enable: whether enable the function
    :param unsigned: whether unsigned input
    :return: array quantized to fixed-point precision
    """
    if enable == 0:
        return x
    power = np.asarray(2. ** qf).astype(np.float64)
    if unsigned == 0:
        clip_val = float(2. ** (qi + qf - 1))
        if floor:
            value = np.floor(x * power)
        else:
            value = np.round(x * power)
        value = np.clip(value, -clip_val, clip_val - 1)  # saturation arithmetic
        result = value / power
    else:
        clip_val = float(2. ** (qi + qf))
        if floor:
            value = np.floor(x * power)
        else:
            value = np.round(x * power)
        value = np.clip(value, 0, clip_val - 1)  # saturation arithmetic
        result = value / power
    return result

# Auto determine minimal possible qi
def quantize_tensor_auto(x, bw, enable):
    if enable == 0:
        return x, 0, 0
    qi = torch.ceil(torch.log2(x)) + 1
    qf = bw - qi
    # if qi < 0 or qf < 0:
    #     print('Quantizing {} to Q({}, {})!'.format(x, qi, qf))
    power = float(2. ** qf)
    clip_val = float(2. ** (qi + qf - 1) - 1)
    value = torch.round(x * power)
    value = torch.clamp(value, -clip_val, clip_val)  # saturation arithmetic
    value = value / power
    return value, qi, qf

def quantize_rnn(net, qi, qf, enable):
    for name, param in net.named_parameters():
        if 'rnn' in name and 'swap' not in name:
            # print("Quantizing Parameter: ", name)
            param.features = quantize_tensor(param.features, qi, qf, enable)
    return net

def quantize_net(net, qi, qf, enable):
    for name, param in net.named_parameters():
        param.features = quantize_tensor(param.features, qi, qf, enable)
    return net

def hard_quantize_tensor(x, m, n, enable=0):
    """
    :param x: input tensor
    :param m: number of integer bits before the decimal point
    :param n: number of fraction bits after the decimal point
    :return: tensor quantized to fixed-point precision
    """

    if not enable:
        return x

    power = 2. ** n
    max_val = 2. ** (m + n - 1) - 1
    value = x * power
    value = torch.floor(value)  # rounding
    value = torch.clamp(value, -max_val, max_val)  # saturation arithmetic
    value = value / power
    return value

def get_temporal_sparsity(list_layer, seq_len, threshold):
    # Evaluate Sparsity
    num_zeros = 0
    num_elems = 0
    # print(seq_len.size())
    # Iterate through layers
    for layer in list_layer:
        all_delta_vec = layer.transpose(0, 1)
        all_delta_vec = torch.abs(all_delta_vec)  # Take absolute values of all delta vector elements
        for i, delta_vec in enumerate(all_delta_vec):
            seq = delta_vec[:seq_len[i], :]
            zero_mask = seq < threshold
            num_zeros += torch.sum(zero_mask)
            num_elems += torch.numel(zero_mask)
    sparsity = float(num_zeros) / float(num_elems)
    return sparsity

# def targeted_dropout(x, gamma, alpha, epoch):
#     """
#     :param x: input tensor
#     :param m: number of integer bits before the decimal point
#     :param n: number of fraction bits after the decimal point
#     :return: tensor quantized to fixed-point precision
#     """
#     import math
#     torch.manual_seed(epoch)
#     torch.cuda.manual_seed_all(epoch)
#     n_dims = len(x.size())
#     n_elements = x.numel()
#     drop_part = math.ceil(n_elements * gamma)
#     weight_vec = x.view(-1)
#     weight_vec_abs = torch.abs(weight_vec)
#     sorted, indices = torch.sort(weight_vec_abs)
#     # print(sorted)
#     drop_indices = indices[0:drop_part]
#     drop_rand_mask = torch.rand(drop_indices.size(0)).cuda()
#     drop_mask = torch.ones(drop_indices.size(0)).cuda()
#     drop_mask = drop_mask.masked_fill_(drop_rand_mask <= alpha, 0)
#     weight_vec[drop_indices] *= drop_mask
#     if n_dims == 1:
#         weight = weight_vec
#     elif n_dims > 1:
#         weight = torch.reshape(weight_vec, (x.size(0), x.size(1)))
#
#     return weight

def aligned_cbwdrop(x, gamma, alpha, num_pe):
    """
    :param x: input tensor (weight matrix)
    :param m: number of integer bits before the decimal point
    :param n: number of fraction bits after the decimal point
    :return: tensor quantized to fixed-point precision
    """
    import math

    n_dims = len(x.size())
    n_rows = x.size(0)

    if n_dims == 1:
        n_cols = 1
    elif n_dims > 1:
        n_cols = x.size(1)
    else:
        raise RuntimeError('Input tensor must have at least 1 dimension.')

    # Split and shuffle weight matrix
    # for i in tqdm(range(0, num_pe), total=num_pe, unit="Submat"):
    for i in range(0, num_pe):
        if n_dims == 1:
            pe_work = x[np.arange(i, n_rows, num_pe)]
            pe_work_abs = torch.abs(pe_work)
            drop_part = math.floor(pe_work_abs.shape[0] * gamma)
            _, indices = torch.sort(pe_work_abs, dim=0)
            drop_indices = indices[0:drop_part]
            drop_rand_mask = torch.rand(drop_indices.size(0)).cuda()
            pe_work[drop_indices] = pe_work[drop_indices].masked_fill_(drop_rand_mask <= alpha, 0)
            x[np.arange(i, n_rows, num_pe)] = pe_work
        elif n_dims > 1:
            pe_work = x[np.arange(i, n_rows, num_pe), :]
            pe_work_abs = torch.abs(pe_work)
            drop_part = math.floor(pe_work_abs.shape[0] * gamma)
            _, indices = torch.sort(pe_work_abs, dim=0)
            drop_indices = indices[0:drop_part, :]
            drop_rand_mask = torch.rand(drop_indices.size(0), drop_indices.size(1)).cuda()
            for j in range(0, n_cols):
                pe_work[drop_indices[:, j], j] = pe_work[drop_indices[:, j], j].masked_fill_(
                    drop_rand_mask[:, j] <= alpha, 0)
            x[np.arange(i, n_rows, num_pe), :] = pe_work
    return x

# def cbtd(x, gamma, alpha):
#     """
#     x: input tensor
#     gamma: target sparsity
#     alpha: drop probability
#     mode: 0 - normal; 1 -
#     """
#     import math
#     device = x.device
#     n_dims = len(x.size())
#     n_rows = x.size(0)
#
#     if n_dims == 1:
#         n_cols = 1
#     elif n_dims == 2:
#         n_cols = x.size(1)
#     else:
#         raise RuntimeError('Input tensor must have 1 or 2 dimensions.')
#
#     x_abs = torch.abs(x)
#     drop_part = math.floor(n_rows * gamma)
#     _, indices = torch.sort(x_abs, dim=0)
#     if n_dims == 1:
#         drop_indices = indices[0:drop_part]
#         drop_rand_mask = torch.rand(drop_indices.size(0)).cuda()
#         x[drop_indices] = x[drop_indices].masked_fill_(drop_rand_mask <= alpha, 0)
#     elif n_dims > 1:
#         drop_indices = indices[0:drop_part, :]
#         drop_rand_mask = torch.rand(drop_indices.size(0), drop_indices.size(1), device=device)
#         for j in range(0, n_cols):
#             x[drop_indices[:, j], j] = x[drop_indices[:, j], j].masked_fill_(
#                 drop_rand_mask[:, j] <= alpha, 0)
#     return x

def castToFixed(x, m, n):
    """
    :param x: input numpy array
    :param m: number of integer bits before the decimal point
    :param n: number of fraction bits after the decimal point
    :return: tensor quantized to fixed-point precision
    """
    power = 2. ** n
    max_val = 2. ** (m + n) - 1
    value = x * power
    value = np.round(value)  # rounding
    value = np.clip(value, -max_val - 1, max_val)  # saturation arithmetic
    return value

def interleave(x, width):
    """
    PyTorch Function to interleave a
    :param x:
    :param width:
    :return:
    """

    x_dim0 = x.size(0)
    num_dim = len(x.size())
    interleaved = []
    for i in range(0, width):
        idx_curr = torch.arange(i, x_dim0, width).long()
        if num_dim == 1:
            interleaved.append(x[idx_curr])
        elif num_dim == 2:
            interleaved.append(x[idx_curr, :])
    interleaved = torch.stack(interleaved)
    return interleaved

def get_array_work(workload, num_array_pe, num_array):
    n_col = workload.size(0)
    hw_n_col = int(num_array_pe * np.ceil(float(n_col) / float(num_array_pe)))
    num_pad_col = int(hw_n_col - n_col)
    workload_pad_col = torch.zeros(num_pad_col)
    act_workload = torch.cat((workload, workload_pad_col), dim=0)

    array_col_work = torch.zeros(int(num_array), int(hw_n_col / num_array))
    array_work = torch.zeros(int(num_array))
    for i in range(0, int(num_array)):
        array_col_work[i] = act_workload[np.arange(i, hw_n_col, num_array)]
        array_work[i] = torch.sum(array_col_work[i])

    return array_work, array_col_work

def cast_fp_to_fxp(x, m, n):
    """
    Converts a floating point number into a fixed point number

    :param x: input numpy array
    :param m: number of integer bits before the decimal point
    :param n: number of fraction bits after the decimal point
    :return: tensor quantized to fixed-point precision
    """
    power = 2. ** n
    max_val = 2. ** (m + n)
    value = x * power
    value = np.round(value)  # rounding
    value = np.clip(value, -max_val, max_val - 1).astype(np.int32)  # saturation arithmetic
    return value

def cast_fxp_to_hex(x, data_type):
    x = np.asarray(x, dtype=data_type)
    return x

# def cast_signed_to_unsigned(x, bw):
#     temp = x
#     idx_below_zero = np.argwhere(temp<0)
#     abs_below_zero = np.abs(temp[idx_below_zero])
#     max_val = 2**bw
#     unsigned_temp = max_val - abs_below_zero
#     temp[idx_below_zero] = unsigned_temp
#     return temp

def cast_signed_to_unsigned(x, bw):
    temp = x
    if temp.ndim == 1:
        temp = np.expand_dims(temp, axis=-1)
    # print("Temp Size: ", temp.shape)
    for i in range(0, temp.shape[-1]):
        # print("0")
        idx_below_zero = np.argwhere(temp[:, i] < 0)
        # print("idx_below_zero", idx_below_zero.shape)
        abs_below_zero = np.abs(temp[idx_below_zero, i])
        # print("abs_below_zero", abs_below_zero.shape)
        max_val = 2 ** bw
        unsigned_temp = max_val - abs_below_zero
        temp[idx_below_zero, i] = unsigned_temp
    return temp

def gen_sim_file(file_path, pdict):
    """
    Generate Modelsim Simulation Files for NN parameters

    :param pdict:    Dictionary from which to save numpy arrays of NN parameters in fp
    :param qi:       number of integer bits before the decimal point in fxp
    :param qf:       number of integer bits after the decimal point in fxp
    :return:         mat_param_fxp - parameter in fixed-point numbers
    """

    # Iterate over params in dict
    pdict_out = {}
    for key, value in pdict.items():

        # Get Inputs
        param = value[0]
        qi = value[1]
        qf = value[2]
        idx = value[3]
        idx_bw = value[4]

        if param.ndim == 1:
            param = np.expand_dims(param, axis=-1)
        if idx is not False:
            if idx.ndim == 1:
                idx = np.expand_dims(idx, axis=-1)

        # Get unsigned variable type
        if idx is not False:
            bit_width = qi + qf + idx_bw
        else:
            bit_width = qi + qf

        if 0 < bit_width <= 8:
            var_type = 'uint8'
        elif 8 < bit_width <= 16:
            var_type = 'uint16'
        elif 16 < bit_width <= 32:
            var_type = 'uint32'
        else:
            raise ValueError('Bit width of parameter must not be larger than 32...')

        # Convert fp to fxp
        if idx is not False:
            # param = quantizeTensor(param, qi, qf, 1) # Remove if necessary
            # fxp_param = param
            fxp_param = cast_fp_to_fxp(param, qi, qf)
            fxp_param_unsigned = cast_signed_to_unsigned(fxp_param, qi + qf)
            fxp_param_shift = fxp_param_unsigned * (2 ** idx_bw)
            param = fxp_param_shift + idx
        else:
            param = cast_fp_to_fxp(param, qi, qf)

        # Get param dimension
        dim0 = param.shape[0]
        dim1 = param.shape[1]

        # Create file
        try:
            f = open(file_path + '/' + key + '.txt', 'w')
        except IOError as err:
            print("I/O error({0}): {1}".format(err.errno, err.strerror))

        # Write parameters column-wisely
        param_hex = cast_fxp_to_hex(param, var_type)
        for i in range(0, dim1):
            for j in range(0, dim0):
                if var_type == 'uint8':
                    f.write('%2.2X\n' % param_hex[j][i])
                elif var_type == 'uint16':
                    f.write('%4.4X\n' % param_hex[j][i])
                elif var_type == 'uint32':
                    f.write('%8.8X\n' % param_hex[j][i])
                else:
                    f.write('%X\n' % param_hex[j][i])
        f.close()

def gen_clib(file_path, file_name, pdict, num_layer=None, input_size=None, hid_size=None):
    """
    Generate C Library for NN parameters

    :param filename: name of the C Library
    :param pdict:    Dictionary from which to save numpy arrays of NN parameters in fp
    :param qdict:    Dictionary from which to save properties of corresponding variables in pdict in tuples
        Elements of each tuple in qdict in order:
        - quantize
        - qi:        number of integer bits before the decimal point in fxp
        - qf:        number of integer bits after the decimal point in fxp
        - var_class: 0 - Constant | 1 - General Tensors | 2 - NN Parameter |
        - print_row_size
    :return:         mat_param_fxp - parameter in fixed-point numbers
    """
    print_row_size = 20
    # Create file
    try:
        f_h = open(file_path + '/' + file_name + '.h', 'w')
    except IOError as err:
        print("I/O error({0}): {1}".format(err.errno, err.strerror))
    try:
        f = open(file_path + '/' + file_name + '.c', 'w')
    except IOError as err:
        print("I/O error({0}): {1}".format(err.errno, err.strerror))

    # Write macro definitions
    macro_name = file_name.upper() + '_H'
    f_h.write('#ifndef ' + macro_name + '\n' + '#define ' + macro_name + '\n')
    f.write('#include "' + file_name + '.h"\n')

    # Iterate over params in dict
    for key, value in pdict.items():
        # Get Inputs
        param = np.asarray(value[0])
        qi = value[1]
        qf = value[2]
        idx = value[3]
        idx_bw = value[4]
        param_type = value[5]

        # Get unsigned variable type
        if idx is not False:
            bit_width = qi + qf + idx_bw
        else:
            bit_width = qi + qf

        if 0 < bit_width <= 8:
            var_type = 'char'
            nn_var_type = 'unsigned char'
            np_var_type = 'uint8'
        elif 8 < bit_width <= 16:
            var_type = 'short'
            nn_var_type = 'unsigned short'
            np_var_type = 'uint16'
        elif 16 < bit_width <= 32:
            var_type = 'long'
            nn_var_type = 'unsigned long'
            np_var_type = 'uint32'
        else:
            raise ValueError('Bit width of parameter must not be larger than 32...')

        # Convert fp to fxp
        if param_type != 0:
            if idx is not False:
                # param = quantizeTensor(param, qi, qf, 1) # Remove if necessary
                fxp_param = cast_fp_to_fxp(param, qi, qf)
                fxp_param_unsigned = cast_signed_to_unsigned(fxp_param, qi + qf)
                fxp_param_shift = fxp_param_unsigned * (2 ** idx_bw)
                param = fxp_param_shift + idx
                param = cast_fxp_to_hex(param, np_var_type)
                param = param
            else:
                param = cast_fp_to_fxp(param, qi, qf)

        if param_type == 0:  # Constant
            f_h.write('#define ' + key.upper() + ' %d\n' % param)
        elif param_type == 1:  # General Tensor
            if param.ndim == 1:
                param = np.expand_dims(param, axis=0)

            # Get param dimension
            size = param.size
            dim0 = param.shape[0]
            dim1 = param.shape[1]

            # Write variable declaration
            f_h.write('/*\n'
                      ' * Var Type: General Matrix \n'
                      ' * Var Name: %s' % key + '\n'
                                                ' * Bit Width: %d' % bit_width + '\n'
                                                                                ' * Dimension: (%d, %d)' % (
                          dim0, dim1) + '\n'
                                        ' */\n'
                      )
            f_h.write('#define ' + key.upper() + '_NUM_ROWS %d\n' % dim0)
            f_h.write('#define ' + key.upper() + '_NUM_COLS %d\n' % dim1)
            f_h.write('#define ' + key.upper() + '_MAT_SIZE %d\n' % size)
            f_h.write('extern const %s %s[' % (var_type, key) + key.upper() + '_MAT_SIZE];\n')
            f.write('const %s %s[' % (var_type, key) + key.upper() + '_MAT_SIZE] = {\n')

            # Write parameters column-wisely
            for i in range(0, dim1):
                for j in range(0, dim0):
                    if var_type == 'float':
                        f.write('%.20f' % param[j][i])
                    else:
                        f.write('%d' % param[j][i])
                    if not (i >= dim1 - 1 and j >= dim0 - 1):
                        f.write(',')
                    if (i * dim0 + j) % print_row_size == print_row_size - 1:
                        f.write('\n')
            f.write('};\n\n')

        elif param_type == 2:  # NN Parameter
            if param.ndim == 1:
                param = np.expand_dims(param, axis=0)

            # Get param dimension
            size = param.size
            dim0 = param.shape[0]
            dim1 = param.shape[1]

            # Write variable declaration
            f_h.write('/*\n'
                      ' * Var Type: NN Parameter Matrix \n'
                      ' * Var Name: %s' % key + '\n'
                                                ' * Bit Width: %d' % bit_width + '\n'
                                                                                ' * Dimension: (%d, %d)' % (
                          dim0, dim1) + '\n'
                                        ' */\n'
                      )

            f_h.write('#define ' + key.upper() + '_NUM_ROWS %d\n' % dim0)
            f_h.write('#define ' + key.upper() + '_NUM_COLS %d\n' % dim1)
            f_h.write('#define ' + key.upper() + '_MAT_SIZE %d\n' % size)
            f_h.write('extern const %s %s[' % (
                nn_var_type, key) + key.upper() + '_MAT_SIZE] __attribute__ ((aligned (%d)));\n' % bit_width)
            f.write('const %s %s[' % (nn_var_type, key) + key.upper() + '_MAT_SIZE] = {\n')
            if num_layer is not None:
                f_h.write('#define ' + key.upper() + '_NUM_LAYERS %d\n' % num_layer)
            if input_size is not None:
                f_h.write('#define ' + key.upper() + '_INP_SIZE %d\n' % input_size)
            if hid_size is not None:
                f_h.write('#define ' + key.upper() + '_HID_SIZE %d\n' % hid_size)
            # if 'rnn' in key:
            #     f_h.write('extern const %s %s[RNN_PARAM_Hweight_tensor_SIZE] __attribute__ ((aligned (%d)));\n' % (var_type, key, bit_width))
            #     f.write('const %s %s[RNN_PARAM_Hweight_tensor_SIZE] = {\n' % (var_type, key))
            # else:
            #     f_h.write('extern const %s %s[%d];\n' % (var_type, key, size))
            #     f.write('const %s %s[%d] = {\n' % (var_type, key, size))

            # Write parameters column-wisely
            for i in range(0, dim1):
                for j in range(0, dim0):
                    f.write('%d' % param[j][i])
                    if not (i >= dim1 - 1 and j >= dim0 - 1):
                        f.write(',')
                    if (i + j) % print_row_size == print_row_size - 1:
                        f.write('\n')
            f.write('};\n\n')

    f_h.write('#endif')
    f.close()
    f_h.close()


from main import util

class Angle:
    def __init__(self, center, point1, point2):
        self.center = center
        self.point1 = point1
        self.point2 = point2
        self.vector1, self.vector2, self.angle = self.compute_angle()

    def compute_angle(self):

        dim = len(self.center)
        vector1 = []
        vector2 = []
        for i in range(dim):
            vector1.append(self.point1[i] - self.center[i])
            vector2.append(self.point2[i] - self.center[i])

        if self.point1 == self.center or self.point2 == self.center:
            angle = 0
        else:
            angle = util.calculate_vector_angle(vector1, vector2)
        return vector1, vector2, angle

    def find(self, point_x, point_y):
        if self.point1 == point_x and self.point2 == point_y:
            return True
        if self.point1 == point_y and self.point2 == point_x:
            return True

        return False

    def get_other_point(self, point):
        if self.point1 == point:
            return self.point2
        elif self.point2 == point:
            return self.point1
        else:
            return None


class GameEngine:
    def get_divisors(self, num):
        divisors = {}
        counter = 1
        for i in range(2, num):
            if num % i == 0:
                divisors[counter] = i
                counter += 1

        return divisors

    def is_prime(self, divisors):
        return len(divisors) == 0


import argparse
import os
from corescore.masks import CoreImageProcessor

def process_images(image_dir, labels):
    """Create masks for labelled images.
    For now, merge rock fragment labels from both core boxes"""
    coreProcessor = CoreImageProcessor("Images",
                                        labels=labels,
                                        merge_fragment_labels=True)
    for image in coreProcessor.core_types:
        mask_file = coreProcessor.processImage(image)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--images', help="Directory of images to make masks for")
    parser.add_argument('--labels', help="JSON file with labels for images")
    args = parser.parse_args()

    images = args.images
    if not images:
        images = 'Images/train'
    labels = args.labels
    if not labels:
        labels = 'Core_labels.json'

    process_images(images, labels)


"""
TODO: copy readme?
      remove haplotypes stuff
      remove blast mode
      remove db stuff?

This is where the magic happens!

It is meant to be able to run locally and in pbs_runner.py there is also specific support for pbs cluster systems.

___Overview___
The pipeline is divided into 4 parts each having it's own .py file.
I   -  Preperation (data_preperation.py)
II  -  Processing (processing.py)
III -  Aggregation (aggregation.py)
IV  -  Visual Summary (summarize.py)

I - Data Prepeation
Input: directory containing fastq/gz files or a directory containing such directories.
Output: fastq files in sizes ready for efficient processing.

Depending on currently availble RAM and CPUs or given values in -mm / --max_memory and -cc / --cpu_count the script
will divide the files into an efficient number of files to later run the processing script on.
If given an merge_opposing it will also merge the forward and backward reads in corresponding fastq files.

II - Processing
Input: fastq files from part I and a reference fasta file.
Output: called_bases - every base called and its attributes.
        ignored_bases - every base ignored, its attributes and why it was ignored.
        suspicious_reads - every suspicious read and why it is suspicious.
        ignored_reads - every read ignored and why.
        read_counter - every read and a number describing how many times it was aligned by blast.
        blast files - alignment output files.

This is where the main logic of the pipeline happens and the runner runs this on each of the fastq files in parallel.
Basically it first aligns the reads to the reference with blast and then goes over every read and nucleotide and decides
whether to filter them out or leave them in.

III - Aggregation
Input: directory containing output of basecalling of stage II
Output: concatenations - concatenated outputs of stage II
        read_id_prefix_file - a json file containing a dictionary of read prefixes to save memory in the other files.
        mutation_read_list - a file describing in which reads each mutation appeared.
        freqs - a frequencies file describing the different alleles and their frequencies for each position.
        consensus_aligned_to_ref - a fasta file of the majority frequency derived from the freqs file aligned to the original reference
        consensus - a fasta file of the majority frequency derived from the freqs file
        read_counter - a file counting how many alignments were called for each read.

IV - Haplotype Inference
Input: directory containing output of stage III
Output: linked_mutations - file containing pairs of mutatations, their pvalues and their frequencies.
        stretches - file containing the aggregations of the stretches by their frequencies.
"""
import argparse
import getpass
import json
import os
import shutil
import concurrent.futures
import Bio
import multiprocessing as mp
import pandas as pd
from Bio import pairwise2, SeqIO
from datetime import datetime
from data_preparation import prepare_data
from summarize import graph_summary, create_stats_file
from processing import process_fastq
from aggregation import aggregate_processed_output, create_freqs_file, create_mutation_read_list_file
from logger import pipeline_logger
from utils import get_files_in_dir, get_sequence_from_fasta, get_mp_results_and_report, create_consensus_file, \
    get_files_by_extension, concatenate_files_by_extension, get_config, md5_dir, md5_file

def parallel_process(processing_dir, fastq_files, reference_file, quality_threshold, task, evalue, dust, num_alignments,
                     soft_masking, perc_identity, mode, reads_overlap):
    with concurrent.futures.ProcessPoolExecutor() as executor:
        future_tasks = {executor.submit(process_fastq, fastq_file, reference=reference_file, output_dir=processing_dir,
                                        quality_threshold=quality_threshold, task=task, evalue=evalue, dust=dust,
                                        num_alignments=num_alignments, soft_masking=soft_masking,
                                        perc_identity=perc_identity, mode=mode, reads_overlap=reads_overlap): fastq_file
                        for fastq_file in fastq_files}

        for future in concurrent.futures.as_completed(future_tasks):
            future.result()

def set_filenames(output_dir):
    filenames = {"freqs_file_path": os.path.join(output_dir, 'freqs.tsv'),
                 "linked_mutations_path": os.path.join(output_dir, 'linked_mutations.tsv'),
                 "mutation_read_list_path": os.path.join(output_dir, 'mutation_read_list.tsv'),
                 "stretches": os.path.join(output_dir, 'stretches.tsv'),
                 "blast_file": os.path.join(output_dir, 'blast.tsv'),
                 "read_counter_file": os.path.join(output_dir, 'read_counter.tsv'),
                 "summary_graphs": os.path.join(output_dir, 'summary.png'),
                 "processing_dir": os.path.join(output_dir, "processing"),
                 'linked_mutations_dir': os.path.join(output_dir, "linked_mutations"),
                 'data_dir': os.path.join(output_dir, "data")}
    filenames['basecall_dir'] = os.path.join(filenames['processing_dir'], 'basecall')
    os.makedirs(filenames['data_dir'], exist_ok=True)
    os.makedirs(filenames["basecall_dir"], exist_ok=True)
    return filenames

def create_consensus_and_check_alignment_with_ref(reference_file, align_to_ref, min_coverage, iteration_data_dir, basecall_dir,
                                                  iteration_counter, min_frequency):
    reference = get_sequence_from_fasta(reference_file)
    freqs_file_path = os.path.join(iteration_data_dir, f"freqs_{iteration_counter}.tsv")
    called_bases_files = get_files_by_extension(basecall_dir, "called_bases")
    create_freqs_file(called_bases_files=called_bases_files, output_path=freqs_file_path, reference=reference)
    if align_to_ref == "Y":
        consensus_path = os.path.join(iteration_data_dir, f"consensus_aligned_to_ref_{iteration_counter}.fasta")
        align_to_ref = True
    else:
        consensus_path = os.path.join(iteration_data_dir, f"consensus_{iteration_counter}.fasta")
        align_to_ref = False
    create_consensus_file(freqs_file=freqs_file_path, min_frequency=min_frequency,
                          min_coverage=min_coverage, output_file=consensus_path, align_to_ref=align_to_ref)
    consensus = get_sequence_from_fasta(consensus_path)
    #TODO: get helpdesk to create the right environment and remove this crap..
    alignment = pairwise2.align.globalxx(consensus, reference)[0]
    if float(Bio.__version__) > 1.76:
        alignment_score = alignment.score / max(len(consensus), len(reference))
        alignment = alignment.seqA
    else:
        alignment_score = int(alignment[2]) / max(len(consensus), len(reference))
        alignment = alignment[0]
    return alignment_score, alignment

def get_consensus_path(basecall_iteration_counter, align_to_ref, iteration_data_dir):
    consensus = os.path.join(iteration_data_dir, "consensus")
    if align_to_ref == "Y":
        consensus += "_aligned_to_ref"
    consensus += f"_{basecall_iteration_counter}.fasta"
    return consensus

def update_meta_data(output_dir, status, db_path, params=None):
    # TODO: write running time in meta_data
    json_file = os.path.join(output_dir, 'meta_data.json')
    if params is not None:
        meta_data = params
        del meta_data['filenames']
        del meta_data['log']
        meta_data['username'] = getpass.getuser()
        meta_data['start_time'] = datetime.now().strftime('%Y-%m-%d-%H:%M')
    else:
        with open(json_file) as read_handle:
            meta_data = json.load(read_handle)
    meta_data['status'] = status
    with open(json_file, 'w') as write_handle:
        json.dump(meta_data, write_handle)
    build_db(db_path)

def build_db(db_path):
    # TODO: get user feedback on what db functionality should be
    os.makedirs(db_path, exist_ok=True)
    outputs = [f.path for f in os.scandir(db_path) if f.is_dir()]
    db_rows = []
    for directory in outputs:
        json_file = os.path.join(directory, "meta_data.json")
        if os.path.isfile(json_file):
            with open(json_file) as read_handle:
                meta_data = json.load(read_handle)
            db_rows.append(meta_data)
    db = pd.DataFrame.from_dict(db_rows)
    db.to_csv(os.path.join(db_path, 'db.tsv'), sep='\t')

def create_dir_also_if_exists(output_dir):
    try:
        os.makedirs(output_dir, exist_ok=False)
    except FileExistsError:
        i = 2
        while i < 1000:
            try:
                output_dir = f"{output_dir}_{i}"
                os.makedirs(output_dir, exist_ok=False)
            except FileExistsError:
                i += 1
    return output_dir

def assign_output_dir(db_path, alias=None):
    now = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
    output_dir_name = now
    if alias:
        output_dir_name = alias + "_" + output_dir_name
    output_dir = os.path.join(db_path, output_dir_name)
    output_dir = create_dir_also_if_exists(output_dir)
    return output_dir

def process_data(align_to_ref, dust, evalue, fastq_files, log, max_basecall_iterations, min_frequency,
                 min_coverage, mode, num_alignments, overlapping_reads, output_dir, perc_identity, processing_dir,
                 quality_threshold, reference_file, soft_masking, task, basecall_dir):
    alignments = []
    for basecall_iteration_counter in range(1, max_basecall_iterations + 1):
        log.info(f"Processing fastq files iteration {basecall_iteration_counter}/{max_basecall_iterations}")
        parallel_process(processing_dir=processing_dir, fastq_files=fastq_files, reference_file=reference_file,
                         quality_threshold=quality_threshold, task=task, evalue=evalue, dust=dust, mode=mode,
                         num_alignments=num_alignments, soft_masking=soft_masking, perc_identity=perc_identity,
                         reads_overlap=overlapping_reads)
        iteration_data_dir = os.path.join(output_dir, 'iteration_data')
        os.makedirs(iteration_data_dir, exist_ok=True)
        alignment_score, alignment = create_consensus_and_check_alignment_with_ref(reference_file=reference_file,
                                                                                iteration_counter=basecall_iteration_counter,
                                                                                basecall_dir=basecall_dir,
                                                                                align_to_ref=align_to_ref,
                                                                                iteration_data_dir=iteration_data_dir,
                                                                                min_coverage=min_coverage,
                                                                                min_frequency=min_frequency)
        alignments.append(alignment)
        log.info(f'Iteration alignment score: {round(alignment_score, 4)}')
        if alignment_score == 1:
            break
        consensus_path = get_consensus_path(basecall_iteration_counter=basecall_iteration_counter,
                                            align_to_ref=align_to_ref, iteration_data_dir=iteration_data_dir)
        reference_file = consensus_path
    return reference_file, alignments, basecall_iteration_counter

def validate_input(output_dir, input_dir, reference_file, mode):
    if mode != 'RefToSeq' and mode != 'SeqToRef':
        raise Exception("blast mode must be either RefToSeq or SeqToRef! ")
    if os.path.exists(output_dir):
        list_dir = os.listdir(output_dir)
        if len(list_dir) > 0:
            if not ((len(list_dir) == 1) and (list_dir[0] == 'pbs_logs')):
                raise Exception("output_dir must be path to a new or empty directory!")
    if not os.path.isfile(reference_file):
        raise Exception("reference_file must exist!")
    else:
        with open(reference_file, "r") as handle:
            fasta = SeqIO.parse(handle, "fasta")
            if not any(fasta):
                raise Exception("reference_file must be of type fasta!")
    if not os.path.isdir(input_dir):
        raise Exception("Input_dir must exist!")
    files_fasta = get_files_by_extension(input_dir, "fastq")
    files_fastagz = get_files_by_extension(input_dir, "fastq.gz")
    if len(files_fasta) == 0 and len(files_fastagz) == 0:
        raise Exception("Could not find files ending with '.fastq' or 'fastq.gz' in input_dir !")

def try_to_rmtree(path_to_delete, retry_attempts=5):
    # this is a hack to avoid a bug rmtree throws about a non-empty dir sometimes...
    if retry_attempts:
        try:
            if os.path.isdir(path_to_delete):
                shutil.rmtree(path_to_delete)
        except:
            try_to_rmtree(path_to_delete, retry_attempts-1)

def remove_unnecessary_files(dirs_to_remove, files_to_remove):
    for file_path in files_to_remove:
        os.remove(file_path)
    for dir_path in dirs_to_remove:
        try_to_rmtree(dir_path)

def runner(input_dir, reference_file, output_dir, max_basecall_iterations, min_coverage, db_comment,
           quality_threshold, task, evalue, dust, num_alignments, soft_masking, perc_identity, mode,
           align_to_ref, cleanup, min_frequency, cpu_count, overlapping_reads,
           db_path, max_memory):
    if not db_path:
        db_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'db')
    if not output_dir:
        output_dir = assign_output_dir(db_path)
    validate_input(output_dir, input_dir, reference_file, mode)
    log = pipeline_logger(logger_name='AccuNGS-Runner', log_folder=output_dir)
    try:
        filenames = set_filenames(output_dir=output_dir)
        if not cpu_count:
            cpu_count = mp.cpu_count()
        input_dir_hash = md5_dir(input_dir)
        reference_file_hash = md5_file(reference_file)
        params = locals().copy()
        update_meta_data(params=params, output_dir=output_dir, status='Setting up...', db_path=db_path)
        log.debug(f"runner params: {params}")  # TODO: why does this contain status..?
        log.info("Preparing data")
        update_meta_data(output_dir=output_dir, status='Preparing data...', db_path=db_path)
        prepare_data(input_dir=input_dir, output_dir=filenames['data_dir'], overlapping_reads=overlapping_reads,
                     cpu_count=cpu_count, max_memory=max_memory)
        data_files = get_files_in_dir(filenames['data_dir'])
        fastq_files = [file_path for file_path in data_files if "fastq.part_" in os.path.basename(file_path)]
        log.info(f"Processing {len(fastq_files)} fastq files.")
        update_meta_data(output_dir=output_dir, status='Processing data...', db_path=db_path)
        reference_file, alignments, iterations = process_data(
            align_to_ref=align_to_ref, dust=dust, min_frequency=min_frequency,
            evalue=evalue, fastq_files=fastq_files, log=log, soft_masking=soft_masking,
            max_basecall_iterations=max_basecall_iterations, min_coverage=min_coverage,
            mode=mode, num_alignments=num_alignments, overlapping_reads=overlapping_reads,
            output_dir=output_dir, perc_identity=perc_identity, reference_file=reference_file,
            processing_dir=filenames['processing_dir'], quality_threshold=quality_threshold,
            task=task, basecall_dir=filenames['basecall_dir'])
        log.info("Aggregating processed fastq files outputs...")
        last_freqs = os.path.join(output_dir, 'iteration_data', f'freqs_{iterations}.tsv')
        shutil.copy(last_freqs, filenames['freqs_file_path'])
        aggregate_processed_output(input_dir=filenames['processing_dir'], output_dir=output_dir,
                                   min_coverage=min_coverage, min_frequency=min_frequency, cleanup=cleanup)
        create_stats_file(output_dir, filenames, alignments)
        log.info("Generating graphs...")
        graph_summary(freqs_file=filenames['freqs_file_path'], blast_file=filenames['blast_file'],
                      read_counter_file=filenames['read_counter_file'],
                      output_file=filenames['summary_graphs'], min_coverage=min_coverage)  # TODO: drop low quality mutations?
        log.info(f"Most outputs are ready in {output_dir} !")
        if cleanup == "Y":
            dirs_to_remove = [filenames['basecall_dir'], filenames['data_dir']]
            files_to_remove = [filenames['blast_file']]
            log.info(f"Removing intermediary files to save space...")
            remove_unnecessary_files(dirs_to_remove, files_to_remove)
        update_meta_data(output_dir=output_dir, status='Done', db_path=db_path)
        log.info(f"Done!")
    except Exception as e:
        log.exception(e)
        update_meta_data(output_dir=output_dir, status="Failed! see logs for details.", db_path=db_path)

def create_runner_parser():
    # TODO: dynamic defaults?
    parser = argparse.ArgumentParser(description="Note: Default parameter values are retrieved from config.ini in your "
                                                 "installation directory.")
    parser.add_argument("-i", "--input_dir", required=True,
                        help="Path to directory containing fastq/gz files or sub directories containg fastq/gz files.")
    parser.add_argument("-o", "--output_dir", help="A directory for output files. "
                                                   "If none is given will put it in the db")
    parser.add_argument("-r", "--reference_file", required=True, help="Full path to reference file (including "
                                                                      "extension) of type fasta to align against.")
    parser.add_argument("-m", "--max_basecall_iterations", type=int,
                        help="Number of times to rerun with previous consensus as the new reference before giving up.")
    parser.add_argument("-or", "--overlapping_reads",
                        help="Y/N/P, run pipeline with, without, or with partial overlapping reads. Y- merge opposing "
                             "reads in the same directory and drop non overlapping areas of the reads. "
                             "P - Merge opposing reads but keep non overlapping areas."
                             "N - No merge, assume reads are independent. "
                             "Y & P assume 2 fastq/gz files in each sub directory of the input_dir.")
    parser.add_argument("-bt", "--blast_task", help="blast's task parameter")
    parser.add_argument("-be", "--blast_evalue", help="blast's e value parameter", type=float)
    parser.add_argument("-bd", "--blast_dust", help="blast's dust parameter")
    parser.add_argument("-bn", "--blast_num_alignments", type=int, help="blast's num_alignments parameter")
    parser.add_argument("-bp", "--blast_perc_identity", type=int, help="blast's perc_identity parameter")
    parser.add_argument("-bs", "--blast_soft_masking", help="blast's soft_masking parameter")
    parser.add_argument("-bm", "--blast_mode", help="RefToSeq or SeqToRef")  # TODO: docs
    parser.add_argument("-qt", "--quality_threshold", type=int,
                        help="phred score must be higher than this to be included")
    parser.add_argument("-mc", "--min_coverage", type=int,
                        help="positions with less than this coverage will be substituted by Ns in the consensus")
    parser.add_argument("-mf", "--min_frequency", type=float,
                        help="positions with less than this frequency will be substituted by Ns in the consensus")
    parser.add_argument("-ar", "--align_to_ref", help="Y/N, generate consensus aligned to the original reference")
    parser.add_argument("-c", "--cleanup", help="Remove intermediary files when done in order to save space")
    parser.add_argument("-cc", "--cpu_count", help="max number of cpus to use (None means all)", type=int)
    parser.add_argument("-db", "--db_path", help='path to db directory')
    parser.add_argument("-dbc", "--db_comment", help='comment to store in db')
    parser.add_argument("-mm", "--max_memory", help='limit memory usage to this many megabytes '
                                                    '(None would use available memory when starting to run)')
    parser.add_argument("-ch", "--calculate_haplotypes", help='Y/N, Run pipeline including calculating haplotypes')
    return parser

if __name__ == "__main__":
    parser = create_runner_parser()
    parser_args = vars(parser.parse_args())
    args = dict(get_config()['runner_defaults'])
    args.update({key: value for key, value in parser_args.items() if value is not None})
    runner(input_dir=args['input_dir'], output_dir=args['output_dir'], reference_file=args['reference_file'],
           max_basecall_iterations=int(args['max_basecall_iterations']), overlapping_reads=args['overlapping_reads'],
           quality_threshold=int(args['quality_threshold']), task=args['blast_task'], max_memory=args['max_memory'],
           evalue=float(args['blast_evalue']), dust=args['blast_dust'],
           num_alignments=int(args['blast_num_alignments']),
           mode=args['blast_mode'], perc_identity=float(args['blast_perc_identity']), cpu_count=args['cpu_count'],
           min_coverage=int(args['min_coverage']), db_comment=args['db_comment'],
           soft_masking=args['blast_soft_masking'], min_frequency=float(args['min_frequency']),
           cleanup=args['cleanup'], align_to_ref=args['align_to_ref'], db_path=args['db_path'])


#!/usr/bin/python3

if __name__ == '__main__':
    import sys

    # Convert the str type to int type.

    try:

        List_Num = [int(item) for item in sys.argv[1:]]

    except:

        print('0')

    # Then Convert the list an iterable data type --> iter()

    Iter_Num = iter(List_Num)
    print("{result_sum}".format(result_sum=sum(Iter_Num)))


text_file = open("input.txt", "r")
lines = text_file.read().split('\n')
text_file.close()

year = 2020

def isInArr(n):
  print(n)
  if(n == None):
    return False

  if str(n[1]) in lines:
    return True

  return False

def getCounterpart(x):
  target = year - int(x)

  counterparts = next(
    filter(
      lambda x: isInArr(x) == True,
      map(lambda x: getSubCounterparts(target, x), lines)
    ),
    None
  )

  if counterparts == None:
    return None

  return [ int(x), counterparts[0], counterparts[1] ]

def getSubCounterparts(target, y):
  return [ int(y), target - int(y) ]

target = next(
  filter(
    lambda x: isInArr(x) == True,
    map(getCounterpart, lines)
  ),
  None
)

print("We found:")
print(target)
print("Sum:")
print(str(target[ 0] + target[1] + target[2]))
print("Product:")
print(str(target[0] * target[1] * target[2]))


from utils_graph_learning import multi_class_accuracy
import torch.nn as nn
import numpy as np

def process_arguments(args):

    # ###### choose the substructures: usually loaded from networkx,
    # ###### except for 'all_simple_graphs' where they need to be precomputed,
    # ###### or when a custom edge list is provided in the input by the user

    # define if degree is going to be used as a feature and when (for each layer or only at initialization)

    # args['degree_as_tag'] = [args['degree_as_tag']] + [False for _ in range(args['num_layers']-1)]

    # define if existing features are going to be retained when the degree is used as a feature
    # args['retain_features'] = [args['retain_features']] + [True for _ in range(args['num_layers']-1)]

    # replicate d_out dimensions if the rest are not defined (msg function, mlp hidden dimension, encoders, etc.)
    # and repeat hyperparams for every layer
    if args['d_msg'] == -1:
        args['d_msg'] = [None for _ in range(args['num_layers'])]
    elif args['d_msg'] is None:
        args['d_msg'] = [args['d_out'] for _ in range(args['num_layers'])]
    else:
        args['d_msg'] = [args['d_msg'] for _ in range(args['num_layers'])]

    if args['d_h'] is None:
        args['d_h'] = [[args['d_out']] * (args['num_mlp_layers'] - 1) for _ in range(args['num_layers'])]
    else:
        args['d_h'] = [[args['d_h']] * (args['num_mlp_layers'] - 1) for _ in range(args['num_layers'])]

    args['d_out_edge_encoder'] = [args['d_out'] for _ in range(args['num_layers'])]

    args['d_out_node_encoder'] = args['d_out']

    args['d_out_id_embedding'] = args['d_out']

    args['d_out_degree_embedding'] = args['d_out']

    # repeat hyperparams for every layer
    args['d_out'] = [args['d_out'] for _ in range(args['num_layers'])]

    args['train_eps'] = [False for _ in range(args['num_layers'])]

    if len(args['final_projection']) == 1:
        args['final_projection'] = [args['final_projection'][0] for _ in range(args['num_layers'])] + [True]

    args['bn'] = [args['bn'] for _ in range(args['num_layers'])]
    args['dropout_features'] = [args['dropout_features'] for _ in range(args['num_layers'])] + [args['dropout_features']]

    # loss function & metrics
    if args['loss_fn'] == 'CrossEntropyLoss':
        assert args['regression'] is False, "Can't use Cross-Entropy loss in regression."
        loss_fn = nn.CrossEntropyLoss()
    elif args['loss_fn'] == 'BCEWithLogitsLoss':
        assert args['regression'] is False, "Can't use binary Cross-Entropy loss in regression."
        loss_fn = nn.BCEWithLogitsLoss()
    elif args['loss_fn'] == 'MSELoss':
        loss_fn = nn.MSELoss()
    elif args['loss_fn'] == 'L1Loss':
        loss_fn = nn.L1Loss()
    else:
        raise NotImplementedError

    if args['prediction_fn'] == 'multi_class_accuracy':
        assert args['regression'] is False, "Can't use Classification Accuracy metric in regression."
        prediction_fn = multi_class_accuracy
    elif args['prediction_fn'] == 'MSELoss':
        prediction_fn = nn.MSELoss(reduction='sum')
    elif args['prediction_fn'] == 'L1Loss':
        prediction_fn = nn.L1Loss(reduction='sum')
    elif args['prediction_fn'] == 'None':
        prediction_fn = None
    else:
        raise NotImplementedError

    if args['regression']:
        perf_opt = np.argmin
    else:
        perf_opt = np.argmax

    return args, loss_fn, prediction_fn, perf_opt


'''
Final Project - Web Startup
userController.py
- JaeYoung Chang (jchang5)
- Maggie Farrell (mfarre22)
'''

import cherrypy
import re, json
from user_library import _user_database

class UserController(object):
    def __init__(self, udb=None):
        if udb is None:
            self.udb = _user_database()
        else:
            self.udb = udb

    def POST_USER(self, username):
        output = {'result': 'success'}

        data = json.loads(cherrypy.request.body.read().decode('utf-8'))

        self.udb.post_user(username, data)

        return json.dumps(output)

    def DELETE_USER(self, username):
        output = {'result': 'success'}

        self.udb.delete_user(username)
        return json.dumps(output)

    def GET_USERS(self):
        output = {'result' : 'success'}

        try:
            result = self.udb.get_users()
            if result is not None:
                output.update(result)
            else:
                output['result'] = 'error'
                output['message'] = 'users not found'
        except Exception as ex:
            output['result'] = 'error'
            output['message'] = str(ex)
        return json.dumps(output)


def run(instructions):
    """Return the accumulator value, or None if a loop is encountered."""
    accumulator = 0
    address = 0  # The address of current instruction.
    visited = set()  # The addresses of visited instructions.

    while address < len(instructions):
        if address in visited:
            return None
        visited.add(address)
        op, arg = instructions[address].split()
        if op == 'nop':
            address += 1
        elif op == 'acc':
            address += 1
            accumulator += int(arg)
        else:  # op == 'jmp'
            address += int(arg)

    return accumulator

def fix(instructions):
    """A generator to yield all fixed programs."""
    replace = dict(nop='jmp', jmp='nop')
    for i, line in enumerate(instructions):
        if line[:3] in replace:
            line = replace[line[:3]] + line[3:]
            yield instructions[:i] + [line] + instructions[i+1:]

def fix_and_run(instructions):
    """Fix program until it works. Then return the accumulator value."""
    for program in fix(instructions):
        accumulator = run(program)
        if accumulator:
            return accumulator

if __name__ == '__main__':
    with open('../data/day08.txt') as f:
        instructions = [line.strip() for line in f]
    print(fix_and_run(instructions))


"""
# Sync Water API with USGS Sites

[USGS Site Web Service](https://waterservices.usgs.gov/rest/Site-Service.html)

[Understanding Output](https://waterservices.usgs.gov/rest/Site-Service.html#Understanding)

URL

- https://waterservices.usgs.gov/nwis/site/?format=rdb&stateCd={state}&period=P52W&siteType=LK,ST&siteStatus=all&hasDataTypeCd=iv,aw

## USGS Fields

| Field                 | Desc                              |
| --------------------- | --------------------------------- |
| agency_cd             | Agency                            |
| site_no               | Site identification number        |
| station_nm            | Site name                         |
| site_tp_cd            | Site type                         |
| dec_lat_va            | Decimal latitude                  |
| dec_long_va           | Decimal longitude                 |
| coord_acy_cd          | Latitude-longitude accuracy       |
| dec_coord_datum_cd    | Decimal Latitude-longitude datum  |
| alt_va                | Altitude of Gage/land surface     |
| alt_acy_va            | Altitude accuracy                 |
| alt_datum_cd          | Altitude datum                    |
| huc_cd                | Hydrologic unit code              |

"""
import csv
from datetime import datetime, timedelta
from io import StringIO
from textwrap import dedent

import helpers.water as water
import requests
from airflow import DAG
from airflow.operators.python import PythonOperator

# Default arguments
default_args = {
    "owner": "airflow",
    "depends_on_past": False,
    "start_date": (datetime.utcnow() - timedelta(days=2)).replace(minute=0, second=0),
    "catchup_by_default": False,
    "email_on_failure": False,
    "email_on_retry": False,
    "retries": 1,
    "retry_delay": timedelta(minutes=10),
    # 'max_active_runs':1,
    # 'concurrency':4,
}

with DAG(
    default_args=default_args,
    dag_id="a2w_sync_usgs_sites",
    tags=["a2w", "usgs"],
    schedule="@daily",
    doc_md=dedent(__doc__),
) as dag:

    ########################################
    def fetch_and_write(state_abbrev):

        url = f"https://waterservices.usgs.gov/nwis/site/?format=rdb&stateCd={state_abbrev}&period=P52W&siteType=LK,ST&siteStatus=all&hasDataTypeCd=iv,aw"
        # print(url)
        r = requests.get(url)
        # print(r.text)
        buff = StringIO(r.text)
        reader = csv.reader(buff, delimiter="\t")

        prepped_data = []
        keys = []
        result = []

        horizontal_datum = {"NAD83": 4269}
        vertical_datum = {"UNKNOWN": 0, "COE1912": 1, "NGVD29": 2, "NAVD88": 3}

        # Cleanup data before parse fields
        # Store header and actual data rows in new variable
        for idx, line in enumerate(reader):

            # only look at the header and data lines
            if line[0].strip() in ["agency_cd", "USGS"]:
                prepped_data.append(line)

        for idx, line in enumerate(prepped_data):

            if idx == 0:
                # this is the header
                keys = line
            else:
                # Build each line (object) by setting the keys and values
                _line = {}
                for i, k in enumerate(keys):
                    _line[k] = line[i].strip()
                result.append(_line)

        state_sites = []

        # Each line (after cleanup) in the results from the USGS API call represents a site
        for line in result:

            site = {}
            site["provider"] = "usgs"
            site["datatype"] = "usgs-site"
            site["code"] = line["site_no"].strip()
            site["state"] = state_abbrev

            # Geometry
            # --------
            geom = {}
            geom["type"] = "Point"
            try:
                geom["coordinates"] = [
                    float(line["dec_long_va"].strip()),
                    float(line["dec_lat_va"].strip()),
                ]
            except:
                geom["coordinates"] = [0, 0]
            site["geometry"] = geom

            # Attributes
            # ---------
            attributes = {}
            attributes["station_name"] = line["station_nm"].strip()
            attributes["site_type"] = line["site_tp_cd"].strip()
            site["attributes"] = attributes

            # site["site_number"] = line["site_no"].strip()
            # site["name"] = line["station_nm"].replace("'", "").strip()
            # site["state_abbrev"] = state_abbrev
            # site["elevation"] = (
            #     float(line["alt_va"].strip()) if line["alt_va"].strip() != "" else None
            # )
            # try:
            #     site["horizontal_datum_id"] = horizontal_datum[
            #         line["dec_coord_datum_cd"]
            #     ]
            # except:
            #     site["horizontal_datum_id"] = 4269
            # site["huc"] = (
            #     f"{line['huc_cd'].strip()}" if line["huc_cd"].strip() != "" else None
            # )
            # try:
            #     site["vertical_datum_id"] = vertical_datum[line["alt_datum_cd"]]
            # except:
            #     site["vertical_datum_id"] = vertical_datum["UNKNOWN"]

            state_sites.append(site)

        # POST to the water API
        # water.sync_usgs_sites(state_sites)
        water_hook = water.WaterHook(method="POST")
        resp = water_hook.request(
            endpoint="/providers/usgs/locations",
            json=state_sites,
        )

    states = water.get_states()

    for state in states:
        fetch_task_id = f"fetch_and_write_{state}"
        get_usgs_sites_task = PythonOperator(
            task_id=fetch_task_id,
            python_callable=fetch_and_write,
            op_kwargs={
                "state_abbrev": state,
            },
        )


import json
import os
import datetime

from model import db
import dao
from utils import *

from flask import Flask, request

app = Flask(__name__)
db_filename = "backend-template.db"

app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///%s" % db_filename
app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
app.config["SQLALCHEMY_ECHO"] = True

db.init_app(app)
with app.app_context():
    db.create_all()

@app.route("/")
def welcome():
    return "Welcome to Backend Template! :)"

# USER ROUTES ----------------------------------------------
@app.route("/users/", methods=["GET"])
def get_users():
    return success_response([user.serialize() for user in dao.get_users()])

@app.route("/users/", methods=["POST"])
def create_user():
    body = json.loads(request.data)
    name = body.get("name")
    email = body.get("email")
    username = body.get("username")

    optional_user, err = dao.create_user(name, email, username)
    if not optional_user:
        return failure_response(err)
    return success_response(optional_user.serialize(), 201)

@app.route("/users/<int:user_id>/", methods=["GET"])
def get_user(user_id):
    user, err = dao.get_user_by_id(user_id)
    if not user:
        return failure_response(err)
    return success_response(user.serialize())

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 5000))
    app.run(host="0.0.0.0", port=port)


# -*- coding: utf-8 -*-
"""
    CreatedDate: 2022-08-04
    FileName   : numberOf2sInRange-1602.py
    Author     : Honghe
    Descreption:
"""

class Solution:
    def numberOf2sInRange(self, n: int) -> int:
        """
        数学规律:其实会发现10、100、1000、10000包含的2的个数呈现的趋势
        :param n:
        :return:
        """
        nums = {
            0:0,
            1:1,
            2:20,
            3:300,
            4:4000,
            5:50000,
            6:600000,
            7:7000000,
            8:80000000,
            9:900000000
        }
        res = 0
        n = str(n)
        length = len(n)
        for i in range(len(n)):
            cur = int(n[i])
            if cur==0:
                continue
            if cur==1:
                res += nums[length-i-1]
            elif cur>2:
                res += nums[length-i-1]*cur+10**(length-i-1)
            else:
                res += nums[length-i-1]*2+int(n[i+1:])+1 if i+1<length else 1
        return res

if __name__ == '__main__':
    sol = Solution()
    n=1000000000
    print(sol.numberOf2sInRange(n))


import torch
from torch import nn, Tensor
import torch.nn.functional as F

from .utils import check_images, check_targets

class RCNN(nn.Module):
    """
    Abstract RCNN
    """
    def __init__(self, transform, backbone, rpn, rcnn_heads):
        super(RCNN, self).__init__()
        self.transform = transform
        self.backbone = backbone
        self.rpn = rpn
        self.rcnn_heads = rcnn_heads

    def forward(self, images, targets=None):
        """
        :param images: list[Tensor], raw images
        :param targets: list[dict[Tensor]], ground truth boxes and corresponding labels
                dict[Tensor]: targets['boxes']=tensor((N,4)), targets['labels']=tensor((N,))
        :return: detection results and losses
        """

        check_images(images)
        if self.training and not check_targets(targets):
            raise ValueError("Invalid targets, but I'm too lazy to tell the details.")

        # step 1: image and targets transform
        images, targets = self.transform(images, targets)

        # step 2: feed into convolutional backbone, generating feature map
        features = self.backbone(images.tensors)

        # step 3: feed features into rpn to generate proposals
        proposals, rpn_losses = self.rpn(images, targets, features)

        # step 4: feed proposals into rcnn head to get final prediction results
        detections, rcnn_losses = self.rcnn_heads(features, proposals, images.image_sizes, targets)

        # step 5: post process, e.g. transform box and image scales
        detections = self.transform.postprocess(detections, images.image_sizes, images.original_image_sizes)

        # losses in training
        losses = {}
        losses.update(rpn_losses)
        losses.update(rcnn_losses)

        return detections, losses

class RCNNHeadMLP(nn.Module):
    """
    RCNN fully connected layers
    """
    def __init__(self, in_channels, out_channels):
        super(RCNNHeadMLP, self).__init__()
        self.fc6 = nn.Linear(in_channels, out_channels)
        self.fc7 = nn.Linear(out_channels, out_channels)

    def forward(self, x):
        x = x.flatten(start_dim=1)
        x = F.relu(self.fc6(x))
        x = F.relu(self.fc7(x))
        return x

class RCNNPredictor(nn.Module):
    """
    RCNN predict head
    """
    def __init__(self, in_channels, num_classes):
        super(RCNNPredictor, self).__init__()
        self.cls_score = nn.Linear(in_channels, num_classes)
        self.bbox_reg = nn.Linear(in_channels, num_classes * 4)

    def forward(self, x):
        x = x.flatten(start_dim=1)
        cls = self.cls_score(x)
        reg = self.bbox_reg(x)

        return cls, reg