import os
import re
import gzip
import shutil
import pickle
import logging
import subprocess
import multiprocessing as mp
import plotly.graph_objects as go
from math import log
from tqdm import tqdm
from itertools import chain
from collections import defaultdict
from plotly.subplots import make_subplots
from mycotools.acc2gff import grab_gff_acc
from mycotools.lib.biotools import list2gff, gff2list, fa2dict, dict2fa
from cloci.lib import input_parsing

# NEED to add PCA output
# NEED locID output and clans
# NEED gbk output option
# NEED to add merge abutting clusters feature
# NEED annotation for gcfs resume and hlgs independent

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


def mk_subplots(labels, axes, axes_labels, alpha=0.6):

    colors = [
        [240, 163, 255],
        [0, 117, 220],
        [153, 63, 0],
        [76, 0, 92],
        [25, 25, 25],
        [0, 92, 49],
        [43, 206, 72],
        [255, 204, 153],
        [128, 128, 128],
        [148, 255, 181],
        [143, 124, 0],
        [157, 204, 0],
        [194, 0, 136],
        [0, 51, 128],
        [255, 164, 5],
        [255, 168, 187],
        [66, 102, 0],
        [255, 0, 16],
        [94, 241, 242],
        [0, 153, 143],
        [224, 255, 102],
        [116, 10, 255],
        [153, 0, 0],
        [255, 255, 128],
        [255, 255, 0],
        [255, 80, 5],
    ]
    #    random.shuffle(colors)
    pl_clr = [
        "rgba(" + ", ".join([str(y) for y in x]) + ", " + str(alpha) + ")"
        for x in colors
    ]
    fig = make_subplots(
        rows=len(axes) - 1, cols=len(axes) - 1, start_cell="bottom-left"
    )

    clr_count = 0
    for i0, axis0 in enumerate(axes):
        for i1, axis1 in enumerate(axes[i0 + 1 :]):
            fig.add_trace(
                go.Scatter(
                    mode="markers",
                    x=axis0,
                    y=axis1,
                    marker=dict(color=pl_clr[clr_count], size=5, opacity=alpha),
                    text=labels,
                ),
                row=i0 + 1,
                col=i0 + i1 + 1,
            )
            fig.update_xaxes(title_text=axes_labels[i0], row=i0 + 1, col=i0 + i1 + 1)
            fig.update_yaxes(
                title_text=axes_labels[i0 + i1 + 1], row=i0 + 1, col=i0 + i1 + 1
            )
            clr_count += 1

    return fig


def mk_3d(labels, axes, axes_labels, alpha=0.6):

    fig = go.Figure()
    fig.add_trace(
        go.Scatter3d(
            mode="markers",
            x=axes[0],
            y=axes[1],
            z=axes[2],
            marker=dict(
                color="rgba(163, 25, 169," + str(alpha) + ")", size=5, opacity=alpha
            ),
            text=labels,
            showlegend=False,
        )
    )
    fig.update_layout(
        scene=dict(
            xaxis_title=axes_labels[0],
            yaxis_title=axes_labels[1],
            zaxis_title=axes_labels[2],
        )
    )  # ,
    #                    width=700,
    #                   margin=dict(r=20, b=10, l=10, t=10))

    return fig


def extract_sig_clus(clus_scores_dict, hgx2loc, top_hgxs, ome2i):
    """Create a hash to seed retrieving clusters based on
    their entry in hgx2loc."""

    sig_clus = defaultdict(dict)
    for top in top_hgxs:
        hgx, gcf, omeIs = top[0], top[1], top[2]
        for gene in hgx2loc[hgx]:
            ome = gene[: gene.find("_")]
            if ome2i[ome] in omeIs:
                if gene not in sig_clus[ome]:
                    sig_clus[ome][gene] = [[set(hgx), gcf]]
                else:
                    sig_clus[ome][gene].append([set(hgx), gcf])

    return dict(sig_clus)


def hash_clusters(gff_path):
    """sliding window of size clusplusminus, identify sequences that may have
    hgs that may be significant, grab their window, and check for a significant
    higher order og combo in the locus"""
    gff_list = gff2list(gff_path)
    cds_dict = input_parsing.compileCDS(
        gff_list, os.path.basename(gff_path).replace(".gff3", "")
    )
    return cds_dict


def merge_clusters(clus_out, cds_dict):
    """because we have significant clusters with borders, we can merge each
    cluster and assume that at a certain number of hgs in a combo will never
    happen via random sampling; therefore, their significance in the
    microsynteny adjustment is sufficient for justification for merging"""

    comp_clus, change = defaultdict(list), False
    for scaf, clusters in clus_out.items():
        gene_list = cds_dict[scaf]
        while clusters:  # while there are clusters
            any_intersect, toDel = False, [0]  # we haven't found intersection
            clus0 = clusters[0]  # grab the first tuple
            loc0 = clus0[0]  # grab the genes in the cluster
            low_i0 = gene_list.index(loc0[0])
            hi_i0 = gene_list.index(loc0[-1])
            for i, clus1 in enumerate(clusters[1:]):  # look at all other clusters
                loc1 = clus1[0]  # grab their loci
                low_i1 = gene_list.index(loc1[0])
                hi_i1 = gene_list.index(loc1[-1])
                if low_i1 - hi_i0 == 1 and False:
                    any_intersect, change = True, True  # we found overlap in
                    # clus_out and in this specific cluster comparison
                    #                   hgx = tuple(sorted(list(set(clus0[0]).union(set(clus1[0])))))
                    # obtain the higher order og combination as the union of the
                    # two cluster sets, then sort it and store as a tuple
                    comp_clus[scaf].append(
                        [
                            #                        hgx,
                            loc0 + loc1,
                            clus0[1].union(clus1[1]),  # , merge_x
                        ]
                    )
                    # append the og combo and the union of the two loci's genes
                    # protein order is random
                    toDel.append(i)  # we can delete the intersected locus as well
                    # as the first
                elif low_i0 - hi_i1 == 1 and False:
                    any_intersect, change = True, True
                    comp_clus[scaf].append([loc1 + loc0, clus0[1].union(clus1[1])])
                    toDel.append(i)
                elif set(loc0).intersection(set(loc1)) and False:
                    any_intersect, change = True, True
                    comp_clus[scaf].append(
                        [set(loc0).union(set(loc1)), clus0[1].union(clus1[1])]
                    )
                    comp_clus[scaf][-1][0] = sorted(
                        comp_clus[scaf][-1][0], key=cds_dict[scaf].index
                    )
                    toDel.append(i)
            if not any_intersect:  # if there is not any intersection in the locus
                comp_clus[scaf].append(
                    clus0
                )  # , clus0[2]]) # simply append the original
                # result
            toDel.sort(reverse=True)  # sort the todel from highest to lowest to
            # not mess up order when indexing
            for i in toDel:
                clusters.pop(i)

    return comp_clus, change


