"""Deterministic symmetry experiment, not an initialization benchmark."""
import math,json
def forward(p):
    h=[1/(1+math.exp(-(p['w'+str(i)]+p['b'+str(i)]))) for i in range(2)]
    out=sum(p['v'+str(i)]*h[i] for i in range(2))+p['c']
    return h,out,.5*(out-1)**2
def gradients(p):
    h,out,_=forward(p);e=out-1;g={'c':e}
    for i in range(2):
        s=str(i);g['v'+s]=e*h[i];g['w'+s]=g['b'+s]=e*p['v'+s]*h[i]*(1-h[i])
    return g
def experiment(p):
    rows=[]
    for step in range(3):
        before=p.copy();h,out,loss=forward(p);g=gradients(p)
        for key in p:
            plus=p.copy();minus=p.copy();plus[key]+=1e-6;minus[key]-=1e-6
            numeric=(forward(plus)[2]-forward(minus)[2])/2e-6
            assert abs(numeric-g[key])<1e-8
        p={key:p[key]-.1*g[key] for key in p}
        rows.append({'step':step+1,'before':before,'hidden':h,'output':out,'loss':loss,'gradients':g,'after':p.copy()})
    return rows
if __name__=='__main__':
    start={'w0':.2,'w1':.2,'b0':0.,'b1':0.,'v0':.5,'v1':.5,'c':0.}
    symmetric=experiment(start.copy());asym_start={**start,'w1':-.2};asymmetric=experiment(asym_start)
    for r in symmetric:
        assert r['hidden'][0]==r['hidden'][1]
        for prefix in ('w','b','v'):assert r['after'][prefix+'0']==r['after'][prefix+'1']
    assert asymmetric[0]['hidden'][0]!=asymmetric[0]['hidden'][1]
    assert asymmetric[0]['gradients']['v0']!=asymmetric[0]['gradients']['v1']
    print(json.dumps({'input':1,'target':1,'learning_rate':.1,'symmetric':symmetric,'asymmetric':asymmetric,'xavier_uniform_gain1_fanin4_fanout2_bound':math.sqrt(6/6),'he_normal_relu_fanin4_theoretical_std':math.sqrt(2/4),'random_sampling':False,'PyTorch_run':False},indent=2))