"""Two-client, two-round scalar simulation; no distributed runtime or privacy mechanism."""
from fractions import Fraction as F
import json
CLIENTS={'A':[F(0),F(2)],'B':[F(4)]*6}
def text(x):return str(x)
def local_update(w,ys):
    gradient=sum(w-y for y in ys)/len(ys)
    return w-F(1,2)*gradient,gradient
def demo():
    w=F(0);rounds=[];all_values=sum(CLIENTS.values(),[])
    for number in [1,2]:
        messages=[];weights=[]
        for name,ys in CLIENTS.items():
            updated,gradient=local_update(w,ys);n=len(ys);weights.append((updated,n))
            messages.append({'client':name,'server_to_client':{'scalar_weight':text(w)},'local_examples':n,'local_gradient':text(gradient),'client_to_server':{'updated_weight':text(updated),'example_count':n}})
        weighted=sum(value*n for value,n in weights)/sum(n for value,n in weights)
        unweighted=sum(value for value,n in weights)/len(weights)
        pooled,gradient=local_update(w,all_values)
        assert weighted==pooled
        rounds.append({'round':number,'server_before':text(w),'messages':messages,'weighted_server_after':text(weighted),'unweighted_alternative':text(unweighted),'pooled_one_step_reference':text(pooled),'weighted_equals_pooled_only_for_this_one_step_setup':True})
        w=weighted
    assert [r['weighted_server_after'] for r in rounds]==['13/8','39/16']
    assert [r['unweighted_alternative'] for r in rounds]==['5/4','33/16']
    return {'client_training_values':{k:[text(x) for x in v] for k,v in CLIENTS.items()},'initial_weight':'0','learning_rate':'1/2','local_steps_per_round':1,'loss':'one half mean squared error for constant predictor w','rounds':rounds,'one_process_simulation':True,'actual_devices_or_network':False,'secure_aggregation_or_DP':False,'TFF_executed':False}
if __name__=='__main__':print(json.dumps(demo(),ensure_ascii=False,indent=2))