def clus2scaf(cds_dict, clusters):

    gene2scaf, index_map = {}, {}
    for scaf, genes in cds_dict.items():
        index_map[scaf] = {}
        for i, gene in enumerate(genes):
            index_map[scaf][gene] = i
            gene2scaf[gene] = scaf

    clus_out = defaultdict(list)
    for loc, gcf in clusters:
        gene0 = loc[0]
        scaf = gene2scaf[gene0]
        ord_loc = sorted(loc, key=lambda pair: index_map[scaf][pair])
        #        hgs = []
        #       for gene in ord_loc:
        #          try:
        #             hgs.append(gene2hg[gene])
        #        except KeyError:
        #           hgs.append('')
        #  clus_info = [hgs, ord_loc, {gcf}]

        clus_out[scaf].append([ord_loc, gcf])

    sorted_clus_out = {
        k: sorted(v, key=lambda x: index_map[k][x[0][0]]) for k, v in clus_out.items()
    }

    return sorted_clus_out, index_map


def write_clusters(clusters, ome, out_file, gff_path, gene2hg, clusplusminus=10):
    cds_dict = input_parsing.compileCDS(
        gff2list(gff_path), os.path.basename(gff_path).replace(".gff3", "")
    )

    clus_out, scaf2gene2i = clus2scaf(cds_dict, clusters)
    merge_out = defaultdict(list)
    for scaf, locs in clus_out.items():
        for loc, gcf in locs:
            merge_out[scaf].append([loc, {gcf}])
    #  change = True
    #    while change: # so long as there is a single merge
    #       clus_out, change = merge_clusters(clus_out, cds_dict) # attempt to merge clusters
    #    clus_out = sorted(clus_out, key = lambda x: len(x[0]), reverse = True)
    # sort the clusters by the size of the OG combination tuple highest to
    # lowest

    for scaf, clusters in merge_out.items():
        mapping = {v: i for i, v in enumerate(cds_dict[scaf])}
        clusters = sorted(clusters, key=lambda pair: mapping[pair[0][0]])
        for i, clus in enumerate(clusters):
            clus[0] = sorted(clus[0], key=cds_dict[scaf].index)
            clusHGs = []
            for gene in clus[0]:
                try:
                    clusHGs.append(gene2hg[gene])
                except KeyError:  # gene not in an OG/is a singleton
                    clusHGs.append("")
            clus.append(clusHGs)
            #         scaf = clus[2]
            #        scaf = clus[3][:20]
            #       if clus[3] != scaf:
            #            scaf = clus[3][:20] + '..'
            #            count = 0
            #           while scaf + '_' + str(count) in scafs:
            #              count += 1
            name = scaf + "_" + str(i)
            #            scafs.add(name)
            clus.append(name)
        clus_out[scaf] = clusters

    out_genes = []
    with open(out_file, "w") as out:
        out.write("#name\thgs\tgenes\tgcfs\n")
        for scaf, clusters in clus_out.items():
            for clus in clusters:
                #           clus[2] = sorted(clus[2])
                hgs = ",".join([str(x) for x in clus[2]])  # og1,og2,...ogn
                genes = ",".join(clus[0])  # prot1,prot2,prot3,...protn
                gcfs = ";".join([str(x) for x in sorted(clus[1])])
                out.write(clus[3] + "\t" + hgs + "\t" + genes + "\t" + gcfs + "\n")
                out_genes.append(clus[0])

    return ome, out_genes


def read_ome_hlgs(ome, ome_hlg_file):
    out_genes = []
    try:
        with open(ome_hlg_file, "r") as raw:
            for line in raw:
                if not line.startswith("#"):
                    #        out.write('#name\thgs\tgenes\tgcfs\n')
                    data = line.rstrip().split()
                    genes = data[2].split(",")
                    out_genes.append(genes)
        return ome, out_genes
    except FileNotFoundError:
        return ome, []


def output_ann_fas(ann_dir, ome, ome_genes, prot_path):
    faa = fa2dict(prot_path)
    genes = list(chain(*ome_genes))
    try:
        out_faa = {k: faa[k] for k in genes}
        with open(ann_dir + ome + ".faa", "w") as out:
            out.write(dict2fa(out_faa))
        return ome, True
    except KeyError:
        logger.error("\t\t" + ome + " accessions modified from input")
        return ome, False


