"""Hypergeometric and ordered-hypergeometric tests used by ActivePathways."""
import numpy as np
from scipy.stats import hypergeom


def hypergeometric(counts):
    """One-sided hypergeometric p-value for a 2x2 contingency table (alternative='greater')."""
    counts = np.asarray(counts)
    if np.any(counts < 0):
        raise ValueError("counts contains negative values. Something went very wrong.")
    m = int(counts[0, 0] + counts[1, 0])  # total annotated genes (white balls)
    n = int(counts[0, 1] + counts[1, 1])  # total non-annotated genes (black balls)
    k = int(counts[0, 0] + counts[0, 1])  # genes drawn
    x = int(counts[0, 0])                 # annotated genes drawn
    # phyper(x-1, m, n, k, lower.tail=FALSE) == hypergeom.sf(x-1, m+n, m, k)
    return float(hypergeom.sf(x - 1, m + n, m, k))


def ordered_hypergeometric(genelist, background, annotations):
    """Walk down genelist and return the smallest hypergeometric p-value.

    Returns {"p_val": float, "ind": int} where ind is the 1-based prefix length
    that achieved the minimum p-value (matching the R API).
    """
    annotations_set = set(annotations)
    background_set = set(background)
    # Indices (0-based) of genelist entries that belong to the annotation set.
    which_in = [i for i, g in enumerate(genelist) if g in annotations_set]
    if not which_in:
        return {"p_val": 1.0, "ind": 1}

    # Build the contingency table for the prefix ending at the first annotation hit.
    first = which_in[0]
    gl_prefix = list(genelist[: first + 1])
    gl_set = set(gl_prefix)
    cl = [g for g in background_set if g not in gl_set]
    genelist0 = len(gl_prefix) - 1                                      # in prefix, not in annotations
    complement1 = sum(1 for g in cl if g in annotations_set)            # outside prefix, annotated
    complement0 = len(cl) - complement1                                 # outside prefix, unannotated
    counts = np.array([[1, complement1], [genelist0, complement0]], dtype=np.int64)
    scores = [hypergeometric(counts)]

    if len(which_in) == 1:
        return {"p_val": scores[0], "ind": first + 1}

    # Incrementally update counts as we extend the prefix to subsequent annotation hits.
    for i in range(1, len(which_in)):
        diff = which_in[i] - which_in[i - 1]
        # i+1 annotated genes are now within the prefix (i is 0-based loop var).
        counts[0, 0] = i + 1
        counts[1, 0] = counts[1, 0] + diff - 1
        counts[0, 1] = counts[0, 1] - 1
        counts[1, 1] = counts[1, 1] - diff + 1
        scores.append(hypergeometric(counts))

    # Match the R behaviour of returning the *last* index that achieves the minimum.
    min_score = min(scores)
    last_min = max(i for i, s in enumerate(scores) if s == min_score)
    return {"p_val": min_score, "ind": which_in[last_min] + 1}
