import logging
import numpy as np
import os
import pandas as pd
import shutil

from argparse import Namespace
from concurrent.futures import ProcessPoolExecutor
from statistics import median, mean
from random import sample

from scarap.callers import *
from scarap.computers import *
from scarap.helpers import * 
from scarap.pan import *
from scarap.readerswriters import *
from scarap.utils import *

def run_pan(args):
    if "species" in args and not args.species is None:
        run_pan_hier(args)
    else:
        run_pan_nonhier(args)

def run_pan_nonhier(args):

    pangenomefout = os.path.join(args.outfolder, "pangenome.tsv")
    if os.path.isfile(pangenomefout):
        logging.info("existing pangenome detected - moving on")
        return()
    
    faafins = read_fastapaths(args.faa_files)
    
    if args.method in ["O-B", "O-D"]:

        logging.info("creating orthofinder directory")
        dir_orthofinder = os.path.join(args.outfolder, "orthofinder")
        os.makedirs(dir_orthofinder, exist_ok = True)

        logging.info("running orthofinder")
        logfile = os.path.join(args.outfolder, "orthofinder.log")
        
        if args.method == "O-B":
            engine = "blast"
        else:
            engine = "diamond"
        run_orthofinder(faafins, dir_orthofinder, logfile, args.threads, 
            engine)
    
        logging.info("creating tidy pangenome file")
        pangenome = read_pangenome_orthofinder(dir_orthofinder)
        write_tsv(pangenome, pangenomefout)
        
    else:
      
        logging.info(f"pangenome will be constructed with the {args.method} "
            "strategy")
        logging.info(f"for alignments of more than {args.max_align} sequences, "
            f"{args.min_reps} to {args.max_reps} representative sequences will "
            "be used")
        infer_pangenome(faafins, args.method, args.min_reps, args.max_reps,
            args.max_align, args.outfolder, args.threads)