def output_res_faas(genes_in, prot_paths, ann_dir, cpus=1):
    with mp.Pool(processes=cpus) as pool:
        res = pool.starmap(
            output_ann_fas,
            tqdm(
                (
                    (ann_dir, ome, ome_genes, prot_paths[ome])
                    for ome, ome_genes in genes_in.items()
                ),
                total=len(genes_in),
            ),
        )

    return set(ome for ome, passing in res if not passing)


def call_hmmsearch(ome, pfam, ann_dir, faa_file, threads=4):
    out_file = ann_dir + ome + "_pfam.out"
    if not os.path.isfile(out_file):
        hmmsearch = subprocess.call(
            [
                "hmmsearch",
                "--cpu",
                str(threads),
                "--domtblout",
                out_file + ".tmp",
                pfam,
                faa_file,
            ],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        os.rename(out_file + ".tmp", out_file)
    else:
        hmmsearch = 0
    return ome, hmmsearch


def run_hmmsearch(pfam, passing_omes, ann_dir, cpus, db=None, threads=4):
    """Manage running hmmsearch for Pfam annotations. Default runs just the
    Pfam hits, inputting a db runs for the full genome"""

    with mp.Pool(processes=round(cpus / 2)) as pool:
        if db:
            hmm_res = pool.starmap(
                call_hmmsearch,
                tqdm(
                    (
                        (ome, pfam, ann_dir, db[ome]["faa"], threads)
                        for ome in list(passing_omes)
                    ),
                    total=len(passing_omes),
                ),
            )
        else:
            hmm_res = pool.starmap(
                call_hmmsearch,
                tqdm(
                    (
                        (ome, pfam, ann_dir, f"{ann_dir}{ome}.faa", threads)
                        for ome in list(passing_omes)
                    ),
                    total=len(passing_omes),
                ),
            )

    failed_omes = []
    for ome, exit_code in hmm_res:
        if exit_code:
            logger.error("\t\tERROR: " + ome + " failed hmmsearch")
            failed_omes.append(ome)

    return failed_omes


def call_iprscan(ome, iprscan, ann_dir, threads=4):
    out_file = ann_dir + ome + "_ipr.out"
    faa_file = ann_dir + ome + ".faa"
    if not os.path.isfile(out_file):
        ipr_code = subprocess.call(
            [
                iprscan,
                "-i",
                faa_file,
                "-cpu",
                str(threads),
                "-o",
                out_file + ".tmp",
                "-f",
                "TSV",
            ],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
    else:
        ipr_code = 0
    return ome, ipr_code


def run_ipr_scan(iprscan, passing_omes, ann_dir, cpus):
    """Initialize IPR scan for all genomes"""
    threads = 4
    with mp.Pool(processes=round(cpus / 2)) as pool:
        ipr_res = pool.starmap(
            call_iprscan,
            tqdm(
                ((ome, iprscan, ann_dir, threads) for ome in list(passing_omes)),
                total=len(passing_omes),
            ),
        )

    failed_omes = []
    for ome, exit_code in ipr_res:
        if exit_code:
            logger.error("\t\t" + ome + " failed InterProScan")
            failed_omes.append(ome)

    return failed_omes


def parse_hmm_res(ome, ann_dir, evalue, threshold, max_hits=None):
    ome_out = ann_dir + ome + "_pfam.out"
    #    lineComp = re.compile(r'([^ ]+)')
    res = defaultdict(list)
    with open(ome_out, "r") as raw:
        for line in raw:
            if not line.startswith("#"):
                #                data = lineComp.findall(line.rstrip())
                data = line.rstrip().split()
                hit_perc = (int(data[18]) - int(float(data[17]))) / int(data[5])
                if hit_perc > threshold:
                    if float(data[6]) < evalue:  # and float(data[7]) > score:
                        # pfam, name, cov_perc, evalue, bit, algn_start, algn_end
                        res[data[0]].append(
                            (
                                data[4],
                                data[3],
                                hit_perc,
                                float(data[21]),
                                int(data[17]),
                                int(data[18]),
                            )
                        )
    ome_res = {}
    for gene, gene_hits in res.items():
        gene_hits = sorted(gene_hits, key=lambda x: x[4], reverse=True)
        if max_hits:
            gene_hits = gene_hits[:max_hits]
        todel, hits = [], set()
        for i, hit in enumerate(gene_hits):
            if hit[0] in hits:
                todel.append(i)
            hits.add(hit)
        for i in reversed(todel):
            del gene_hits[i]
        ome_res[gene] = gene_hits

    return ome, ome_res


def compile_hmm_res(
    ann_dir, passing_omes, evalue=0.001, threshold=0.5, cpus=1, max_hits=None
):

    with mp.Pool(processes=cpus) as pool:
        omebyome_res = pool.starmap(
            parse_hmm_res,
            tqdm(
                ((ome, ann_dir, evalue, threshold) for ome in list(passing_omes)),
                total=len(passing_omes),
            ),
        )
    ome_res = {}
    for ome, res in omebyome_res:
        ome_res[ome] = res

    return ome_res


def parse_ipr_res(ome, ann_dir, evalue, threshold):
    # NOT BUILT YET
    ome_out = ann_dir + ome + "_ipr.out"
    lineComp, res = re.compile(r"([^ ]+)"), defaultdict(list)
    with open(ome_out, "r") as raw:
        for line in raw:
            if not line.startswith("#"):
                data = lineComp.findall(line.rstrip())
                hit_perc = (int(data[15]) - int(float(data[14]))) / int(data[5])
                if hit_perc > threshold:
                    if float(data[6]) < evalue:  # and float(data[7]) > score:
                        # pfam, name, cov_perc, evalue, bit, algn_start, algn_end
                        res[data[0]].append(
                            (
                                data[4],
                                data[3],
                                hit_perc,
                                float(data[6]),
                                float(data[7]),
                                int(data[16]),
                                int(data[17]),
                            )
                        )
    ome_res = {}
    for gene, gene_hits in res.items():
        gene_hits = sorted(gene_hits, key=lambda x: x[3], reverse=True)
        todel, hits = [], set()
        for i, hit in enumerate(gene_hits):
            if hit[0] in hits:
                todel.append(i)
            hits.add(hit)
        for i in reversed(todel):
            del gene_hits[i]
        ome_res[gene] = gene_hits

    return ome, ome_res


def compile_ipr_res(ann_dir, passing_omes, evalue=0.01, threshold=0.5, cpus=1):

    with mp.Pool(processes=cpus) as pool:
        omebyome_res = pool.starmap(
            parse_ipr_res,
            tqdm(
                ((ome, ann_dir, evalue, threshold) for ome in list(passing_omes)),
                total=len(passing_omes),
            ),
        )
    ome_res = {}
    for ome, res in omebyome_res:
        ome_res[ome] = res

    return ome_res


def ann_mngr(
    genes_list,
    prot_paths,
    wrk_dir,
    pfam,
    ipr_scan,
    evalue=0.0001,
    threshold=0.5,
    cpus=1,
):

    ann_res = None
    ann_dir = wrk_dir + "ann/"
    if not os.path.isdir(ann_dir):
        os.mkdir(ann_dir)

    failed_omes = output_res_faas(genes_list, prot_paths, ann_dir, cpus=cpus)
    passing_omes = set(genes_list.keys()).difference(failed_omes)

    if pfam:
        logger.info("\tHmmsearch'ing Pfam database")
        hmm_failed_omes = run_hmmsearch(pfam, passing_omes, ann_dir, cpus)
        passing_omes = passing_omes.difference(set(hmm_failed_omes))
        failed_omes = failed_omes.union(set(hmm_failed_omes))
        logger.info("\tParsing hmmsearch output")
        ann_res = compile_hmm_res(ann_dir, passing_omes, evalue, threshold, cpus)
    elif ipr_scan:
        logger.info("\tRunning InterProScan")
        ipr_failed_omes = run_ipr_scan(ipr_scan, passing_omes, ann_dir, cpus)
        passing_omes = passing_omes.difference(set(ipr_failed_omes))
        failed_omes = failed_omes.union(set(ipr_failed_omes))
        logger.info("\tParsing InterProScan output")
        pass

    return ann_res, failed_omes


def annotate_clusters(
    genes_list, ome, ome_dir, gene2hg, prefix="hlg", ann_res={}  # gff_path, prot_path,
):

    if not os.path.isfile(f"{ome_dir}{ome}/{prefix}.tsv"):
        return

    with open(ome_dir + ome + f"/{prefix}.tsv.tmp", "w") as out:
        out.write("#name\thgs\tgenes\tgcfs\tpfams\n")
        with open(ome_dir + ome + f"/{prefix}.tsv", "r") as raw:
            for line in raw:
                if not line.startswith("#"):
                    clus_d = line.split()
                    genes = clus_d[2].split(",")
                    anns_p = [ann_res[x] if x in ann_res else None for x in genes]
                    anns = [
                        "|".join([y[1] if y else "" for y in x]) if x else ""
                        for x in anns_p
                    ]
                    ann_str = ";".join(anns)
                    out.write("\t".join(clus_d + [ann_str]) + "\n")
    shutil.move(f"{ome_dir}{ome}/{prefix}.tsv.tmp", f"{ome_dir}{ome}/{prefix}.tsv")


#    return ome#, svg_dict
#    return ome, clus_gffs, clus_fas, svg_dict


def output_hgxs(hgx2dist, hgx2omes, hgx2i, i2ome, out_dir):
    maxhgxd = max(hgx2dist.values())
    minhgxd = min(hgx2dist.values())

    denom = maxhgxd - minhgxd
    hgx_output = []
    for hgx, omes in hgx2omes.items():
        hgx_id = hgx2i[hgx]
        hgx_output.append(
            [
                ",".join([str(x) for x in hgx]),
                hgx_id,
                (log(hgx2dist[hgx]) - minhgxd) / denom,
                ",".join([i2ome[x] for x in omes]),
            ]
        )  # HGxs at this stage are not segregated into groups
    hgx_output = sorted(hgx_output, key=lambda x: x[3], reverse=True)
    with gzip.open(out_dir + "hgxs.tsv.gz", "wt") as out:
        out.write("#hgs\thgx_id\tnrm_log_tmd\tomes")  # \tpdd\tomes')
        for entry in hgx_output:
            out.write("\n" + "\t".join([str(x) for x in entry]))


def write_hlgs_txt_wpos(
    hlg_hgxs,
    hlg_omes,
    logg2d,
    hlg2clan,
    omes2patch,
    hlg2apds,
    hgx2omes2gcl,
    hgx2omes2id,
    hgx2omes2pos,
    i2ome,
    out_dir,
    group_name,
):
    hlg_output = []
    for i, omesc in hlg_omes.items():
        hlg_hgx = hlg_hgxs[i]
        if i in hlg2apds:
            apds = hlg2apds[i]
        else:
            apds = "na"
        try:
            hlg_output.append(
                [
                    ",".join([str(x) for x in hlg_hgx]),
                    i,
                    hlg2clan[i],
                    logg2d[omesc],
                    omes2patch[omesc],
                    hgx2omes2gcl[hlg_hgx][omesc],
                    hgx2omes2id[hlg_hgx][omesc],
                    hgx2omes2pos[hlg_hgx][omesc],
                    100
                    * (
                        1 - (hgx2omes2id[hlg_hgx][omesc] / hgx2omes2pos[hlg_hgx][omesc])
                    ),
                    apds,
                    ",".join([str(i2ome[x]) for x in omesc]),  # ,
                    #        dnds_dict[hgx][0], dnds_dict[hgx][1], str(dnds_dict[hgx][2]),
                ]
            )
        except ZeroDivisionError:
            hlg_output.append(
                [
                    ",".join([str(x) for x in hlg_hgx]),
                    i,
                    hlg2clan[i],
                    logg2d[omesc],
                    omes2patch[omesc],
                    hgx2omes2gcl[hlg_hgx][omesc],
                    hgx2omes2id[hlg_hgx][omesc],
                    hgx2omes2pos[hlg_hgx][omesc],
                    0,
                    apds,
                    ",".join([str(i2ome[x]) for x in omesc]),  # ,
                    #        dnds_dict[hgx][0], dnds_dict[hgx][1], str(dnds_dict[hgx][2]),
                ]
            )
        except KeyError:
            try:
                hlg_output.append(
                    [
                        ",".join([str(x) for x in hlg_hgx]),
                        i,
                        hlg2clan[i],
                        logg2d[omesc],
                        omes2patch[omesc],
                        hgx2omes2gcl[hlg_hgx][omesc],
                        hgx2omes2id[hlg_hgx][omesc],
                        "na",
                        "na",
                        apds,
                        ",".join([str(i2ome[x]) for x in omesc]),  # ,
                        #        dnds_dict[hgx][0], dnds_dict[hgx][1], str(dnds_dict[hgx][2]),
                    ]
                )
            except KeyError:
                logger.error(f"missing HGx/omes from results: {hlg_hgx}, {omesc}")

    hlg_output = sorted(hlg_output, key=lambda x: x[3], reverse=True)
    with gzip.open(out_dir + f"{group_name}s.tsv.gz", "wt") as out:
        out.write(
            f"#hgs\t{group_name}\thlc\tnrm_log_tmd\tpds"
            + "\tgcl\tmmi\tmmp\tcsb\tapds\tomes"
        )  # \tmmp\tomes') #+ \
        #'selection_coef\tmean_dnds\tog_dnds\t' + \
        #   'total_dist'
        #            )
        for entry in hlg_output:
            out.write("\n" + "\t".join([str(x) for x in entry]))
    return hlg_output


def write_hlgs_txt_wopos(
    hlg_hgxs,
    hlg_omes,
    logg2d,
    hlg2clan,
    omes2patch,
    hgx2omes2gcl,
    hgx2omes2id,
    i2ome,
    out_dir,
    group_name,
):
    hlg_output = []
    for i, hlg_hgx in hlg_hgxs.items():
        omesc = hlg_omes[i]
        hlg_output.append(
            [
                ",".join([str(x) for x in hlg_hgx]),
                i,
                hlg2clan[i],
                logg2d[omesc],
                omes2patch[omesc],
                hgx2omes2gcl[hlg_hgx][omesc],
                hgx2omes2id[hlg_hgx][omesc],
                ",".join([str(i2ome[x]) for x in omesc]),  # ,
                #        dnds_dict[hgx][0], dnds_dict[hgx][1], str(dnds_dict[hgx][2]),
            ]
        )

    hlg_output = sorted(hlg_output, key=lambda x: x[3], reverse=True)
    with gzip.open(out_dir + f"{group_name}s.tsv.gz", "wt") as out:
        out.write(
            f"#hgs\t{group_name}\thlc\tnrm_log_tmd\tpds" + "\tgcl\tmmi\tomes"
        )  # \tmmp\tomes') #+ \
        #'selection_coef\tmean_dnds\tog_dnds\t' + \
        #   'total_dist'
        #            )
        for entry in hlg_output:
            out.write("\n" + "\t".join([str(x) for x in entry]))
    return hlg_output


def calc_logg2d(hlg_omes, omes2dist):
    logg2d_prep = {}
    for omesc in hlg_omes.values():
        logg2d_prep[omesc] = log(omes2dist[omesc])
    maxhlgd = max(logg2d_prep.values())  # max observed, even of those truncated/removed
    minhlgd = min(logg2d_prep.values())
    denom = maxhlgd - minhlgd
    logg2d = {k: (v - minhlgd) / denom for k, v in logg2d_prep.items()}
    return logg2d


def threshold_hlg_by_dist(logg2d, dist_thresh, hlgs, hlg_hgxs, hlg_omes, omes2dist):
    if dist_thresh:
        gcfs, gcf_omes, gcf_hgxs = {}, {}, {}
        for hlg, locs in hlgs.items():
            hgxc = hlg_hgxs[hlg]
            omesc = hlg_omes[hlg]
            if logg2d[omesc] >= dist_thresh:
                gcfs[hlg] = locs
                gcf_omes[hlg] = omesc
                gcf_hgxs[hlg] = hgxc
        return gcfs, gcf_omes, gcf_hgxs
    else:
        return hlgs, hlg_omes, hlg_hgxs


def threshold_hlg_wpos(
    gcl_thresh,
    patch_thresh,
    id_perc,
    pos_perc,
    csb_thresh,
    hlgs,
    hlg_hgxs,
    hlg_omes,
    omes2patch,
    hgx2omes2gcl,
    hgx2omes2id,
    hgx2omes2pos,
):

    if any(x > 0 for x in [gcl_thresh, patch_thresh, id_perc, pos_perc, csb_thresh]):
        gcfs, gcf_omes, gcf_hgxs = {}, {}, {}
        for hlg, locs in hlgs.items():
            check = False  # have we added a new list
            hgxc = hlg_hgxs[hlg]
            omesc = hlg_omes[hlg]
            try:
                patch = omes2patch[omesc]
                gcl = hgx2omes2gcl[hgxc][omesc]
                id_, pos = hgx2omes2id[hgxc][omesc], hgx2omes2pos[hgxc][omesc]
                if pos != 0 and id_ != 0:
                    csb = 1 - (id_ / pos)
                else:
                    csb = 0
                if (
                    gcl >= gcl_thresh
                    and patch >= patch_thresh
                    and id_ >= id_perc
                    and pos >= pos_perc
                    and csb >= csb_thresh
                ):
                    gcfs[hlg] = locs
                    gcf_omes[hlg] = omesc
                    gcf_hgxs[hlg] = hgxc
            except TypeError:
                if (
                    hgx2omes2gcl[hgxc][omesc] >= gcl_thresh
                    and omes2patch[omesc] >= patch_thresh
                    and hgx2omes2id[hgxc][omesc] >= id_perc
                ):
                    gcfs[hlg] = locs
                    gcf_omes[hlg] = omesc
                    gcf_hgxs[hlg] = hgxc
            except KeyError:
                logger.warning(f"\tmissing from proxies and removed: {omesc, hgxc}")
        return gcfs, gcf_omes, gcf_hgxs
    else:
        return hlgs, hlg_omes, hlg_hgxs


def threshold_hlg_wopos(
    gcl_thresh,
    patch_thresh,
    id_perc,
    hlgs,
    hlg_hgxs,
    hlg_omes,
    omes2patch,
    hgx2omes2gcl,
    hgx2omes2id,
):
    if any(x > 0 for x in [gcl_thresh, patch_thresh, id_perc]):
        logger.info("\tApplying thresholds")
        logger.info("\t\t" + str(len(hlgs)) + " HLGs before")
        gcfs, gcf_omes, gcf_hgxs = {}, {}, {}
        for hlg, locs in hlgs.items():
            hgxc = hlg_hgxs[hlg]
            omesc = hlg_omes[hlg]
            if (
                hgx2omes2gcl[hgxc][omesc] >= gcl_thresh
                and omes2patch[omesc] >= patch_thresh
                and hgx2omes2id[hgxc][omesc] >= id_perc
            ):
                gcfs[hlg] = locs
                gcf_omes[hlg] = omesc
                gcf_hgxs[hlg] = hgxc

        return gcfs, gcf_omes, gcf_hgxs
    else:
        return hlgs, hlg_omes, hlg_hgxs


def output_thresholds(
    logg2d,
    dist_thresh,
    gcl_thresh,
    patch_thresh,
    id_perc,
    pos_perc,
    csb_thresh,
    hlgs,
    hlg_hgxs,
    hlg_omes,
    omes2dist,
    omes2patch,
    hlg2clan,
    hgx2omes2gcl,
    hgx2omes2id,
    hgx2omes2pos,
    hlg2apds,
    out_dir,
    i2ome,
):
    logger.info(f"\tFiltering HLGs for GCFs")
    logger.info("\tApplying thresholds")
    logger.info("\t\t" + str(len(hlgs)) + " HLGs before")
    hlgs, hlg_omes, hlg_hgxs = threshold_hlg_by_dist(
        logg2d, dist_thresh, hlgs, hlg_hgxs, hlg_omes, omes2dist
    )
    if hgx2omes2pos:
        gcfs, gcf_omes, gcf_hgxs = threshold_hlg_wpos(
            gcl_thresh,
            patch_thresh,
            id_perc,
            pos_perc,
            csb_thresh,
            hlgs,
            hlg_hgxs,
            hlg_omes,
            omes2patch,
            hgx2omes2gcl,
            hgx2omes2id,
            hgx2omes2pos,
        )
        gcf_output = write_hlgs_txt_wpos(
            gcf_hgxs,
            gcf_omes,
            logg2d,
            hlg2clan,
            omes2patch,
            hlg2apds,
            hgx2omes2gcl,
            hgx2omes2id,
            hgx2omes2pos,
            i2ome,
            out_dir,
            "gcf",
        )

    else:
        gcfs, gcf_omes, gcf_hgxs = threshold_gcf_wopos(
            gcl_thresh,
            patch_thres,
            id_perc,
            gcfs,
            gcf_hgxs,
            gcf_omes,
            omes2patch,
            hgx2omes2gcl,
            hgx2omes2id,
        )
        gcf_output = write_hlgs_txt_wopos(
            gcf_hgxs,
            gcf_omes,
            logg2d,
            hlg2clan,
            omes2patch,
            hgx2omes2gcl,
            hgx2omes2id,
            i2ome,
            out_dir,
            "gcf",
        )

    logger.info(f"\t\t{len(gcfs)} GCFs pass")

    return gcf_output, gcfs, gcf_omes, gcf_hgxs


def output_figures(hlg_output, prefix, out_dir):
    # legacy, need to remove unless dn/ds is reimplemented
    #    if dnds_dict:
    #       axes = [[],[],[],[],[], []]
    #  else:
    axes = [[], [], [], [], [], []]
    labels = []

    try:
        for entry in hlg_output:
            labels.append(["HLG:", str(entry[1]) + " | Omes: " + entry[-1]])
            axes[0].append(entry[5])
            axes[1].append(entry[4])
            axes[2].append(entry[3])
            axes[3].append(entry[6])
            axes[4].append(entry[7])
            axes[5].append(entry[8])
    except ValueError:  # no positives
        labels = []
        axes = [[], [], [], []]
        for entry in hlg_output:
            labels.append(["HLG:", str(entry[1]) + " | Omes: " + entry[-1]])
            axes[0].append(entry[5])
            axes[1].append(entry[4])
            axes[2].append(entry[3])
            axes[3].append(entry[6])

    logger.info(f"\tOutputting {prefix.upper()} plots")
    if len(axes) == 4:
        axes_labels = ["PDS", "GCL", "nl_TMD", "MMI"]
    else:
        axes_labels = ["PDS", "GCL", "nl_TMD", "MMI", "MMP", "CSB"]

    fig = mk_subplots(labels, axes, axes_labels, alpha=0.6)
    fig.write_html(out_dir + f"{prefix}_proxies.html")
    fig = mk_3d(labels, axes[:3], axes_labels[:3], alpha=0.7)
    fig.write_html(out_dir + f"{prefix}s.html")


def quick_thresh(ome, ome_dir, hlgs):
    hlg_f = ome_dir + ome + "/hlg.tsv"
    gcf_f = ome_dir + ome + "/gcf.tsv"
    with open(hlg_f, "r") as raw, open(gcf_f, "w") as out:
        for line in raw:
            if line.startswith("#"):
                out.write(line)
            else:
                hlg = int(line.rstrip().split()[3])
                if hlg in hlgs:
                    out.write(line)


def threshold_gcf_quick(
    db,
    out_dir,
    ome_dir,
    dist_thresh,
    gcl_thresh,
    patch_thresh,
    id_perc,
    pos_perc,
    csb_thresh,
    cpus=1,
):

    ome2hlgs = defaultdict(set)
    logger.info("\tLoading HLGs file")
    with gzip.open(out_dir + "hlgs.tsv.gz", "rt") as raw, gzip.open(
        out_dir + "gcfs.tsv.gz", "wt"
    ) as out:
        for line in raw:
            if line.startswith("#"):
                out.write(line)
            else:
                d = line.rstrip().split()
                hlg = int(d[1])
                tmd = float(d[3])
                pds = float(d[4])
                gcl = float(d[5])
                mmi = float(d[6])
                mmp = float(d[7])
                csb = float(d[8]) / 100
                omes = d[9].split(",")
                if (
                    tmd >= dist_thresh
                    and gcl >= gcl_thresh
                    and pds >= patch_thresh
                    and mmi >= id_perc
                    and mmp >= pos_perc
                    and csb >= csb_thresh
                ):
                    for ome in omes:
                        ome2hlgs[ome].add(hlg)
                    out.write(line)

    logger.info("\tExtracting GCFs")
    with mp.Pool(processes=cpus) as pool:
        pool.starmap(
            quick_thresh, ((ome, ome_dir, hlgs) for ome, hlgs in ome2hlgs.items())
        )


def threshold_gcf_bypass(
    db,
    out_dir,
    wrk_dir,
    i2ome,
    gene2hg,
    dist_thresh,
    gcl_thresh,
    patch_thresh,
    id_perc,
    pos_perc,
    csb_thresh,
    ipr_path,
    pfam_path,
    cpus=1,
):

    logger.info("\tLoading data structures")
    with open(wrk_dir + "hlgs.pickle", "rb") as raw:
        hlgs, hlg_omes, hlg_hgxs, hlg2clan = pickle.load(raw)

    with open(wrk_dir + "pds.pickle", "rb") as in_pick:
        omes2patch = pickle.load(in_pick)

    with open(wrk_dir + "gcl.pickle", "rb") as pickin:
        hgx2omes2gcl = pickle.load(pickin)
    with open(wrk_dir + "mmi.pickle", "rb") as pickin:
        hgx2omes2id = pickle.load(pickin)
    if os.path.isfile(wrk_dir + "mmp.pickle"):
        with open(wrk_dir + "mmp.pickle", "rb") as pickin:
            hgx2omes2pos = pickle.load(pickin)
    else:
        hgx2omes2pos = None
    if os.path.isfile(wrk_dir + "apds.pickle"):
        with open(wrk_dir + "apds.pickle", "rb") as pickin:
            hlg2apds = pickle.load(pickin)
    else:
        hlg2apds = {}

    with open(wrk_dir + "omes2dist.pickle", "rb") as raw:
        omes2dist = pickle.load(raw)

    logg2d = calc_logg2d(hlg_omes, omes2dist)
    gcf_output, gcfs, gcf_omes, gcf_hgxs = output_thresholds(
        logg2d,
        dist_thresh,
        gcl_thresh,
        patch_thresh,
        id_perc,
        pos_perc,
        csb_thresh,
        hlgs,
        hlg_hgxs,
        hlg_omes,
        omes2dist,
        omes2patch,
        hlg2clan,
        hgx2omes2gcl,
        hgx2omes2id,
        hgx2omes2pos,
        hlg2apds,
        out_dir,
        i2ome,
    )
    output_figures(gcf_output, "gcf", out_dir)

    with mp.Pool(processes=cpus) as pool:
        hlg_gene_res = pool.starmap(
            read_ome_hlgs, ((ome, f"{out_dir}ome/hlg.tsv") for ome in i2ome)
        )
    hlg_genes = {}
    for ome, ome_hlg_gene in hlg_gene_res:
        if ome_hlg_gene:
            hlg_genes[ome] = ome_hlg_gene

    gcf_genes = write_ome_output("gcf", gcfs, out_dir, db, cpus, gene2hg)
    annotation_mngr(
        hlg_genes,
        db,
        wrk_dir,
        pfam_path,
        ipr_path,
        gene2hg,
        out_dir,
        prefix="hlg",
        cpus=cpus,
    )
    annotation_mngr(
        gcf_genes,
        db,
        wrk_dir,
        pfam_path,
        ipr_path,
        gene2hg,
        out_dir,
        prefix="gcf",
        cpus=cpus,
    )


def write_ome_output(prefix, hlgs, out_dir, db, cpus, gene2hg):
    logger.info(f"\tWriting {prefix.upper()} files")
    write_clus_cmds = []
    ome2clusters = defaultdict(list)
    for hlg, loci in hlgs.items():
        for loc in loci:
            ome = loc[0][: loc[0].find("_")]
            ome2clusters[ome].append([loc, hlg])

    ome_dir = out_dir + "ome/"
    if not os.path.isdir(ome_dir):
        os.mkdir(ome_dir)
    for ome, clusters in ome2clusters.items():
        if not os.path.isdir(ome_dir + ome):
            os.mkdir(ome_dir + ome)
        gff = db[ome]["gff3"]
        out_file = ome_dir + ome + f"/{prefix}.tsv"
        write_clus_cmds.append([clusters, ome, out_file, gff, gene2hg])
    with mp.Pool(processes=cpus) as pool:
        out_genes = pool.starmap(
            write_clusters, tqdm(write_clus_cmds, total=len(write_clus_cmds))
        )
        pool.close()
        pool.join()
    genes = {x[0]: x[1] for x in out_genes if x[1]}
    return genes


def annotation_mngr(
    hlg_genes, db, wrk_dir, pfam_path, ipr_path, gene2hg, out_dir, prefix="hlg", cpus=1
):
    logger.info("\tSearching")
    prot_paths = {}
    for ome in hlg_genes:
        prot_paths[ome] = db[ome]["faa"]

    ann_res, failedOmes = ann_mngr(
        hlg_genes,
        prot_paths,
        wrk_dir,
        pfam_path,
        ipr_path,
        evalue=0.0001,
        threshold=0.5,
        cpus=cpus,
    )
    logger.info("\tAnnotating")
    annotate_clusters_cmds = []
    for ome, genes in hlg_genes.items():
        if ome in failedOmes:
            continue
        gff_path = db[ome]["gff3"]
        pro_path = db[ome]["faa"]
        try:
            annotate_clusters_cmds.append(
                [
                    genes,  # gff_path, pro_path,
                    ome,
                    out_dir + "ome/",
                    gene2hg,
                    prefix,
                    ann_res[ome],
                ]
            )
        except KeyError:
            pass
    #             annotate_clusters_cmds.append([genes, #gff_path, pro_path,
    #                                  ome, out_dir + 'ome/', gene2hg,
    #                                prefix])

    with mp.Pool(processes=cpus) as pool:
        pool.starmap(
            annotate_clusters,
            tqdm(annotate_clusters_cmds, total=len(annotate_clusters_cmds)),
        )
        pool.close()
        pool.join()


def output_hlgs(
    db,
    wrk_dir,
    hlgs,
    hlg_omes,
    i2ome,
    out_dir,
    hlg_hgxs,
    omes2dist,
    omes2patch,
    hgx2omes2gcl,
    hgx2omes2id,
    hgx2omes2pos,
    gene2hg,
    plusminus,
    ome2i,
    hlg2clan,
    hlg2apds,
    dist_thresh,
    gcl_thresh,
    patch_thresh,
    id_perc,
    pos_perc,
    csb_thresh,
    ipr_path=None,
    pfam_path=None,
    dnds_dict={},
    cpus=1,
):

    logger.info("\tWriting cluster scores")

    logg2d = calc_logg2d(hlg_omes, omes2dist)
    if hgx2omes2pos:
        hlg_output = write_hlgs_txt_wpos(
            hlg_hgxs,
            hlg_omes,
            logg2d,
            hlg2clan,
            omes2patch,
            hlg2apds,
            hgx2omes2gcl,
            hgx2omes2id,
            hgx2omes2pos,
            i2ome,
            out_dir,
            "hlg",
        )
    else:
        hlg_output = write_hlgs_txt_wopos(
            hlg_hgxs,
            hlg_omes,
            logg2d,
            hlg2clan,
            omes2patch,
            hgx2omes2gcl,
            hgx2omes2id,
            i2ome,
            out_dir,
            "hlg",
        )

    if dist_thresh or gcl_thresh or patch_thresh or id_perc or pos_perc:
        gcf_output, gcfs, gcf_omes, gcf_hgxs = output_thresholds(
            logg2d,
            dist_thresh,
            gcl_thresh,
            patch_thresh,
            id_perc,
            pos_perc,
            csb_thresh,
            hlgs,
            hlg_hgxs,
            hlg_omes,
            omes2dist,
            omes2patch,
            hlg2clan,
            hgx2omes2gcl,
            hgx2omes2id,
            hgx2omes2pos,
            hlg2apds,
            out_dir,
            i2ome,
        )
    else:
        gcf_output, gcfs, gcf_omes, gcf_hgxs = None, hlgs, hlg_omes, hlg_hgxs

    output_figures(hlg_output, "hlg", out_dir)
    if gcf_output:
        output_figures(gcf_output, "gcf", out_dir)

    prefix = "hlg"
    hlg_genes = write_ome_output("hlg", hlgs, out_dir, db, cpus, gene2hg)
    if gcf_output:
        #        prefix = 'gcf'
        gcf_genes = write_ome_output("gcf", gcfs, out_dir, db, cpus, gene2hg)

    if pfam_path or ipr_path:
        logger.info("X. Annotating clusters")

    annotation_mngr(
        hlg_genes,
        db,
        wrk_dir,
        pfam_path,
        ipr_path,
        gene2hg,
        out_dir,
        prefix="hlg",
        cpus=cpus,
    )
    if dist_thresh or gcl_thresh or patch_thresh or id_perc or pos_perc:
        annotation_mngr(
            hlg_genes,
            db,
            wrk_dir,
            pfam_path,
            ipr_path,
            gene2hg,
            out_dir,
            prefix="gcf",
            cpus=cpus,
        )
