"""Current-coded integer IF hidden layer and analog spike-count readout; no KNP claim."""
import numpy as np
CAP=4.0
SCALE=4096

def normalize(image):
    """Bbox>16 to20px bilinear, integer center-of-mass shift to13.5; sharedJSalgorithm."""
    image=np.asarray(image,dtype=np.uint8).reshape(28,28)
    ys,xs=np.nonzero(image>16)
    if len(xs)<2: return np.zeros((28,28),np.uint8)
    x0,x1=int(xs.min()),int(xs.max()); y0,y1=int(ys.min()),int(ys.max())
    cw,ch=x1-x0+1,y1-y0+1; ratio=20/max(cw,ch)
    nw=max(1,int(np.floor(cw*ratio+.5))); nh=max(1,int(np.floor(ch*ratio+.5)))
    xx=np.clip((np.arange(nw)+.5)*cw/nw-.5,0,cw-1)
    yy=np.clip((np.arange(nh)+.5)*ch/nh-.5,0,ch-1)
    lx=xx.astype(int); ly=yy.astype(int); hx=np.minimum(lx+1,cw-1); hy=np.minimum(ly+1,ch-1)
    ax=xx-lx; ay=yy-ly; crop=image[y0:y1+1,x0:x1+1].astype(float)
    patch=(crop[ly[:,None],lx]*(1-ax)+crop[ly[:,None],hx]*ax)*(1-ay[:,None])+(crop[hy[:,None],lx]*(1-ax)+crop[hy[:,None],hx]*ax)*ay[:,None]
    out=np.zeros((28,28),np.uint8); ox=(28-nw)//2; oy=(28-nh)//2
    out[oy:oy+nh,ox:ox+nw]=np.floor(patch+.5).astype(np.uint8)
    mass=float(out.sum())
    if mass:
        gridy,gridx=np.mgrid[:28,:28]; dx=int(np.floor(13.5-(out*gridx).sum()/mass+.5)); dy=int(np.floor(13.5-(out*gridy).sum()/mass+.5))
        shifted=np.zeros_like(out); sx0=max(0,-dx); sx1=min(28,28-dx); sy0=max(0,-dy); sy1=min(28,28-dy)
        shifted[sy0+dy:sy1+dy,sx0+dx:sx1+dx]=out[sy0:sy1,sx0:sx1]; out=shifted
    return out

def augment(x,rng):
    a=x.reshape(-1,28,28).copy(); n=len(a)
    # Shift eachimage independently; label and samplemembership neverchange.
    dx=rng.integers(-2,3,n); dy=rng.integers(-2,3,n); p=np.pad(a,((0,0),(2,2),(2,2)))
    a=p[np.arange(n)[:,None,None],np.arange(28)[None,:,None]+(2-dy)[:,None,None],np.arange(28)[None,None,:]+(2-dx)[:,None,None]]
    mode=rng.random(n); p=np.pad(a,((0,0),(1,1),(1,1)))
    for pick,operation in [(mode<.10,np.maximum),(mode>.90,np.minimum)]:
        if not pick.any(): continue
        neighbors=[p[pick,1:29,1:29],p[pick,:28,1:29],p[pick,2:30,1:29],p[pick,1:29,:28],p[pick,1:29,2:30]]
        a[pick]=operation.reduce(neighbors)
    return a.reshape(n,784).astype(np.float32)/255

def quantized(params):
    w1,b1,w2,b2=params
    # Integer drives all fit exactJSNumber arithmetic; weights smallenoughint16.
    assert max(abs(w1).max(),abs(w2).max())*SCALE<32767
    return np.floor(w1*SCALE+.5).astype(np.int16),np.floor(b1*SCALE*255+.5).astype(np.int32),np.floor(w2*SCALE+.5).astype(np.int16),np.floor(b2*SCALE+.5).astype(np.int32)

def predict(x,params,steps=64,batch=512,return_scores=False):
    w1,b1,w2,b2=quantized(params); threshold=int(CAP*SCALE*255); results=[]
    # Driveinteger dot evaluatedfloat64: exactinteger sum(<2^53), matchingJS.
    for first in range(0,len(x),batch):
        drive=np.clip(x[first:first+batch].astype(np.float64)@w1.astype(np.float64)+b1,0,threshold)
        # Algebraic count of finiteIFsteps; browser independently executessteps.
        counts=np.floor(drive*steps/threshold).astype(np.int32)
        scores=counts.astype(np.float64)@w2.astype(np.float64)*CAP/steps+b2
        results.append(scores)
    scores=np.concatenate(results)
    return scores if return_scores else scores.argmax(1)