"""P-value merging methods (Fisher, Brown, Stouffer, Strube and their directional variants)."""
import warnings

import numpy as np
import pandas as pd
from scipy.stats import chi2, norm

# Methods that integrate directional fold-change information.
_DIRECTIONAL_METHODS = {"Fisher_directional", "DPM", "Stouffer_directional", "Strube_directional"}
_NON_DIRECTIONAL_METHODS = {"Fisher", "Brown", "Stouffer", "Strube"}
_VALID_METHODS = _NON_DIRECTIONAL_METHODS | _DIRECTIONAL_METHODS

# p-values smaller than this floor are clamped to avoid numerical underflow.
_PVAL_FLOOR = 1e-300


def _to_array(scores):
    """Normalize the scores argument to (array, names_or_None, is_matrix)."""
    if isinstance(scores, pd.DataFrame):
        return scores.to_numpy(dtype=float), (list(scores.index), list(scores.columns)), True
    if isinstance(scores, dict):
        names = list(scores.keys())
        return np.asarray(list(scores.values()), dtype=float), names, False
    arr = np.asarray(scores)
    if arr.ndim == 2:
        return arr.astype(float), (None, None), True
    if arr.ndim == 1:
        return arr.astype(float), None, False
    raise ValueError("scores must be a matrix or vector")


def _validate_scores(arr):
    if arr.dtype.kind not in "fiu":
        raise ValueError("scores must be numeric")
    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) | (arr > 1)):
        raise ValueError("All values in scores must be in [0,1]")


def _floor_pvalues(arr):
    if arr.size and arr.min() < _PVAL_FLOOR:
        warnings.warn(f"p-values smaller than {_PVAL_FLOOR} are replaced with {_PVAL_FLOOR}")
        arr = np.where(arr < _PVAL_FLOOR, _PVAL_FLOOR, arr)
    return arr


def _validate_direction(scores_arr, scores_meta, is_matrix, scores_direction, constraints_vector, method):
    """Mirror the R-side validation for directional inputs; returns coerced numpy arrays."""
    if (scores_direction is None) ^ (constraints_vector is None):
        raise ValueError("Both scores_direction and constraints_vector must be provided")
    if scores_direction is None:
        return None, 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")

    if isinstance(scores_direction, pd.DataFrame):
        sd = scores_direction.to_numpy(dtype=float)
        sd_index, sd_columns = list(scores_direction.index), list(scores_direction.columns)
    else:
        sd = np.asarray(scores_direction)
        sd_index = sd_columns = None
    if not is_matrix and sd.ndim == 2:
        raise ValueError("scores and scores_direction must be the same data type")
    if is_matrix and sd.ndim != 2:
        raise ValueError("scores and scores_direction must be the same data type")
    if sd.dtype.kind not in "fiu":
        raise ValueError("scores_direction must be numeric")
    if np.any(np.isnan(sd)):
        raise ValueError("scores_direction cannot contain missing values, we recommend replacing NA with 0 or removing")
    if method in _NON_DIRECTIONAL_METHODS:
        raise ValueError("Only DPM, Fisher_directional, Stouffer_directional, and Strube_directional methods support directional integration")

    if is_matrix:
        if sd.shape[0] != scores_arr.shape[0]:
            raise ValueError("scores and scores_direction must have the same number of rows")
        scores_index, scores_columns = scores_meta if scores_meta else (None, None)
        if scores_index is not None and sd_index is not None:
            if set(sd_index) != set(scores_index):
                raise ValueError("scores_direction gene names must match scores genes")
            if list(scores_index) != list(sd_index):
                raise ValueError("scores genes should be in the same order as scores_direction genes")
        if scores_columns is None or sd_columns is None:
            raise ValueError("column names must be provided to scores and scores_direction")
        if set(sd_columns) != set(scores_columns):
            raise ValueError("scores_direction column names must match scores column names")
        if len(cv) != sd.shape[1]:
            raise ValueError("constraints_vector should have the same number of entries as columns in scores_direction")
        zero_cols = np.where(cv == 0)[0]
        if zero_cols.size and not np.all(sd[:, zero_cols] == 0):
            raise ValueError("scores_direction entries must be set to 0's for columns that do not contain directional information")
    else:
        if len(cv) != sd.size:
            raise ValueError("constraints_vector should have the same number of entries as scores_direction")
        if sd.size != scores_arr.size:
            raise ValueError("scores_direction should have the same number of entries as scores")
        zero_idx = np.where(cv == 0)[0]
        if zero_idx.size and not np.all(sd[zero_idx] == 0):
            raise ValueError("scores_direction entries that do not contain directional information must be set to 0's")
    return sd.astype(float), cv.astype(float)


