import json
import numpy as np


def center(x):
    return x - x.mean(0, keepdims=True) - x.mean(1, keepdims=True) + x.mean()


def sinkhorn(cost, eps, s, r, iters):
    logk = -cost / eps
    lu = np.zeros(len(s)); lv = np.zeros(len(r))
    for _ in range(iters):
        lu = np.log(s) - np.logaddexp.reduce(logk + lv[None, :], axis=1)
        lv = np.log(r) - np.logaddexp.reduce(logk.T + lu[None, :], axis=1)
    return np.exp(lu[:, None] + logk + lv[None, :])


def inverse_cost(w, eps, delta=0.0):
    return -eps * center(np.log(w + delta))


def math_checks(seed=7):
    rng = np.random.default_rng(seed)
    m, n, eps = 9, 7, 0.7
    c = rng.normal(size=(m, n))
    s = rng.random(m); s /= s.sum()
    r = rng.random(n); r /= r.sum()
    w = sinkhorn(c, eps, s, r, 300)
    d = np.linalg.norm(center(c))
    exact = np.linalg.norm(center(c) - inverse_cost(w, eps)) / d
    marg = max(np.max(abs(w.sum(1)-s)), np.max(abs(w.sum(0)-r)))
    g = rng.normal(size=m)[:, None] + rng.normal(size=n)[None, :]
    wg = sinkhorn(c + g, eps, s, r, 300)
    gauge_plan = np.linalg.norm(w-wg) / np.linalg.norm(w)
    gauge_inverse = np.linalg.norm(inverse_cost(wg, eps)-inverse_cost(w, eps)) / d
    convergence = []
    for it in [3, 5, 10, 30, 100]:
        wi = sinkhorn(c, eps, s, r, it)
        convergence.append({
            'iters': it,
            'marginal_error': float(max(np.max(abs(wi.sum(1)-s)), np.max(abs(wi.sum(0)-r)))),
            'recovery_error': float(np.linalg.norm(center(c)-inverse_cost(wi, eps))/d),
            'min_w': float(wi.min())})
    floor = []
    for delta in [0.0, 1e-12, 1e-8, 1e-5, 1e-3]:
        floor.append({'delta': delta, 'recovery_error': float(np.linalg.norm(center(c)-inverse_cost(w, eps, delta))/d)})
    return {'exact_recovery_error': float(exact), 'marginal_error': float(marg),
            'gauge_plan_relative_change': float(gauge_plan),
            'gauge_inverse_relative_change': float(gauge_inverse),
            'convergence': convergence, 'floor': floor, 'min_w_300': float(w.min())}


def attention_benchmark(seed=11, trials=300):
    rng = np.random.default_rng(seed)
    n, d, eps = 12, 16, 0.5
    methods = ['row_softmax', 'sinkhorn', 'sinkhorn_inverse_penalty']
    scores = {x: [] for x in methods}; entropy = {x: [] for x in methods}
    for _ in range(trials):
        q = rng.normal(size=(n, d)); k = rng.normal(size=(n, d)); v = rng.normal(size=(n, d))
        target = np.argmax(q @ k.T / np.sqrt(d), axis=1)
        raw = q @ k.T / np.sqrt(d)
        row = np.exp(raw-raw.max(1,keepdims=True)); row /= row.sum(1,keepdims=True)
        w = sinkhorn(-raw, eps, np.ones(n)/n, np.ones(n)/n, 30)
        # The inverse-cost consistency term is evaluated, not optimized: it is zero
        # for a plan generated by the same kernel, isolating geometry rather than tuning.
        inv = inverse_cost(w, eps)
        structured = center(raw)
        penalty = np.mean((structured - inv)**2)
        for name, a in [('row_softmax', row), ('sinkhorn', w), ('sinkhorn_inverse_penalty', w)]:
            pred = np.argmax(a, axis=1)
            scores[name].append(np.mean(pred == target))
            entropy[name].append(float(-np.mean(np.sum(a*np.log(a+1e-30),1))))
    return {m: {'retrieval_accuracy': float(np.mean(scores[m])),
                'entropy': float(np.mean(entropy[m]))} for m in methods}


if __name__ == '__main__':
    out = {'math': math_checks(), 'attention': attention_benchmark()}
    print(json.dumps(out, indent=2))