"""Theil's U (uncertainty coefficient) direccional entre dos columnas categoricas.

U(a|b) mide cuanta incertidumbre de `a` se elimina conociendo `b`, normalizado a
[0, 1]. Es ASIMETRICO: theils_u(a, b) != theils_u(b, a) en general, lo que lo
distingue de medidas simetricas como Cramer's V y permite detectar dependencias
direccionales (p.ej. ciudad -> pais).
"""

import math
from collections import Counter


def _entropy(counts: list) -> float:
    """Entropia de Shannon (base natural) de una lista de conteos.

    Args:
        counts: conteos por categoria (enteros >= 0).

    Returns:
        entropia en nats; 0.0 si no hay observaciones.
    """
    total = sum(counts)
    if total == 0:
        return 0.0
    h = 0.0
    for c in counts:
        if c > 0:
            p = c / total
            h -= p * math.log(p)
    return h


def theils_u(a: list, b: list) -> float:
    """Theil's U direccional U(a|b): incertidumbre de `a` explicada por `b`.

    Calcula la fraccion de la entropia de la distribucion marginal de `a` que se
    elimina al condicionar sobre los valores de `b`. Es una medida de asociacion
    ASIMETRICA en [0, 1]:

    - U(a|b) = 1.0 -> conocer `b` determina por completo `a`.
    - U(a|b) = 0.0 -> `b` no aporta nada sobre `a` (independencia).

    Las entropias usan la misma base (logaritmo natural), por lo que la base se
    cancela en el cociente y el resultado es independiente de ella.

    Args:
        a: columna categorica objetivo (cuya incertidumbre se mide).
        b: columna categorica condicionante (el conocimiento que se aporta).
            Ambas listas se emparejan por indice; los pares con algun None se
            descartan antes de calcular.

    Returns:
        Theil's U(a|b) como float en [0.0, 1.0]. Devuelve 0.0 (nunca None ni
        excepcion) si hay menos de 2 pares validos o si H(a) == 0 (es decir, `a`
        ya es constante y no hay incertidumbre que eliminar).
    """
    # Empareja por indice y descarta pares con algun None.
    pairs = [
        (av, bv)
        for av, bv in zip(a, b)
        if av is not None and bv is not None
    ]
    if len(pairs) < 2:
        return 0.0

    # H(a): entropia de la distribucion marginal de a.
    a_counts = Counter(av for av, _ in pairs)
    h_a = _entropy(list(a_counts.values()))
    if h_a == 0.0:
        return 0.0

    # H(a|b) = suma_b p(b) * H(a | b=valor).
    by_b: dict = {}
    for av, bv in pairs:
        by_b.setdefault(bv, Counter())[av] += 1
    total = len(pairs)
    h_a_given_b = 0.0
    for bv, a_sub in by_b.items():
        p_b = sum(a_sub.values()) / total
        h_a_given_b += p_b * _entropy(list(a_sub.values()))

    u = (h_a - h_a_given_b) / h_a
    # Clampa a [0, 1] para absorber errores de redondeo en coma flotante.
    if u < 0.0:
        return 0.0
    if u > 1.0:
        return 1.0
    return u