#!/usr/bin/env python3
"""Pack encoder 1's extracted weights (engine/rr/enc1_ref.py --extract) for engine/rr/enc1_m.comp.

usage: enc1_pack.py WEIGHTS.npz OUT.bin OUT.h

The shader runs every product transposed (D^T = W^T X^T: weights are the A operand, activations the B operand), so a
result's rows are output channels and its columns tokens, and the result is the next product's B operand in place.
The rows come out in the order of W's columns, so every weight's columns are permuted here into the k order NVIDIA's
next product contracts them in (engine/rr/enc1_ref.py: q8(x)[:, inter(n)]), and the f16 vectors that act on those
rows likewise. Keys and queries are token slots: slot s is token inter(96)[s] (the prologue's key list), so the
attention bias is stored transposed in slot order.

  A tiles     W (K x N) -> W^T 16x16 tiles [mt][kt], row-major (m, k): 256 * (mt * K/16 + kt) bytes
  B tiles     the V projection runs X Wv (A = tokens): tiles [kt][nt], row-major (k, n)
  f16         vectors as is; the attention bias as 96 x 64 [key slot][query slot]
"""
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from enc1_ref import inter  # noqa: E402

I32, I64, I96 = inter(32), inter(64), inter(96)


def a_tiles(w):
    k, n = w.shape
    return np.ascontiguousarray(w.T).reshape(n // 16, 16, k // 16, 16).transpose(0, 2, 1, 3).reshape(-1)


def b_tiles(w):
    k, n = w.shape
    return w.reshape(k // 16, 16, n // 16, 16).transpose(0, 2, 1, 3).reshape(-1)


CUBIC_MAX = 0x3873   # binary16 bits of 0.55615234375, the clamp of the scaled score


def act_table():
    """the MLP's hidden activation and its FP8 conversion, q8(act(h)), for every binary16 h, indexed by its bits: one
    byte load in place of the six binary16 operations and the conversion (the shaders' ACTQ)"""
    from enc1_ref import act
    from rrmath import q8
    h = np.arange(65536, dtype=np.uint32).astype(np.uint16).view(np.float16)
    with np.errstate(all='ignore'):
        return np.asarray(q8(act(h).astype(np.float64)), np.uint8)


def cubic_table():
    """the cubic softmax numerator's bits (enc1_ref.softmax_cubic) for every binary16 score, indexed by its bits (the
    scaling and the clamp included; NaN scores take NVIDIA's min / max, which return the other operand)"""
    x = np.arange(65536, dtype=np.uint32).astype(np.uint16).view(np.float16)
    with np.errstate(invalid='ignore', over='ignore'):
        m = (x * np.float16(0.017181396484375)).astype(np.float16)
    lim = np.float16(0.55615234375)
    m = np.where(np.isnan(m), lim, np.clip(m, -lim, lim)).astype(np.float16)

    def fma16(a, b, c):
        return (a.astype(np.float64) * b.astype(np.float64) + c.astype(np.float64)).astype(np.float16)
    t = fma16(m, -m, np.full_like(m, np.float16(0.92724609375)))
    u = fma16(m, t, np.full_like(m, np.float16(1.375)))
    return ((u.view(np.uint16).astype(np.uint32) << 5) & 0x7FE0).astype(np.uint16)


def main():
    src, out_bin, out_h = sys.argv[1:4]
    w = dict(np.load(src))
    e = {}
    for h in range(2):
        e[f'wk{h}'] = ('A', np.vstack([w[f'wk{h}a'], w[f'wk{h}b']])[:, I32])
        e[f'wv{h}'] = ('B', np.vstack([w[f'wv{h}a'], w[f'wv{h}b']])[:, I32])
        e[f'wo{h}'] = ('A', w[f'wo{h}'][:, I64])
        e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(I64, I96)].T)
    for b in range(8):
        e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32])
        e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32])
        e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, I64])
    e['wm'], e['bm'] = ('A', w['wm'][:, I96]), ('H', w['bm'][I96])
    e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh'])
    e['gain0'] = ('H', w['gain0'])
    for k in ('table7', 'gain14', 'table16'):
        e[k] = ('H', w[k][I64])
    e['cubic'] = ('U', cubic_table())
    blob, defines = bytearray(), []
    for name, (kind, a) in e.items():
        while len(blob) % 256:
            blob.append(0)
        off = len(blob)
        if kind == 'A':
            blob += a_tiles(a).tobytes()
        elif kind == 'B':
            blob += b_tiles(a).tobytes()
        elif kind == 'U':
            blob += np.ascontiguousarray(a, np.uint16).tobytes()
        else:
            blob += np.ascontiguousarray(a, np.float16).tobytes()
        defines.append(f'#define W_{name.upper()} {off}u   // {kind} {a.shape}')
    open(out_bin, 'wb').write(bytes(blob))
    open(out_h, 'w').write('// generated by engine/rr/enc1_pack.py\n' + '\n'.join(defines) + '\n')
    print(f'{len(blob)} bytes, {len(e)} entries')


if __name__ == '__main__':
    main()