# -----------------------------------------------------------------------------
# Per-row scalar combiners (operate on a single row of p-values)
# -----------------------------------------------------------------------------
def fishers_method(p_values):
    """Fisher's chi-squared statistic: sum(-2 * log(p))."""
    return float(np.sum(-2.0 * np.log(p_values)))


def fishers_directional(p_values, scores_direction, constraints_vector):
    """Directional Fisher statistic; zero-direction tests are pooled non-directionally."""
    p_values = np.asarray(p_values, dtype=float)
    sd = np.asarray(scores_direction, dtype=float)
    cv = np.asarray(constraints_vector, dtype=float)
    # Replicate R's directionality = constraints * sign(scores_direction); /0 -> NaN flags non-directional entries.
    with np.errstate(divide="ignore", invalid="ignore"):
        directionality = cv * (sd / np.abs(sd))
    directional_mask = ~np.isnan(directionality)
    chisq_directional = abs(-2.0 * np.sum(np.log(p_values[directional_mask]) * directionality[directional_mask]))
    chisq_nondirectional = abs(-2.0 * np.sum(np.log(p_values[~directional_mask])))
    return float(chisq_directional + chisq_nondirectional)


def stouffers_method(p_values):
    """Stouffer's Z-statistic for a row of p-values."""
    p_values = np.asarray(p_values, dtype=float)
    z_values = norm.ppf(p_values / 2.0)
    return float(np.sum(z_values) / np.sqrt(len(p_values)))


def stouffers_directional(p_values, scores_direction, constraints_vector):
    """Directional Stouffer Z-statistic; non-directional entries are summed separately."""
    p_values = np.asarray(p_values, dtype=float)
    sd = np.asarray(scores_direction, dtype=float)
    cv = np.asarray(constraints_vector, dtype=float)
    k = len(p_values)
    with np.errstate(divide="ignore", invalid="ignore"):
        directionality = cv * (sd / np.abs(sd))
    directional_mask = ~np.isnan(directionality)
    z_directional = abs(np.sum(norm.ppf(p_values[directional_mask] / 2.0) * directionality[directional_mask]))
    z_nondirectional = abs(np.sum(norm.ppf(p_values[~directional_mask] / 2.0)))
    return float((z_directional + z_nondirectional) / np.sqrt(k))


# -----------------------------------------------------------------------------
# Covariance helpers (Brown / Strube)
# -----------------------------------------------------------------------------
def _pop_var(x):
    """Population variance (R's var * (n-1)/n), ignoring NaNs."""
    x = np.asarray(x, dtype=float)
    x = x[~np.isnan(x)]
    if x.size == 0:
        return 0.0
    return float(np.var(x, ddof=1) * (x.size - 1) / x.size)


def _pop_sd(x):
    return float(np.sqrt(_pop_var(x)))


