"""Entropy/KL support handling and a stable softmax gradient check. Python 3.10+."""
from __future__ import annotations
import math


def distribution(values):
    """Pre: iterable. Post: finite normalized nonnegative tuple, or ValueError."""
    values = tuple(values)
    if not values or any(not math.isfinite(x) or x < 0 for x in values) or not math.isclose(sum(values), 1):
        raise ValueError("expected a nonnegative distribution summing to 1")
    return values


def cross_entropy(p, q):
    """Pre: equal-length distributions. Post: cross-entropy in nats, including +infinity."""
    p, q = distribution(p), distribution(q)
    if len(p) != len(q):
        raise ValueError("distributions must have matching support dimensions")
    if any(pi > 0 and qi == 0 for pi, qi in zip(p, q)):
        return math.inf
    return -sum(pi*math.log(qi) for pi, qi in zip(p, q) if pi)


def kl(p, q):
    """Pre: equal-length distributions. Post: KL(p||q) in nats."""
    p, q = distribution(p), distribution(q)
    return cross_entropy(p, q) - cross_entropy(p, p)


def softmax(logits):
    """Pre: finite nonempty logits. Post: normalized probabilities using a stable shift."""
    if not logits or any(not math.isfinite(z) for z in logits):
        raise ValueError("logits must be finite and nonempty")
    maximum = max(logits)
    shifted = [math.exp(z-maximum) for z in logits]
    total = sum(shifted)
    return [z/total for z in shifted]


def classification_loss(logits, target):
    """Pre: valid class index. Post: log-sum-exp cross-entropy in nats."""
    if not 0 <= target < len(logits):
        raise ValueError("target outside class range")
    maximum = max(logits)
    return maximum + math.log(sum(math.exp(z-maximum) for z in logits)) - logits[target]


def demo():
    p, q = [.5, .25, .25], [.25, .25, .5]
    print("H(p), CE(p,q), KL(p||q) in nats:",
          cross_entropy(p, p), cross_entropy(p, q), kl(p, q))
    z = [1002., 1001., 1000.]
    probabilities = softmax(z)
    print("Stable softmax:", [round(x, 6) for x in probabilities])
    print("CE gradient at target 0:", [round(x-(i==0), 6) for i, x in enumerate(probabilities)])
    print("KL([1,0] || [.9,.1]):", kl([1, 0], [.9, .1]))
    print("Reverse KL:", kl([.9, .1], [1, 0]))


def self_test():
    count = 0
    # Invariant: prior shifts preserve probabilities/gradients. Bound: remaining shift cases.
    for shift in (0, 1000, -1000):
        logits = [2+shift, 1+shift, shift]
        probabilities = softmax(logits)
        check(all(abs(a-b) < 1e-12 for a, b in zip(probabilities, softmax([2, 1, 0]))), True)
        for i in range(3):
            plus, minus = list(logits), list(logits)
            plus[i] += 1e-5
            minus[i] -= 1e-5
            numerical = (classification_loss(plus, 0)-classification_loss(minus, 0))/2e-5
            check(abs(numerical-(probabilities[i]-(i==0))) < 1e-7, True)
            count += 1
        count += 1
    check(kl([.9, .1], [1, 0]), math.inf)
    check(abs(kl((x for x in (.2, .8)), [.2, .8])) < 1e-12, True)
    return count + 2

def check(actual, expected):
    """Require an expected result; this check also runs under python -O."""
    if actual != expected:
        raise AssertionError(f"expected {expected!r}, got {actual!r}")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        print(f"{self_test()} checks passed")
    else:
        demo()