import json, math
from pathlib import Path
import numpy as np

SEED = 17

def dyadic_rectangles(n):
    out=[]
    L=int(round(math.log2(n)))
    for ah in range(L+1):
        h=2**ah
        for aw in range(L+1):
            w=2**aw
            for y in range(0,n,h):
                for x in range(0,n,w): out.append((y,x,h,w))
    return out

def contains(a,b):
    y,x,h,w=a; yy,xx,hh,ww=b
    return y<=yy and x<=xx and y+h>=yy+hh and x+w>=xx+ww

def antichain(rects):
    kept=[]
    for r in rects:
        if not any(contains(k,r) or contains(r,k) for k in kept): kept.append(r)
    return kept

def halo(mask,threshold=.5):
    n=mask.shape[0]; out=np.zeros_like(mask,dtype=bool)
    for y,x,h,w in dyadic_rectangles(n):
        if mask[y:y+h,x:x+w].mean()>threshold: out[y:y+h,x:x+w]=True
    return out

def ancestors(r, n):
    y,x,h,w=r; hs=[]; vs=[]
    for _ in range(int(round(math.log2(n)))+1):
        hs.append((y,x,h,w)); nw=min(n,2*w); x=(x//nw)*nw; w=nw
    y,x,h,w=r
    for _ in range(int(round(math.log2(n)))+1):
        vs.append((y,x,h,w)); nh=min(n,2*h); y=(y//nh)*nh; h=nh
    return hs,vs

def in_halo_rect(oh,r):
    y,x,h,w=r
    return bool(oh[y:y+h,x:x+w].all())

def depth(r, oh, n):
    ha,va=ancestors(r,n)
    e1=sum(in_halo_rect(oh,a) for a in ha)-1
    e2=sum(in_halo_rect(oh,a) for a in va)-1
    return max(0,e1),max(0,e2)

def build_H(mask,rects,s):
    n=mask.shape[0]; oh=halo(mask); H=np.zeros((n,n),float); info=[]
    for r in rects:
        y,x,h,w=r; e1,e2=depth(r,oh,n)
        alpha=(e1+1)**(-s)*(e2+1)**(-(1-s))
        H[y:y+h,x:x+w]+=alpha
        info.append((r,e1,e2,alpha))
    return H,info

def metrics(H,mask):
    v=H[mask]
    return {'mean':float(v.mean()),'p99':float(np.quantile(v,.99)),'max':float(v.max()),'exp_mean':float(np.exp(.25*v).mean())}

def selected_rects(n, seed, count=28):
    rs=dyadic_rectangles(n); rng=np.random.default_rng(seed)
    weights=np.array([1/(r[2]*r[3])**.35 for r in rs]); weights/=weights.sum()
    idx=rng.choice(len(rs),size=min(count*4,len(rs)),replace=False,p=weights)
    return [rs[i] for i in idx]

def router_demo(n=16):
    mask=np.ones((n,n),bool); rs=selected_rects(n,101,32)
    hb,ib=build_H(mask,rs,.5); ra=antichain(rs); ha,ia=build_H(mask,ra,.5)
    return {'candidate_count':len(rs),'antichain_count':len(ra),'baseline':metrics(hb,mask),'idea':metrics(ha,mask)}

def main():
    n=16; mask=np.ones((n,n),bool); base=selected_rects(n,22,40)
    sym=[]
    for s in [.1,.2,.3,.5,.7,.8,.9]:
        H1,_=build_H(mask,base,s); H2,_=build_H(mask,base,1-s)
        sym.append({'s':s,'relative_symmetry_error':float(np.max(np.abs(H1-H2))/(np.max(H1)+1e-12)), 'mean':float(H1.mean())})
    H,info=build_H(mask,antichain(base),.5); scale=[]
    for lam in [.25,.5,1.,2.,4.]:
        HH=lam*H; scale.append({'lambda':lam,'mean_H':float(HH.mean()),'log_exp_penalty':float(np.log(np.exp(.25*HH).mean()))})
    directional=[]
    for s in [.01,.1,.25,.5,.75,.9,.99]:
        Hs,_=build_H(mask,antichain(base),s)
        directional.append({'s':s,'p99':float(np.quantile(Hs[mask],.99)),'max':float(Hs.max())})
    result={'seed':SEED,'grid':n,'symmetry_sweep':sym,'lambda_sweep':scale,'directional_sweep':directional,'router_demo':router_demo(n)}
    Path('results.json').write_text(json.dumps(result,indent=2)); print(json.dumps(result,indent=2))

if __name__=='__main__': main()