"""Merge two ActivePathways analyses (e.g. directional vs non-directional) into one set of files."""
import numpy as np
import pandas as pd

from .gmt import GMT, read_gmt, write_gmt


def merge_results(
    enriched_pathways,
    enriched_pathways_directional,
    gmt_file,
    output_prefix="Aggregated",
    col_colors=None,
    tests=None,
):
    """Combine two ActivePathways result tables for joint Cytoscape visualisation.

    Writes <output_prefix>_pathways.txt, _subgroups.txt, and _pathways.gmt.
    Returns the aggregated pathways table.
    """
    gmt = read_gmt(gmt_file)
    if tests is None:
        gene_cols = [c for c in enriched_pathways.columns if c.startswith("Genes_")]
        tests = [c[len("Genes_"):] for c in gene_cols] + ["combined"]
    if not tests:
        raise ValueError("Tests parameter must be provided (e.g., ['rna', 'protein', 'combined']) or NULL")

    tests_to_check = [t for t in tests if t != "combined"]
    expected_cols = {f"Genes_{t}" for t in tests_to_check}
    for df in (enriched_pathways, enriched_pathways_directional):
        if not expected_cols.issubset(df.columns):
            raise ValueError("All tests (except 'combined') must exist as columns with 'Genes_' prefix in both enriched_pathways and enriched_pathways_directional")

    if col_colors is None:
        from .cytoscape import _rainbow
        col_colors = _rainbow(len(tests))
    elif len(col_colors) != len(tests):
        raise ValueError("col_colors must have the same length as tests")

    # Concatenate the two result tables and aggregate by term taking the minimum adjusted p-value.
    combined = pd.concat([
        enriched_pathways[["term_id", "term_name", "adjusted_p_val", "evidence"]],
        enriched_pathways_directional[["term_id", "term_name", "adjusted_p_val", "evidence"]],
    ], ignore_index=True)

    pathways_txt = (
        combined.groupby(["term_id", "term_name"], as_index=False, sort=False)["adjusted_p_val"]
        .min()
    )
    pathways_txt.to_csv(f"{output_prefix}_pathways.txt", sep="\t", index=False)

    # Build the binary evidence matrix: which tests supported each unique term.
    term_ids = list(pathways_txt["term_id"])
    indicator = np.zeros((len(term_ids), len(tests)), dtype=int)
    for term_idx, term_id in enumerate(term_ids):
        ev_lists = combined.loc[combined["term_id"] == term_id, "evidence"]
        evidence_values = set()
        for ev in ev_lists:
            if isinstance(ev, str):
                evidence_values.update([e.strip() for e in ev.split(",")])
            elif isinstance(ev, (list, tuple, np.ndarray, pd.Series)):
                evidence_values.update(ev)
        for j, t in enumerate(tests):
            if t in evidence_values:
                indicator[term_idx, j] = 1

    col_significance = pd.DataFrame(indicator, columns=list(tests))
    col_significance.insert(0, "term_id", term_ids)

    # Annotate which terms were lost / gained / shared between methods.
    set_orig = set(enriched_pathways["term_id"])
    set_dir = set(enriched_pathways_directional["term_id"])
    impact = []
    for tid in term_ids:
        if tid in set_orig and tid not in set_dir:
            impact.append(1)  # lost when adding directionality
        elif tid in set_dir and tid not in set_orig:
            impact.append(2)  # gained
        else:
            impact.append(0)  # shared
    col_significance["directional_impact"] = impact

    instruct_str = (
        f'piechart: attributelist="{",".join(tests)}"'
        f' colorlist="{",".join(col_colors)}" showlabels=FALSE'
    )
    col_significance["instruct"] = instruct_str
    col_significance.to_csv(f"{output_prefix}_subgroups.txt", sep="\t", index=False)

    # Subset the source GMT to the union of pathways found by either run.
    filtered_gmt = GMT([gmt[tid] for tid in term_ids if tid in gmt])
    write_gmt(filtered_gmt, f"{output_prefix}_pathways.gmt")

    return pathways_txt


def merge_gmt(gmt_file, term_ids):
    """Read gmt_file and return a GMT containing only the specified term_ids."""
    gmt = read_gmt(gmt_file)
    return GMT([gmt[tid] for tid in term_ids if tid in gmt])
