#!/usr/bin/env python3
"""Losslessly repack official weights using the existing fused tensor names.

INT4 codes and BF16 group scales remain unchanged. Execution precision is chosen
by the experimental runtime, not baked into this source-weight package.
"""
import argparse
import hashlib
import json
import mmap
from pathlib import Path
import struct
import zlib

import torch
from safetensors import safe_open
from pack_model import HEADER, ENTRY, align, pad_to
from quantize_model import source_linear_names
from official_int4_reference import unpack4, REVISION

p=argparse.ArgumentParser(description=__doc__)
p.add_argument('--template',type=Path,required=True)
p.add_argument('--official',type=Path,required=True)
p.add_argument('--output',type=Path,required=True)
a=p.parse_args()
if a.output.exists(): raise FileExistsError(a.output)
torch.set_num_threads(4)
index=json.loads((a.official/'model.safetensors.index.json').read_text())['weight_map']
handles={s:safe_open(a.official/s,framework='pt') for s in set(index.values())}
def get(name): return handles[index[name]].get_tensor(name)
with a.template.open('rb') as src, mmap.mmap(src.fileno(),0,access=mmap.ACCESS_READ) as data:
    header=list(HEADER.unpack_from(data)); count=header[4]; table=header[7]; strings=header[8]; payload=header[10]
    entries=[]
    for i in range(count):
        e=list(ENTRY.unpack_from(data,table+i*ENTRY.size))
        name=data[strings+e[0]:strings+e[0]+e[1]].decode()
        entries.append((name,e))
    a.output.parent.mkdir(parents=True,exist_ok=True)
    with a.output.open('w+b') as out:
        out.write(data[:payload]); written=[]; linear_count=0
        for name,e in entries:
            raw=None
            if e[4]==3:
                base=name.removesuffix('.weight'); names=source_linear_names(base)
                k,n=e[8:10]
                if '.mlp.experts.' in base:
                    codes=[]
                    for key in names:
                        shape=tuple(get(key+'_shape').tolist())
                        codes.append(unpack4(get(key+'_packed'),shape))
                    q=torch.cat(codes).t().contiguous().reshape(-1).numpy().view('uint8') & 15
                    raw=(q[::2] | (q[1::2]<<4)).tobytes()
                    e[2]=5;e[5]=4;e[6]=2;e[7]=32
                    assert len(raw)==k*n//2
                else:
                    weight=torch.cat([get(key) for key in names])
                    assert weight.dtype==torch.bfloat16 and weight.shape[1]==k and weight.shape[0]<=n
                    if weight.shape[0]<n:weight=torch.nn.functional.pad(weight,(0,0,0,n-weight.shape[0]))
                    raw=weight.contiguous().view(torch.uint16).numpy().tobytes()
                    e[2]=1;e[5]=0;e[6]=0;e[7]=0
                linear_count+=1
                if linear_count%256==0:print(f'packed {linear_count} linears',flush=True)
            elif name.endswith('.scales') and '.mlp.experts.' in name:
                base=name.removesuffix('.scales')
                scales=torch.cat([get(key+'_scale') for key in source_linear_names(base)])
                raw=scales.contiguous().view(torch.uint16).numpy().tobytes()
                e[2]=1;e[3]=2;e[5]=4;e[6]=0;e[7]=32;e[8:12]=[scales.shape[0],scales.shape[1],0,0]
            if raw is None:raw=data[e[12]:e[12]+e[13]]
            pad_to(out,align(out.tell()));e[12]=out.tell();e[13]=len(raw);out.write(raw)
            e[14]=e[15]=0;e[22]=hashlib.sha256(raw).digest();written.append(ENTRY.pack(*e))
        pad_to(out,align(out.tell()));header[11]=out.tell()-payload;header[12]=out.tell();header[13]=bytes.fromhex(REVISION)
        header[6]|=0x100;header[14]=0;header[14]=zlib.crc32(HEADER.pack(*header)) & 0xffffffff
        out.seek(0);out.write(HEADER.pack(*header));out.seek(table)
        for entry in written:out.write(entry)
    print(json.dumps({'output':str(a.output),'bytes':a.output.stat().st_size,'linears':linear_count,'source_revision':REVISION,'quantization':'unchanged official INT4 group32 and high precision weights'}),flush=True)