#!/usr/bin/env python3
"""Pack a Swin stage's extracted weights (engine/rr/swin_ref.py --extract) for engine/rr/swin_m.comp.

usage: swin_pack.py STAGE WEIGHTS.npz OUT.bin OUT.h

enc1_pack.py's layout for any stage of engine/rr/swin_ref.STAGES: transposed products, every weight's columns in the
k order of the product that consumes its result (slot j holds natural channel inter(n)[j]); the expand's 4C columns
are four sub-positions of C channels, each permuted the same way. OUT.h also defines the stage's shape for the shader
(KIND, C, CX, NH).
"""
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 swin_ref import STAGES  # noqa: E402


def main():
    stage, src, out_bin, out_h = sys.argv[1:5]
    kind, C, CX, NH, _ = STAGES[stage]
    w = dict(np.load(src))
    I32, I96, IC = inter(32), inter(96), inter(C)
    e = {}
    if kind == 'dec':
        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'wk{h}'] = ('A', w[f'wk{h}'][:, I32])
        e[f'wv{h}'] = ('B', w[f'wv{h}'][:, I32])
        e[f'wo{h}'] = ('A', w[f'wo{h}'][:, IC])
        e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(inter(64), I96)].T)
    for b in range(C // 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}'][:, IC])
    if kind == 'enc' and CX:
        e['wm'], e['bm'] = ('A', w['wm'][:, inter(CX)]), ('H', w['bm'][inter(CX)])
    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][IC])
    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 {0 if kind == "enc" else 1}', f'#define C {C}', f'#define CX {CX}', f'#define NH {NH}']
    if NH < 2:
        defines.append('#define W_WK1 (W_WK0 + 1u)')
    if C // 8 < 2:
        defines.append('#define W_W1_1 (W_W1_0 + 1u)')
    open(out_bin, 'wb').write(bytes(blob))
    open(out_h, 'w').write(f'// generated by engine/rr/swin_pack.py for {stage}\n' + '\n'.join(shape + defines) + '\n')
    print(f'{stage}: {len(blob)} bytes, {len(e)} entries')


if __name__ == '__main__':
    main()