"""Main pathway enrichment driver - Python port of the R ActivePathways function."""
import warnings

import numpy as np
import pandas as pd

from ._pvalue_correction import p_adjust
from .gmt import GMT, Pathway, is_gmt, make_background, read_gmt
from .merge_p import merge_p_values
from .statistical_tests import ordered_hypergeometric


def _scores_to_dataframe(scores, name="scores"):
    """Coerce array/DataFrame input to a DataFrame with string row and column labels."""
    if isinstance(scores, pd.DataFrame):
        df = scores.copy()
    else:
        arr = np.asarray(scores)
        if arr.ndim != 2:
            raise ValueError(f"{name} must be a numeric matrix")
        df = pd.DataFrame(arr)
    if df.shape[0] == 0 or df.shape[1] == 0:
        raise ValueError(f"{name} must be a numeric matrix")
    if not np.issubdtype(df.to_numpy().dtype, np.number):
        raise ValueError(f"{name} must be a numeric matrix")
    return df


def _validate_inputs(
    scores,
    scores_direction,
    constraints_vector,
    cutoff,
    significant,
    merge_method,
    custom_colors,
    color_palette,
    color_integrated_only,
):
    """Validate all inputs, mirroring the original R checks."""
    arr = scores.to_numpy()
    if np.any(np.isnan(arr)):
        raise ValueError("scores cannot contain missing values, we recommend replacing NA with 1 or removing")
    if np.any(arr < 0) or np.any(arr > 1):
        raise ValueError("All values in scores must be in [0,1]")
    if scores.index.duplicated().any():
        raise ValueError("scores matrix contains duplicated genes - rownames must be unique")

    if (scores_direction is None) ^ (constraints_vector is None):
        raise ValueError("Both scores_direction and constraints_vector must be provided")
    if scores_direction is not None:
        cv = np.asarray(constraints_vector)
        if cv.dtype.kind not in "fiu":
            raise ValueError("constraints_vector must be a numeric vector")
        if not np.all(np.isin(cv, [1, -1, 0])):
            raise ValueError("constraints_vector must contain the values: 1, -1 or 0")
        sd_arr = scores_direction.to_numpy()
        if not np.issubdtype(sd_arr.dtype, np.number):
            raise ValueError("scores_direction must be a numeric matrix")
        if np.any(np.isnan(sd_arr)):
            raise ValueError("scores_direction cannot contain missing values, we recommend replacing NA with 0 or removing")
        if scores.shape[0] != scores_direction.shape[0]:
            raise ValueError("scores and scores_direction must have the same number of rows")
        if not set(scores_direction.index).issubset(set(scores.index)):
            raise ValueError("scores_direction gene names must match scores genes")
        if list(scores.index) != list(scores_direction.index):
            raise ValueError("scores genes should be in the same order as scores_direction genes")
        if scores.columns is None or scores_direction.columns is None:
            raise ValueError("column names must be provided to scores and scores_direction")
        if not set(scores_direction.columns).issubset(set(scores.columns)):
            raise ValueError("scores_direction column names must match scores column names")
        if len(cv) != scores_direction.shape[1]:
            raise ValueError("constraints_vector should have the same number of entries as columns in scores_direction")
        if merge_method in {"Fisher", "Brown", "Stouffer", "Strube"}:
            raise ValueError("Only DPM, Fisher_directional, Stouffer_directional, and Strube_directional methods support directional integration")
        zero_cols = np.where(cv == 0)[0]
        if zero_cols.size and not np.all(sd_arr[:, zero_cols] == 0):
            raise ValueError("scores_direction entries must be set to 0's for columns that do not contain directional information")

    for label, val in (("cutoff", cutoff), ("significant", significant)):
        if val is None or not np.isscalar(val):
            raise ValueError(f"length({label}) == 1 is not TRUE")
        if not isinstance(val, (int, float, np.integer, np.floating)) or isinstance(val, bool):
            raise ValueError(f"is.numeric({label}) is not TRUE")
        if val < 0 or val > 1:
            raise ValueError(f"{label} must be a value in [0,1]")

    if custom_colors is not None:
        if not (isinstance(custom_colors, (list, tuple, np.ndarray, pd.Series, dict)) and
                all(isinstance(c, str) for c in (custom_colors.values() if isinstance(custom_colors, dict) else custom_colors))):
            raise ValueError("colors must be provided as a character vector")
        n_colors = len(custom_colors)
        if n_colors != scores.shape[1]:
            raise ValueError("incorrect number of colors is provided")
        if isinstance(custom_colors, dict):
            if not all(name in scores.columns for name in custom_colors.keys()):
                raise ValueError("names() of the custom colors vector should match the scores column names")
    if custom_colors is not None and color_palette is not None:
        raise ValueError("Both custom_colors and color_palette are provided. Specify only one of these parameters for node coloring.")

    if not isinstance(color_integrated_only, str):
        raise ValueError("color must be provided as a character vector")


