"""Generate the four output files used by Cytoscape's EnrichmentMap app."""
import numpy as np
import pandas as pd

from .gmt import GMT, write_gmt

# Default qualitative palettes (similar to RColorBrewer); ``None`` means fall back to a rainbow.
_BREWER_PALETTES = {
    "Set1": ["#E41A1C", "#377EB8", "#4DAF4A", "#984EA3", "#FF7F00", "#FFFF33", "#A65628", "#F781BF", "#999999"],
    "Set2": ["#66C2A5", "#FC8D62", "#8DA0CB", "#E78AC3", "#A6D854", "#FFD92F", "#E5C494", "#B3B3B3"],
    "Set3": ["#8DD3C7", "#FFFFB3", "#BEBADA", "#FB8072", "#80B1D3", "#FDB462", "#B3DE69", "#FCCDE5", "#D9D9D9", "#BC80BD", "#CCEBC5", "#FFED6F"],
    "Pastel1": ["#FBB4AE", "#B3CDE3", "#CCEBC5", "#DECBE4", "#FED9A6", "#FFFFCC", "#E5D8BD", "#FDDAEC", "#F2F2F2"],
    "Pastel2": ["#B3E2CD", "#FDCDAC", "#CBD5E8", "#F4CAE4", "#E6F5C9", "#FFF2AE", "#F1E2CC", "#CCCCCC"],
    "Dark2": ["#1B9E77", "#D95F02", "#7570B3", "#E7298A", "#66A61E", "#E6AB02", "#A6761D", "#666666"],
    "Accent": ["#7FC97F", "#BEAED4", "#FDC086", "#FFFF99", "#386CB0", "#F0027F", "#BF5B17", "#666666"],
    "Paired": ["#A6CEE3", "#1F78B4", "#B2DF8A", "#33A02C", "#FB9A99", "#E31A1C", "#FDBF6F", "#FF7F00", "#CAB2D6", "#6A3D9A", "#FFFF99", "#B15928"],
}


def _rainbow(n):
    """Generate ``n`` evenly-spaced HSV colours (R's ``grDevices::rainbow`` analogue)."""
    import colorsys
    return ["#{:02X}{:02X}{:02X}".format(*[int(255 * c) for c in colorsys.hsv_to_rgb(i / max(n, 1), 1.0, 1.0)])
            for i in range(n)]


def _resolve_colors(tests, color_palette, custom_colors, color_integrated_only):
    """Determine the per-test colour list, mirroring the R logic order of precedence."""
    if color_palette is None and custom_colors is None:
        col_colors = _rainbow(len(tests))
    elif custom_colors is not None:
        if isinstance(custom_colors, dict):
            # Reorder named colours to match the test order.
            col_colors = [custom_colors[name] for name in tests if name in custom_colors]
            # Insert the integrated colour at the position of "combined".
            if "combined" in tests:
                col_colors.insert(tests.index("combined"), color_integrated_only)
        else:
            col_colors = list(custom_colors)
            if "combined" in tests:
                col_colors.insert(tests.index("combined"), color_integrated_only)
    else:
        if color_palette not in _BREWER_PALETTES:
            raise ValueError("palette must be from the RColorBrewer package")
        palette = _BREWER_PALETTES[color_palette]
        col_colors = palette[: len(tests)]
    # The "combined" slot must always show ``color_integrated_only``.
    if "combined" in tests:
        col_colors[tests.index("combined")] = color_integrated_only
    return col_colors


def prepare_cytoscape(
    terms,
    gmt,
    cytoscape_file_tag,
    col_significance,
    color_palette=None,
    custom_colors=None,
    color_integrated_only="#FFFFF0",
):
    """Write pathways.txt, subgroups.txt, pathways.gmt, and legend.pdf.

    terms must contain term_id, term_name, and adjusted_p_val columns.
    """
    if col_significance is None:
        terms.to_csv(f"{cytoscape_file_tag}pathways.txt", sep="\t", index=False)
        write_gmt(gmt, f"{cytoscape_file_tag}pathways.gmt")
        return

    # The first two columns of col_significance are 'term_id' and 'evidence'; the rest are Genes_<col>.
    gene_cols = [c for c in col_significance.columns if c.startswith("Genes_")]
    tests = [c[len("Genes_"):] for c in gene_cols] + ["combined"]

    # Build a 0/1 indicator matrix marking which test(s) supported each pathway.
    indicator = np.zeros((len(col_significance), len(tests)), dtype=int)
    for row_idx, ev in enumerate(col_significance["evidence"]):
        if isinstance(ev, str):
            ev = [ev]
        for t in ev:
            if t in tests:
                indicator[row_idx, tests.index(t)] = 1

    out = pd.DataFrame(indicator, columns=tests)
    out.insert(0, "term_id", col_significance["term_id"].to_numpy())

    col_colors = _resolve_colors(tests, color_palette, custom_colors, color_integrated_only)
    instruct_str = (
        f'piechart: attributelist="{",".join(tests)}"'
        f' colorlist="{",".join(col_colors)}" showlabels=FALSE'
    )
    out["instruct"] = instruct_str

    terms.to_csv(f"{cytoscape_file_tag}pathways.txt", sep="\t", index=False)
    out.to_csv(f"{cytoscape_file_tag}subgroups.txt", sep="\t", index=False)
    write_gmt(gmt, f"{cytoscape_file_tag}pathways.gmt")
    _write_legend(tests, col_colors, f"{cytoscape_file_tag}legend.pdf")


def _write_legend(tests, col_colors, filename):
    """Render a simple legend PDF (matplotlib imported lazily)."""
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError:  # pragma: no cover - exercised only when matplotlib is missing
        return

    n = len(tests)
    max_label_width = max(len(t) for t in tests) * 0.05 + 1
    fig_height = 5.5 if n > 20 else (n * 0.25 + 1)
    fig_width = np.ceil(n / 20) * max_label_width if n > 20 else max_label_width
    fig, ax = plt.subplots(figsize=(max(fig_width, 2), max(fig_height, 1.5)))
    ax.axis("off")
    handles = [plt.Rectangle((0, 0), 1, 1, color=c) for c in col_colors]
    ax.legend(handles, tests, title="Contribution", loc="center", frameon=False)
    fig.tight_layout()
    fig.savefig(filename, format="pdf")
    plt.close(fig)
