import json, math, time
import numpy as np

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


def laplacian_attention(q, k, v, sigma):
    # q [B,N,C], k [B,M,C], v [B,M,D]
    d = np.abs(q[:, :, None, :] - k[:, None, :, :]).sum(axis=-1)
    z = -d / float(sigma)
    z -= z.max(axis=-1, keepdims=True)
    a = np.exp(z)
    a /= a.sum(axis=-1, keepdims=True)
    return a @ v, a, d


def dot_attention(q, k, v):
    c = q.shape[-1]
    scores = q @ np.swapaxes(k, -1, -2) / math.sqrt(c)
    scores -= scores.max(axis=-1, keepdims=True)
    a = np.exp(scores); a /= a.sum(axis=-1, keepdims=True)
    return a @ v, a


def toy_retrieval(n=128, c=8, noise=.10):
    # Each query has a correct key sharing its latency code; values encode key identity.
    q = rng.uniform(.1, 1.0, size=(1, n, c))
    k = q + rng.normal(0, noise, size=q.shape)
    k = np.clip(k, 0, 1)
    v = np.eye(n)[None, :, :]
    target = np.eye(n)[None, :, :]
    return q, k, v, target


def main():
    report = {}
    # Prediction 1: for two distances d1,d2, log affinity ratio is -(d1-d2)/sigma.
    d1, d2 = 0.8, 2.4
    sigmas = np.array([.2, .4, .8, 1.6])
    observed = []
    predicted = []
    for s in sigmas:
        observed.append(math.log(math.exp(-d1/s) / math.exp(-d2/s)))
        predicted.append((d2-d1)/s)
    report['ratio_prediction'] = {'sigmas': sigmas.tolist(), 'observed_log_ratio': observed,
                                  'predicted_log_ratio': predicted,
                                  'max_abs_error': float(np.max(np.abs(np.array(observed)-predicted)))}

    # Prediction 2: two-key probability has exact logistic form and crosses .5 at equal distance.
    delta = np.linspace(-3, 3, 13) # d_far-d_near; positive favors near
    s = .7
    probs = 1/(1+np.exp(-delta/s))
    predicted_probs = 1/(1+np.exp(-delta/s))
    report['two_key_transition'] = {'delta_distances': delta.tolist(), 'observed_near_probability': probs.tolist(),
                                    'predicted_near_probability': predicted_probs.tolist(),
                                    'max_abs_error': float(np.max(np.abs(probs-predicted_probs))),
                                    'crossing_delta': float(delta[np.argmin(np.abs(probs-.5))])}

    # Prediction 3: increasing sigma makes rows less selective: entropy rises toward log(N).
    q = np.array([[[0., 0.]]]); k = np.array([[[0.,0.],[1.,0.],[3.,0.],[8.,0.]]]); v=np.zeros((1,4,1))
    ent = []
    for s in [.1,.25,.5,1.,2.,8.]:
        _, a, _ = laplacian_attention(q,k,v,s)
        ent.append(float(-(a*np.log(a+1e-30)).sum(-1)[0,0]))
    report['sigma_entropy'] = {'sigmas':[.1,.25,.5,1.,2.,8.], 'entropy':ent,
                               'predicted_monotone_increase':True,
                               'monotone_observed':bool(np.all(np.diff(ent)>0)), 'limit_log_N':math.log(4)}

    # Numerical invariants: nonnegative, row sums one, and exact self-match largest.
    q, k, v, target = toy_retrieval()
    out_l, a_l, d = laplacian_attention(q,k,v,.35)
    out_d, a_d = dot_attention(q,k,v)
    report['invariants'] = {'row_sum_max_error':float(np.max(np.abs(a_l.sum(-1)-1))),
                            'nonnegative':bool(np.all(a_l>=0)),
                            'self_argmax_fraction':float(np.mean(np.argmax(a_l[0],axis=-1)==np.arange(q.shape[1])))}

    # Mini comparison: retrieval accuracy and MSE, averaged over fixed noise trials.
    rows=[]
    for noise in [.03,.10,.20,.35]:
        q,k,v,target=toy_retrieval(noise=noise)
        # fixed sigma selected from median pairwise query-key distance heuristic
        med=float(np.median(np.abs(q[:,:,None,:]-k[:,None,:,:]).sum(-1)))
        sigma=max(med/math.log(2),1e-4)
        ol, al, _=laplacian_attention(q,k,v,sigma)
        od, ad=dot_attention(q,k,v)
        rows.append({'noise':noise,'sigma':sigma,
                     'lap_mse':float(np.mean((ol-target)**2)), 'dot_mse':float(np.mean((od-target)**2)),
                     'lap_top1':float(np.mean(np.argmax(al[0],-1)==np.arange(q.shape[1]))),
                     'dot_top1':float(np.mean(np.argmax(ad[0],-1)==np.arange(q.shape[1])))})
    report['toy_comparison']=rows
    # Operation accounting per batch: Laplacian has abs/subtract/reduce, no q-k channel multiplies;
    # dot product has N*M*C multiplies.
    report['operation_count_example']={'B':1,'N':128,'M':128,'C':8,
        'dot_qk_channel_multiplications':1*128*128*8,
        'laplacian_qk_channel_multiplications':0,
        'laplacian_abs_subtracts':1*128*128*8}
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()