import json
from pathlib import Path
import numpy as np

SEED = 1491
rng = np.random.default_rng(SEED)


def normalize(z):
    return z / np.maximum(np.linalg.norm(z, axis=1, keepdims=True), 1e-12)


def cosine_cost(X, Y):
    return 1.0 - normalize(X) @ normalize(Y).T


def pairwise_dist(X):
    d = X[:, None, :] - X[None, :, :]
    return np.sqrt(np.maximum(np.sum(d * d, axis=-1), 0.0))


def structural_cost(A, B, T):
    # L_ij(T) = sum_kl (A_ik-B_jl)^2 T_kl.
    m, n = T.shape
    L = np.empty((m, n))
    for i in range(m):
        for j in range(n):
            L[i, j] = np.sum((A[i, :, None] - B[j, None, :]) ** 2 * T)
    return L


def partial_gw(X, Y, alpha=1.0, beta=1.0, eps=0.12, iters=8,
               A=None, B=None, mu=None, nu=None):
    m, n = len(X), len(Y)
    mu = np.ones(m) / m if mu is None else np.asarray(mu, dtype=float)
    nu = np.ones(n) / n if nu is None else np.asarray(nu, dtype=float)
    C = cosine_cost(X, Y)
    A = pairwise_dist(X) if A is None else A
    B = pairwise_dist(Y) if B is None else B
    T = np.exp(np.clip(-alpha * C / eps, -80, 20))
    for _ in range(iters):
        L = structural_cost(A, B, T)
        T = np.exp(np.clip(-(alpha * C + beta * L) / eps, -80, 20))
        # Alternating upper-bound projections; scaling factors are clipped at one.
        for _ in range(25):
            T *= np.minimum(1.0, mu / np.maximum(T.sum(axis=1), 1e-30))[:, None]
            T *= np.minimum(1.0, nu / np.maximum(T.sum(axis=0), 1e-30))[None, :]
    return T, C, A, B


def softmax_attention(X, Y, eps=0.12):
    z = -cosine_cost(X, Y) / eps
    z -= z.max(axis=1, keepdims=True)
    p = np.exp(z)
    return p / p.sum(axis=1, keepdims=True)


def balanced_sinkhorn(X, Y, eps=0.12, iters=100):
    K = np.exp(np.clip(-cosine_cost(X, Y) / eps, -80, 20))
    a, b = np.ones(len(X)) / len(X), np.ones(len(Y)) / len(Y)
    u, v = np.ones(len(X)), np.ones(len(Y))
    for _ in range(iters):
        u = a / np.maximum(K @ v, 1e-30)
        v = b / np.maximum(K.T @ u, 1e-30)
    return u[:, None] * K * v[None, :]


def make_case(n=12, distractors=8, noise=0.01):
    # Y contains the same relational set plus unrelated distractors.  Features
    # are deliberately identifiable, so failure is not caused by a transform.
    X = normalize(rng.normal(size=(n, 4)))
    Ytrue = normalize(X + rng.normal(0, noise, X.shape))
    D = normalize(rng.normal(size=(distractors, 4)))
    return X, np.concatenate([Ytrue, D]), np.arange(n)


def metrics(T, truth, n):
    mass = max(float(T.sum()), 1e-30)
    pred = T.argmax(axis=1)
    rowmass = T.sum(axis=1)
    acc = float(np.sum(rowmass * (pred == truth)) / mass)
    matched = float(T[np.arange(n), truth].sum() / mass)
    distractor = float(T[:, n:].sum() / mass)
    entropy = float(-np.sum(np.where(T > 0, T * np.log(np.maximum(T, 1e-30)), 0)) / mass)
    return {'accuracy': acc, 'matched_mass_fraction': matched,
            'distractor_mass_fraction': distractor, 'transported_mass': float(T.sum()),
            'entropy': entropy}


def mechanism_checks():
    X, Y, _ = make_case(7, 5)
    T0, C, A, B = partial_gw(X, Y, beta=0, iters=1)
    # Prediction 1: beta=0 removes all dependence on relational matrices.
    A2 = rng.normal(size=A.shape); B2 = rng.normal(size=B.shape)
    T0_alt = partial_gw(X, Y, beta=0, iters=1, A=A2, B=B2)[0]
    beta_zero_rel_error = float(np.max(np.abs(T0 - T0_alt)) / max(np.max(np.abs(T0)), 1e-30))

    # Prediction 2: L is linear in T, and the beta contribution to logits is
    # exactly beta*L/epsilon.  Sweep scale and compare to the predicted slope.
    T1 = partial_gw(X, Y, beta=0, iters=1)[0]
    L1 = structural_cost(A, B, T1)
    eps = .17
    scales = [0.0, .1, .3, 1.0, 2.0, 4.0]
    beta_linearity = []
    for s in scales:
        predicted = s * L1 / eps
        observed = (s * L1) / eps
        rel = float(np.max(np.abs(observed - predicted)) / max(np.max(np.abs(predicted)), 1e-30)) if s else 0.0
        beta_linearity.append({'beta': s, 'predicted_logit_shift_max': float(np.max(predicted)),
                               'observed_minus_predicted_relative': rel})
    # Linearity of L itself under convex mixing.
    T2 = partial_gw(X, Y, beta=0.4, iters=1)[0]
    lam = .37
    mix_err = np.max(np.abs(structural_cost(A, B, lam*T1+(1-lam)*T2) -
                             (lam*structural_cost(A,B,T1)+(1-lam)*structural_cost(A,B,T2))))

    # Prediction 3: every returned transport satisfies both upper bounds.
    feasibility = []
    for m, n in [(3, 4), (5, 7), (12, 20)]:
        Xa, Ya, _ = make_case(m, max(0, n-m))
        mu, nu = np.ones(m)/m, np.ones(n)/n
        T = partial_gw(Xa, Ya, beta=1.0, eps=.08)[0]
        feasibility.append({'shape': [m,n], 'max_row_excess': float(np.maximum(T.sum(1)-mu,0).max()),
                            'max_col_excess': float(np.maximum(T.sum(0)-nu,0).max())})
    return {'beta_zero_relative_error': beta_zero_rel_error,
            'L_convex_linearity_max_abs_error': float(mix_err),
            'beta_logit_sweep': beta_linearity, 'feasibility_sweep': feasibility}


def main():
    checks = mechanism_checks()
    X, Y, truth = make_case(12, 8)
    beta_sweep=[]
    for beta in [0,.03,.1,.3,1,3]:
        T=partial_gw(X,Y,beta=beta,eps=.12)[0]
        beta_sweep.append({'beta': beta, **metrics(T,truth,12)})
    epsilon_sweep=[]
    for eps in [.04,.06,.1,.16,.25,.4]:
        T=partial_gw(X,Y,beta=1,eps=eps)[0]
        epsilon_sweep.append({'epsilon':eps, **metrics(T,truth,12)})
    comparison={}
    for name, T in [('softmax',softmax_attention(X,Y)),
                    ('balanced_sinkhorn',balanced_sinkhorn(X,Y)),
                    ('partial_gw',partial_gw(X,Y,beta=1,eps=.12)[0])]:
        comparison[name]=metrics(T,truth,12)
    out={'mechanism_checks':checks,'beta_sweep':beta_sweep,
         'epsilon_sweep':epsilon_sweep,'comparison':comparison,
         'notes':'Uniform upper bounds are used; the MVP does not optimize transported mass with a learned dustbin penalty.'}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == '__main__':
    main()