"""Assemble a self-contained public release directory (code + reports) next to its transformer.

usage: package_release.py fp8|nvfp4
Copies runtime/CLI modules to the release root and reproducibility sources to source/, writes
LICENSE/Notice/requirements. README.md is written separately from measured results. Does not publish.
"""
import json, shutil, sys
from pathlib import Path
from common import ROOT, BASE

fmt = sys.argv[1]
SRC = ROOT / 'src'
OUT = ROOT / fmt / 'release'
RUNTIME = {'fp8': ['fp8_runtime.py'], 'nvfp4': ['nvfp4_runtime.py', 'dynamic_scale.py', 'fp8_runtime.py']}[fmt]
COMMON = ['acceleration.py', 'vae_fast.py', 'transparency.py', 'prompt_cache.py']
SOURCE = ['fast_decode.py', 'verify_video.py', 'final_round.sh', 'speed_round3.sh', 'assemble_mixed.py', 'hybrid.py', 'test_split.py', 'micro_attn.py', 'micro_vae.py', 'micro_tune.py', 'common.py', 'prompts.py', 'smoke.py', 'verify_base.py', 'calibrate.py', 'hessians.py', 'hessian_store.py', 'gptq.py',
          'weight_quant.py', 'quantize_fp8.py', 'quantize_nvfp4.py', 'evaluate.py', 'metrics.py', 'benchmark.py', 'speedlab.py',
          'video.py', 'package_release.py', 'test_gptq.py', 'test_vae.py']
for f in RUNTIME + COMMON:
    if (SRC / f).exists():
        shutil.copy(SRC / f, OUT / f)
gen = (SRC / 'generate_template.py').read_text()
fp8_revision = json.loads((ROOT / 'fp8' / 'publication.json').read_text())['revision'] if fmt == 'nvfp4' else 'not-used-by-fp8-release'
gen = (gen.replace('RUNTIME_MODULE', RUNTIME[0][:-3]).replace('RUNTIME_NAME', fmt.upper())
          .replace('DEFAULT_FP8_FIRST_STEPS', '1').replace('FP8_REVISION_PLACEHOLDER', fp8_revision))
(OUT / 'generate.py').write_text(gen)
(OUT / 'source').mkdir(exist_ok=True)
for f in SOURCE:
    if (SRC / f).exists():
        shutil.copy(SRC / f, OUT / 'source' / f)
shutil.copy(BASE / 'LICENSE', OUT / 'LICENSE')
(OUT / 'Notice').write_text(
    'Qwen is licensed under the Qwen RESEARCH LICENSE AGREEMENT, Copyright (c) 2026 Hangzhou Tongyi Laboratory Technology Co., Ltd. '
    'All Rights Reserved.\n\nBuilt with Qwen.\nProCreations modified the Qwen-Image-2.1-Turbo transformer weights through calibrated '
    f'{fmt.upper()} quantization on 2026-10-09. The custom runtime, acceleration and calibration/evaluation scripts were added by '
    'ProCreations. Non-commercial research and evaluation only under the included upstream license.\n')
req = ['torch==2.14.0  # CUDA 13.0 build: pip install torch==2.14.0 --index-url https://download.pytorch.org/whl/cu130', 'torchvision==0.29.0', 'transformers==5.17.0', 'accelerate==1.15.0', 'triton==3.8.0', 'safetensors==0.8.0',
       'huggingface-hub', 'pillow==12.3.0', 'sentencepiece==0.2.2', 'diffusers==0.41.0']
if fmt == 'nvfp4':
    req.append('flashinfer-python==0.7.1')
(OUT / 'requirements.txt').write_text('\n'.join(req) + '\n# Optional, faster attention (--attention sage/auto): SageAttention 2.2.0 built from source (tested commit d1a57a5;\n# Torch 2.14 needs a C++20 host compiler), e.g. pip install --no-build-isolation git+https://github.com/thu-ml/SageAttention.git@d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5\n')
(OUT / 'requirements-evaluation.txt').write_text('-r requirements.txt\nlpips==0.1.4\nscikit-image==0.26.0\nimageio-ffmpeg\n')
shutil.copy(ROOT / 'base_integrity.json', OUT / 'reports' / 'base_integrity.json')
print('PACKAGED', OUT)