#!/usr/bin/env python3
"""Re-export only calibrated linear assets; reuse baseline raw tensors/islands."""
import argparse
import copy
import json
from pathlib import Path
import time

import torch
from quantize_model import AssetWriter, TensorSource, source_linear_names

p=argparse.ArgumentParser(description=__doc__)
p.add_argument('--source',type=Path,required=True)
p.add_argument('--baseline-manifest',type=Path,required=True)
p.add_argument('--scales',type=Path,required=True)
p.add_argument('--output',type=Path,required=True)
a=p.parse_args()
torch.set_num_threads(4)
baseline=a.baseline_manifest.resolve()
manifest=json.loads(baseline.read_text())
assert manifest.get('flags')==1, 'baseline must be a complete model'
writer=AssetWriter(a.output.resolve(),a.scales)
source=TensorSource(a.source.resolve())
start=time.monotonic()
count=0
for item in manifest['tensors']:
    if item['role']!='linear_weight':
        continue
    name=item['name'].removesuffix('.weight')
    tensors=[source.tensor(n) for n in source_linear_names(name)]
    n=sum(t.shape[0] for t in tensors)
    writer.linear(name,tensors,layer=item.get('layer',0xFFFFFFFF),
                  expert=item.get('expert',0xFFFFFFFF),pad_n=n!=item['dims'][1],
                  required_k_splits=item['flags'])
    count+=1
    if count%200==0:
        print(json.dumps({'exported_linears':count,'seconds':time.monotonic()-start}),flush=True)
replacement={t['name']:t for t in writer.entries}
result=copy.deepcopy(manifest)
result['tensors']=[]
for item in manifest['tensors']:
    if item['name'] in replacement:
        result['tensors'].append(replacement.pop(item['name']))
    else:
        original=copy.deepcopy(item)
        for key in ('path','aux_path'):
            if original.get(key):
                original[key]=str((baseline.parent/original[key]).resolve())
                assert Path(original[key]).is_file()
        result['tensors'].append(original)
assert not replacement, 'candidate unexpectedly added tensors'
result['calibration']={'method':'activation_diagonal_weighted_scale_search_v1',
                       'scales_sha256':writer.calibration_sha256,
                       'preserve_baseline_k_partitions':True}
(a.output/'package-manifest.json').write_text(json.dumps(result,indent=2)+'\n')
print('PASS: complete candidate manifest; baseline assets unchanged',flush=True)