def transform_data(dat):
    """Empirical transformation used inside Brown's covariance estimation."""
    dat = np.asarray(dat, dtype=float)
    # If all values are equal, the covariance is the zero matrix; return as-is to avoid division by 0.
    if np.allclose(dat.min(), dat.max()):
        return dat
    dvm = np.nanmean(dat)
    dvsd = _pop_sd(dat)
    s = (dat - dvm) / dvsd
    # Empirical CDF over the standardized values, evaluated at each standardized point.
    sorted_s = np.sort(s)
    ranks = np.searchsorted(sorted_s, s, side="right")
    ecdf_vals = ranks / len(s)
    return -2.0 * np.log(ecdf_vals)


def calculate_covariances(data_matrix):
    """Compute the covariance matrix for Brown's/DPM methods via row-wise transform_data."""
    data_matrix = np.asarray(data_matrix, dtype=float)
    # Apply transform_data per row -> result keeps shape (n_tests, n_items); each row is a variable.
    transformed = np.apply_along_axis(transform_data, 1, data_matrix)
    return np.cov(transformed, ddof=1)


# -----------------------------------------------------------------------------
# Matrix combiners (Brown, DPM, Strube)
# -----------------------------------------------------------------------------
def _brown_scaling(cov_matrix):
    """Compute Brown's degrees-of-freedom and scale factor from a covariance matrix."""
    n = cov_matrix.shape[1]
    expected = 2 * n
    # Sum the strict lower-triangle covariances and double them (symmetric matrix).
    cov_sum = 2.0 * np.sum(np.tril(cov_matrix, k=-1))
    var = (4 * n) + cov_sum
    sf = var / (2.0 * expected)
    df = (2.0 * expected ** 2) / var
    if df > 2 * n:
        df = 2 * n
        sf = 1.0
    return df, sf


def browns_method(p_values, data_matrix=None, cov_matrix=None):
    """Brown's method for a matrix of p-values (rows = items, cols = tests)."""
    if data_matrix is None and cov_matrix is None:
        raise ValueError("Either data_matrix or cov_matrix must be supplied")
    if data_matrix is not None and cov_matrix is not None:
        warnings.warn("Both data_matrix and cov_matrix were supplied. Ignoring data_matrix")
    if cov_matrix is None:
        cov_matrix = calculate_covariances(data_matrix)
    df, sf = _brown_scaling(np.asarray(cov_matrix))
    fisher_chisq = np.array([fishers_method(row) for row in p_values])
    return chi2.sf(fisher_chisq / sf, df=df)


def dpm(p_values, scores_direction, constraints_vector, data_matrix=None, cov_matrix=None):
    """Directional Brown's method (DPM)."""
    if data_matrix is None and cov_matrix is None:
        raise ValueError("Either data_matrix or cov_matrix must be supplied")
    if data_matrix is not None and cov_matrix is not None:
        warnings.warn("Both data_matrix and cov_matrix were supplied. Ignoring data_matrix")
    if cov_matrix is None:
        cov_matrix = calculate_covariances(data_matrix)
    df, sf = _brown_scaling(np.asarray(cov_matrix))
    fisher_chisq = np.array([
        fishers_directional(p_values[i], scores_direction[i], constraints_vector)
        for i in range(p_values.shape[0])
    ])
    return chi2.sf(fisher_chisq / sf, df=df)


def _strube_scaling(p_values):
    """Compute Strube's normalisation factor ``sqrt(k) / sqrt(sum |cor|)``."""
    # Use 'complete.obs' equivalent: pandas .corr() ignores NaNs pairwise; here we have no NaNs.
    cor_mtx = np.corrcoef(p_values, rowvar=False)
    cor_mtx = np.nan_to_num(np.abs(cor_mtx), nan=0.0)
    k = p_values.shape[1]
    return np.sqrt(k) / np.sqrt(cor_mtx.sum())


def strubes_method(p_values):
    """Strube's method - Stouffer's method adjusted for between-test correlation."""
    p_values = np.asarray(p_values, dtype=float)
    stouffer_z = np.array([stouffers_method(row) for row in p_values])
    adjusted_z = stouffer_z * _strube_scaling(p_values)
    return 2.0 * norm.cdf(-1.0 * np.abs(adjusted_z))


