MSTreesV2.py

# POC 
# Don't use in production

from __future__ import print_function
import numpy as np, networkx as nx, argparse
from numba import jit
from ete3 import Tree
from subprocess import Popen, PIPE
import sys, os, tempfile, platform, re, tempfile
import time
from datetime import datetime
import multiprocessing
from multiprocessing import Pool
from tqdm import tqdm
import warnings
from numba.core.errors import NumbaPendingDeprecationWarning
warnings.filterwarnings("ignore", category=NumbaPendingDeprecationWarning)

base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

params = dict(
    edmonds_Linux = os.path.join(base_dir, 'binaries', 'edmonds-linux'),
)

start_time = time.time()

@jit(nopython=True)
def contemporary(a,b,c, n_loci) :
    a[0], a[1] = max(min(a[0], n_loci-0.5), 0.5), max(min(a[1], n_loci-0.5), 0.5);
    b, c = max(min(b, n_loci-0.5), 0.5), max(min(c, n_loci-0.5), 0.5)
    if b >= a[0] + c and b >= a[1] + c :
        return False
    elif b == c :
        return True
    s11, s12 = np.sqrt(1-a[0]/n_loci), (2*n_loci - b - c)/2/np.sqrt(n_loci*(n_loci-a[0]))
    v = 1-((n_loci-a[1])*(n_loci-c)/n_loci+(n_loci-b))/2/n_loci
    s21, s22 = 1+a[1]*v/(b-2*n_loci*v), 1+c*v/(b-2*n_loci*v)

    p1 = a[0]*np.log(1-s11*s11) + (n_loci-a[0])*np.log(s11*s11) + (b+c)*np.log(1-s11*s12) + (2*n_loci-b-c)*np.log(s11*s12)
    p2 = a[1]*np.log(1-s21) + (n_loci-a[1])*np.log(s21) + b*np.log(1-s21*s22) + (n_loci-b)*np.log(s21*s22) + c*np.log(1-s22) + (n_loci-c)*np.log(s22)
    return p1 >= p2

