"""Co-ocurrencia de ausencias: matriz de correlacion de Pearson entre mascaras de nulos.

Funcion pura del grupo eda, nucleo del capitulo de missingness. Recibe la mascara
binaria de ausencias de una tabla (1 = falta, 0 = presente, alineada por fila) y
mide hasta que punto las columnas faltan juntas. Para cada par de columnas con
varianza en su ausencia calcula la correlacion de Pearson entre los vectores
binarios, mas las cifras de solapamiento de conjuntos (co-missing, either-missing,
Jaccard). Compone la funcion atomica `pearson` del registry; no reimplementa la
correlacion. Lectura defensiva; NUNCA lanza.
"""

from datascience import pearson


def missingness_correlation(null_mask, top_k=20) -> dict:
    """Correlacion de co-ocurrencia de ausencias entre columnas.

    Args:
        null_mask: dict {col: [int 0/1, ...]} alineado por fila (1 = el valor
            falta en esa fila). Todas las listas se asumen de la misma longitud.
        top_k: numero maximo de pares a devolver, ordenados por |corr| desc.

    Returns:
        dict con:
          - columns: columnas con varianza en la ausencia (al menos un 1 y al
            menos un 0), en orden de entrada.
          - matrix: matriz len(columns) x len(columns) de correlacion de Pearson
            entre las mascaras binarias, diagonal 1.0.
          - pairs: lista de hasta top_k pares (i<j) ordenados por |corr| desc.
            Cada par: {a, b, corr, co_missing, either_missing, jaccard}.
          - n_excluded: numero de columnas con algun nulo pero sin varianza
            (constantes en la ausencia: siempre presentes o siempre ausentes).
          - excluded_cols: lista de esas columnas (en orden de entrada).

        Si hay menos de 2 columnas con varianza, columns/matrix/pairs van vacios
        pero n_excluded/excluded_cols se rellenan igualmente. NUNCA lanza.
    """
    # Salida base, defensiva ante entradas no-dict.
    result = {
        "columns": [],
        "matrix": [],
        "pairs": [],
        "n_excluded": 0,
        "excluded_cols": [],
    }

    if not isinstance(null_mask, dict) or not null_mask:
        return result

    varying = []          # columnas con varianza en la ausencia
    varying_vecs = []     # sus vectores binarios saneados (floats 0.0/1.0)
    excluded_cols = []    # columnas con nulos pero sin varianza (constantes)

    for col, raw in null_mask.items():
        if not isinstance(raw, (list, tuple)):
            continue
        # Sanea a 0/1: cualquier valor truthy distinto de 0 cuenta como ausencia.
        vec = [1 if bool(v) else 0 for v in raw]
        if not vec:
            continue
        ones = sum(vec)
        zeros = len(vec) - ones
        if ones > 0 and zeros > 0:
            varying.append(col)
            varying_vecs.append([float(v) for v in vec])
        elif ones > 0:
            # Tiene nulos pero todos (constante en la ausencia): sin varianza.
            excluded_cols.append(col)
        # ones == 0 -> columna siempre presente, sin nulos: no se cuenta como
        # excluida (no aporta ausencia al analisis de co-ocurrencia).

    result["n_excluded"] = len(excluded_cols)
    result["excluded_cols"] = excluded_cols

    n = len(varying)
    if n < 2:
        return result

    result["columns"] = list(varying)

    # Matriz de correlacion de Pearson, diagonal 1.0.
    matrix = [[0.0] * n for _ in range(n)]
    for i in range(n):
        matrix[i][i] = 1.0
    for i in range(n):
        for j in range(i + 1, n):
            r = pearson(varying_vecs[i], varying_vecs[j])
            matrix[i][j] = r
            matrix[j][i] = r
    result["matrix"] = matrix

    # Pares con cifras de solapamiento de conjuntos.
    pairs = []
    for i in range(n):
        vi = varying_vecs[i]
        for j in range(i + 1, n):
            vj = varying_vecs[j]
            co_missing = 0
            either_missing = 0
            for a, b in zip(vi, vj):
                a_miss = a != 0.0
                b_miss = b != 0.0
                if a_miss and b_miss:
                    co_missing += 1
                if a_miss or b_miss:
                    either_missing += 1
            jaccard = co_missing / either_missing if either_missing > 0 else 0.0
            pairs.append({
                "a": varying[i],
                "b": varying[j],
                "corr": matrix[i][j],
                "co_missing": co_missing,
                "either_missing": either_missing,
                "jaccard": jaccard,
            })

    pairs.sort(key=lambda p: abs(p["corr"]), reverse=True)
    result["pairs"] = pairs[:top_k] if top_k is not None and top_k >= 0 else pairs

    return result
