import re
import os
import sys
import gzip
import shutil
import hashlib
import logging
import multiprocessing as mp
from tqdm import tqdm
from ast import literal_eval
from cogent3 import PhyloNode, load_tree
from collections import defaultdict
from itertools import combinations
from datetime import datetime
from mycotools.acc2fa import dbmain as acc2fa
from mycotools.lib.biotools import gff2list, dict2fa, fa2dict
from mycotools.lib.kontools import format_path, collect_files, collect_dirs

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def init_log(log_file, log_dict):
    with open(log_file, "w") as out:
        out.write(
            f'mtdb\t{log_dict["mtdb"]}\n'
            + f'focal_genes\t{log_dict["focal_genes"]}\n'
            + f'microsynt_constraint\t{log_dict["microsynt_constraint"]}\n'
            + f'microsynt_tree\t{log_dict["microsynt_tree"]}\n'
            + f'hg_file\t{log_dict["hg_file"]}\n'
            + f'dist_type\t{log_dict["dist_type"]}\n'
            + f'uniq_sp\t{log_dict["uniq_sp"]}\n'
            + f'plusminus\t{log_dict["plusminus"]}\n'
            + f'hgp_percentile\t{log_dict["hgp_percentile"]}\n'
            + f'hgx_percentile\t{log_dict["hgx_percentile"]}\n'
            + f'aligner\t{log_dict["aligner"]}\n'  #                + f'hlg_percentile\t{log_dict["hlg_percentile"]}\n' \
            + f'gene_id\t{log_dict["gene_id"]}\n'
            + f'hlg_id\t{log_dict["hlg_id"]}\n'
            + f'hlg_sim\t{log_dict["hlg_sim"]}\n'
            + f'orig_hlg_id\t{log_dict["hlg_id"]}\n'
            + f'topology_merge\t{log_dict["topology_merge"]}\n'
            + f'min_topology_sim\t{log_dict["min_topology_sim"]}\n'
            + f'domain_inflation\t{log_dict["domain_inflation"]}\n'
            + f'hlg_inflation\t{log_dict["hlg_inflation"]}\n'
            + f'tuning\t{log_dict["tuning"]}\n'
            + f'id_percent\t{log_dict["id_percent"]}\n'
            + f'pos_percent\t{log_dict["pos_percent"]}\n'
            + f'iipr_percent\t{log_dict["iipr_percent"]}\n'
            + f'patch_threshold\t{log_dict["patch_threshold"]}\n'
            + f'gcl_threshold\t{log_dict["gcl_threshold"]}\n'
            + f'dist_threshold\t{log_dict["dist_threshold"]}\n'
            + f'partition\t{log_dict["partition"]}\n'
            + f'null_samples\t{log_dict["null_samples"]}\n'
            + f'n50\t{log_dict["n50"]}\n'
            + f'hg_dir\t{log_dict["hg_dir"]}\n'
            + f'hgx_dir\t{log_dict["hgx_dir"]}\n'
            + f'ipr_path\t{log_dict["ipr_path"]}\n'
            + f'pfam_path\t{log_dict["pfam_path"]}'
        )


