#!/usr/bin/env python3
"""Pack a preset F stage's weights, read from the user's nvngx_dlssd.dll (engine/rr/rr2_ref.dll_weights), for
engine/rr/rr2_m.comp.

usage: rr2_pack.py STAGE NVNGX_DLSSD_DLL OUT.bin OUT.h

swin_pack.py's layout: transposed products as 16 x 16 tiles, every tensor's outputs in slot order (slot j holds the
natural index inter(j)); a merge's or the expand's inputs are sub-positions of C channels each. OUT.h also defines
the stage's shape for the shader (KIND, C, NH, NB, CM, CL, GUIDE).
"""
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
from enc1_pack import a_tiles, b_tiles, cubic_table, act_table  # noqa: E402
from rr2_ref import STAGES, EXPAND, dll_weights  # noqa: E402


def main():
    stage, dll, out_bin, out_h = sys.argv[1:5]
    w = dll_weights(open(dll, 'rb').read(), stage)
    if stage in EXPAND:
        C, CL, NH, NB = EXPAND[stage]
        kind, CM, guide = 2, 0, 0
    else:
        C, c_up, CM, NH, NB = STAGES[stage]
        kind, CL = (4 if stage == 'enc0' else 5 if stage == 'dec0_1' else 3 if stage.startswith('dec') else 1 if c_up else 0), 0
        guide = 0 if kind >= 3 else 1
    I32, I64, IC = inter(32), inter(64), inter(C)
    e = {}
    if kind == 1:
        e['w1'], e['b1'] = ('A', w['w1'][:, IC]), ('H', w['b1'][IC])
    if kind == 4:   # encoder 0's embedding branches
        for b in range(2):
            e[f'ewd{b}'] = ('A', w['ewd'][b][:, I32])
            e[f'ebd{b}'] = ('H', w['ebd'][b][I32])
            e[f'ewu{b}'] = ('A', w['ewu'][b][:, IC])
        e['ebu'] = ('H', w['ebu'][IC])
    if kind == 2:
        I4C = [C * (j // C) + IC[j % C] for j in range(4 * C)]
        e['we'], e['be'] = ('A', w['we'][:, I4C]), ('H', w['be'][I4C])
    for h in range(NH):
        e[f'wqk{h}'] = ('A', w['wqk'][h][:, I32])
        e[f'wv{h}'] = ('B', w['wv'][h][:, I32])
        e[f'wo{h}'] = ('A', w['wo'][h][:, IC])
        e[f'bias{h}'] = ('H', w['bias'][h][np.ix_(I64, I64)])   # [key slot][query slot]; symmetric
    e['bo'] = ('H', w['bo'][IC])
    for b in range(NB):
        e[f'wd{b}'] = ('A', w['wd'][b][:, I32])
        e[f'bd{b}'] = ('H', w['bd'][b][I32])
        e[f'wu{b}'] = ('A', w['wu'][b][:, IC])
    e['bu'] = ('H', w['bu'][IC])
    if CM:
        e['wm'], e['bm'] = ('A', w['wm'][:, inter(CM)]), ('H', w['bm'][inter(CM)])
    if 'wh' in w:
        e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh'])
    if kind == 5:   # dec0's output tail: tanh of every binary16 value, correctly rounded, by its bits
        with np.errstate(all='ignore'):
            x = np.arange(65536, dtype=np.uint32).astype(np.uint16).view(np.float16).astype(np.float64)
            e['tanh'] = ('U', np.tanh(x).astype(np.float16).view(np.uint16))
    e['cubic'] = ('U', cubic_table())
    e['act'] = ('Y', act_table())
    blob, defines = bytearray(), []
    for name, (kd, a) in e.items():
        while len(blob) % 256:
            blob.append(0)
        off = len(blob)
        if kd == 'A':
            blob += a_tiles(a).tobytes()
        elif kd == 'B':
            blob += b_tiles(a).tobytes()
        elif kd == 'Y':
            blob += np.ascontiguousarray(a, np.uint8).tobytes()
        elif kd == '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   // {kd} {a.shape}')
    shape = [f'#define KIND {kind}', f'#define C {C}', f'#define NH {NH}', f'#define NB {NB}', f'#define CM {CM}',
             f'#define CL {CL}', f'#define GUIDE {guide}']
    open(out_bin, 'wb').write(bytes(blob))
    open(out_h, 'w').write(f'// generated by engine/rr/rr2_pack.py for {stage}\n' + '\n'.join(shape + defines) + '\n')
    print(f'{stage}: {len(blob)} bytes, {len(e)} entries')


if __name__ == '__main__':
    main()