def strubes_directional(p_values, scores_direction, constraints_vector):
    """Directional Strube's method."""
    p_values = np.asarray(p_values, dtype=float)
    stouffer_z = np.array([
        stouffers_directional(p_values[i], scores_direction[i], constraints_vector)
        for i in range(p_values.shape[0])
    ])
    adjusted_z = stouffer_z * _strube_scaling(p_values)
    return 2.0 * norm.cdf(-1.0 * np.abs(adjusted_z))


# -----------------------------------------------------------------------------
# Public dispatcher
# -----------------------------------------------------------------------------
def merge_p_values(scores, method="Fisher", scores_direction=None, constraints_vector=None):
    """Merge a vector or matrix of p-values into one combined p-value per row.

    scores can be a 1-D list/array or a 2-D matrix/DataFrame (one column per test).
    Returns a scalar for vector input or a Series indexed by gene for matrix input.
    """
    arr, meta, is_matrix = _to_array(scores)
    _validate_scores(arr)
    if isinstance(meta, list) and not is_matrix:
        names = meta
    elif is_matrix and isinstance(meta, tuple):
        names = meta[0]
    else:
        names = None

    if method not in _VALID_METHODS:
        raise ValueError(
            "Only Fisher, Brown, Stouffer and Strube methods are currently supported for non-directional analysis. "
            "And only DPM, Fisher_directional, Stouffer_directional, and Strube_directional are supported for directional analysis"
        )
    if method in _DIRECTIONAL_METHODS and scores_direction is None:
        raise ValueError("scores_direction and constraints_vector must be provided for directional analyses")

    sd_arr, cv_arr = _validate_direction(arr, meta, is_matrix, scores_direction, constraints_vector, method)

    if not is_matrix:
        if method in {"Brown", "Strube", "DPM", "Strube_directional"}:
            raise ValueError("Brown's, DPM, Strube's, and Strube_directional methods cannot be used with a single list of p-values")
        arr = _floor_pvalues(arr)
        if method == "Fisher":
            return float(chi2.sf(fishers_method(arr), df=2 * len(arr)))
        if method == "Fisher_directional":
            return float(chi2.sf(fishers_directional(arr, sd_arr, cv_arr), df=2 * len(arr)))
        if method == "Stouffer":
            return float(2.0 * norm.cdf(-1.0 * abs(stouffers_method(arr))))
        if method == "Stouffer_directional":
            return float(2.0 * norm.cdf(-1.0 * abs(stouffers_directional(arr, sd_arr, cv_arr))))

    # Matrix path: a single-column matrix has nothing to merge.
    if arr.shape[1] == 1:
        out = arr[:, 0]
        return pd.Series(out, index=names) if names is not None else out
    arr = _floor_pvalues(arr)

    if method == "Fisher":
        merged = np.array([chi2.sf(fishers_method(row), df=2 * arr.shape[1]) for row in arr])
    elif method == "Fisher_directional":
        merged = np.array([
            chi2.sf(fishers_directional(arr[i], sd_arr[i], cv_arr), df=2 * arr.shape[1])
            for i in range(arr.shape[0])
        ])
    elif method == "Brown":
        merged = browns_method(arr, data_matrix=arr.T)
    elif method == "DPM":
        merged = dpm(arr, sd_arr, cv_arr, data_matrix=arr.T)
    elif method == "Stouffer":
        merged = np.array([2.0 * norm.cdf(-1.0 * abs(stouffers_method(row))) for row in arr])
    elif method == "Stouffer_directional":
        merged = np.array([
            2.0 * norm.cdf(-1.0 * abs(stouffers_directional(arr[i], sd_arr[i], cv_arr)))
            for i in range(arr.shape[0])
        ])
    elif method == "Strube":
        merged = strubes_method(arr)
    else:  # Strube_directional
        merged = strubes_directional(arr, sd_arr, cv_arr)

    return pd.Series(merged, index=names) if names is not None else merged