def read_log(log_file, log_dict):
    log_res = {}
    with open(log_file, "r") as raw:
        for line in raw:
            key = line[: line.find("\t")]
            res = line[line.find("\t") + 1 :].rstrip()
            if key not in log_dict:  # old feature
                log_res[key] = False
                continue
            elif key == "orig_hlg_id":
                if round(float(res) * 100) <= round(float(log_dict["hlg_id"]) * 100):
                    log_res[key] = True
                else:
                    log_dict["orig_hlg_id"] = log_dict["hlg_id"]
                    log_res[key] = False
            elif key == "hgx_percentile":
                if round(float(res) * 100) > round(
                    float(log_dict["hgx_percentile"]) * 100
                ):
                    log_res["hgx_percentile_I"] = False
                    log_res["hgx_percentile_II"] = False
                elif round(float(res) * 100) != round(
                    float(log_dict["hgx_percentile"]) * 100
                ):
                    log_res["hgx_percentile_I"] = True
                    log_res["hgx_percentile_II"] = False
                else:
                    log_res["hgx_percentile_I"] = True
                    log_res["hgx_percentile_II"] = True
            elif key == "hlg_id":
                continue
            elif key == "hlg_sim":
                if res.lower() != str(log_dict[key]).lower():
                    log_res[key] = False
                else:
                    log_res[key] = True
            elif res != str(log_dict[key]):
                log_res[key] = False
                if key == "tuning":
                    # if you're tuning and weren't before
                    if log_dict["tuning"]:
                        # then the old inflation is irrelevant
                        log_res["domain_inflation"] = False
                        log_res["hlg_inflation"] = False
                    # elif you're not tuning, were before, and inflations are the same
                    elif log_res["domain_inflation"] and log_res["hlg_inflation"]:
                        log_res["tuning"] = True
                        log_dict["tuning"] = literal_eval(res)
                    # elif the domain inflation is the same and the hlg is not
            #                   else:
            #                        log_res['hlg_inflation'] = False
            else:
                log_res[key] = True
                if key == "tuning" and log_dict["tuning"]:
                    if inflation_rnd1:
                        log_res["domain_inflation"] = True
                    else:
                        inflation_rnd2 = None
                        log_res["domain_inflation"] = False
                    if inflation_rnd2:
                        log_res["hlg_inflation"] = True
                    else:
                        log_res["hlg_inflation"] = False
            # if there's an inflation, tuning is complete
            if key == "domain_inflation":
                if log_dict["tuning"]:
                    inflation_rnd1 = literal_eval(res)
                else:  # else the inflation is the inflation in the log_dict
                    inflation_rnd1 = log_dict["domain_inflation"]
            elif key == "hlg_inflation":
                # if you're tuning
                if log_dict["tuning"] and log_res["hlg_inflation"]:
                    inflation_rnd2 = literal_eval(res)
                else:
                    inflation_rnd2 = log_dict["hlg_inflation"]

    init_discrep = []
    try:
        if not log_res["mtdb"]:
            init_discrep.append("-d")
            log_res["n50"] = False
    except KeyError:
        init_discrep.append("-d")
    try:
        if not log_res["partition"]:
            init_discrep.append("-nr/-np")
            log_res["null_samples"] = False
    except KeyError:
        init_discrep.append("-nr/-np")
    try:
        if not log_res["focal_genes"]:
            init_discrep.append("-f")
            log_res["microsynt_tree"] = False
    except KeyError:
        init_discrep.append("-f")
    try:
        if not log_res["microsynt_constraint"]:
            init_discrep.append("-c")
            log_res["microsynt_tree"] = False
    except KeyError:
        init_discrep.append("-c")
    try:
        if not log_res["n50"]:
            init_discrep.append("--n50")
            log_res["plusminus"] = False
    except KeyError:
        init_discrep.append("--n50")
    try:
        if not log_res["dist_type"]:
            init_discrep.append("-mmd")
            log_res["uniq_sp"] = False
    except KeyError:
        init_discrep.append("-mmd")
    try:
        if not log_res["uniq_sp"]:
            init_discrep.append("-u")
            log_res["null_samples"] = False
    except KeyError:
        init_discrep.append("-u")
    try:
        if not log_res["plusminus"]:
            init_discrep.append("-w")
            log_res["null_samples"] = False
    except KeyError:
        init_discrep.append("-w")
    try:
        if not log_res["null_samples"]:
            init_discrep.append("-ns")
            log_res["hgp_percentile"] = False
    except KeyError:
        init_discrep.append("-ns")
    try:
        if not log_res["hgp_percentile"]:
            init_discrep.append("-hp")
            log_res["hgx_percentile_I"] = False
    except KeyError:
        init_discrep.append("-hp")
    try:
        if not log_res["hgx_percentile_I"]:
            init_discrep.append("-xp")
            log_res["hgx_percentile_II"] = False
    except KeyError:
        init_discrep.append("-xp")
    try:
        if not log_res["hgx_percentile_II"]:
            init_discrep.append("-xp")
            log_res["topology_merge"] = False
    except KeyError:
        init_discrep.append("-xp")
    try:
        if not log_res["topology_merge"]:
            log_res["min_topology_sim"] = False
            init_discrep.append("-tm")
    except KeyError:
        init_discrep.append("-tm")
    try:
        if not log_res["min_topology_sim"]:
            init_discrep.append("-ts")
            log_res["hlg_sim"] = False
    except KeyError:
        init_discrep.append("-ts")
    try:
        if not log_res["aligner"]:
            init_discrep.append("-a")
            log_res["gene_id"] = False
    except KeyError:
        log_res["aligner"] = False
        init_discrep.append("-a")
    try:
        if not log_res["gene_id"]:
            init_discrep.append("-mg")
            log_res["hlg_sim"] = False
    except KeyError:
        init_discrep.append("-mg")
    try:
        if not log_res["hlg_sim"]:
            init_discrep.append("-s")
            log_res["orig_hlg_id"] = False
    except KeyError:
        init_discrep.append("-s")
    try:
        if not log_res["orig_hlg_id"]:
            init_discrep.append("-ml")
            log_res["domain_inflation"] = False
    except KeyError:
        init_discrep.append("-ml")
    try:
        if not log_res["domain_inflation"]:
            init_discrep.append("-I2/-T")
            log_res["hlg_inflation"] = False
    except KeyError:
        init_discrep.append("-I1/-T")
        logger.error(
            "corrupted log.txt."
            + "\nIf not rectified, future runs will completely overwrite the current"
        )
        sys.exit(149)
    try:
        if not log_res["hlg_inflation"]:
            init_discrep.append("-I2/-T")
    #            log_res['hlg_percentile'] = False
    except KeyError:
        init_discrep.append("-I2/-T")
        logger.error("corrupted log.txt\nFuture runs may overwrite the current")
        sys.exit(149)
    #    try:
    #       if not log_res['hlg_percentile']:
    #          init_discrep.append('-fp')
    # except KeyError:
    #    init_discrep.append('-fp')

    # stop gap for legacy data
    if "hg_dir" not in log_res:
        log_res["hg_dir"] = True
    if "hgx_dir" not in log_res:
        log_res["hgx_dir"] = True

    return log_res, inflation_rnd1, inflation_rnd2, init_discrep