def enrichment_analysis(genelist, gmt, background):
    """Run an ordered hypergeometric test for each GMT term.

    Returns a DataFrame with columns: term_id, term_name, adjusted_p_val,
    term_size, and overlap (list of genes).
    """
    rows = []
    for term in gmt:
        result = ordered_hypergeometric(genelist, background, term.genes)
        # ``ind`` is 1-based; slice the gene list and intersect with the term's genes.
        prefix = list(genelist[: result["ind"]])
        annot = set(term.genes)
        overlap = [g for g in prefix if g in annot]
        rows.append({
            "term_id": term.id,
            "term_name": term.name,
            "adjusted_p_val": result["p_val"],
            "term_size": len(term.genes),
            "overlap": overlap,
        })
    return pd.DataFrame(rows, columns=["term_id", "term_name", "adjusted_p_val", "term_size", "overlap"])


def column_significance(scores, gmt, background, cutoff, significant, correction_method, pvals):
    """Compute per-column enrichment evidence for each term.

    Runs enrichment independently per column; marks overlap genes as NaN for
    non-significant terms. Evidence column lists which columns drove significance.
    """
    dt = pd.DataFrame({"term_id": gmt.ids()})
    dt["evidence"] = None  # placeholder; populated after column passes are done
    for col in scores.columns:
        col_scores = scores[col]
        kept = col_scores[col_scores <= cutoff].sort_values()
        ranked = list(kept.index)
        res = enrichment_analysis(ranked, gmt, background)
        adj = p_adjust(res["adjusted_p_val"].to_numpy(), correction_method)
        # Replace overlap with NaN where the term is not significant for this column.
        overlaps = res["overlap"].tolist()
        masked = [ov if a <= significant else np.nan for ov, a in zip(overlaps, adj)]
        dt[col] = masked

    # Build the per-row evidence list: which columns marked the term significant; if none did,
    # mark as 'combined' if the merged p-value passed, else 'none'.
    ev_columns = list(scores.columns)
    evidence = []
    for i, term_id in enumerate(dt["term_id"]):
        contributing = [c for c in ev_columns if not _is_na(dt.at[i, c])]
        if contributing:
            evidence.append(contributing)
        else:
            evidence.append(["combined"] if pvals[i] <= significant else ["none"])
    dt["evidence"] = evidence
    # Rename per-column evidence columns to the R-style "Genes_<col>" prefix.
    dt = dt.rename(columns={c: f"Genes_{c}" for c in ev_columns})
    return dt


def _is_na(value):
    """Return True if value is a scalar NaN (lists containing NaN count as not NaN)."""
    if isinstance(value, list):
        return False
    try:
        return bool(pd.isna(value))
    except (TypeError, ValueError):
        return False


def active_pathways(
    scores,
    gmt,
    background=None,
    geneset_filter=(5, 1000),
    cutoff=0.1,
    significant=0.05,
    merge_method="Fisher",
    correction_method="holm",
    cytoscape_file_tag=None,
    color_palette=None,
    custom_colors=None,
    color_integrated_only="#FFFFF0",
    scores_direction=None,
    constraints_vector=None,
):
    """Integrative pathway enrichment for multivariate omics p-value matrices.

    Mirrors the R ActivePathways function. Returns a DataFrame of significant
    pathways, or None when no terms pass the significance threshold.
    """
    scores_df = _scores_to_dataframe(scores)
    # Ensure row/column labels are strings so downstream string ops behave consistently.
    scores_df.columns = [str(c) for c in scores_df.columns]
    scores_df.index = [str(i) for i in scores_df.index]

    scores_direction_df = None
    if scores_direction is not None:
        scores_direction_df = _scores_to_dataframe(scores_direction, name="scores_direction")
        scores_direction_df.columns = [str(c) for c in scores_direction_df.columns]
        scores_direction_df.index = [str(i) for i in scores_direction_df.index]

    _validate_inputs(scores_df, scores_direction_df, constraints_vector, cutoff, significant,
                     merge_method, custom_colors, color_palette, color_integrated_only)

    if isinstance(gmt, str):
        gmt = read_gmt(gmt)
    if not is_gmt(gmt):
        raise TypeError("gmt must be a GMT object or a path to a GMT file")
    if len(gmt) == 0:
        raise ValueError("No pathways in gmt made the geneset_filter")

    if background is None:
        background = make_background(gmt)
    if not (isinstance(background, (list, tuple, np.ndarray, pd.Index)) and
            all(isinstance(b, str) for b in background)):
        raise ValueError("background must be a character vector")
    background = list(background)

    if geneset_filter is not None:
        if not (isinstance(geneset_filter, (list, tuple, np.ndarray)) and len(geneset_filter) == 2):
            if isinstance(geneset_filter, (list, tuple, np.ndarray)) and len(geneset_filter) != 2:
                raise ValueError("geneset_filter must be length 2")
            raise ValueError("geneset_filter must be a numeric vector")
        gf = list(geneset_filter)
        for g in gf:
            if g is None:
                continue
            if not isinstance(g, (int, float, np.integer, np.floating)) or isinstance(g, bool):
                raise ValueError("geneset_filter must be a numeric vector")
            if g < 0:
                raise ValueError("geneset_filter limits must be positive")

    return _run_pipeline(
        scores_df, scores_direction_df, constraints_vector, gmt, background,
        geneset_filter, cutoff, significant, merge_method, correction_method,
        cytoscape_file_tag, color_palette, custom_colors, color_integrated_only,
    )