def add_args() :
    n_cores = multiprocessing.cpu_count()

    parser = argparse.ArgumentParser(description='For details, see "https://github.com/achtman-lab/GrapeTree/blob/master/README.md".\nIn brief, GrapeTree generates a NEWICK tree to the default output (screen) \nor a redirect output, e.g., a file. ', formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument('--profile', '-p', dest='fname', help='An input filename of a file containing MLST.', required=False)
    parser.add_argument('--distance', '-d', dest='distance', help='An inpute filename containing distance matrix.', required=False)
    parser.add_argument('--matrix', '-x', dest='matrix_type', help='Specify the matrix type as either "symmetric" or "asymmetric" when a distance matrix is provided', required=False)
    parser.add_argument('--chunk_size', '-c', dest='chunk_size', help='Chunk size', required=True)
    parser.add_argument('--n_proc', '-n', dest='number_of_processes', help='Number of CPU processes in parallel use. Default is half of available cores.', type=int, default=n_cores // 2)
    parser.add_argument('--sep', '-s', dest='sep', help='Separator for the input file. Default is tab.', default='\t')
    parser.add_argument('--keep_files', '-k', dest='keep_files', help='Keep the files generated by the script. Default is True.', action='store_true')

    args = parser.parse_args()
    
    if not args.fname and not args.distance:
        parser.error("Error: Either --profile or --distance must be provided.")
        
    if args.number_of_processes > n_cores - 1:
        parser.error("Error: n_proc cannot be greater than the total number of cores - 1.")
        
    if args.distance and not args.matrix_type:
        parser.error("Error: --matrix is required when --distance is selected")
    
    if args.matrix_type and args.matrix_type not in ['asymmetric', 'symmetric']:
        parser.error("Error: --matrix must be either 'asymmetric' or 'symmetric'")
    
    args.profile = args.fname
    args.n_proc = args.number_of_processes
    args.handle_missing = 'pair_delete'
    args.method = 'MSTree'
    
    if args.profile:
        args.matrix_type = 'asymmetric'
    
    args.heuristic = 'harmonic'
    args.branch_recraft = True
    
    if args.distance:
        args.rows, args.columns = count_lines_and_columns(args.distance)
        print(f"[info] {args.distance} has {args.rows} rows and {args.columns} columns.")
    else:
        args.rows, args.columns = count_lines_and_columns(args.profile)
        print(f"[info] {args.profile} has {args.rows} rows and {args.columns} columns.")

    args.dtype = np.float32
    
    if args.chunk_size is None:
        if args.rows <= 1000:
            args.chunk_size = 500
        elif args.rows <= 20000:
            args.chunk_size = 5000
        else:
            args.chunk_size = 10000
    else:
        args.chunk_size = int(args.chunk_size)
    
    print(f"[info] The chunk size is set to {args.chunk_size}.")

    return args.__dict__

def count_lines_and_columns(filename, sep='\t'):
        with open(filename, 'r') as f:
            line_count = sum(1 for _ in f) - 1
            f.seek(0)
            column_count = len(next(f).split(sep)) - 1
        return line_count, column_count
    
def parallel_distance(callup) :
    func, index_range = callup
    res = np.memmap(params['dist_file'], dtype=params['dtype'], mode='r+', shape=(params['rows'], params['rows']))
    eval('distance_matrix.'+func)(res[:, index_range[0]:index_range[1]], index_range)
    res.flush()

class distance_matrix(object) :
    @staticmethod
    def get_distance(func):
        int_time = time.time()
        print(f"[info] get_distance method started...")
        n_profile = params['rows']
        n_proc = min(int(params['n_proc']), n_profile)
        res = np.memmap(params['dist_file'], dtype=params['dtype'], mode='w+', shape=(params['rows'], params['rows']))
        indices = np.array([[n_profile * v / n_proc + 0.5, n_profile * (v + 1) / n_proc + 0.5] for v in np.arange(n_proc, dtype=float)], dtype=int)
        if n_proc > 1:
            with Pool(n_proc) as pool:
                for _ in tqdm(pool.imap(parallel_distance, [[func, idx] for idx in indices]), total=len(indices), desc="Calculating distances"):
                    pass
        else:
            for _ in tqdm([parallel_distance([func, [0, n_profile]])], total=1, desc="Calculating distances"):
                pass

        res.flush()
        print(f"[info] get_distance method finished in {time.time() - int_time} seconds.")
        return res

    @staticmethod
    def asymmetric(distances, index_range=None) :
        profiles = np.memmap(params['prof_file'], dtype=params['dtype'], mode='r', shape=(params['rows'], params['columns']))
        if index_range is None :
            index_range = [0, profiles.shape[0]]
        presences = (profiles > 0)
        for i2, id in enumerate(np.arange(*index_range)) :
            profile, presence = profiles[id], presences[id]
            diffs = np.sum(((profiles != profile) & presence), axis=1) * float(presence.size)/np.sum(presence)
            distances[:, i2] = diffs
        return distances

    # TODO: Implement chunking for this section
    @staticmethod
    def symmetric_link(links) :
        int_time = time.time()
        print(f"[info] symmetric_link method started...")
        profiles = np.memmap(params['prof_file'], dtype=params['dtype'], mode='r', shape=(params['rows'],params['columns']))
        presences = (profiles > 0)
        sl = [ [ s, t, np.sum((profiles[s] != profiles[t]) & presences[s] & presences[t]) ] \
                 for s, t, d in links ]
        print(f"[info] symmetric_link method finished in {time.time() - int_time} seconds.")
        return sl
        
    @staticmethod
    def harmonic(n_str):
        int_time = time.time()
        print(f"[info] harmonic method started...")
        dist = np.memmap(params['dist_file'], dtype=params['dtype'], mode='r', shape=(params['rows'], params['rows']))
        chunk_size = params['chunk_size']
        n_rows, _ = dist.shape
        weights = np.zeros(n_rows, dtype=params['dtype'])
        for start_row in range(0, n_rows, chunk_size):
            end_row = min(start_row + chunk_size, n_rows)
            dist_chunk = dist[start_row:end_row, :]
            harmonic_sums = np.sum(1.0 / (dist_chunk + 0.1), axis=1)
            weights[start_row:end_row] = n_rows / harmonic_sums
        cw = np.vstack([-np.array(n_str, dtype=params['dtype']), weights])
        weights[np.lexsort(cw)] = np.arange(n_rows, dtype=params['dtype']) / n_rows
        print(f"[info] harmonic method finished in {time.time() - int_time} seconds.")
        return weights

class methods(object):
    
    @staticmethod
    def _symmetric(weight, **params) :
        int_time = time.time()
        print(f"[info] _symmetric method started...")
        dist = np.memmap(params['dist_file'], dtype=params['dtype'], mode='r', shape=(params['rows'], params['rows']))
        dist_temp = np.memmap(tempfile.NamedTemporaryFile(delete=True), dtype=params['dtype'], mode='w+', shape=(params['rows'], params['rows']))
        
        def minimum_spanning_tree(dist):
            n_node = dist.shape[0]
            nodes = np.arange(n_node)
            ng = {n:[n] for n in nodes}
            edges = np.array([ [x, y, dist[x, y]] for y in np.arange(n_node) for x in np.arange(y) ])
            edges = edges[np.argsort(edges.T[2])].astype(int)
            mst = []
            for m, e in enumerate(edges) :
                if nodes[e[0]] == nodes[e[1]] :
                    continue
                mst.append(e.tolist())
                if nodes[e[0]] > nodes[e[1]] :
                    s, e = nodes[e[1]], nodes[e[0]]
                else :
                    s, e = nodes[e[0]], nodes[e[1]]
                nodes[ng[e]] = s
                ng[s].extend(ng.pop(e))
            return mst
        
        n_rows = params['rows']
        chunk_size = params['chunk_size']

        for start in range(0, n_rows, chunk_size):
            end = min(start + chunk_size, n_rows)
            dist_temp[start:end, :] = np.round(dist[start:end, :], 0) + weight[start:end].reshape([weight.size, 1])
            
        for start in range(0, n_rows, chunk_size):
            end = min(start + chunk_size, n_rows)
            dist_temp[start:end, start:end] = np.where(np.eye(end - start, dtype=bool), 0, dist_temp[start:end, start:end])
        
        for start in range(0, n_rows, chunk_size):
            end = min(start + chunk_size, n_rows)
            dist_chunk = dist_temp[start:end, :]
            dist_temp[start:end, :] = np.minimum(dist_chunk, dist_chunk.T)
        
        try:
            g = nx.Graph(dist_temp)
            ms = nx.minimum_spanning_tree(g)
            res =  [[d[0], d[1], int(d[2]['weight'])] for d in ms.edges(data=True)]
            print(f"[info] _symmetric method finished in {time.time() - int_time} seconds.")
            return res
        except :
            res = minimum_spanning_tree(dist)
            print(f"[info] _symmetric method finished in {time.time() - int_time} seconds.")
            return res
        
    # TODO: Implement chunking for this section
    @staticmethod
    def _asymmetric(weight, **params) :
        int_time = time.time()
        print(f"[info] _asymmetric method started...")
        dist = np.memmap(params['dist_file'], dtype=params['dtype'], mode='r', shape=(params['rows'], params['rows']))
        int_dtype = np.int32
        def get_shortcut(dist, weight, cutoff=20):
            if dist.shape[0] < 3000 :
                cutoff = 2
            elif dist.shape[0] < 10000 :
                cutoff = 5
            elif dist.shape[0] < 30000 :
                cutoff = 10
            link = np.array(np.where(dist < (cutoff+1) ), dtype=int_dtype)
            link = link.T[weight[link[0]] < weight[link[1]]].T
            link = link.T[np.lexsort(link)]
            return link[np.unique(link.T[1], return_index=True)[1]].astype(int_dtype)
        try:
            presence = np.arange(weight.shape[0], dtype=int_dtype)
            shortcuts = get_shortcut(dist, weight)
            for (s, t, d) in shortcuts :
                dist[s, dist[s] > dist[t]] = dist[t, dist[s] > dist[t]]
            presence[shortcuts.T[1]] = -1
            dist = dist.T[presence >= 0].T[presence >= 0]
            presence = presence[presence >=0]
            weight2 = weight[presence]
            dist = np.round(dist, 0) + weight2.reshape([weight2.size, -1])
            np.fill_diagonal(dist, 0.0)
            dist_file = params['tempfix'] + '.dist.list'
            with open(dist_file, 'w') as fout :
                for d in dist :
                    fout.write('{0}\n'.format('\t'.join(['{0:.5f}'.format(dd) for dd in (d+(1.-0.000005)) ])))
            del d, dist
            mstree = Popen([params['edmonds_' + platform.system()], dist_file], stdout=PIPE).communicate()[0]
            if isinstance(mstree, bytes) :
                mstree = mstree.decode('utf8')
            mstree = np.array([ br.strip().split() for br in mstree.strip().split('\n')], dtype=float).astype(int)
            assert mstree.size > 0
            mstree.T[2] -= 1
            mstree.T[:2] = presence[mstree.T[:2]]
            m = mstree.tolist() + shortcuts.tolist()
            print(f"[info] _asymmetric method finished in {time.time() - int_time} seconds.")
            return m
        except Exception as e:
            print("[error] Edmonds' algorithm crashed :( .")
            sys.exit(1)

    @staticmethod
    def _branch_recraft(branches, weights) :
        int_time = time.time()
        print(f"[info] _branch_recraft method started...")
        dist = np.memmap(params['dist_file'], dtype=params['dtype'], mode='r+', shape=(params['rows'], params['rows']))
        n_loci = params['columns']
        group_id, groups, childrens = {b:b for br in branches for b in br[:2]}, \
            {b:[b] for br in branches for b in br[:2]}, \
            {b:[] for br in branches for b in br[:2]}
        branches = sorted(branches, key=lambda br:[dist[br[0], br[1]]] + sorted([weights[br[0]], weights[br[1]]]))
        i = 0
        while i < len(branches) :
            src, tgt, brlen = branches[i]
            sources, targets = groups[group_id[src]], groups[group_id[tgt]]
            tried = {}
            if len(sources) > 1 :
                for w, d, s in sorted(zip(weights[sources], dist[sources, tgt], sources))[:3] :
                    if s == src : break
                    if d < 1.5*dist[src, tgt] :
                        if contemporary([dist[s, src], dist[src, s]], d, dist[src, tgt], n_loci) :
                            tried[src], src = s, s
                            break
                while src not in tried :
                    tried[src] = src
                    mid_nodes = sorted([[weights[s], dist[s,tgt], s] for s in childrens[src] if s not in tried and dist[s,tgt] < 2*dist[src, tgt]])
                    for w, d, s in mid_nodes :
                        if d < dist[src, tgt] :
                            if not contemporary([dist[src, s], dist[s, src]], dist[src, tgt], d, n_loci) :
                                tried[src], src = s, s
                                break
                        elif w < weights[src] :
                            if contemporary([dist[s, src], dist[src, s]], d, dist[src, tgt], n_loci) :
                                tried[src], src = s, s
                                break
                        tried[s] = src
            if len(targets) > 1 :
                for w, d, t in sorted(zip(weights[targets], dist[src, targets], targets))[:3] :
                    if t == tgt : break
                    if d < 1.5*dist[src, tgt] :
                        if contemporary([dist[t, tgt], dist[tgt, t]], d, dist[src, tgt], n_loci) :
                            tried[tgt], tgt = t, t
                            break
                while tgt not in tried :
                    tried[tgt] = tgt
                    mid_nodes = sorted([[weights[t], dist[src,t], t] for t in childrens[tgt] if t not in tried and dist[src, t] < 2*dist[src, tgt]])
                    for w, d, s in mid_nodes :
                        if d < dist[src, tgt] :
                            if not contemporary([dist[tgt, t], dist[t, tgt]], dist[src, tgt], d, n_loci) :
                                tried[tgt], tgt = t, t
                                break
                        elif w < weights[tgt] :
                            if contemporary([dist[t, tgt], dist[tgt, t]], d, dist[src, tgt], n_loci) :
                                tried[tgt], tgt = t, t
                                break
                        tried[t] = tgt
            brlen = dist[src, tgt]
            branches[i] = [src, tgt, brlen]
            if i >= len(branches) - 1 or branches[i+1][2] >= brlen:
                tid = group_id[tgt]
                for t in targets :
                    group_id[t] = group_id[src]
                groups[group_id[src]].extend(groups.pop(tid, []))
                childrens[src].append(tgt)
                childrens[tgt].append(src)
                i += 1
            else :
                branches[i:] = sorted(branches[i:], key=lambda br:br[2])
                
        print(f"[info] _branch_recraft method finished in {time.time() - int_time} seconds.")
        return branches

    @staticmethod
    def _network2tree(branches):
        int_time = time.time()
        print(f"[info] _network2tree method started...")
        names = np.memmap(params['names_file'], dtype='<U100', mode='r', shape=(params['rows'],))
        branches.sort(key=lambda x:x[2], reverse=True)
        branch = []
        in_use = {branches[0][0]:1}
        while len(branches) :
            remain = []
            for br in branches :
                if br[0] in in_use :
                    branch.append(br)
                    in_use[br[1]] = 1
                elif br[1] in in_use :
                    branch.append([br[1], br[0], br[2]])
                    in_use[br[0]] = 1
                else :
                    remain.append(br)
            branches = remain

        tre = Tree()
        nodeFinder = {}

        tre.name = branch[0][0]
        nodeFinder[tre.name] = tre
        for src, tgt, dif in branch :
            node = nodeFinder[src]
            child = node.add_child(name=tgt, dist=dif)
            nodeFinder[child.name] = child
        for node in tre.traverse('postorder') :
            if not node.is_leaf() :
                name = node.name
                node.name = ''
                node.add_child(name=names[name], dist=0.)
            else :
                node.name = names[node.name]
        print(f"[info] _network2tree method finished in {time.time() - int_time} seconds.")
        return tre

    @staticmethod
    def MSTree(embeded, matrix_type='asymmetric', heuristic='harmonic', branch_recraft=True, handle_missing='pair_delete', **params) :
        int_time = time.time()
        print(f"[info] MSTree method started...")
        names = np.memmap(params['names_file'], dtype='<U100', mode='r', shape=(params['rows'],))
        if (params['profile']):
            distance_matrix.get_distance(matrix_type)
        weight = eval('distance_matrix.'+heuristic)([len(embeded[n]) for n in names])
        tree = eval('methods._'+matrix_type)(weight, **params)
        if branch_recraft :
            tree = methods._branch_recraft(tree, weight)
        if (params['profile']):
            if matrix_type != 'blockwise' :
                tree = distance_matrix.symmetric_link(tree)
        tree = methods._network2tree(tree)
        print(f"[info] MSTree method finished in {time.time() - int_time} seconds.")
        return tree
    
# TODO: to chunkify this method
def nonredundant(names, profiles) :
    int_time = time.time()
    print(f"[info] nonredundant method started...")
    names = np.memmap(params['names_file'], dtype='<U100', mode='r', shape=(params['rows'],))
    profiles = np.memmap(params['prof_file'], dtype='<U10', mode='r', shape=(params['rows'], params['columns']))
    #encoded_profile = np.array([np.unique(p, return_inverse=True)[1]+1 for p in profiles.T]).T
    encoded_profile = np.array([np.unique(p, return_inverse=True)[1] + 1 for p in profiles.T], dtype=np.int32).T
    encoded_profile[ (profiles == '0') | (profiles == 'N') | (profiles == '-')] = 0
    sorted_profiles = np.lexsort(encoded_profile.T)
    names = names[sorted_profiles]
    profiles = encoded_profile[sorted_profiles]
    presence = (np.sum(profiles > 0, 1) > 0)
    names, profiles = names[presence], profiles[presence]
    uniqueness = np.concatenate([[1], np.sum(np.diff(profiles, axis=0) != 0, 1) > 0])
    embeded = {names[0]:[]}
    embeded_group = embeded[names[0]]
    for n, u in zip(names, uniqueness) :
        if u == 0 :
            embeded_group.append(n)
        else :
            embeded[n] = [n]
            embeded_group = embeded[n]
    uniqueness_indexes = uniqueness > 0
    names = names[uniqueness_indexes]
    profiles = profiles[uniqueness_indexes]
    chunk_size = params['chunk_size']
    new_names = np.memmap(params['names_file'], dtype=names.dtype, mode='w+', shape=(len(names),))
    for start in range(0, len(names), chunk_size):
        end = min(start + chunk_size, len(names))
        new_names[start:end] = names[start:end]
    new_names.flush()
    new_profiles = np.memmap(params['prof_file'], dtype=profiles.dtype, mode='w+', shape=profiles.shape)
    for start in range(0, len(profiles), chunk_size):
        end = min(start + chunk_size, len(profiles))
        new_profiles[start:end, :] = profiles[start:end, :]
    new_profiles.flush()
    print(f"[info] nonredundant method finished in {time.time() - int_time} seconds.")
    return new_names, new_profiles, embeded

def nonredundant_dist():
    int_time = time.time()
    print(f"[info] nonredundant_dist method started...")

    chunk_size = params['chunk_size']
    total_rows = params['rows']
    total_columns = params['columns']

    names = np.memmap(params['names_file'], dtype='<U100', mode='r', shape=(total_rows,))
    dist = np.memmap(params['dist_file'], dtype=params['dtype'], mode='r', shape=(total_rows, total_columns))
    
    unique_rows_set = set()
    uniqueness_list = []

    for start in range(0, total_rows, chunk_size):
        end = min(start + chunk_size, total_rows)
        chunk_dist = dist[start:end]
        chunk_unique_rows, chunk_uniqueness = np.unique(chunk_dist, axis=0, return_inverse=True)
        for row in chunk_unique_rows:
            unique_rows_set.add(tuple(row))
        uniqueness_list.extend(chunk_uniqueness + len(unique_rows_set) - len(chunk_unique_rows))

    unique_rows = np.array(list(unique_rows_set))
    uniqueness = np.array(uniqueness_list)

    embeded_group = {}
    for name, unique_idx in zip(names, uniqueness):
        unique_row = unique_rows[unique_idx].tobytes()
        if unique_row not in embeded_group:
            embeded_group[unique_row] = [name]
        else:
            embeded_group[unique_row].append(name)

    embeded = {value[0]: value for value in embeded_group.values()}
    new_names_list = list(embeded.keys())
    new_names = np.array(new_names_list)
    
    name_to_index = {name: idx for idx, name in enumerate(names)}
    new_indices = [name_to_index[name] for name in new_names]
    
    new_dist_memmap = np.memmap(f"{params['dist_file']}_temp", dtype=params['dtype'], mode='w+', shape=(len(new_names), len(new_names)))

    for i in range(0, len(new_indices), chunk_size):
        for j in range(0, len(new_indices), chunk_size):
            chunk_indices_i = new_indices[i:i + chunk_size]
            chunk_indices_j = new_indices[j:j + chunk_size]
            new_dist_memmap[i:i + len(chunk_indices_i), j:j + len(chunk_indices_j)] = dist[np.ix_(chunk_indices_i, chunk_indices_j)]
    new_dist_memmap.flush()

    os.remove(params['dist_file'])
    os.rename(f"{params['dist_file']}_temp", params['dist_file'])
    
    names_memmap = np.memmap(f"{params['names_file']}_temp", dtype='<U100', mode='w+', shape=(len(new_names),))
    for start in range(0, len(new_names), chunk_size):
        end = min(start + chunk_size, len(new_names))
        names_memmap[start:end] = new_names[start:end]
    names_memmap.flush()
    
    os.remove(params['names_file'])
    os.rename(f"{params['names_file']}_temp", params['names_file'])

    print(f"[info] nonredundant_dist method finished in {time.time() - int_time} seconds.")
    return new_dist_memmap.shape[0], new_dist_memmap.shape[1], embeded

def process_distance_chunk(chunk, start_idx):
    names = np.memmap(params['names_file'], dtype='<U100', mode='r+', shape=(params['rows'],))
    dist = np.memmap(params['dist_file'], dtype=params['dtype'], mode='r+', shape=(params['rows'], params['rows']))
    lines = chunk.strip().split('\n')
    for i, line in enumerate(lines):
        parts = line.split('\t')
        sample = parts[0]
        distances = np.array(parts[1:], dtype=params['dtype'])
        names[start_idx + i] = sample
        dist[start_idx + i, :len(distances)] = distances

def load_distance_matrix():
    names = np.memmap(params['names_file'], dtype='<U100', mode='r+', shape=(params['rows'],))
    dist = np.memmap(params['dist_file'], dtype=params['dtype'], mode='r+', shape=(params['rows'], params['rows']))
    start_idx = 0
    with open(params['distance'], 'r') as file:
        file.readline()
        while True:
            chunk = ''.join([file.readline() for _ in range(params['chunk_size'])])
            if not chunk.strip():
                break
            process_distance_chunk(chunk, start_idx)
            start_idx += len(chunk.strip().split('\n'))
    names.flush()
    dist.flush()

def profiles_generator(file_path, chunk_size, sep):
    with open(file_path, 'r') as fin:
        header = fin.readline().strip().split(sep)
        allele_cols = [id for id, col in enumerate(header) if id > 0 and not col.startswith('#') and not col.lower() in {'st_id', 'st'}]
        while True:
            names_chunk = []
            profiles_chunk = []
            for _ in range(chunk_size):
                line = fin.readline().strip()
                if not line:
                    break
                part = line.split('\t')
                name = re.sub(r'[\(\)\ \,\"\';]', '_', part[0])
                names_chunk.append(name)
                profile = [part[i].upper() for i in allele_cols]
                profiles_chunk.append(profile)
            if not profiles_chunk:
                break
            yield names_chunk, profiles_chunk
            
def backend(**args) :
    global params
    params.update(args)
    
    temp_file = tempfile.NamedTemporaryFile(delete=True, dir='.')
    params['tempfix'] = temp_file.name
    if params['profile']:
        params['prof_file'] = f'{temp_file.name}.prof.npy'
    params['names_file'] = f'{temp_file.name}.names.npy'
    params['dist_file'] = f'{temp_file.name}.dist.npy'
    params['nwk_file'] = f'{temp_file.name}.nwk'
    
    if params['distance']:
        int_time = time.time()
        print(f"[info] The names file will be saved in {params['names_file']}.")
        print(f"[info] The distance file will be saved in {params['dist_file']}.")
        print(f"[info] The distance file for edmonds will be saved in {temp_file.name}.dist.list")
        print(f"[info] Processing {params['distance']} in chunks...")
        names = np.memmap(params['names_file'], dtype='<U100', mode='w+', shape=(params['rows'],))
        dist = np.memmap(params['dist_file'], dtype=params['dtype'], mode='w+', shape=(params['rows'], params['rows']))
        load_distance_matrix()
        rows, columns, embeded = nonredundant_dist()
        params['rows'] = rows
        params['columns'] = columns
    else:
        print(f"[info] The profile file will be saved in {params['prof_file']}.")
        print(f"[info] The names file will be saved in {params['names_file']}.")
        print(f"[info] The distance file will be saved in {params['dist_file']}.")
        print(f"[info] The distance file for edmonds will be saved in {temp_file.name}.dist.list")
       
            
        int_time = time.time()
        print(f"[info] Processing {params['profile']} in chunks...")
        
        profiles = np.memmap(params['prof_file'], dtype='<U10', mode='w+', shape=(params['rows'], params['columns']))
        names = np.memmap(params['names_file'], dtype='<U100', mode='w+', shape=(params['rows'],))
        chunk_count = 0
        
        total_chunks = (params['rows'] + params['chunk_size'] - 1) // params['chunk_size']
        for names_chunk, profiles_chunk in tqdm(profiles_generator(params['fname'], params['chunk_size'], params['sep']), total=total_chunks, desc="Processing chunks"):
            profiles_chunk = np.array(profiles_chunk)
            names_chunk = np.array(names_chunk)
            start = chunk_count * params['chunk_size']
            end = start + len(profiles_chunk)
            profiles[start:end] = profiles_chunk
            names[start:end] = names_chunk
            chunk_count += 1

        profiles.flush()
        names.flush()
        print(f"[info] Processing finished in {time.time() - int_time} seconds.")
        
        names, profiles, embeded = nonredundant(names, profiles)
        print(f"[info] New shape of profiles: {profiles.shape}")
        print(f"[info] New shape of names: {names.shape}")
        params['rows'] = profiles.shape[0]
        params['columns'] = profiles.shape[1]
        
    dist = np.memmap(params['dist_file'], mode='r', dtype=np.float32, shape=(params['rows'],params['rows']))
    tre = eval('methods.' + params['method'])(embeded, **params)
    maxDist = 0.
    for node in tre.iter_descendants() :
        if node.dist > maxDist: maxDist = node.dist
    if maxDist > 3 :
        for node in tre.iter_descendants('postorder') :
            if node.dist < 0.1 and node.dist > 0 :
                for s in node.get_sisters() :
                    s.dist += node.dist
                node.dist = 0
    for leaf in tre.get_leaves() :
        embeded_group = embeded[leaf.name]
        if len(embeded_group) > 1 :
            leaf.name = ''
            for n in embeded_group :
                leaf.add_child(name=n, dist=0.)

    if not params['keep_files']:
        print(f"[info] Removing temporary files...")
        try:
            if 'prof_file' in params and os.path.isfile(params['prof_file']):
                os.unlink(params['prof_file'])
                
            if 'names_file' in params and os.path.isfile(params['names_file']):
                os.unlink(params['names_file'])
            
            if 'dist_file' in params and os.path.isfile(params['dist_file']):
               os.unlink(params['dist_file'])
            
            if os.path.isfile(f"{temp_file.name}.dist.list"):
                os.unlink(f"{temp_file.name}.dist.list")
                
        except Exception as e:
            print(f"[error] Error in removing temporary files...")
        
    with open(params['nwk_file'], 'w') as f:
        f.write(tre.write(format=1).replace("'", ""))
        
    print(f"[info] The nwk file will is saved in {params['nwk_file']}")
        
if __name__ == '__main__' :
    print("#########################################################")
    print("####       !!! DON'T USE IN PRODUCTION !!!       ########")
    print("#########################################################")
    print(f"[info] MSTreesV2 started at {datetime.now()}...")
    print("#########################################################")
    backend(**add_args())
    print(f"[info] Process completed in {time.time() - start_time} seconds.")
    print("#########################################################")
    print(f"[info] MSTreesV2 finished at {datetime.now()}...")
    print("#########################################################")