import argparse
import json
import re
import time
from pathlib import Path

import torch
from safetensors import safe_open
from transformers import AutoTokenizer
from packed_int4 import PackedLinear, load_int4

parser = argparse.ArgumentParser()
parser.add_argument('--folder', required=True)
parser.add_argument('--source', required=True)
parser.add_argument('--reload-only', action='store_true')
args = parser.parse_args()
folder, source = Path(args.folder), Path(args.source)
torch.set_num_threads(4)
torch.manual_seed(20261020)
model = load_int4(folder)
tokenizer = AutoTokenizer.from_pretrained(folder, trust_remote_code=True, local_files_only=True)
probe = tokenizer('The sum of 17 and 23 is', return_tensors='pt').to('cuda')
with torch.inference_mode():
    actual = model(**probe, use_cache=False, logits_to_keep=1).logits.float().cpu()
assert torch.isfinite(actual).all()
if args.reload_only:
    expected = torch.load(folder/'reference_probe.pt', weights_only=True)
    error = (actual-expected).abs().max().item()
    assert torch.equal(actual, expected), error
    validation = json.loads((folder/'validation.json').read_text())
    validation['fresh_process_reload_max_logit_error'] = error
    validation['complete'] = True
    (folder/'validation.json').write_text(json.dumps(validation, indent=2, ensure_ascii=False)+'\n')
    print(json.dumps({'stage': 'reload_verified', 'max_logit_error': error}), flush=True)
    raise SystemExit(0)
torch.save(actual, folder/'reference_probe.pt')
manifest = json.loads((folder/'edgerock_quantization.json').read_text())
decoded = len(manifest['modules'])
assert sum(isinstance(module, PackedLinear) for module in model.modules()) == decoded
print(json.dumps({'stage': 'packed_model_loaded', 'quantized_modules': decoded}), flush=True)
index = json.loads((source/'model.safetensors.index.json').read_text())['weight_map']
protected = 0
for shard in sorted(folder.glob('model*.safetensors')):
    with safe_open(shard, framework='pt', device='cpu') as checkpoint:
        ordinary = [name for name in checkpoint.keys() if name in index and
                    name.removesuffix('.weight') not in manifest['modules']]
        for original_name in sorted({index[name] for name in ordinary}):
            with safe_open(source/original_name, framework='pt', device='cpu') as original:
                for name in ordinary:
                    if index[name] == original_name:
                        candidate = checkpoint.get_tensor(name)
                        reference = original.get_tensor(name)
                        assert candidate.dtype == reference.dtype and torch.equal(candidate, reference), name
                        protected += 1
assert protected > 0
cases = [
    ('multiply', 'Calculate 17 * 23. Give only the answer.', r'\b391\b'),
    ('equation', 'Solve x: 3*x + 7 = 31. Give only the answer.', r'\b8\b'),
    ('memory_update', 'The active project code is COBALT-7391.\n'
        + 'The archive contains ordinary observations about the weather and landscape. '*120
        + '\nUpdate: the active project code is now AMBER-5286. The earlier code is obsolete.\n'
        + 'Additional notes record the temperature, wind, and clouds each afternoon. '*160
        + '\nWhat is the active project code now? Give only that code.', r'AMBER-5286'),
]
rows = []
for name, prompt, expected in cases:
    text = tokenizer.apply_chat_template([{'role': 'user', 'content': prompt}],
                                         tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors='pt', add_special_tokens=False).to('cuda')
    started = time.perf_counter()
    torch.cuda.reset_peak_memory_stats()
    with torch.inference_mode():
        output = model.generate(**inputs, max_new_tokens=2048, do_sample=False, use_cache=True)
    torch.cuda.synchronize()
    answer = tokenizer.decode(output[0, inputs.input_ids.shape[1]:], skip_special_tokens=True)
    final = answer.split('</think>')[-1].strip()
    row = {
        'name': name, 'prompt': prompt, 'prompt_tokens': inputs.input_ids.shape[1],
        'output_tokens': output.shape[1]-inputs.input_ids.shape[1],
        'output': answer, 'final_answer': final,
        'matched_expected': bool(re.search(expected, final)),
        'finished_by_eos': int(output[0, -1]) in (248046, 248044),
        'seconds': time.perf_counter()-started, 'peak_allocated_gib': torch.cuda.max_memory_allocated()/2**30,
    }
    rows.append(row)
    print(json.dumps({'stage': 'generation', **row}, ensure_ascii=False), flush=True)
validation = {
    'complete': False, 'format': manifest['method'].upper()+' W4A16', 'unquantized_tensors_exact': protected,
    'packed_linear_modules': decoded,
    'generations': rows, 'all_generation_checks_passed': all(row['matched_expected'] and row['finished_by_eos'] for row in rows),
    'runtime': 'GPU packed INT4 linear weights, fused Triton BF16-activation matmul',
    'hardware': torch.cuda.get_device_name(0), 'max_new_tokens': 2048, 'do_sample': False,
}
(folder/'validation.json').write_text(json.dumps(validation, indent=2, ensure_ascii=False)+'\n')
assert validation['all_generation_checks_passed'], 'Generation checks did not pass; publication has not been started.'