import argparse
import copy
import json
import shutil
import time
from pathlib import Path

import modelopt.torch.quantization as mtq
import torch
from modelopt.torch.export import export_hf_checkpoint
from safetensors import safe_open
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

parser = argparse.ArgumentParser()
parser.add_argument('--source', required=True)
parser.add_argument('--output', required=True)
args = parser.parse_args()
source, output = Path(args.source), Path(args.output)
output.mkdir(parents=True, exist_ok=False)
torch.set_num_threads(4)
torch.manual_seed(20261020)
config = AutoConfig.from_pretrained(source, trust_remote_code=True, local_files_only=True)
config.backend = 'torch'
config.prefill_chunk_size = 512
model = AutoModelForCausalLM.from_pretrained(
    source, config=config, trust_remote_code=True, local_files_only=True,
    dtype=torch.bfloat16, device_map={'': 0}, attn_implementation='sdpa',
).eval().requires_grad_(False)
tokenizer = AutoTokenizer.from_pretrained(source, trust_remote_code=True, local_files_only=True)
tokenizer.chat_template = (Path(__file__).parent/'chat_template.jinja').read_text()
model.generation_config.eos_token_id = [248046, 248044]
model.generation_config.pad_token_id = tokenizer.pad_token_id
modules = {}
recipe = copy.deepcopy(mtq.W4A16_NVFP4_CFG)
for name, module in model.named_modules():
    if not isinstance(module, torch.nn.Linear):
        continue
    if name != 'lm_head' and min(module.weight.shape) >= 128 and '.admission_' not in name:
        assert module.in_features % 16 == 0 and module.out_features % 16 == 0, name
        modules[name] = list(module.weight.shape)
    else:
        recipe['quant_cfg'].append({'quantizer_name': name+'.*', 'enable': False})
assert modules
probe = tokenizer('The sum of 17 and 23 is', return_tensors='pt').to('cuda')
started = time.perf_counter()
print(json.dumps({'stage': 'quantizing', 'modules': len(modules), 'format': 'W4A16_NVFP4'}), flush=True)
with torch.inference_mode():
    mtq.quantize(model, recipe, forward_loop=lambda current: current(**probe, use_cache=False, logits_to_keep=1))
    quantized_probe = model(**probe, use_cache=False, logits_to_keep=1).logits.float().cpu()
assert torch.isfinite(quantized_probe).all()
torch.save(quantized_probe, output.parent/'quantizer_probe.pt')
print(json.dumps({'stage': 'exporting'}), flush=True)
export_hf_checkpoint(model, export_dir=output, save_modelopt_state=False, max_shard_size='4GB')
tokenizer.save_pretrained(output)
for name in ('configuration_edgerock.py', 'modeling_edgerock.py', 'LICENSE', 'LICENSE-CODE', 'NOTICE'):
    shutil.copy2(source/name, output/name)
for name in ('load_nvfp4.py', 'quantize.py'):
    shutil.copy2(Path(__file__).parent/name, output/name)
shutil.copy2(Path(__file__).parent/'chat_template.jinja', output/'chat_template.jinja')
hf_config = json.loads((output/'hf_quant_config.json').read_text())
assert hf_config['quantization']['quant_algo'] == 'W4A16_NVFP4', hf_config
manifest = {
    'format': 'NVFP4', 'recipe': 'W4A16_NVFP4', 'weight_encoding': 'E2M1, two values per byte, low nibble first',
    'block_size': 16, 'block_scale_dtype': 'float8_e4m3fn', 'global_scale_dtype': 'float32',
    'activation_dtype': 'bfloat16', 'scale_layout': 'row-major', 'modules': modules,
    'embedding_and_lm_head': 'unquantized, tied',
    'preserved': ['admission networks', 'small recurrent gates', 'normalization', 'history capacities and update rules'],
}
(output/'nvfp4_config.json').write_text(json.dumps(manifest, indent=2)+'\n')
packed = {}
for shard in sorted(output.glob('model*.safetensors')):
    with safe_open(shard, framework='pt', device='cpu') as handle:
        for name, shape in modules.items():
            if name+'.weight' not in handle.keys():
                continue
            weight = handle.get_tensor(name+'.weight')
            scale = handle.get_tensor(name+'.weight_scale')
            global_scale = handle.get_tensor(name+'.weight_scale_2')
            assert weight.dtype == torch.uint8 and list(weight.shape) == [shape[0], shape[1]//2], name
            assert scale.dtype == torch.float8_e4m3fn and scale.numel() == shape[0]*shape[1]//16, name
            assert global_scale.dtype == torch.float32 and global_scale.numel() == 1, name
            assert torch.isfinite(scale.float()).all() and global_scale.item() > 0, name
            packed[name] = {'shape': shape, 'weight_bytes': weight.numel(), 'scale_bytes': scale.numel()+4}
assert set(packed) == set(modules)
receipt = {
    'format': 'W4A16_NVFP4', 'quantized_modules': len(packed), 'modules': packed,
    'conversion_seconds': time.perf_counter()-started,
    'weight_bytes': sum(path.stat().st_size for path in output.glob('model*.safetensors')),
    'gpu': torch.cuda.get_device_name(0), 'modelopt_version': '0.47.0',
}
(output/'conversion.json').write_text(json.dumps(receipt, indent=2)+'\n')
print(json.dumps({'stage': 'conversion_complete', 'modules': len(packed), 'weight_bytes': receipt['weight_bytes']}), flush=True)