from microbit import *
import struct,gc,time,random,math
from array import array
uart.init(115200)
# BitFlow MoE v4 - micro:bit 推理
# DIM=40 FFN=80 3专家 词表149 (~58K int8权重)
VS=149;D=40;CT=48;FD=80;NE=3;NL=2;RD=32
RB=2+NE*12;NB=RB+4;HI=NB+2
# 加载权重
wf=open("weights.bin","rb")
nt=struct.unpack('<H',wf.read(2))[0]
LEN=list(struct.unpack('<%dH'%nt,wf.read(2*nt)))
SC=list(struct.unpack('<%df'%nt,wf.read(4*nt)))
W=wf.read()
wf.close()
del nt,wf;gc.collect()
OFF=[];o=0
for l in LEN:OFF.append(o);o+=l
del o;gc.collect()
# 单专家状态 (NL层, 路由器选一个专家)
S=[array('f',bytes(4*D*D)) for _ in range(NL)]
z=[array('f',bytes(4*D)) for _ in range(NL)]
# 工作缓冲区
bx=array('f',bytes(4*D));bh=array('f',bytes(4*D))
bq=array('f',bytes(4*D));bk=array('f',bytes(4*D))
bv=array('f',bytes(4*D));bgd=array('f',bytes(4*D))
bnum=array('f',bytes(4*D));by=array('f',bytes(4*D))
bout=array('f',bytes(4*D));bff=array('f',bytes(4*FD))
bpr=array('f',bytes(4*VS));bcnt=array('i',bytes(4*VS))
brh=array('f',bytes(4*RD));brl=array('f',bytes(4*NE))
blg=array('f',bytes(4*VS))
bsum=array('f',bytes(4*D));brx=array('f',bytes(4*D));bct=[0]
gc.collect()
# 预解量化小权重
def dqS(i):
 s=SC[i];o=OFF[i];n=LEN[i];r=array('f',bytes(4*n))
 for j in range(n):
  v=W[o+j]
  if v>=128:v-=256
  r[j]=v*s
 return r
SMW=[]
for e in range(NE):
 b0=2+e*12
 SMW.append(dqS(b0));SMW.append(dqS(b0+1));SMW.append(dqS(b0+2))
 SMW.append(dqS(b0+8));SMW.append(dqS(b0+9))
rf1b=dqS(RB+1);rf2b=dqS(RB+3)
nfw=dqS(NB);nfb=dqS(NB+1)
del dqS;gc.collect()
# 融合解量化+矩阵乘 (省去DB大缓冲区)
def dqmv(i,ni,no,x,out):
 s=SC[i];o=OFF[i]
 for r in range(no):
  ss=0.0;b=r*ni
  for c in range(ni):
   v=W[o+b+c]
   if v>=128:v-=256
   ss+=v*s*x[c]
  out[r]=ss
def ln(w,b,x,o):
 m=0.0
 for i in range(D):m+=x[i]
 m/=D;v=0.0
 for i in range(D):t=x[i]-m;v+=t*t
 v/=D;iv=1.0/math.sqrt(v+1e-5)
 for i in range(D):o[i]=w[i]*(x[i]-m)*iv+b[i]
def e1(x):return x+1.0 if x>=0 else math.exp(x)
def sg(x):
 if x>=0:e=math.exp(-x);return 1.0/(1.0+e)
 e=math.exp(x);return e/(1.0+e)
def ge(x):
 c=0.7978845608
 return x*sg(2.0*c*(x+0.044715*x*x*x))
def sm(l,t=0.3):
 m=-1e9
 for i in range(VS):
  v=l[i]-bcnt[i]*2.0
  if v>m:m=v
 s=0.0
 for i in range(VS):
  v=l[i]-bcnt[i]*2.0
  e=math.exp((v-m)/t)
  bpr[i]=e;s+=e
 r=random.random()*s;c=0.0
 for i in range(VS):
  c+=bpr[i]
  if c>=r:return i
 return VS-1