def rm_old_data(log_res, out_dir, wrk_dir):

    todel, tosave, save_dir = [], [], None
    if not log_res["dist_type"] or not log_res["uniq_sp"]:
        if os.path.isfile(out_dir + "working/omes2dist.pickle"):
            os.remove(out_dir + "working/omes2dist.pickle")
    if not log_res["null_samples"]:
        if os.path.isdir(wrk_dir + "null/"):
            shutil.rmtree(wrk_dir + "null/")
    if not log_res["hgp_percentile"]:
        seed_file = out_dir + "hgps.tsv.gz"
        seed_arr = wrk_dir + "microsynt.npy"
        if os.path.isfile(seed_file):
            tosave.append(seed_file)
        if os.path.isfile(seed_arr):
            os.remove(seed_arr)
        clus_pickle = wrk_dir + "hgx2loc.pickle"
        if os.path.isfile(clus_pickle):
            os.remove(clus_pickle)
    if not log_res["hgx_percentile_I"]:
        skip_files = collect_files(wrk_dir + "null/", "*")
        skip_files = [x for x in skip_files if os.path.basename(x).startswith("skip.")]
        for skip_file in skip_files:
            os.remove(skip_file)
        done_file = [
            x
            for x in collect_files(f"{wrk_dir}null/", "txt")
            if os.path.basename(x).startswith("done")
        ]
        if done_file:
            os.remove(done_file[0])
    if not log_res["hgx_percentile_II"]:
        hlg_dir = wrk_dir + "hlg/"
        if os.path.isdir(hlg_dir):
            todel.append(hlg_dir)
        groups = ["hgcs", "gene2hgx"]
        for group in groups:
            if os.path.isfile(f"{wrk_dir}{group}.pickle"):
                os.remove(f"{wrk_dir}{group}.pickle")
    if not log_res["min_topology_sim"]:
        groups = ["group.III"]
        clan_data = ["hg.json.gz", "hgx.json.gz", "json.gz"]
        for file_ in clan_data:
            if os.path.isfile(f"{wrk_dir}clan2loci.{file_}"):
                os.remove(f"{wrk_dir}clan2loci.{file_}")
        for group in groups:
            if os.path.isfile(f"{wrk_dir}{group}.pickle"):
                os.remove(f"{wrk_dir}{group}.pickle")
        todel.append(f"{wrk_dir}hlg/")
    if not log_res["aligner"]:
        hgx_dir = f"{wrk_dir}hgx/"
        logger.warning(
            "\talignment software changed, old alignments retained:\n\t\t{hgx_dir}"
        )
    #       shutil.rmtree(hgx_dir)
    if not log_res["orig_hlg_id"]:
        hlg_dir = f"{wrk_dir}hlg/"
        adj_rows = collect_files(hlg_dir, "tmp.w")
        adj_rows.extend(collect_files(hlg_dir, "tmp.r"))
        adj_rows.extend(collect_files(hlg_dir, "tmp"))
        row_file = wrk_dir + "hlg/loci.adj"
        mci_file = wrk_dir + "hlg/loci.mci"
        mcl_rows = wrk_dir + "hlg/mcl_rows.tsv"
        for f in [row_file, mci_file, mcl_rows]:
            if os.path.isfile(f):
                todel.append(f)

    if not log_res["hlg_inflation"]:
        hlgs_file = f"{wrk_dir}hlgs.pickle"
        if os.path.isfile(hlgs_file):
            os.remove(hlgs_file)
        mcl_res = f"{wrk_dir}hlg/loci.clus"
        tosave.extend([mcl_res, mcl_res + ".tmp"])
        hgx_files = ["gcl.pickle", "mmi.pickle", "mmp.pickle"]
        for file_ in hgx_files:
            if os.path.isfile(f"{wrk_dir}{file_}"):
                os.remove(f"{wrk_dir}{file_}")

        #    if not log_res['hlg_percentile']:
        hlg_output = out_dir + "hlgs.tsv.gz"
        gcf_output = out_dir + "gcfs.tsv.gz"
        kern_file = out_dir + "hgxs.tsv.gz"
        ome_dir = out_dir + "ome/"
        ann_dir = wrk_dir + "ann/"
        tosave.extend([hlg_output, gcf_output])
        if os.path.isdir(ome_dir):
            count = 0
            save_dir = f"{out_dir}run{count}/"
            while os.path.isdir(save_dir):
                count += 1
                save_dir = f"{out_dir}run{count}/"
            os.mkdir(save_dir)
            os.mkdir(save_dir + "working/")
            os.mkdir(save_dir + "working/hlg/")
            os.mkdir(save_dir + "working/hlg/lg")
            shutil.move(ome_dir, save_dir)
            if os.path.isfile(kern_file):
                shutil.copy(kern_file, save_dir + os.path.basename(kern_file))
            if os.path.isfile(hlg_output):
                shutil.move(hlg_output, save_dir + os.path.basename(hlg_output))
            if os.path.isfile(out_dir + "hlgs.html"):
                shutil.move(out_dir + "hlgs.html", save_dir + "hlgs.html")
            if os.path.isfile(out_dir + "hlg_metrics.html"):
                shutil.move(out_dir + "hlg_metrics.html", save_dir + "hlg_metrics.html")
            if os.path.isfile(gcf_output):
                shutil.move(gcf_output, save_dir + os.path.basename(clus_file))
            if os.path.isfile(out_dir + "gcfs.html"):
                shutil.move(out_dir + "gcfs.html", save_dir + "gcfs.html")
            if os.path.isfile(out_dir + "gcf_metrics.html"):
                shutil.move(out_dir + "gcf_metrics.html", save_dir + "gcf_metrics.html")

            shutil.move(out_dir + "log.txt", save_dir + "log.txt")
            for f in collect_files(f"{wrk_dir}hlg/", "*"):
                shutil.copy(f, f"{save_dir}working/hlg/{os.path.basename(f)}")
            for f in collect_files(f"{wrk_dir}hlg/lg/", "*"):
                shutil.copy(f, f"{save_dir}working/hlg/lg/{os.path.basename(f)}")

    #        else:
    #           if os.path.isfile(clus_file):
    #              os.remove(clus_file)

    #        if os.path.isdir(hmm_dir):
    #           shutil.rmtree(hmm_dir)
    #      if os.path.isdir(f'{out_dir}net/'):
    #         shutil.move(f'{out_dir}net/', f'{save_dir}net/')

    if not log_res["hg_dir"]:
        if os.path.isdir(f"{wrk_dir}hg/"):
            shutil.rmtree(f"{wrk_dir}hg/")
    if not log_res["hgx_dir"]:
        if os.path.isdir(f"{wrk_dir}hgx/"):
            shutil.rmtree(f"{wrk_dir}hgx/")

    if save_dir:
        for f in tosave:
            try:
                shutil.move(f, save_dir + os.path.basename(f))
            except FileNotFoundError:
                continue
    else:
        for f in tosave:
            try:
                os.remove(f)
            except FileNotFoundError:
                continue
    for f in todel:
        if os.path.isdir(f):
            shutil.rmtree(f)
        elif os.path.isfile(f):
            os.remove(f)

    if not log_res["hg_file"]:
        shutil.rmtree(wrk_dir)
        os.mkdir(wrk_dir)
        for key in log_res:
            log_res[key] = False