def _run_pipeline(
    scores_df,
    scores_direction_df,
    constraints_vector,
    gmt,
    background,
    geneset_filter,
    cutoff,
    significant,
    merge_method,
    correction_method,
    cytoscape_file_tag,
    color_palette,
    custom_colors,
    color_integrated_only,
):
    """Filter inputs, merge p-values, run enrichment, and optionally write Cytoscape files."""
    contribution = scores_df.shape[1] > 1
    if not contribution:
        warnings.warn("scores matrix contains only one column. Column contributions will not be calculated")

    # Drop genes that are not part of the user-supplied background.
    background_set = set(background)
    orig_length = scores_df.shape[0]
    scores_df = scores_df.loc[[g for g in scores_df.index if g in background_set]]
    if scores_direction_df is not None:
        scores_direction_df = scores_direction_df.loc[[g for g in scores_direction_df.index if g in background_set]]
    if scores_df.shape[0] == 0:
        raise ValueError("scores does not contain any genes in the background")
    if scores_df.shape[0] < orig_length:
        warnings.warn(f"{orig_length - scores_df.shape[0]} rows were removed from scores because they are not found in the background")

    # If the background lacks some GMT genes, intersect each term's gene list with it.
    all_gmt_genes = set()
    for term in gmt:
        all_gmt_genes.update(term.genes)
    if not all_gmt_genes.issubset(background_set):
        new_terms = []
        for term in gmt:
            kept = [g for g in term.genes if g in background_set]
            if kept:
                new_terms.append(Pathway(id=term.id, name=term.name, genes=kept))
        gmt = GMT(new_terms)

    if geneset_filter is not None:
        orig_n = len(gmt)
        lo, hi = geneset_filter[0], geneset_filter[1]
        kept = []
        for term in gmt:
            n = len(term.genes)
            if lo is not None and not (isinstance(lo, float) and np.isnan(lo)) and n < lo:
                continue
            if hi is not None and not (isinstance(hi, float) and np.isnan(hi)) and n > hi:
                continue
            kept.append(term)
        gmt = GMT(kept)
        if len(gmt) == 0:
            raise ValueError("No pathways in gmt made the geneset_filter")
        if len(gmt) < orig_n:
            warnings.warn(f"{orig_n - len(gmt)} terms were removed from gmt because they did not make the geneset_filter")

    # Merge per-gene p-values across columns to get a single ranked list.
    merged = merge_p_values(scores_df, merge_method, scores_direction_df, constraints_vector)
    if isinstance(merged, np.ndarray):
        merged = pd.Series(merged, index=scores_df.index)
    elif np.isscalar(merged):
        merged = pd.Series([merged], index=scores_df.index)
    merged = merged[merged <= cutoff]
    if merged.empty:
        raise ValueError("No genes made the cutoff")
    ordered_scores = list(merged.sort_values().index)

    # Pathway-level enrichment using the ranked gene list.
    res = enrichment_analysis(ordered_scores, gmt, background)
    res["adjusted_p_val"] = p_adjust(res["adjusted_p_val"].to_numpy(), correction_method)

    significant_mask = res["adjusted_p_val"] <= significant
    if not significant_mask.any():
        warnings.warn("No significant terms were found")
        return None

    if contribution:
        sig_cols = column_significance(scores_df, gmt, background, cutoff, significant,
                                       correction_method, res["adjusted_p_val"].to_numpy())
        # Drop the term_id column from sig_cols when joining to avoid duplication.
        res = pd.concat([res.reset_index(drop=True), sig_cols.drop(columns="term_id").reset_index(drop=True)], axis=1)
    else:
        sig_cols = None

    if cytoscape_file_tag is not None and significant_mask.any():
        from .cytoscape import prepare_cytoscape  # local import to keep matplotlib optional
        sig_terms = res.loc[significant_mask, ["term_id", "term_name", "adjusted_p_val"]].reset_index(drop=True)
        sig_gmt_terms = [t for t, keep in zip(gmt, significant_mask) if keep]
        sig_gmt = GMT(sig_gmt_terms)
        sig_cols_subset = sig_cols.loc[significant_mask].reset_index(drop=True) if sig_cols is not None else None
        prepare_cytoscape(sig_terms, sig_gmt, cytoscape_file_tag, sig_cols_subset,
                          color_palette, custom_colors, color_integrated_only)

    return res.loc[significant_mask].reset_index(drop=True)