def run_pan_hier(args):
  
    logging.info("creating genome table with species and faa paths")
    genometbl1 = read_species(args.species)
    genometbl1.species = [sp.replace(" ", "_") for sp in genometbl1.species]
    faapaths = read_fastapaths(args.faa_files)
    genomes = [filename_from_path(path) for path in faapaths]
    genometbl2 = pd.DataFrame({"faapath": faapaths, "genome": genomes})
    genometbl = pd.merge(genometbl1, genometbl2)
    
    logging.info("creating the necessary subfolders")
    speciespansdio = os.path.join(args.outfolder, "speciespangenomes")
    pseudogenomesdio = os.path.join(args.outfolder, "pseudogenomes")
    pseudogenomesfio = os.path.join(args.outfolder, "pseudogenomes.tsv")
    pseudopandio = os.path.join(args.outfolder, "pseudopangenome")
    pseudopanfio = os.path.join(args.outfolder, "pseudopangenome.tsv")
    pangenomefout = os.path.join(args.outfolder, "pangenome.tsv")
    os.makedirs(speciespansdio, exist_ok = True)
    os.makedirs(pseudogenomesdio, exist_ok = True)
    os.makedirs(pseudopandio, exist_ok = True)

    nonhier_args = dict(
        threads=args.threads,
        method=args.method,
        max_align=args.max_align,
        max_reps=args.max_reps,
        min_reps=args.min_reps
    )
    
    logging.info("PHASE 1: inferring species-level pangenomes")
    for species, genomesubtbl in genometbl.groupby("species"):
        speciespanfio = os.path.join(speciespansdio, species + ".tsv")
        if os.path.exists(speciespanfio):
            logging.info(f"existing pangenome found for {species}")
            continue
        logging.info(f"started inferring pangenome of {species}")
        dout = os.path.join(speciespansdio, species)
        os.makedirs(dout, exist_ok = True)
        faapaths_sub = genomesubtbl["faapath"].values.tolist()
        faapathsfio = os.path.join(dout, "faapaths.txt")
        write_lines(faapaths_sub, faapathsfio)
        run_pan_nonhier(Namespace(faa_files = faapathsfio, outfolder = dout,
            **nonhier_args))
        shutil.move(os.path.join(dout, "pangenome.tsv"), speciespanfio)
        shutil.rmtree(dout)
    
    logging.info("PHASE 2: constructing pseudogenome per species")
    n_species = genometbl["species"].nunique()
    pseudogenomes = [None] * n_species
    for ix, speciespanfio in enumerate(listpaths(speciespansdio)):
        species = filename_from_path(speciespanfio)
        logging.info(f"constructing pseudogenome for {species}")
        speciespan = read_genes(speciespanfio)
        tmpdio = os.path.join(args.outfolder, "temp")
        pseudogenomes[ix] = create_pseudogenome(speciespan, faapaths, tmpdio)
        pseudogenomes[ix]["species"] = species
    pseudogenomes = pd.concat(pseudogenomes)
    write_tsv(pseudogenomes, pseudogenomesfio)
    
    logging.info("PHASE 3: inferring pseudopangenome")
    pseudogenomes = pseudogenomes.rename(columns = {"species": "orthogroup"})
    gather_orthogroup_sequences(pseudogenomes, faapaths, pseudogenomesdio)
    run_pan_nonhier(Namespace(faa_files = pseudogenomesdio, 
        outfolder = pseudopandio, **nonhier_args))
    shutil.move(os.path.join(pseudopandio, "pangenome.tsv"), pseudopanfio)
    shutil.rmtree(pseudogenomesdio)
    shutil.rmtree(pseudopandio)
    
    logging.info("PHASE 4: inflating pseudopangenome with species pangenomes")
    pangenomes = [None] * len(os.listdir(speciespansdio))
    for ix, speciespanfio in enumerate(listpaths(speciespansdio)):
        pangenomes[ix] = read_genes(speciespanfio)
        pangenomes[ix]["species"] = filename_from_path(speciespanfio)
    pangenomes = pd.concat(pangenomes)
    pangenomes = pangenomes.rename(columns = {"orthogroup": "speciesfam"})
    pseudopan = read_genes(pseudopanfio)
    pseudopan = pseudopan.rename(columns = {"genome": "species"})
    speciesfamtbl = pd.merge(pseudopan, pangenomes, on = ["gene", "species"], 
        how = "left")
    speciesfamtbl = speciesfamtbl[["speciesfam", "species", "orthogroup"]]
    pangenome = pd.merge(pangenomes, speciesfamtbl, 
        on = ["speciesfam", "species"])
    pangenome = pangenome[["gene", "genome", "orthogroup"]]
    write_tsv(pangenome, pangenomefout)