def log_check(log_dict, log_path, out_dir, wrk_dir, flag=True):

    gcf_thresh = False
    if not os.path.isfile(log_path):
        log_res = {x: False for x in log_dict}
        try:
            rm_old_data(log_res, out_dir, wrk_dir)
        except KeyError:
            logger.warning("missing log, rerunning without deleting")
        init_log(log_path, log_dict)
    log_res, inflation_1, inflation_2, init_discrep = read_log(log_path, log_dict)
    if any(not log_res[x] for x in log_res):
        failed = set([x for x, v in log_res.items() if not v])
        for skip in [
            "patch_threshold",
            "gcl_threshold",
            "iipr_percent",
            "id_percent",
            "pos_percent",
            "dist_threshold",
            "ipr_path",
            "pfam_path",
        ]:
            if skip in failed:
                failed.remove(skip)
        if failed:
            try:
                if not flag:
                    logger.info(
                        f"\n{init_discrep[0]} changed; removing data to checkpoint"
                    )
                    rm_old_data(log_res, out_dir, wrk_dir)
                    init_log(log_path, log_dict)
                else:
                    logger.error(
                        "-n not called and incompatible parameters: \
                            \n\t"
                        + init_discrep[0]
                    )
                    sys.exit(15)
            # irrelevant missing discrepancy
            except IndexError:
                init_log(log_path, log_dict)
        else:
            if os.path.isdir(out_dir + "ome/"):
                tsvs = set(collect_files(out_dir + "ome/", "tsv", recursive=True))
                # NEED an intelligent resume
                ome_dirs = [
                    f"{out_dir}ome/{x}/"
                    for x in os.listdir(out_dir + "ome/")
                    if os.path.isdir(f"{out_dir}ome/{x}/")
                ]
                if all(f"{x}hlg.tsv" in tsvs for x in ome_dirs) or os.path.isfile(
                    out_dir + "hlgs.tsv.gz"
                ):
                    logger.info("\nBypassing to filtration")
                    gcf_thresh = True
                else:
                    init_log(log_path, log_dict)
            else:
                init_log(log_path, log_dict)
    elif os.path.isdir(out_dir + "ome/"):
        tsvs = set(collect_files(out_dir + "ome/", "tsv", recursive=True))
        ome_dirs = [
            f"{out_dir}ome/{x}/"
            for x in os.listdir(out_dir + "ome/")
            if os.path.isdir(f"{out_dir}ome/{x}/")
        ]
        if all(f"{x}hlg.tsv" in tsvs for x in ome_dirs):
            logger.info("\nBypassing to filtration")
            gcf_thresh = True

    return log_res, inflation_1, inflation_2, gcf_thresh


