"""Three-token self-attention permutation experiment; no learned model."""
import math,json
X=[[1.,0.,0.,0.],[0.,1.,0.,0.],[0.,0.,1.,0.]]
PERM=[2,0,1]
def pe(pos):return [math.sin(pos),math.cos(pos),math.sin(pos/100),math.cos(pos/100)]
def add(a,b):return [x+y for x,y in zip(a,b)]
def attention(rows):
    outputs=[]
    for q in rows:
        scores=[sum(x*y for x,y in zip(q,k))/2 for k in rows]
        peak=max(scores);weights=[math.exp(s-peak) for s in scores];total=sum(weights);weights=[w/total for w in weights]
        outputs.append([sum(w*v[j] for w,v in zip(weights,rows)) for j in range(4)])
    return outputs
def align(permuted):return [permuted[PERM.index(i)] for i in range(3)]
def max_difference(a,b):return max(abs(x-y) for rowa,rowb in zip(a,b) for x,y in zip(rowa,rowb))
def experiment():
    P=[pe(i) for i in range(3)];permuted=[X[i] for i in PERM]
    plain=attention(X);plain_permuted=align(attention(permuted))
    with_position=attention([add(x,p) for x,p in zip(X,P)])
    with_position_permuted=align(attention([add(x,p) for x,p in zip(permuted,P)]))
    # A simple pooled sum before attention is NOT a valid order witness here.
    sum1=[sum(x[j]+p[j] for x,p in zip(X,P)) for j in range(4)]
    sum2=[sum(x[j]+p[j] for x,p in zip(permuted,P)) for j in range(4)]
    assert max(abs(a-b) for a,b in zip(sum1,sum2))<1e-12
    assert max_difference(plain,plain_permuted)<1e-12
    assert max_difference(with_position,with_position_permuted)>.01
    assert pe(0)==[0.,1.,0.,1.]
    return {'symbolic_token_vectors_A_B_C':X,'sinusoidal_position_vectors':P,'new_sequence_original_indices':PERM,'all_comparison_rows_aligned_to':['A','B','C'],'Q_K_V':'all identity transforms on supplied rows','scale_sqrt_d':2,'causal_mask':False,'attention_without_PE':plain,'permuted_without_PE_aligned':plain_permuted,'attention_with_PE':with_position,'permuted_with_PE_aligned':with_position_permuted,'no_PE_max_difference':max_difference(plain,plain_permuted),'with_PE_max_difference':max_difference(with_position,with_position_permuted),'pooled_sum_before_attention_equal':True,'learned_weights_or_real_tokenizer':False}
if __name__=='__main__':print(json.dumps(experiment(),indent=2))