"""Sampled weight diff of DeepSeek-V4.1-Flash vs DeepSeek-V4.1-Flash-1017 over HTTP range reads.

usage: diff.py DIR   (DIR has headers.json from headers.py; writes DIR/diff.json)
For every sampled tensor it reads the same byte windows from both releases and reports
frac (share of elements whose bits changed) and rel (||b - a|| / ||a||, decoded values;
FP4 codes are compared without their block scales, which are sampled as tensors of
their own). Pure Python: no numpy on this machine."""
import array, json, math, os, random, re, sys, threading, time, urllib.error, urllib.request, concurrent.futures as cf

D = sys.argv[1]
H = json.load(open(f'{D}/headers.json'))
A, B = 'DeepSeek-V4.1-Flash', 'DeepSeek-V4.1-Flash-1017'
WIN = 16384
random.seed(1017)


def e4m3(b):
    s, e, m = b >> 7, (b >> 3) & 15, b & 7
    if e == 15 and m == 7:
        return 0.0
    v = (m / 8) * 2 ** -6 if e == 0 else (1 + m / 8) * 2 ** (e - 7)
    return -v if s else v


E4M3 = [e4m3(i) for i in range(256)]
E2M1 = [0, .5, 1, 1.5, 2, 3, 4, 6, -0., -.5, -1, -1.5, -2, -3, -4, -6]


