"""Course 13 CPU studio. Python 3.12+, standard library only for planar training.
Reference and learner modes are separate. No learner text is evaluated.
Use --check to check reference math; --mode learner intentionally needs implementation.
The richer MuJoCo bridge remains in controller-simulator.zip with pinned requirements.
"""
import argparse
import hashlib
import sys
import json
import math
from pathlib import Path

PARAMS = 110

def randoms(seed=913):
    while True:
        seed = (1664525 * seed + 1013904223) & 0xffffffff
        yield seed / 4294967296

def initial():
    r = randoms()
    return [.15 * (2 * next(r) - 1) for _ in range(PARAMS)]

def forward(w, x):
    h = [math.tanh(w[48+j] + sum(x[i]*w[i*8+j] for i in range(6))) for j in range(8)]
    y = [math.tanh(w[104+k] + sum(h[j]*w[56+j*6+k] for j in range(8))) for k in range(6)]
    return h, y

def examples():
    rows=[]
    for i in range(8):
        side = 1 if i%2 else -1
        goal = [.75*side, (i%4-1.5)*.12]
        p = [-side*.65, (i%4-1.5)*.22]
        episode=[]
        for _ in range(20):
            a=[max(-1,min(1,3*(goal[j]-p[j]))) for j in range(2)]
            episode.append((p[:],a))
            p=[p[j]+.08*a[j] for j in range(2)]
        for t,(p,_) in enumerate(episode):
            x=p+goal+[int(side==1),int(side==-1)]
            targets=sum([episode[min(t+k,19)][1] for k in range(3)],[])
            rows.append((x,targets))
    return rows

def loss_gradient(w, rows):
    g=[0.]*PARAMS
    loss=0.
    for x,target in rows:
        h,y=forward(w,x)
        dh=[0.]*8
        for k in range(6):
            error=y[k]-target[k]
            loss+=error*error/2
            dy=error*(1-y[k]*y[k])
            g[104+k]+=dy
            for j in range(8):
                g[56+j*6+k]+=dy*h[j]
                dh[j]+=dy*w[56+j*6+k]
        for j in range(8):
            dz=dh[j]*(1-h[j]*h[j])
            g[48+j]+=dz
            for i in range(6):
                g[i*8+j]+=dz*x[i]
    n=6*len(rows)
    return loss/n,[v/n for v in g]

def learner_gradient(weights, rows):
    # TODO: implement the same objective using your own forward/backward code.
    # Do not call loss_gradient. Preserve duplicate contributions and simultaneous updates.
    raise NotImplementedError("Complete learner_gradient, then explain a finite-difference test")

def checks(gradient=loss_gradient):
    w=initial()
    # Public miniature case separate from the demonstration trajectories.
    rows=[([.2,-.4,.7,.1,1,0],[.3,-.2,.1,.4,-.3,.2])]
    _,g=gradient(w,rows)
    assert len(g)==110
    for i in range(110):
        a=w[:]; b=w[:]; a[i]+=1e-5; b[i]-=1e-5
        finite=(loss_gradient(a,rows)[0]-loss_gradient(b,rows)[0])/2e-5
        assert abs(finite-g[i])<1e-8, (i,finite,g[i])
    # Independent known value: all zero weights -> outputs zero, bias gradient = -target/6.
    _,zero=gradient([0.]*110,rows)
    assert abs(zero[104]+.3/6)<1e-12
    assert abs(zero[105]-.2/6)<1e-12
    assert all(v==0 for v in zero[:104])
    print("110 gradients and independent zero-network oracle passed; public software checks only")

def main():
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('--mode',choices=['reference','learner'],default='reference')
    p.add_argument('--check',action='store_true')
    p.add_argument('--updates',type=int,default=100)
    p.add_argument('--output',type=Path)
    args=p.parse_args()
    gradient=loss_gradient if args.mode=='reference' else learner_gradient
    if args.check: checks(gradient)
    if args.output:
        if not 0<=args.updates<=400: raise ValueError('Use 0–400 updates')
        weights=initial(); rows=examples()
        for _ in range(args.updates):
            _,g=gradient(weights,rows)
            proposed=[w-.6*d for w,d in zip(weights,g)]
            if len(proposed)!=110 or any(not math.isfinite(w) or abs(w)>8 for w in proposed):
                raise ValueError('Rejected update; preceding checkpoint retained')
            weights=proposed
        with args.output.open('x',encoding='utf8') as f:
            json.dump({'format':'robot-learning-studio-v1','source_sha256':hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),'python_version':sys.version,'provenance':'Original LearnAI teaching code and synthetic data; see expanded-study-NOTICE.txt; no separate open-source license asserted','mode':args.mode,'policy':{'weights':weights,'updates':args.updates},'trainingRows':len(rows),'loss':loss_gradient(weights,rows)[0],'humanReview':'pending','limits':'Authored ideal-contact planar model; no hardware or richer-embodiment result'},f,indent=2,allow_nan=False)
        print('Saved',args.output)

if __name__=='__main__': main()