def run_build(args): 
    
    fin_faapaths = args.faa_files
    fin_pangenome = args.pangenome
    dout = args.outfolder
    core_prefilter = args.core_prefilter
    core_filter = args.core_filter
    max_cores = args.max_core_genes
    threads = args.threads
    
    # define output paths/folders
    dout_ogseqs = os.path.join(dout, "orthogroups")
    dout_alis = os.path.join(dout, "alignments")
    dout_tmp = os.path.join(dout, "tmp")
    fout_cutoffs = os.path.join(dout, "cutoffs.csv")
    fout_hits = os.path.join(dout, "hits.tsv")
    fout_genes = os.path.join(dout, "genes.tsv")

    if os.path.isfile(fout_cutoffs):
        logging.info("existing database detected - moving on")
        return()

    logging.info("creating output subfolders")
    for dir in [dout_ogseqs, dout_alis]:
        os.makedirs(dir, exist_ok = True)
        
    logging.info("reading pangenome")
    pangenome = read_genes(fin_pangenome)
    # read separate table with genomes of all genes, in case the pangenome
    # table is not complete or core genes get selected
    fins_faas = read_fastapaths(fin_faapaths)
    genes_genomes = extract_genes(fins_faas)
        
    if core_prefilter != 0:
        logging.info(f"applying core prefilter of {core_prefilter}")
        corefams = determine_corefams(pangenome, core_prefilter)
        pangenome = filter_groups(pangenome, corefams)

    logging.info("gathering sequences of orthogroups")
    gather_orthogroup_sequences(pangenome, fins_faas, dout_ogseqs, 1)
    logging.info(f"gathered sequences for {len(os.listdir(dout_ogseqs))} "
        f"orthogroups")

    logging.info("aligning orthogroups")
    orthogroups = [os.path.splitext(f)[0] for f in os.listdir(dout_ogseqs)]
    fouts_ogseqs = make_paths(orthogroups, dout_ogseqs, ".fasta")
    fouts_alis = make_paths(orthogroups, dout_alis, ".aln")
    run_mafft_parallel(fouts_ogseqs, fouts_alis)
    
    # run profile search (function does its own logging)
    run_profilesearch(fins_faas, fouts_alis, fout_hits, dout_tmp, threads)
    
    logging.info("training score cutoffs for profiles") 
    colnames = ["gene", "profile", "score"]
    hits = pd.read_csv(fout_hits, sep = "\t", names = colnames, 
        usecols = [0, 1, 2])
    hits[["gene", "profile"]] = hits[["gene", "profile"]].\
        applymap(lambda x: x.split(" ")[0])
    cutoffs = train_cutoffs(hits, pangenome)
    
    if core_filter != 0 or max_cores != 0:
        logging.info(f"applying core filter of {core_filter} and maximum "
            f"number of core genes of {max_cores}")
        genes = process_scores(hits, cutoffs, top_profiles = False)
        genes = pd.merge(genes_genomes, genes, how = "right")
        corefams = determine_corefams(genes, core_filter = core_filter,
            max_cores = max_cores)
        hits = hits[hits["profile"].isin(corefams)]
        cutoffs = cutoffs[cutoffs["profile"].isin(corefams)]
        for file in os.listdir(dout_alis):
            if not os.path.splitext(file)[0] in corefams:
                os.remove(os.path.join(dout_alis, file))
        logging.info(f"{len(corefams)} core genes were identified")
    
    logging.info("applying score cutoffs to pangenome")
    genes = process_scores(hits, cutoffs)
    genes = pd.merge(genes_genomes, genes, how = "right")
    
    logging.info("writing output files")
    write_tsv(genes, fout_genes)
    write_tsv(cutoffs, fout_cutoffs)
    
    logging.info("removing temporary files and folders")
    shutil.rmtree(dout_ogseqs)
    os.remove(fout_hits)
    
def run_search(args): 
    
    fin_qpaths = args.faa_files
    din_db = args.db
    dout = args.outfolder
    threads = args.threads
    
    # define output paths/folders
    dout_tmp = os.path.join(dout, "tmp")
    din_alis = os.path.join(din_db, "alignments")
    fin_cutoffs = os.path.join(din_db, "cutoffs.csv")
    fout_hits = os.path.join(dout, "hits.tsv")
    fout_genes = os.path.join(dout, "genes.tsv")

    if os.path.isfile(fout_genes):
        logging.info("existing search results detected - moving on")
        return()
    
    # run profile search (function does its own logging)
    fins_alis = [os.path.join(din_alis, f) for f in os.listdir(din_alis)]
    fins_queries = read_fastapaths(fin_qpaths)
    if os.path.isfile(fout_hits):
        logging.info("existing mmseqs2 hits detected - moving on")
    else:
        run_profilesearch(fins_queries, fins_alis, fout_hits, dout_tmp, threads)
    
    logging.info("reading hits and score cutoffs")
    colnames = ["gene", "profile", "score"]
    hits = pd.read_csv(fout_hits, sep = "\t", names = colnames, 
        usecols = [0, 1, 2])
    hits[["gene", "profile"]] = hits[["gene", "profile"]].\
        applymap(lambda x: x.split(" ")[0])
    colnames = ["profile", "cutoff"]
    cutoffs = pd.read_csv(fin_cutoffs, sep = "\t", names = colnames)
    
    logging.info("applying score cutoffs to hits")
    genes = process_scores(hits, cutoffs)
    logging.info("extracting genome names from faa files")
    genes_genomes = extract_genes(fins_queries, threads = threads)
    logging.info("merging search results and genome names")
    genes = pd.merge(genes_genomes, genes, how = "right")
    
    logging.info("writing output files")
    write_tsv(genes, fout_genes)
    
    logging.info("removing temporary files and folders")
    os.remove(fout_hits)

def run_checkgenomes(args):

    logging.info("checking genomes")
    coregenome = read_genes(args.coregenome)
    genomes = checkgenomes(coregenome)
    write_tsv(genomes, os.path.join(args.outfolder, "genomes.tsv"))