def decode(dtype, raw):
    if dtype == 'BF16':
        u = array.array('H', raw[: len(raw) // 2 * 2])
        return array.array('f', array.array('I', (x << 16 for x in u)).tobytes()).tolist(), u.tolist()
    if dtype == 'F32':
        f = array.array('f', raw[: len(raw) // 4 * 4])
        return f.tolist(), array.array('I', f.tobytes()).tolist()
    if dtype == 'F8_E4M3':
        return [E4M3[x] for x in raw], list(raw)
    if dtype == 'F8_E8M0':
        return [float(x - 127) for x in raw], list(raw)   # exponents: compare in log2
    if dtype == 'I8':                                        # two FP4 (E2M1) codes per byte
        codes = [c for x in raw for c in (x & 15, x >> 4)]
        return [E2M1[c] for c in codes], codes
    raise ValueError(dtype)


CDN, LOCK = {}, threading.Lock()


def get(url, a, b):
    req = urllib.request.Request(url, headers={'Range': f'bytes={a}-{b}', 'User-Agent': 'curl/8.7.1'})
    for i in range(8):
        try:
            with urllib.request.urlopen(req, timeout=120) as r:
                return r.read(), r.geturl()
        except urllib.error.HTTPError as e:
            if e.code not in (429, 500, 502, 503, 504):
                raise
            time.sleep(min(60, 2 ** i + random.random()))
        except Exception:
            time.sleep(min(60, 2 ** i + random.random()))
    raise RuntimeError(f'gave up on {url}')


def cdn(repo, f):
    """The signed CDN URL of a file, resolved once: range reads then skip huggingface.co."""
    key = (repo, f)
    with LOCK:
        if key in CDN and CDN[key][1] > time.time():
            return CDN[key][0]
    _, final = get(f'https://huggingface.co/deepseek-ai/{repo}/resolve/main/{f}', 0, 0)
    with LOCK:
        CDN[key] = (final, time.time() + 1800)
    return final


def rng(repo, t, a, b):
    start = t['base'] + t['off'][0]
    data, _ = get(cdn(repo, t['file']), start + a, start + b - 1)
    return data


def windows(t):
    n = t['off'][1] - t['off'][0]
    if n <= 2 * WIN:
        return [(0, n)]
    return [(int(n * f) // 64 * 64, int(n * f) // 64 * 64 + WIN) for f in (0.25, 0.75)]


def category(name):
    if name.startswith('vision.') or name.startswith('aligner.') or name.startswith('image_'):
        return 'vision encoder'
    if name in ('embed.weight',):
        return 'token embeddings'
    if name in ('head.weight',):
        return 'output head'
    if name.startswith('mtp.'):
        if any(x in name for x in ('markov_head', 'confidence_head', 'main_proj')):
            return 'DSpark heads'
        return 'MTP layers'
    if '.engram.' in name:
        return 'Engram tables' if 'embed' in name else 'Engram projections'
    if '.ffn.experts.' in name:
        return 'routed experts'
    if '.ffn.shared_experts.' in name:
        return 'shared experts'
    if '.ffn.gate.' in name:
        return 'router'
    if '.attn.indexer.' in name or '.attn.compressor.' in name:
        return 'sparse-attention indexer'
    if '.attn.' in name:
        return 'attention'
    if '.hc_' in name:
        return 'hyper-connections'
    if 'norm' in name:
        return 'norms'
    return 'other'


def pick():
    names = list(H[A])
    out = []
    experts = {}
    for n in names:
        m = re.match(r'((?:layers|mtp)\.\d+)\.ffn\.experts\.(\d+)\.', n)
        if m:
            experts.setdefault(m.group(1), set()).add(int(m.group(2)))
    sample_e = {L: set(sorted(es)[::max(1, len(es) // 4)]) for L, es in experts.items()}   # 4 experts per layer
    for n in names:
        m = re.match(r'((?:layers|mtp)\.\d+)\.ffn\.experts\.(\d+)\.', n)
        if m and int(m.group(2)) not in sample_e[m.group(1)]:
            continue
        if n.startswith('vision.blocks.') and not re.match(r'vision\.blocks\.(0|16|31)\.', n):
            continue
        out.append(n)
    return out


def job(name):
    ta, tb = H[A][name], H[B][name]
    rows = []
    for a, b in windows(ta):
        ra, rb = rng(A, ta, a, b), rng(B, tb, a, b)
        va, ca = decode(ta['dtype'], ra)
        vb, cb = decode(tb['dtype'], rb)
        k = min(len(va), len(vb))
        changed = sum(1 for i in range(k) if ca[i] != cb[i])
        num = sum((vb[i] - va[i]) ** 2 for i in range(k) if not (math.isnan(va[i]) or math.isnan(vb[i])))
        den = sum(va[i] ** 2 for i in range(k) if not math.isnan(va[i]))
        rows.append((k, changed, num, den))
    k = sum(r[0] for r in rows)
    m = re.match(r'(?:layers|mtp|vision\.blocks)\.(\d+)\.', name)
    return dict(name=name, cat=category(name), layer=int(m.group(1)) if m else None, dtype=ta['dtype'],
                n=k, frac=sum(r[1] for r in rows) / k if k else 0,
                rel=math.sqrt(sum(r[2] for r in rows) / sum(r[3] for r in rows)) if sum(r[3] for r in rows) else None)


names = pick()
OUT = f'{D}/diff.jsonl'
done = set()
if os.path.exists(OUT):
    done = {json.loads(l)['name'] for l in open(OUT) if l.strip()}
todo = [n for n in names if n not in done]
print(len(names), 'tensors,', len(todo), 'to sample', flush=True)
for repo, f in sorted({(r, H[r][n]['file']) for n in todo for r in (A, B)}):   # resolve once, one at a time
    cdn(repo, f)
print('resolved', len(CDN), 'CDN URLs', flush=True)
with open(OUT, 'a') as out, cf.ThreadPoolExecutor(24) as ex:
    futs = {ex.submit(job, n): n for n in todo}
    for i, f in enumerate(cf.as_completed(futs)):
        try:
            r = f.result()
        except Exception as e:
            print('failed', futs[f], e, flush=True)
            continue
        out.write(json.dumps(r) + '\n'); out.flush()
        if i % 250 == 0:
            print(i, r['name'], round(r['frac'], 4), r['rel'], flush=True)
print('done', sum(1 for _ in open(OUT)))