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


def vdc(m):
    x, f = 0.0, 0.5
    while m:
        x += (m & 1) * f
        m >>= 1
        f *= 0.5
    return x


def sequence(D, n):
    return np.array([min(D-1, int(D*vdc(m))) for m in range(n)], dtype=int)


def dyadic_sets(D):
    # Every node in the fixed balanced binary tree, represented by a leaf interval.
    out = []
    width = D
    while width >= 1:
        for lo in range(0, D, width):
            out.append((lo, lo+width))
        width //= 2
    return out


def endpoint_leaf_error(seq, D):
    counts = np.bincount(seq, minlength=D)
    n = len(seq)
    return float(np.max(np.abs(counts - n / D)))


def discrepancy(seq, D):
    counts = np.zeros(D, dtype=int)
    max_leaf = 0.0
    max_nested = 0.0
    arg = None
    sets = dyadic_sets(D)
    for n, x in enumerate(seq, 1):
        counts[x] += 1
        max_leaf = max(max_leaf, float(np.max(np.abs(counts - n/D))))
        for lo, hi in sets:
            e = abs(int(counts[lo:hi].sum()) - n*(hi-lo)/D)
            if e > max_nested:
                max_nested, arg = e, (n, lo, hi)
    return max_leaf, max_nested, arg


def random_discrepancy(D, n, trials=40, seed=7):
    rng = np.random.default_rng(seed)
    vals, nested = [], []
    sets = dyadic_sets(D)
    for _ in range(trials):
        s = rng.integers(0, D, size=n)
        hist = np.zeros((n, D), dtype=np.int16)
        hist[np.arange(n), s] = 1
        counts = np.cumsum(hist, axis=0)
        vals.append(float(np.max(np.abs(counts[-1] - n / D))))
        mx = 0.0
        k = np.arange(1, n + 1)
        for lo, hi in sets:
            actual = counts[:, lo:hi].sum(axis=1)
            target = k * (hi - lo) / D
            mx = max(mx, float(np.max(np.abs(actual - target))))
        nested.append(mx)
    return float(np.mean(vals)), float(np.mean(nested))


def assignment_trial(D=8, tokens=128, capacity_factor=1.0, seed=11):
    # Synthetic learned router preferences: 70% prefer expert 0, remaining are uniform.
    rng = np.random.default_rng(seed)
    cap = int(math.ceil(tokens / D * capacity_factor))
    pref = np.where(rng.random(tokens) < .70, 0, rng.integers(0,D,size=tokens))
    schedules = {
        'random': rng.integers(0,D,size=tokens),
        'greedy': pref.copy(),
        'vdc': sequence(D,tokens),
    }
    result = {}
    for name, proposed in schedules.items():
        used = np.zeros(D, dtype=int); accepted = 0
        # Assignment takes proposed candidate; capacity overflow is dropped.
        for x in proposed:
            if used[x] < cap:
                used[x] += 1; accepted += 1
        result[name] = {'drop_rate': 1-accepted/tokens,
                        'load_var': float(np.var(used)), 'loads': used.tolist()}
    return result


def main():
    np.set_printoptions(suppress=True)
    # Prediction P1: first D points are a permutation, hence zero discrepancy at n=D.
    # P2: nested discrepancy grows no faster than the supplied logarithmic budget B(D).
    # P3: at every prefix, VDC leaf discrepancy is <= 1 (a stronger empirical prediction
    # for this dyadic construction); random routing should grow with D and prefix length.
    sweep=[]
    for D in [2,4,8,16,32,64]:
        n=4*D
        s=sequence(D,n)
        leaf,nested,arg=discrepancy(s,D)
        B=math.log(D)/(3*math.log(2))+1
        rd,rn=random_discrepancy(D,n)
        sweep.append({'D':D,'vdc_leaf_max':leaf,'vdc_nested_max':nested,
                      'bound_B':B,'random_leaf_mean':rd,'random_nested_mean':rn,
                      'at_n_D_endpoint_error':endpoint_leaf_error(sequence(D,D),D)})
    # Capacity sweep averages fixed seeds, same proposals for each method per trial.
    caps=[]
    for cf in [0.75,1.0,1.25,1.5]:
        agg={k:[] for k in ['random','greedy','vdc']}
        for seed in range(20):
            r=assignment_trial(capacity_factor=cf,seed=seed)
            for k in agg: agg[k].append(r[k])
        caps.append({'capacity_factor':cf, **{k:{'drop_rate':float(np.mean([z['drop_rate'] for z in v])),
              'load_var':float(np.mean([z['load_var'] for z in v]))} for k,v in agg.items()}})
    out={'predictions':[
        'P1: at n=D, VDC visits every leaf exactly once (max leaf error 0).',
        'P2: max nested prefix error is below B(D)=log(D)/(3 log 2)+1.',
        'P3: dyadic VDC leaf prefix error stays <=1, while random routing has larger mean imbalance.'
    ],'sweep':sweep,'capacity_sweep':caps}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()