"""Stage the release (weights, tokenizer, evaluation records, training code) into /job/export/hf and a reproducibility
copy into /job/export/project. README.md is written separately from results.json and reviewed before upload."""
import json, shutil
import numpy as np
from datasets import load_from_disk
from common import *

HF = EXPORT/'hf'; PROJECT = EXPORT/'project'
EVAL = CKPT/'evaluation'; SEL = CKPT/'selection'

def copy(src, dest):
    dest = Path(dest); dest.parent.mkdir(parents=True, exist_ok=True); shutil.copy2(src, dest)

def looks():
    """Frozen candidate plus the disclosed second look; the publication rule was fixed before the second result."""
    out = [json.loads((EVAL/'results.json').read_text()) | {'look': 'frozen by validation NLL', 'file': 'results.json',
                                                           'logits': 'candidate-benchmark-final.npz'}]
    for f in sorted(EVAL.glob('results-*.json')):
        tag = f.stem[len('results-'):]
        out.append(json.loads(f.read_text()) | {'look': 'second look (long-trained), maintainer request', 'file': f.name,
                                                'logits': f'candidate-{tag}-benchmark-final.npz'})
    def key(r):
        o = r['benchmark_default']['overall']
        return (r['correct'], -o['false_approve_count'])
    chosen = out[0]
    for r in out[1:]:
        if key(r) > key(chosen): chosen = r
    return chosen, out


def main():
    results, all_looks = looks()
    cand = Path(results['candidate']['path']); assert sha(cand/'model.safetensors') == results['candidate']['sha256']
    HF.mkdir(parents=True, exist_ok=True)
    for f in ['model.safetensors', 'config.json', 'tokenizer.json']: copy(cand/f, HF/f)
    # transformers 5 writes tokenizer_class "TokenizersBackend", which transformers 4.x and llama.cpp cannot load.
    tc = json.loads((cand/'tokenizer_config.json').read_text())
    tc['tokenizer_class'] = 'PreTrainedTokenizerFast'
    for k in ['backend', 'is_local', 'local_files_only']: tc.pop(k, None)
    tc['model_max_length'] = MAX_LEN
    atomic_json(HF/'tokenizer_config.json', tc)
    cfg = json.loads((HF/'config.json').read_text()); cfg.pop('_name_or_path', None); atomic_json(HF/'config.json', cfg)
    bench = load_from_disk(str(DATA/'benchmark'))
    logits = np.load(EVAL/results['logits'])['logits']
    np.savez(HF/'benchmark_predictions.npz', logits=logits, labels=np.array(bench['labels']), source_row=np.array(bench['source_row']))
    slim = {k: v for k, v in results.items() if k not in ('plan', 'logits', 'file')}
    atomic_json(HF/'eval_results.json', slim)
    atomic_json(HF/'eval'/'benchmark_looks.json', {
        'rule': json.loads((ROOT/'plan.json').read_text()).get('second_look'),
        'published': results['candidate']['name'],
        'looks': [{'candidate': r['candidate']['name'], 'look': r['look'], 'correct': r['correct'],
                   'false_approve_count': r['benchmark_default']['overall']['false_approve_count'],
                   'false_deny_count': r['benchmark_default']['overall']['false_deny_count'],
                   'auroc': r['benchmark_default']['overall']['auroc'],
                   'long_16k_64k_accuracy': r['benchmark_default']['slices']['length']['16k-64k']['accuracy'],
                   'mid_4k_16k_accuracy': r['benchmark_default']['slices']['length']['4k-16k']['accuracy'],
                   'audit_accuracy': r['validation_audit']['overall']['accuracy'],
                   'validation_nll': r['candidate']['nll']} for r in all_looks]})
    label = 'candidate' if results['file'] == 'results.json' else f"candidate-{results['candidate']['name']}"
    for src, dest in [(f'{label}-published-probes', 'candidate-published-probes'), (f'{label}-fresh-probes', 'candidate-fresh-probes'),
                      ('reference-published-probes', 'reference-published-probes'), ('reference-fresh-probes', 'reference-fresh-probes')]:
        copy(EVAL/f'{src}.json', HF/'eval'/f'{dest}.json')
    copy(ROOT/'tool_probes.json', HF/'eval'/'tool_probes.json'); copy(ROOT/'fresh_probes.json', HF/'eval'/'fresh_probes.json')
    copy(ROOT/'run_tool_probes.py', HF/'eval'/'run_tool_probes.py')
    copy(SEL/'summary.json', HF/'eval'/'selection_summary.json')
    copy(CKPT/'pilot_choice.json', HF/'eval'/'lr_pilots.json')
    copy(DATA/'rope_sweep.json', HF/'eval'/'rope_sweep.json')
    copy(DATA/'smoke.json', HF/'eval'/'runtime_smoke.json')
    for run in ['p1_short', 'p2_long']:
        copy(CKPT/run/'complete.json', HF/'eval'/'training'/f'{run}.json')
    copy(DATA/'data_audit.json', HF/'data_audit.json')
    copy(JOB/'work'/'environment.json', HF/'environment.json')
    atomic_json(HF/'source_revisions.json', {k: {x: v.get(x) for x in ['repo', 'revision', 'weights_sha256', 'files'] if x in v} for k, v in sources().items() if k != 'teacher'}
                | {'teacher': {'name': 'Private Auto 3B classifier (unpublished)', 'weights_sha256': TEACHER_SHA}})
    for f in ['common.py', 'download.py', 'prepare.py', 'rope.py', 'smoke.py', 'engine.py', 'train.py', 'choose.py', 'evaluate.py', 'chain.py', 'stage.py', 'card.py', 'plan.json']:
        copy(ROOT/f, HF/'training'/f)
    # Project copy for reproduction (source, runner, logs); no data, checkpoints or caches.
    if PROJECT.exists(): shutil.rmtree(PROJECT)
    shutil.copytree(ROOT, PROJECT/'src', ignore=shutil.ignore_patterns('__pycache__'))
    shutil.copytree(JOB/'work'/'runner', PROJECT/'runner')
    shutil.copytree(JOB/'work'/'state', PROJECT/'tasks')
    shutil.copytree(LOG, PROJECT/'logs')
    copy(JOB/'work'/'environment.json', PROJECT/'environment.json')
    event('staged', files=sorted(str(p.relative_to(HF)) for p in HF.rglob('*') if p.is_file()))

if __name__ == '__main__': main()