def run_checkgroups(args):

    logging.info("checking core orthogroups")
    coregenome = read_genes(args.coregenome)
    orthogroups = checkgroups(coregenome)
    write_tsv(orthogroups, os.path.join(args.outfolder, "orthogroups.tsv"))

def run_filter(args):

    logging.info("reading pangenome")
    pangenome = read_genes(args.pangenome)
    if not args.genomes is None:
        logging.info("filtering genomes")
        genomes = read_lines(args.genomes)
        pangenome = filter_genomes(pangenome, genomes)
    if not args.orthogroups is None:
        logging.info("filtering orthogroups")
        orthogroups = read_lines(args.orthogroups)
        pangenome = filter_groups(pangenome, orthogroups)
    write_tsv(pangenome, os.path.join(args.outfolder, "pangenome.tsv"))

def run_concat(args):
  
    fin_faapaths = args.faa_files
    fin_pangenome = args.pangenome
    dout = args.outfolder
    core_filter = args.core_filter
    max_cores = args.max_core_genes
    fin_ffnpaths = args.ffn_files

    sm_aas_fout = os.path.join(dout, "supermatrix_aas.fasta")
    sm_nucs_fout = os.path.join(dout, "supermatrix_nucs.fasta")
    seqs_aas_dio = os.path.join(dout, "seqs_aas")
    seqs_nucs_dio = os.path.join(dout, "seqs_nucs")
    alis_aas_dio = os.path.join(dout, "alis_aas")
    alis_nucs_dio = os.path.join(dout, "alis_nucs")
    
    # exit if requested supermatrices already exist
    if os.path.isfile(sm_aas_fout) and (fin_ffnpaths is None or \
        os.path.isfile(sm_nucs_fout)):
        logging.info("requested supermatrices already exist - moving on")
        return()
    
    logging.info("reading core genome") 
    coregenome = read_genes(fin_pangenome)
    orthogroups = coregenome["orthogroup"].unique()
    genomes = coregenome["genome"].unique()
    logging.info(f"detected {len(orthogroups)} orthogroups in "
        f"{len(genomes)} genomes")
    
    if core_filter != 0 or max_cores != 0:
        logging.info(f"applying core filter of {core_filter} and maximum "
            f"number of core genes of {max_cores}")
        corefams = determine_corefams(coregenome, core_filter, max_cores)
        coregenome = filter_groups(coregenome, corefams)
        orthogroups = coregenome["orthogroup"].unique()

    logging.info("removing same-genome copies of core genes")
    coregenome = coregenome.drop_duplicates(["genome", "orthogroup"], 
        keep = False)
    
    seqs_aas_fios = make_paths(orthogroups, seqs_aas_dio, ".fasta")
    alis_aas_fios = make_paths(orthogroups, alis_aas_dio, ".aln")
        
    # move on if dir already exists and is not empty
    if os.path.isdir(seqs_aas_dio) and os.listdir(seqs_aas_dio):
      
        logging.info("existing amino acid sequences detected - moving on")
        
    else:
    
        logging.info("gathering amino acid sequences of orthogroups")
        faa_fins = read_fastapaths(fin_faapaths)
        os.makedirs(seqs_aas_dio, exist_ok = True)
        gather_orthogroup_sequences(coregenome, faa_fins, seqs_aas_dio)
        
    if os.path.isdir(alis_aas_dio) and os.listdir(alis_aas_dio):
      
        logging.info("existing amino acid alignments detected - moving on")
        
    else: 

        logging.info("aligning orthogroups on the amino acid level")
        os.makedirs(alis_aas_dio, exist_ok = True)
        run_mafft_parallel(seqs_aas_fios, alis_aas_fios)

    logging.info("concatenating amino acid alignments")
    construct_supermatrix(coregenome, alis_aas_fios, sm_aas_fout)

    if not fin_ffnpaths is None:
      
        seqs_nucs_fios = make_paths(orthogroups, seqs_nucs_dio, ".fasta")
        alis_aas_fios = make_paths(orthogroups, alis_aas_dio, ".aln")
        alis_nucs_fios = make_paths(orthogroups, alis_nucs_dio, ".aln")
        
        if os.path.isdir(seqs_nucs_dio) and os.listdir(seqs_nucs_dio):
          
            logging.info("existing nucleotide sequences detected - moving on")
            
        else:
    
            logging.info("gathering nucleotide sequences of orthogroups")
            ffn_fins = read_fastapaths(fin_ffnpaths)
            os.makedirs(seqs_nucs_dio, exist_ok = True)
            gather_orthogroup_sequences(coregenome, ffn_fins, seqs_nucs_dio)
            
        if os.path.isdir(alis_nucs_dio) and os.listdir(alis_nucs_dio):
          
            logging.info("existing nucleotide alignments detected - moving on")
            
        else:

            logging.info("aligning orthogroups on the nucleotide level")
            os.makedirs(alis_nucs_dio, exist_ok = True)
            reverse_align_parallel(seqs_nucs_fios, alis_aas_fios, 
                alis_nucs_fios)

        logging.info("concatenating nucleotide alignments")
        construct_supermatrix(coregenome, alis_nucs_fios, sm_nucs_fout)
        
    logging.info("removing temporary folders")
    shutil.rmtree(seqs_aas_dio)
    shutil.rmtree(alis_aas_dio)
    try:
        shutil.rmtree(seqs_nucs_dio)
        shutil.rmtree(alis_nucs_dio)
    except FileNotFoundError:
        pass