def compileCDS(gff_list, ome):
    """
    Inputs the gff_list and organism ome code. Compiles CDS entries for loci parsing and outputs:
    cds_dict = {contig: {protein: [[CDSstart, CDSstop]] } }
    Sorts the cds_dict for each protein from smallest to largest
    """

    cds_dict, fail = defaultdict(dict), False
    for entry in gff_list:
        if entry["type"] == "CDS":
            prot_prep_i0 = entry["attributes"].index(";Alias=")  # grab the
            # mycotools accession index
            try:  # grab the end of the Alias by grabbing the index for a ';',
                # that is after the accession tag's occurrence.
                # then add that to the start of the accession tag and add one to
                # remove that semicolon. If there is a value error, the accession
                # is at the end of the attributes (gff col 8 (?))
                prot_prep_i1 = (
                    entry["attributes"][prot_prep_i0 + 1 :].index(";")
                    + prot_prep_i0
                    + 1
                )
            except ValueError:
                prot_prep_i1 = len(entry["attributes"])
            prot = entry["attributes"][prot_prep_i0:prot_prep_i1].replace(";Alias=", "")
            # obtain the protein accession
            if prot not in cds_dict[entry["seqid"]] and prot:  # add the protein
                # to the cds_dict entry
                cds_dict[entry["seqid"]][prot] = []
            elif not prot:  # if there isn't a valid accession it may mean the
                # mycotools curation did not work or the user did not curate
                # correctly
                logger.info(entry["attributes"], prot_prep_i0, prot_prep_i1)
                if not fail:
                    logger.warning("\t" + ome + " has proteins in gff with no Alias")
                fail = True
                continue
            cds_dict[entry["seqid"]][prot].extend(
                [int(entry["start"]), int(entry["end"])]
            )  # add the coordinates

    for contig in cds_dict:
        for prot in cds_dict[contig]:
            cds_dict[contig][prot].sort()  # sort the coordinates of the proteins
            # lowest to highest
        cds_dict[contig] = list(
            sorted(cds_dict[contig].keys(), key=lambda k: cds_dict[contig][k][0])
        )  # sort the proteins in the contigs lowest
        # to highest coordinates

    return dict(cds_dict)


def compileCDS2(gff_list, ome):
    """
    Inputs the gff_list and organism ome code. Compiles CDS entries for loci parsing and outputs:
    cds_dict = {contig: {protein: [[CDSstart, CDSstop]] } }
    Sorts the cds_dict for each protein from smallest to largest
    """

    cds_dict, cds_dict2, fail = defaultdict(dict), {}, False
    for entry in gff_list:
        if entry["type"] == "CDS":
            prot_prep_i0 = entry["attributes"].index(";Alias=")  # grab the
            # mycotools accession index
            try:  # grab the end of the Alias by grabbing the index for a ';',
                # that is after the accession tag's occurrence.
                # then add that to the start of the accession tag and add one to
                # remove that semicolon. If there is a value error, the accession
                # is at the end of the attributes (gff col 8 (?))
                prot_prep_i1 = (
                    entry["attributes"][prot_prep_i0 + 1 :].index(";")
                    + prot_prep_i0
                    + 1
                )
            except ValueError:
                prot_prep_i1 = len(entry["attributes"])
            prot = entry["attributes"][prot_prep_i0:prot_prep_i1].replace(";Alias=", "")
            # obtain the protein accession
            if prot not in cds_dict[entry["seqid"]] and prot:  # add the protein
                # to the cds_dict entry
                cds_dict[entry["seqid"]][prot] = []
                cds_dict2[prot] = []
            elif not prot:  # if there isn't a valid accession it may mean the
                # mycotools curation did not work or the user did not curate
                # correctly
                logger.info(entry["attributes"], prot_prep_i0, prot_prep_i1)
                if not fail:
                    logger.warning("\t" + ome + " has proteins in gff with no Alias")
                fail = True
                continue
            cds_dict[entry["seqid"]][prot].extend(
                [int(entry["start"]), int(entry["end"])]
            )  # add the coordinates
            cds_dict2[prot].append(entry)

    for contig in cds_dict:
        for prot in cds_dict[contig]:
            cds_dict[contig][prot].sort()  # sort the coordinates of the proteins
            # lowest to highest
        cds_dict[contig] = list(
            sorted(cds_dict[contig].keys(), key=lambda k: cds_dict[contig][k][0])
        )  # sort the proteins in the contigs lowest
        # to highest coordinates

    return dict(cds_dict), cds_dict2