# GLA状态更新 (单层)
def gs(e,layer,h):
 b0=2+e*12;dec=SMW[e*5+2]
 dqmv(b0+3,D,D,h,bq)
 dqmv(b0+4,D,D,h,bk)
 dqmv(b0+5,D,D,h,bv)
 dqmv(b0+6,D,D,h,bgd);gc.collect()
 for i in range(D):bq[i]=e1(bq[i]);bk[i]=e1(bk[i])
 for i in range(D):bgd[i]=sg(bgd[i])*dec[i]
 sl=S[layer];zl=z[layer]
 for i in range(D):
  si=i*D;gd=bgd[i];ki=bk[i]
  for j in range(D):sl[si+j]=gd*sl[si+j]+ki*bv[j]
 for i in range(D):zl[i]=bgd[i]*zl[i]+bk[i]
 for j in range(D):bnum[j]=0.0
 for i in range(D):
  si=i*D;qi=bq[i]
  for j in range(D):bnum[j]+=qi*sl[si+j]
 den=1e-6
 for i in range(D):den+=bq[i]*zl[i]
 iv=1.0/den
 for j in range(D):by[j]=bnum[j]*iv
 dqmv(b0+7,D,D,by,bout);gc.collect()
# 专家块 (残差+GLA+FFN)
def bs(e,x,layer):
 b0=2+e*12
 ln(SMW[e*5],SMW[e*5+1],x,bh)
 gs(e,layer,bh)
 for i in range(D):x[i]+=bout[i]
 ln(SMW[e*5+3],SMW[e*5+4],x,bh)
 dqmv(b0+10,D,FD,bh,bff)
 for i in range(FD):bff[i]=ge(bff[i])
 dqmv(b0+11,FD,D,bff,bout);gc.collect()
 for i in range(D):x[i]+=bout[i]
# 路由器 (D→RD→NE)
def route(x):
 dqmv(RB,D,RD,x,brh)
 for i in range(RD):brh[i]=max(0.0,brh[i]+rf1b[i])
 dqmv(RB+2,RD,NE,brh,brl)
 for i in range(NE):brl[i]+=rf2b[i]
 bi=0;bv2=brl[0]
 for i in range(1,NE):
  if brl[i]>bv2:bv2=brl[i];bi=i
 return bi
# 前向: gen_flag=False时跳过head(种子处理加速)
def fw(ctx,gen_flag=True):
 t=min(len(ctx)-1,CT-1);ti=ctx[-1]
 s0=SC[0];o0=OFF[0];s1=SC[1];o1=OFF[1]
 for i in range(D):
  v0=W[o0+ti*D+i]
  if v0>=128:v0-=256
  v1=W[o1+t*D+i]
  if v1>=128:v1-=256
  bx[i]=v0*s0+v1*s1
 gc.collect()
 for i in range(D):bsum[i]+=bx[i]
 bct[0]+=1
 for i in range(D):brx[i]=bsum[i]/bct[0]
 ei=route(brx)
 for layer in range(NL):bs(ei,bx,layer)
 ln(nfw,nfb,bx,bh)
 if gen_flag:
  dqmv(HI,D,VS,bh,blg);gc.collect()
  return blg,ei
 return None,ei
# 重置状态
def rs():
 for l in range(NL):
  for i in range(D*D):S[l][i]=0.0
  for i in range(D):z[l][i]=0.0
 for i in range(D):bsum[i]=0.0
 bct[0]=0
# 生成
def gen(ids,n):
 NL_ch=VS-1
 ids=list(ids)+[NL_ch]
 sl=len(ids);c=list(ids);rs();gc.collect()
 for i in range(VS):bcnt[i]=0
 uart.write("S:%s\r\n"%",".join(str(i) for i in c))
 for i in range(len(c)-1):
  fw(c[:i+1],False);gc.collect()
 for _ in range(n):
  t0=time.ticks_ms()
  lg,ei=fw(c)
  nx=sm(lg)
  dt=time.ticks_diff(time.ticks_ms(),t0)
  uart.write("T:%d %dms E%d\r\n"%(nx,dt,ei))
  bcnt[nx]+=1
  del lg;gc.collect();c.append(nx)
 uart.write("O:%s\r\n"%",".join(str(i) for i in c[sl:]))
gc.collect()
print("READY")
print("free=%d"%gc.mem_free())
while True:
 if uart.any():
  l=uart.readline()
  if l:
   l=l.strip()
   if l[:4]==b"GEN:":
    p=l[4:].split(b"|")
    ids=[int(x) for x in p[0].split(b",") if x.strip()]
    nn=int(p[1]) if len(p)>1 else 10
    if ids:gen(ids,nn)
 sleep(50)