"""Pooling on fabricated token vectors; no transformer/model loaded."""
from fractions import Fraction as F
import math,json
X=[[2,0],[0,2],[2,2],[100,100]];MASK=[1,1,1,0]
def mean_pool(rows,mask):
    if not rows or len(rows)!=len(mask) or not rows[0] or any(len(r)!=len(rows[0]) for r in rows):raise ValueError('rectangular nonempty rows and same-length mask')
    if any(type(m) is not int or m not in (0,1) for m in mask):raise ValueError('binary integer mask')
    if not all(math.isfinite(x) for r in rows for x in r):raise ValueError('finite values')
    count=sum(mask)
    if not count:return None
    return [sum(F(r[j])*m for r,m in zip(rows,mask))/count for j in range(len(rows[0]))]
def unit(v):
    norm=math.sqrt(sum(float(x)**2 for x in v))
    if not norm:return None
    return [float(x)/norm for x in v]
if __name__=='__main__':
    masked=mean_pool(X,MASK);plain=mean_pool(X,[1,1,1,1]);content=mean_pool(X,[0,1,1,0]);cls=X[0]
    assert masked==[F(4,3),F(4,3)] and plain==[26,26] and content==[1,2] and cls==[2,0]
    assert mean_pool(X[:-1]+[[1000000,1000000]],MASK)==masked
    assert mean_pool([[3,0]]+X[1:],MASK)!=masked
    assert mean_pool(X,[0,0,0,0]) is None
    assert all(math.isclose(x,1/math.sqrt(2)) for x in unit(masked))
    print(json.dumps({'masked_mean_including_CLS':[str(x) for x in masked],'unmasked_mean':[str(x) for x in plain],'content_only_mean':[str(x) for x in content],'symbolic_CLS_vector':cls,'L2_after_masked_mean':unit(masked),'real_model_loaded':False},indent=2))
