"""Small dense networks and Adam, shared by PPO and its finite-difference checks."""
import json
import numpy as np
class Dense:
 def __init__(self,model=None,sizes=None,mean=None,std=None,rng=None):
  rng=rng or np.random.default_rng(1)
  if model:
   self.meta=model;self.mean=np.array(model['inputMean'],dtype='f');self.std=np.array(model['inputStd'],dtype='f');self.w=[];self.b=[];width=model['inputDim']
   for layer in model['layers']:
    n=len(layer['bias']);self.w.append(np.array(layer['weight'],dtype='f').reshape(n,width).T.copy());self.b.append(np.array(layer['bias'],dtype='f'));width=n
  else:
   self.meta={};self.mean=mean;self.std=std;self.w=[(rng.standard_normal((a,b))*np.sqrt(1/a)).astype('f') for a,b in zip(sizes,sizes[1:])];self.b=[np.zeros(n,'f') for n in sizes[1:]]
  self.p=[v for pair in zip(self.w,self.b) for v in pair];self.m=[np.zeros_like(v) for v in self.p];self.v=[np.zeros_like(v) for v in self.p];self.steps=0
 def forward(self,x,cache=False):
  a=(x-self.mean)/self.std;acts=[a]
  for i,(w,b) in enumerate(zip(self.w,self.b)):a=a@w+b;a=np.tanh(a) if i<len(self.w)-1 else a;acts.append(a)
  return (a,acts) if cache else a
 def gradients(self,acts,grad):
  result=[]
  for i in range(len(self.w)-1,-1,-1):
   result[0:0]=[acts[i].T@grad,grad.sum(0)]
   if i:grad=(grad@self.w[i].T)*(1-acts[i]**2)
  return result
 def update(self,grads,rate):
  self.steps+=1;norm=np.sqrt(sum(float((g*g).sum()) for g in grads));scale=min(1.,.5/max(1e-8,norm))
  for i,(p,g) in enumerate(zip(self.p,grads)):
   g=g*scale;self.m[i]=.9*self.m[i]+.1*g;self.v[i]=.999*self.v[i]+.001*g*g;p-=rate*(self.m[i]/(1-.9**self.steps))/(np.sqrt(self.v[i]/(1-.999**self.steps))+1e-8)
 def export(self,path,name,metadata):
  model={**self.meta,'name':name,'training':'PPO reinforcement learning after behavior cloning','rl':metadata,'layers':[{'weight':w.T.flatten().tolist(),'bias':b.tolist()} for w,b in zip(self.w,self.b)]}
  path.write_text(json.dumps(model,separators=(',',':')))
def sigmoid(x):return 1/(1+np.exp(-np.clip(x,-30,30)))
def log_prob(mean,z,fire,sd):
 p=sigmoid(mean[:,4]);return -.5*(((z-mean[:,:4])/sd)**2+2*np.log(sd)+np.log(2*np.pi)).sum(1)+fire*np.log(p+1e-8)+(1-fire)*np.log(1-p+1e-8)
def ppo_gradient(mean,z,fire,oldlog,adv,sd,clip=.2):
 logp=log_prob(mean,z,fire,sd);ratio=np.exp(np.clip(logp-oldlog,-15,15));active=np.where(adv>=0,ratio<=1+clip,ratio>=1-clip);factor=-adv*ratio*active/len(adv)
 grad=np.zeros_like(mean);grad[:,:4]=factor[:,None]*(z-mean[:,:4])/(sd**2);grad[:,4]=factor*(fire-sigmoid(mean[:,4]));return grad,float(np.mean(oldlog-logp)),float(np.mean(np.minimum(ratio*adv,np.clip(ratio,1-clip,1+clip)*adv)))
def gae(rewards,values,next_value,dones,gamma=.99,lam=.95):
 adv=np.zeros_like(rewards);last=np.zeros_like(next_value)
 for t in range(len(rewards)-1,-1,-1):
  following=next_value if t==len(rewards)-1 else values[t+1];nonterminal=1-dones[t];delta=rewards[t]+gamma*following*nonterminal-values[t];last=delta+gamma*lam*nonterminal*last;adv[t]=last
 return adv,adv+values

MACRO_HEADS=[(0,11),(11,4),(15,7),(22,9),(31,5),(38,3)]
def macro_probabilities(raw,group_masks):
 result=[]
 for start,width in MACRO_HEADS:
  logits=raw[:,start:start+width].copy()
  if start==22:logits=np.where(group_masks,logits,-1e9)
  p=np.exp(logits-logits.max(1,keepdims=True));p/=p.sum(1,keepdims=True);result.append(p)
 return result
def macro_log_prob(raw,categories,z,masks,sd=.3):
 probabilities=macro_probabilities(raw,masks);log=np.zeros(len(raw),dtype=raw.dtype)
 for i,p in enumerate(probabilities):log+=np.log(p[np.arange(len(p)),categories[:,i]]+1e-8)
 return log-.5*(((z-raw[:,36:38])/sd)**2+2*np.log(sd)+np.log(2*np.pi)).sum(1)
def macro_ppo_gradient(raw,categories,z,masks,oldlog,adv,sd=.3,clip=.2):
 log=macro_log_prob(raw,categories,z,masks,sd);ratio=np.exp(np.clip(log-oldlog,-15,15));active=np.where(adv>=0,ratio<=1+clip,ratio>=1-clip);factor=-adv*ratio*active/len(raw);grad=np.zeros_like(raw)
 for h,((start,width),p) in enumerate(zip(MACRO_HEADS,macro_probabilities(raw,masks))):
  dp=-p;dp[np.arange(len(p)),categories[:,h]]+=1;grad[:,start:start+width]=dp*factor[:,None]
 grad[:,36:38]=factor[:,None]*(z-raw[:,36:38])/sd**2
 return grad,float(np.mean(oldlog-log))