def parseLoci(gff_path, ome_num, gene2hg, plusminus=6):
    """obtain a set of tuples of HG pairs {(OG0, OG1)...}"""

    gff_list = gff2list(gff_path)  # open here to improve pickling
    cds_dict = compileCDS(gff_list, os.path.basename(gff_path).replace(".gff3", ""))
    pairs = []
    for scaf in cds_dict:  # for each contig
        for i, seq in enumerate(cds_dict[scaf]):  # for each gene
            if i < plusminus:  # does the locus start at the contig border?
                locus = cds_dict[scaf][: i + plusminus + 1]
            else:  # or I extend it to the +/- or opposite border
                locus = cds_dict[scaf][i - plusminus : i + plusminus + 1]
            loc_og = []
            for gene in locus:
                try:  # failure should be relatively rare, so faster than `if`
                    loc_og.append(gene2hg[gene])
                except KeyError:  # it's either a singleton or not in an OG
                    pass
            pairs.extend(
                [tuple(sorted(x)) for x in combinations(set(loc_og), 2)]
            )  # don't deal with
            # same OGs - tandem duplications will just naturally cooccur more often

    out_pairs = set(pairs)  # unique pairs of OGs
    return ome_num, out_pairs


def load_seedScores(file_):  # , seed_thresh):

    out_hgs = []
    with gzip.open(file_, "rt") as raw:
        for line in raw:
            if not line.startswith("#"):
                d = line.rstrip().split()
                out_hgs.append((int(d[0]), int(d[1])))

    return out_hgs


def split_tree(tree, name):
    """Prepare tree for rooting"""
    node = tree.get_node_matching_name(name)
    length = (node.length or 0.0) / 2
    parent = node.parent
    # report if the node is at the root
    if not parent:
        return None, None
    parent.children.remove(node)
    node.parent = None
    left = node.unrooted_deepcopy()
    left.name = f"{name}-L"
    right = parent.unrooted_deepcopy()
    right.name = f"{name}-R"
    left.length = length
    right.length = length
    return left, right


def root_tree(tree, outgroup=[], midpoint=False):
    """Root the tree based on outgroup(s)"""
    if midpoint:
        # root at the midpoint
        return tree.root_at_midpoint()
    elif isinstance(outgroup, str) or (
        isinstance(outgroup, list) and len(outgroup) == 1
    ):
        # root via outgroup tip if there is 1 list member or it is a string
        if isinstance(outgroup, list):
            outgroup = outgroup[0]
        if outgroup:
            return tree.rooted_with_tip(outgroup).bifurcating()
        else:
            return tree
    elif len(outgroup > 2):
        raise RootError(
                f"can only root using two tips"
            )
    else:
        # root by determining the MRCA branch as the MRCA that contains as few tips as possible
        # while including those delineated in the outgroup list
        mrca = tree.get_connecting_node(outgroup)
        left, reft = split_tree(tree, mrca.name)
        try:
            return type(tree)(name="root", children=[left, reft])
        except:
            raise RootError(
                f"tree could not be rooted with supplied tips: " + f"{outgroup}"
            )


def compile_tree(i2ome, tree_path, root=[]):
    #    with open(tree_path, 'r') as raw:
    #       nwk = raw.read()

    phylo = load_tree(tree_path)
    phylo.name_unnamed_nodes()
    omes = set(i2ome)
    tips = set(phylo.get_tip_names())
    missing = tips.difference(omes)
    if missing:
        logger.warning(
            f"\t\tremoving {len(missing)} " + "missing tips from microsynteny tree"
        )
        todel = []
        for n in phylo.tips():
            if n.name in missing:
                todel.append(n)
        for n in todel:
            n.parent.remove(n)
            phylo.prune()

    if isinstance(root, list):
        try:
            phylo = root_tree(phylo, outgroup=root)
        except:
            logger.error(f"tree could not be rooted with supplied tips: " + f"{root}")
            sys.exit(449)
        phylo.write(tree_path, with_distances=True)
    else:
        phylo = root_tree(phylo, midpoint=True)

    phylo.reassign_names({v: str(i) for i, v in enumerate(i2ome)})
    return phylo


def write_hg_fa(hg_file, big_dict, genes):
    fa_str = dict2fa({k: big_dict[k] for k in genes})
    with open(f"{hg_file}.tmp", "w") as out:
        out.write(fa_str)
    os.rename(f"{hg_file}.tmp", hg_file)


def output_hg_fas(db, genes, hg_file):
    fa_str = dict2fa(acc2fa(db, genes, error=False, spacer="\t\t", coord_check=False))
    with open(hg_file + ".tmp", "w") as out:
        out.write(fa_str)
    os.rename(hg_file + ".tmp", hg_file)


def cp_files(f0, f1):
    shutil.copy(f0, f1)


def symlink_files(f0, f1):
    os.symlink(f0, f1)


