"""Reproduce the bounded mathematical examples added in EEML v0.25.
Python 3 standard library only. These checks are not empirical validation.
"""
import itertools
import json
import math

results = {}
def record(name, condition, **values):
    if not condition:
        raise AssertionError((name, values))
    results[name] = values or {"passed": True}

def probs(costs, eps):
    lowest = min(costs)
    weights = [math.exp(-(s-lowest)/eps) for s in costs]
    total = sum(weights)
    return [w/total for w in weights]

# Stable log-sum-exp and tied minima, including negative costs in the finite lemma.
max_error = 0.0
for costs in ([2, 2, 5], [-3, 1, 7], [4]):
    for eps in (2, .5, .01, .0001):
        m = min(costs)
        f = m - eps * math.log(sum(math.exp(-(s-m)/eps) for s in costs))
        assert m-eps*math.log(len(costs))-1e-12 <= f <= m+1e-12
        if eps == .0001:
            max_error = max(max_error, abs(f-m))
p = probs([2, 2, 5], .001)
record('finite_epsilon_bridge', max_error < .0002 and p[:2] == [.5, .5], max_limit_error=max_error, tied_minima=p)

# Enumerate all 32 full states independently of the effective expression.
xs = list(itertools.product((-1, 1), repeat=4))
eps = .7
J = .3
full = [sum(math.exp(-(4*abs(J)-J*y*sum(x))/eps) for y in (-1,1)) for x in xs]
seff = [4*abs(J)-eps*math.log(2*math.cosh(J*sum(x)/eps)) for x in xs]
effective = [math.exp(-s/eps) for s in seff]
pmarg = [w/sum(full) for w in full]
peff = [w/sum(effective) for w in effective]
marginal_error = max(abs(a-b) for a,b in zip(pmarg, peff))
# Deleting y removes the interactions and gives a uniform visible distribution.
deleted_error = max(abs(p-1/16) for p in pmarg)
record('hidden_spin_marginal', marginal_error < 1e-14 and deleted_error > .01, exact_error=marginal_error, naive_deletion_error=deleted_error)
# Walsh projection isolates the genuine four-body coefficient.
j = .01
s = [-math.log(2*math.cosh(j*sum(x))) for x in xs]
c4 = sum(v*math.prod(x) for v,x in zip(s,xs))/16
leading = 2*j**4
record('four_way_coupling', c4 > 0 and abs(c4/leading-1) < .002, coefficient=c4, leading_order=leading)

# KL/free-energy identity for arbitrary distributions, including zero mass.
costs = [0, 1, 3]
q = probs(costs, eps)
z = sum(math.exp(-s/eps) for s in costs)
errors=[]
for p in ([.2,.3,.5], [0,1,0], q):
    f = sum(a*s for a,s in zip(p,costs)) + eps*sum(a*math.log(a) for a in p if a)
    kl = sum(a*math.log(a/b) for a,b in zip(p,q) if a)
    errors.append(abs(f+eps*math.log(z)-eps*kl))
    assert kl >= -1e-14
record('relative_entropy_identity', max(errors) < 1e-14, max_error=max(errors))

# Response checked against numerical differentiation of the full distribution.
A=[-1,2,4]; B=[2,0,-1]; h=.2
mean=lambda v,p:sum(a*b for a,b in zip(v,p))
def response_mean(h):return mean(A,probs([s-h*b for s,b in zip(costs,B)],eps))
p=probs([s-h*b for s,b in zip(costs,B)],eps)
cov=mean([a*b for a,b in zip(A,B)],p)-mean(A,p)*mean(B,p)
step=1e-5
derivative=(response_mean(h+step)-response_mean(h-step))/(2*step)
record('static_response', abs(derivative-cov/eps) < 1e-8, finite_difference=derivative, covariance_over_epsilon=cov/eps)

# Full ODE solved numerically, compared with the exact reduced memory relation.
# x'=y, y'=-x-2y, x(0)=1, y(0)=0; solution x=(1+t)e^-t.
def rhs(v):return (v[1],-v[0]-2*v[1])
v=(1.0,0.0); dt=.001
for _ in range(2000):
    k1=rhs(v);k2=rhs(tuple(v[i]+dt*k1[i]/2 for i in (0,1)))
    k3=rhs(tuple(v[i]+dt*k2[i]/2 for i in (0,1)))
    k4=rhs(tuple(v[i]+dt*k3[i] for i in (0,1)))
    v=tuple(v[i]+dt*(k1[i]+2*k2[i]+2*k3[i]+k4[i])/6 for i in (0,1))
t=2.0;n=2000;dx=t/n
integrand=lambda s:math.exp(-2*(t-s))*(1+s)*math.exp(-s)
integ=dx/3*(integrand(0)+integrand(t)+sum((4 if i%2 else 2)*integrand(i*dx) for i in range(1,n)))
record('memory_redescription', abs(v[0]-(1+t)*math.exp(-t)) < 1e-10 and abs(v[1]+integ) < 1e-10, full_x=v[0], reduced_x=(1+t)*math.exp(-t), derivative_residual=abs(v[1]+integ), memoryless_x=1.0)

# Joint consistency cannot be recovered from separately feasible pair models.
states=list(itertools.product((0,1),repeat=3))
constraints=(lambda a,b,c:a==b, lambda a,b,c:b==c, lambda a,b,c:a!=c)
record('joint_consistency', all(any(f(*v) for v in states) for f in constraints) and not any(all(f(*v) for f in constraints) for v in states))

# Euler bookkeeping, conditional on the geometric fixed-point and resolution facts.
chi_torus=sum((-1)**k*math.comb(4,k) for k in range(5))
fixed_points=2**4
chi_quotient=(chi_torus+fixed_points)//2
chi_resolution=chi_quotient+fixed_points*(2-1)
record('kummer_bookkeeping', (chi_torus,chi_quotient,chi_resolution)==(0,8,24), torus=chi_torus, quotient=chi_quotient, resolution=chi_resolution)
print(json.dumps({"scope":"Mathematical example checks only; no empirical validation or new Lean verification", "checks":results}, indent=2))