import json, time
import numpy as np

LAYERS, WIDTH = 4, 2

def instance(seed, K=8):
    r=np.random.default_rng(seed); edges=[]
    for t in range(LAYERS):
        src=range(t*WIDTH,(t+1)*WIDTH) if t else [-1]
        dst=range((t+1)*WIDTH,(t+2)*WIDTH) if t<LAYERS-1 else [-2]
        edges += [(a,b) for a in src for b in dst]
    foc=[]
    for _ in edges:
        l=np.maximum(.03,r.normal(1.2,.25,K)); u=l+r.uniform(.05,.75,K)
        foc.append((l,u,r.dirichlet(np.ones(K))))
    return edges,foc

def routes(edges):
    lookup={p:i for i,p in enumerate(edges)}; out=[]
    for q in np.ndindex(*(WIDTH,)*(LAYERS-1)):
        ns=[-1]+[t*WIDTH+q[t-1] for t in range(1,LAYERS)]+[-2]
        x=np.zeros(len(edges),dtype=np.int8)
        for a,b in zip(ns[:-1],ns[1:]): x[lookup[a,b]]=1
        out.append(x)
    return np.asarray(out)

def compressed(foc,groups):
    # Deliberately conservative surrogate: merged block cost is its mass times
    # max lower endpoint, hence every merge only raises edge cost.
    vals=[]
    for e,(l,u,m) in enumerate(foc):
        vals.append(sum(max(float(np.dot(m[z],l[z])),m[z].sum()*float(np.max(l[z]))) for z in groups[e]))
    return np.asarray(vals)

def merge_groups(groups,e,i,j):
    g=[[z[:] for z in ge] for ge in groups]; g[e][i]+=g[e][j]; del g[e][j]; return g

def route(c,X):
    v=X@c; k=int(np.argmin(v)); return k,float(v[k])

def compress(foc,X,target,mode='regret',lam=.1):
    groups=[[[i] for i in range(len(foc[e][0]))] for e in range(len(foc))]
    c=compressed(foc,groups); x0=X[np.argmin(X@c)]
    while sum(len(z) for ge in groups for z in ge)>target*len(groups):
        best=None
        for e,g in enumerate(groups):
            for i in range(len(g)):
              for j in range(i+1,len(g)):
                ng=merge_groups(groups,e,i,j); nc=compressed(foc,ng); d=nc-c
                if mode=='regret': score=float(d@x0+lam*np.abs(d).sum())
                elif mode=='random': score=float(np.random.default_rng(e+i*31+j).random())
                elif mode=='mass':
                    score=-float(foc[e][2][g[i]].sum()+foc[e][2][g[j]].sum())
                else:
                    l,u,m=foc[e]; ai,bj=g[i],g[j]
                    lo=max(min(l[ai]),min(l[bj])); hi=min(max(u[ai]),max(u[bj]))
                    union=max(max(u[ai]),max(u[bj]))-min(min(l[ai]),min(l[bj]))
                    score=-max(0.,hi-lo)/max(union,1e-12)
                if best is None or score<best[0]: best=(score,ng,nc)
        if best is None: break
        _,groups,c=best
    return groups,c,x0

def theorem_check():
    rng=np.random.default_rng(7); violation=0.; ratios=[]
    for _ in range(8):
        e,f=instance(int(rng.integers(1e9))); X=routes(e); g,ch,_=compress(f,X,3); c=compressed(f,[[[i] for i in range(8)] for _ in f])
        delta=X@(ch-c); star=int(np.argmin(X@c)); hat=int(np.argmin(X@ch))
        regret=float((X[hat]-X[star])@c); bound=float(delta[star])
        violation=max(violation,regret-bound); ratios.append(regret/(bound+1e-12))
    return {'max_bound_violation':float(violation),'mean_regret_bound_ratio':float(np.mean(ratios))}

def main():
    start=time.time(); check=theorem_check(); summary={}
    for K in (4,8,12):
      for target in (1,2,3):
        vals=[]
        for seed in range(3):
            e,f=instance(seed+100,K); X=routes(e); c=compressed(f,[[[i] for i in range(K)] for _ in f]); star,_=route(c,X)
            _,ch,_=compress(f,X,target); hat,_=route(ch,X); d=X@(ch-c)
            vals.append((float((X[hat]-X[star])@c),int(hat!=star),float(d[star]),float(np.sum(ch-c))))
        summary[f'K{K}_to{target}']={'mean_regret':float(np.mean([v[0] for v in vals])),'flip_rate':float(np.mean([v[1] for v in vals])),'mean_bound':float(np.mean([v[2] for v in vals])),'mean_inflation':float(np.mean([v[3] for v in vals]))}
    baselines={}
    for mode in ('regret','jaccard','random'):
        vals=[]
        for seed in range(3):
            e,f=instance(seed+900,8); X=routes(e); c=compressed(f,[[[i] for i in range(8)] for _ in f]); star,_=route(c,X)
            _,ch,_=compress(f,X,2,mode); h,_=route(ch,X); vals.append(((X[h]-X[star])@c,h!=star))
        baselines[mode]={'mean_regret':float(np.mean([v[0] for v in vals])),'flip_rate':float(np.mean([v[1] for v in vals]))}
    print(json.dumps({'theorem':check,'sweep':summary,'baselines':baselines,'seconds':time.time()-start},indent=2))
if __name__=='__main__': main()