def load_ome2fa(db, cpus=1):
    #    big_fa = {}
    #    with mp.Pool(processes = cpus) as pool:
    #       fa_dicts = pool.map(fa2dict,
    #                          tqdm((row['faa'] for row in db.values()),
    #                         total = len(db)))
    #    pool.close()
    #   pool.join()
    #    for res in fa_dicts:
    #       big_fa = {**big_fa, **res}
    fas = []
    for row in tqdm(db.values(), total=len(db)):
        fas.append(fa2dict(row["faa"]))
    big_fa = {}
    for fa in fas:
        for nam, seq in fa.items():
            big_fa[nam] = seq

    return big_fa


def load_ome2fa_mp(db, cpus=1):
    big_fa = {}
    with mp.Pool(processes=cpus) as pool:
        fa_dicts = pool.map(
            fa2dict, tqdm((row["faa"] for row in db.values()), total=len(db))
        )
        pool.close()
        pool.join()
    for res in fa_dicts:
        big_fa = {**big_fa, **res}

    return big_fa


def big_acc2fa(db, hg_dir, hgs, hg2gene, big_fa, cpus=1):
    for hg in tqdm(hgs, total=len(hgs)):
        write_hg_fa(f"{hg_dir}{hg}.faa", big_fa, hg2gene[hg])


def big_acc2fa_mp(db, hg_dir, hgs, hg2gene, big_fa, cpus=1):
    big_dict = mp.Manager().dict(big_fa)
    with mp.Pool(processes=cpus) as pool:
        pool.starmap(
            write_hg_fa,
            tqdm(
                ((f"{hg_dir}{hg}.faa", big_dict, hg2gene[hg]) for hg in hgs),
                total=len(hgs),
            ),
        )
        pool.close()
        pool.join()


def hg_fa_mngr(wrk_dir, hg_dir, hgs, db, hg2gene, cpus=1, low_mem=False):
    new_hg_dir = wrk_dir + "hg/"
    if not os.path.isdir(new_hg_dir):
        os.mkdir(new_hg_dir)
    # extract only used HGs
    hg_files = set(
        [int(os.path.basename(x[:-4])) for x in collect_files(new_hg_dir, "faa")]
    )
    missing_hgs = sorted(set(hgs).difference(hg_files))

    if not missing_hgs:
        return new_hg_dir

    if not hg_dir:
        if low_mem:  # this process is way too slow
            with mp.Pool(processes=cpus) as pool:
                pool.starmap(
                    output_hg_fas,
                    tqdm(
                        (
                            (db, hg2gene[hg], f"{new_hg_dir}{hg}.faa")
                            for hg in missing_hgs
                        ),
                        total=len(missing_hgs),
                    ),
                )
                pool.close()
                pool.join()
        else:
            big_fa = load_ome2fa(db, cpus=cpus)
            big_acc2fa(db, new_hg_dir, missing_hgs, hg2gene, big_fa, cpus)
    else:  # predetermined hg input
        hg_fa_cmds = []
        orthofinder = False
        if not os.path.isfile(f"{hg_dir}{hgs[0]}.faa"):
            digits = len(str(hgs[0]))
            zeros = 7 - digits
            if os.path.isfile(f'{hg_dir}OG{"0" * zeros}{hgs[0]}.fa'):
                orthofinder = True
                ext = ".fa"
            elif os.path.isfile(f"{hg_dir}{hgs[0]}.fa"):
                ext = ".fa"
            elif os.path.isfile(f"{hg_dir}{hgs[0]}.fasta"):
                ext = ".faa"
        else:
            raise FileNotFoundError("HG fastas missing")

        if orthofinder:
            for hg in missing_hgs:
                digits = len(str(hg))
                zeros = 7 - digits
                hg_file = f'{hg_dir}OG{"0" * zeros}{hg}.fa'
                hg_fa_cmds.append((hg_file, f"{new_hg_dir}{hg}.faa"))
        else:
            hg_fa_cmds = [
                (f"{hg_dir}{hg}{ext}", f"{new_hg_dir}{hg}.faa") for hg in missing_hgs
            ]
        copy_hgs = False
        try:
            os.symlink(hg_fa_cmds[0][0], hg_fa_cmds[0][1])
            del hg_fa_cmds[0]
        except:
            copy_hgs = True

        if copy_hgs:
            with mp.Pool(processes=cpus) as pool:
                pool.starmap(cp_files, hg_fa_cmds)
                pool.close()
                pool.join()
        else:
            with mp.Pool(processes=cpus) as pool:
                pool.starmap(symlink_files, hg_fa_cmds)
                pool.close()
                pool.join()
    return new_hg_dir