def run_sample(args):
  
    if "exact" in args and args.exact:
        ali_mode = 3
        logging.info("identity calculation set to exact")
    else:
        ali_mode = 1
        logging.info("identity calculation set to approximate")

    logging.info("creating output subfolders")
    dio_seqs = os.path.join(args.outfolder, "tmp_seqs")
    dio_alis = os.path.join(args.outfolder, "tmp_alis")
    os.makedirs(dio_seqs, exist_ok = True)
    os.makedirs(dio_alis, exist_ok = True)
    
    fout_clusters = os.path.join(args.outfolder, "clusters.tsv")
    fout_seeds = os.path.join(args.outfolder, "seeds.txt")
    fout_identities = os.path.join(args.outfolder, "identities.tsv")
    
    if os.path.isfile(fout_clusters):
        logging.info("existing results detected - moving on")
        return()

    logging.info("reading core genome") 
    core = read_genes(args.pangenome)
    fams = core["orthogroup"].unique()
    genomes = core["genome"].unique()
    logging.info(f"detected {len(fams)} orthogroups in "
        f"{len(genomes)} genomes")
    
    if args.max_genomes == 0:
        args.max_genomes = len(genomes)
        logging.info(f"max_genomes set to {args.max_genomes}")
    
    logging.info("removing same-genome copies of core genes")
    core = core.drop_duplicates(["genome", "orthogroup"], keep = False)
    fams = core["orthogroup"].unique() # in case some fams were fully removed
    
    if os.path.isdir(dio_seqs) and os.listdir(dio_seqs):
      
        logging.info("existing sequences detected - moving on")
        
    else:
    
        logging.info("gathering sequences of orthogroups")
        fins_faas = read_fastapaths(args.fasta_files)
        gather_orthogroup_sequences(core, fins_faas, dio_seqs)
    
    logging.info("creating database for alignments")
    for dir in ["sequenceDB", "logs"]:
        makedirs_smart(f"{dio_alis}/{dir}")
    fio_seqs = f"{dio_alis}/seqs.fasta"
    open(fio_seqs, "w").close()
    for fam in fams:
        with open(f"{dio_seqs}/{fam}.fasta", "r") as hin_fam:
            with open(fio_seqs, "a+") as hout_seqs:
                hout_seqs.write(hin_fam.read())
    run_mmseqs(["createdb", fio_seqs, f"{dio_alis}/sequenceDB/db"], 
        f"{dio_alis}/logs/createdb.log", threads = args.threads)
    os.remove(fio_seqs)
    
    logging.info("initializing empty identity matrix")
    id_m = np.zeros([len(genomes), 0])
    
    logging.info("adding genomes until max_genomes or identity threshold is "
        "reached")
    core_grouped = core.copy().groupby("orthogroup")
    core_grouped = [core_fam for fam, core_fam in core_grouped]
    seeds = []
    while True:
      
        print("|", end = "", flush = True)
      
        # add column of zeros to identity matrix
        id_m = np.hstack((id_m, np.zeros([id_m.shape[0], 1])))
      
        # select seed
        max_ids = np.amax(id_m, 1) # max of each row = max along columns
        for seed in seeds: max_ids[seed] = 1.1 # avoid seeds being reused
        s = np.argmin(max_ids)
        seeds.append(s)
        min_max_ids = max_ids[s]
        seed_genome = genomes[s]
        
        # if sampled genomes close enough to each other: break the loop
        if min_max_ids > args.identity:
            print("")
            logging.info("identity threshold reached")
            # remove last column from identity matrix
            id_m = np.delete(id_m, -1, 1)
            seeds = seeds[:-1]
            break
        
        # prepare prefilter database
        for dir in ["prefDB", "alignmentDB"]:
            makedirs_smart(f"{dio_alis}/{dir}")
        pref = core_grouped.copy()
        for p, fam in enumerate(pref):
            if not seed_genome in fam.genome.tolist():
                pref[p] = None
                continue
            fam["query"] = fam[fam.genome == seed_genome].iloc[0]["gene"]
        pref = pd.concat(pref)
        pref = pref.rename(columns = {"gene": "target"})
        lookup = pd.read_csv(f"{dio_alis}/sequenceDB/db.lookup", sep = "\t", 
            usecols = [0, 1], names = ["id", "sequence"])
        lookup = dict(zip(lookup["sequence"].tolist(), lookup["id"].tolist()))
        pref["query_id"] = [lookup[q] for q in pref["query"].tolist()]
        pref["target_id"] = [lookup[t] for t in pref["target"].tolist()]
        pref = pref[["query_id", "target_id"]]
        write_tsv(pref, f"{dio_alis}/pref.tsv")
        run_mmseqs(["tsv2db", f"{dio_alis}/pref.tsv", f"{dio_alis}/prefDB/db", 
            "--output-dbtype", "7"], f"{dio_alis}/logs/tsv2db.log")
        
        # perform alignments 
        run_mmseqs(["align", f"{dio_alis}/sequenceDB/db",
            f"{dio_alis}/sequenceDB/db", f"{dio_alis}/prefDB/db", 
            f"{dio_alis}/alignmentDB/db", "--alignment-mode", str(ali_mode)], 
            f"{dio_alis}/logs/align.log", threads = args.threads)
        run_mmseqs(["createtsv", f"{dio_alis}/sequenceDB/db", 
            f"{dio_alis}/sequenceDB/db", f"{dio_alis}/alignmentDB/db", 
            f"{dio_alis}/hits.tsv"], f"{dio_alis}/logs/createtsv.log")
        
        # parse alignment results
        hits = pd.read_csv(f"{dio_alis}/hits.tsv", sep = "\t", 
            usecols = [1, 3], names = ["gene", "identity"])
        hits = core.merge(hits, on = "gene", how = "right")
        hits = hits.drop(["gene", "orthogroup"], axis = 1)

        # calculate mean/median identity with seed per genome
        if args.method == "median":
            ids = hits.groupby("genome").aggregate(median)
        elif args.method[:4] == "mean":
            p = args.method[4:]
            if (p == ""): p = "100"
            p = int(p) / 100
            def meanp(l, p):
                start = int(round(((1 - p) / 2) * len(l)))
                stop = int(round((1 - (1 - p) / 2) * len(l)))
                if start == stop: return(median(l))
                return(mean(sorted(l)[start:stop]))
            ids = hits.groupby("genome").aggregate(lambda l: meanp(l, p))
                
        # add mean/median identity to identity matrix
        for g, genome in enumerate(genomes):
            id_m[g, -1] = ids.loc[genome, "identity"]
            
        # if enough sampled genomes: break the loop
        if np.size(id_m, 1) == args.max_genomes:
            print("")
            logging.info("max number of genomes sampled")
            break
            
    logging.info("writing seeds.txt")
    with open(fout_seeds, "a") as hout_seeds:
        for s in seeds: hout_seeds.write(genomes[s] + "\n")
            
    logging.info("writing identities.tsv")
    id_df = pd.DataFrame(id_m)
    id_df.columns = [genomes[s] for s in seeds]
    id_df["genome"] = genomes 
    cols = id_df.columns.tolist()
    cols = [cols[-1]] + cols[:-1]
    id_df = id_df[cols]
    id_df.to_csv(fout_identities, sep = "\t", index = False, header = True)
        
    logging.info("clustering the genomes")
    clusters = np.argmax(id_m, 1)
    genomes_clusters = pd.DataFrame({"genome": genomes, "cluster": clusters})
    logging.info(f"{len(set(clusters))} clusters found")
    
    logging.info("writing clusters.tsv")
    write_tsv(genomes_clusters, fout_clusters)
        
    logging.info("removing temporary folders")
    shutil.rmtree(dio_seqs)
    shutil.rmtree(dio_alis)
    
