"""Informacion mutua entre dos columnas pareadas (relaciones lineales y no lineales).

Funcion pura del grupo eda. Mide la dependencia estadistica general entre dos
columnas (numericas, categoricas o mezcla), capturando relaciones de cualquier
forma -- no solo lineales como Pearson. Es la metrica "general" de la matriz de
asociacion: complementa a `pearson` (solo lineal num-num) y `cramers_v` (solo
cat-cat).
"""

import math
from collections import Counter

import numpy as np
from sklearn.metrics import mutual_info_score


def _discretize(values: list, numeric: bool, bins: int) -> list:
    """Discretiza una columna a etiquetas enteras.

    Columnas numericas -> `bins` cubos por cuantiles (np.digitize sobre los
    bordes de cuantil). Columnas categoricas -> factorizacion valor->id.
    """
    if numeric:
        arr = np.asarray(values, dtype=float)
        # Bordes interiores por cuantiles (excluye 0 y 1 para usar digitize).
        qs = np.linspace(0.0, 1.0, bins + 1)[1:-1]
        if qs.size == 0:
            # bins <= 1 -> todo cae en un unico cubo.
            return [0] * len(arr)
        edges = np.quantile(arr, qs)
        # Bordes unicos: cuantiles repetidos (columnas con poca variacion)
        # colapsan en menos cubos, lo cual es correcto (menos entropia).
        edges = np.unique(edges)
        return list(np.digitize(arr, edges))
    # Categorica: mapa valor -> id entero, en orden de aparicion.
    ids: dict = {}
    out = []
    for v in values:
        if v not in ids:
            ids[v] = len(ids)
        out.append(ids[v])
    return out


def _entropy(labels: list) -> float:
    """Entropia de Shannon (nats) de una secuencia de etiquetas."""
    n = len(labels)
    if n == 0:
        return 0.0
    h = 0.0
    for c in Counter(labels).values():
        p = c / n
        h -= p * math.log(p)
    return h


def mutual_info_columns(
    a: list,
    b: list,
    a_numeric: bool = False,
    b_numeric: bool = False,
    bins: int = 10,
    normalized: bool = True,
) -> float:
    """Informacion mutua entre dos columnas pareadas posicion a posicion.

    Empareja `a` y `b`, descarta los pares donde cualquiera de los dos sea None,
    discretiza cada columna (numericas por cuantiles, categoricas por
    factorizacion) y calcula la informacion mutua. Captura relaciones de
    cualquier forma (lineal o no, num-num, cat-cat, num-cat).

    Args:
        a: lista de valores de la primera columna (None se descarta).
        b: lista de valores pareada con `a` (mismo criterio).
        a_numeric: si True, `a` se discretiza en `bins` cuantiles; si False se
            factoriza como categorica.
        b_numeric: idem para `b`.
        bins: numero de cubos por cuantiles para columnas numericas.
        normalized: si True devuelve la NMI = MI / sqrt(H(a)*H(b)) en [0, 1]
            (1 = dependencia total). Si False devuelve la MI cruda en nats.

    Returns:
        float. NMI en [0, 1] si normalized; MI en nats (>= 0) si no. Devuelve
        0.0 si hay menos de 2 pares validos o si alguna columna discretizada
        tiene entropia 0 (constante) bajo normalized. Nunca None ni excepcion.
    """
    pairs = [
        (x, y)
        for x, y in zip(a, b)
        if x is not None and y is not None
    ]
    if len(pairs) < 2:
        return 0.0

    a_vals = [x for x, _ in pairs]
    b_vals = [y for _, y in pairs]

    a_disc = _discretize(a_vals, a_numeric, bins)
    b_disc = _discretize(b_vals, b_numeric, bins)

    mi = float(mutual_info_score(a_disc, b_disc))

    if not normalized:
        return max(0.0, mi)

    ha = _entropy(a_disc)
    hb = _entropy(b_disc)
    if ha <= 0.0 or hb <= 0.0:
        # Alguna columna es constante -> no hay informacion compartida medible.
        return 0.0

    nmi = mi / math.sqrt(ha * hb)
    # Clampa a [0, 1] por seguridad numerica.
    return max(0.0, min(1.0, nmi))