def sha_tune_file(tune_file):
    with open(tune_file, "r") as raw:
        tune_data = [x.rstrip() for x in raw if not x.startswith("#")]
    tune = {}
    for cluster in tune_data:
        try:
            name, rawgenes, rawomes, rawfalse = cluster.split("\t")
        except ValueError:
            name, rawgenes, rawomes = cluster.split("\t")
            rawfalse = None
        if name in tune:
            logger.error(f"duplicate entry for {name} in tune file")
        if "," in rawgenes:
            genes = [x.rstrip().lstrip() for x in rawgenes.split(",")]
        else:
            genes = [x.rstrip().lstrip() for x in rawgenes.split()]
        if "," in rawomes:
            omes = [x.rstrip().lstrip() for x in rawomes.split(",")]
        else:
            omes = [x.rstrip().lstrip() for x in rawomes.split()]
        if rawfalse:
            if "," in rawfalse:
                false_omes = [x.rstrip().lstrip() for x in rawfalse.split(",")]
            else:
                false_omes = [x.rstrip().lstrip() for x in rawfalse.split()]
        else:
            false_omes = []
        tune[name] = [
            tuple(sorted(set(genes))),
            tuple(sorted(set(omes))),
            tuple(sorted(set(false_omes))),
        ]

    tune_sha = hashlib.sha256(str(sorted(tune)).encode("utf-8")).hexdigest()
    return tune_sha


def init_run(
    db,
    out_dir,
    near_single_copy_genes,
    constraint_path,
    tree_path,
    hg_file,
    plusminus,
    seed_perc,  # clus_perc,
    hgx_perc,
    aligner,
    id_perc,
    pos_perc,
    iipr_perc,
    patch_thresh,
    gcl_thresh,
    dist_thresh,
    samples,
    n50thresh,
    flag,
    min_gene_id,
    min_hlg_id,
    inflation_rnd1,
    inflation_rnd2,
    simfun,
    tune_file,
    dist_type,
    uniq_sp,
    partition,
    min_topology_sim,
    topology_merge,
    hg_dir,
    hgx_dir,
    ipr_path,
    pfam_path,
):

    wrk_dir = out_dir + "working/"
    if not os.path.isdir(wrk_dir):
        os.mkdir(wrk_dir)

    nul_dir = wrk_dir + "null/"
    if not os.path.isdir(nul_dir):
        os.mkdir(nul_dir)

    log_path = out_dir + "log.txt"
    if not tree_path:
        tree_path = out_dir + "microsynt.newick"

    if tune_file:
        tune_sha = sha_tune_file(tune_file)
    else:
        tune_sha = None

    log_dict = {
        "mtdb": hashlib.sha256(str(sorted(db.keys())).encode("utf-8")).hexdigest(),
        "focal_genes": hashlib.sha256(
            ",".join(sorted(near_single_copy_genes)).encode("utf-8")
        ).hexdigest(),
        "microsynt_constraint": format_path(constraint_path),
        "microsynt_tree": tree_path,
        "hg_file": hg_file,
        "plusminus": plusminus,
        "dist_type": dist_type,
        "uniq_sp": bool(uniq_sp),
        "hgp_percentile": seed_perc,
        "hgx_percentile": hgx_perc,
        "aligner": aligner,
        "gene_id": min_gene_id,
        "orig_hlg_id": None,
        "hlg_id": min_hlg_id,
        "hlg_sim": str(simfun).lower(),
        #        'hlg_percentile': clus_perc,
        "id_percent": id_perc,
        "pos_percent": pos_perc,
        "iipr_percent": iipr_perc,
        "patch_threshold": patch_thresh,
        "dist_threshold": dist_thresh,
        "gcl_threshold": gcl_thresh,
        "domain_inflation": inflation_rnd1,
        "hlg_inflation": inflation_rnd2,
        "tuning": tune_sha,
        "partition": partition,
        "topology_merge": topology_merge,
        "min_topology_sim": min_topology_sim,
        "null_samples": samples,
        "n50": n50thresh,
        "hg_dir": hg_dir,
        "hgx_dir": hgx_dir,
        "ipr_path": ipr_path,
        "pfam_path": pfam_path,
    }

    log_res, inflation_1, inflation_2, gcf_bypass = log_check(
        log_dict, log_path, out_dir, wrk_dir, flag
    )

    # symlink files on an individual basis so the new dir is writeable without
    # propagating new changes back to the reference dir
    hg_dir, hgx_dir = wrk_dir + "hg/", wrk_dir + "hgx/"

    existing_hg = set(os.path.basename(x) for x in collect_files(hg_dir, "faa"))
    existing_hgx = set(os.path.basename(x) for x in collect_files(hgx_dir, "out"))
    if log_dict["hg_dir"]:
        if not os.path.isdir(hg_dir):
            os.mkdir(wrk_dir + "hg/")
        prev_hg = set(
            os.path.basename(x) for x in collect_files(log_dict["hg_dir"], "faa")
        )
        missing_hg = prev_hg.difference(existing_hg)
        for i in list(missing_hg):
            os.symlink(log_dict["hg_dir"] + i, wrk_dir + "hg/" + i)
    if log_dict["hgx_dir"]:
        if not os.path.isdir(hgx_dir):
            os.mkdir(wrk_dir + "hgx/")
        prev_hgx = set(
            os.path.basename(x) for x in collect_files(log_dict["hgx_dir"], "out")
        )
        missing_hgx = prev_hgx.difference(existing_hgx)
        for i in list(missing_hgx):
            os.symlink(log_dict["hgx_dir"] + i, wrk_dir + "hgx/" + i)

    return wrk_dir, nul_dir, inflation_1, inflation_2, gcf_bypass