def run_fetch(args):

    logging.info("creating subfolder for fastas")
    dout_seqs = os.path.join(args.outfolder, "fastas")
    os.makedirs(dout_seqs, exist_ok = True)

    logging.info("reading genes") 
    genes = read_genes(args.pangenome)
    fams = genes["orthogroup"].unique()
    genomes = genes["genome"].unique()
    logging.info(f"detected {len(fams)} orthogroups in "
        f"{len(genomes)} genomes")
        
    logging.info("gathering sequences of orthogroups")
    fins_fastas = read_fastapaths(args.fasta_files)
    gather_orthogroup_sequences(genes, fins_fastas, dout_seqs)
    
def run_core(args):
  
    fin_faapaths = args.faa_files
    dout = args.outfolder
    method = args.method
    min_reps = args.min_reps
    max_reps = args.max_reps
    max_align = args.max_align
    seeds = args.seeds
    core_prefilter = args.core_prefilter
    core_filter = args.core_filter
    max_cores = args.max_core_genes
    threads = args.threads
    
    # define paths
    dout_seedpan = os.path.join(dout, "seedpan")
    dout_seedcore = os.path.join(dout, "seedcore")
    fout_seedpaths = os.path.join(dout, "seedpaths.txt")
    fout_nonseedpaths = os.path.join(dout, "nonseedpaths.txt")
    fout_seedpan = os.path.join(dout_seedpan, "pangenome.tsv")
    fout_genes_core = os.path.join(dout_seedcore, "genes.tsv")
    fout_genes = os.path.join(dout, "genes.tsv")
    
    # make output subfolders 
    for dir in [dout_seedpan, dout_seedcore]:
        os.makedirs(dir, exist_ok = True)

    logging.info("selecting random seed genomes")
    fins_faas = read_fastapaths(fin_faapaths)
    if not os.path.isfile(fout_seedpaths):
        fins_seeds = sample(fins_faas, seeds)
        fins_nonseeds = [fin for fin in fins_faas if not fin in fins_seeds]
        write_tsv(pd.DataFrame({"path": fins_seeds}), fout_seedpaths)
        write_tsv(pd.DataFrame({"path": fins_nonseeds}), fout_nonseedpaths)

    logging.info("STEP 1 - inferring pangenome of seed genomes")
    args_pan = Namespace(faa_files = fout_seedpaths, outfolder = dout_seedpan,
        method = method, min_reps = min_reps, max_reps = max_reps, 
        max_align = max_align, threads = threads)
    run_pan(args_pan)
    
    logging.info("STEP 2 - building database of seed core genes and searching "
        "in seed faas")
    args_build = Namespace(faa_files = fout_seedpaths, pangenome = fout_seedpan,
        outfolder = dout_seedcore, core_prefilter = core_prefilter,
        core_filter = core_filter, max_core_genes = max_cores, 
        threads = threads)
    run_build(args_build)
    
    if os.stat(fout_nonseedpaths).st_size == 0:
        logging.info("all faa files are seeds - skipping search in non-seeds")
        shutil.copy(fout_genes_core, fout_genes)
        return()

    logging.info("STEP 3 - identifying core orthogroups in non-seed genomes")
    args_search = Namespace(faa_files = fout_nonseedpaths, db = dout_seedcore,
        outfolder = dout, threads = threads)
    run_search(args_search)
    
    logging.info("adding search results in seeds to search results of "
        "non-seeds")
    with open(fout_genes, "a") as hout_genes:
        with open(fout_genes_core) as hout_genes_core:
            for line in hout_genes_core:
                hout_genes.write(line)
    
