"""Reload saved ARTI cases and compare uninterrupted / resumed continuation.

Run after reading_qa.py all. Seed 17, both languages. The extra three updates
are continuation diagnostics and never enter benchmark effect scores.
"""
import argparse
import json
from pathlib import Path
import tempfile
from huggingface_hub import snapshot_download
import reading_qa as q
from arti.training_run import ARTITrainingRunConfig


def verify(work):
    base=Path(snapshot_download(q.BASE,revision=q.REVISION,cache_dir=work/'hub-cache',local_files_only=True))
    root=Path(tempfile.mkdtemp(prefix='lifecycle-',dir=work))
    receipts,count={},0
    for lang in ('squad','cmrc'):
        folder=work/lang/'arti-17'
        rows=json.loads((work/lang/'heldout.json').read_text(encoding='utf8'))
        predictor=q.CausalLMPredictor.from_pretrained(folder/'native-run/model',device='cuda',base_model=base,local_files_only=True)
        answers,seconds=q.generate(predictor.model,predictor.tokenizer,rows)
        expected=json.loads((folder/'predictions.json').read_text(encoding='utf8'))
        assert [a['id'] for a in answers]==[a['id'] for a in expected]
        assert [a['prediction'] for a in answers]==[a['prediction'] for a in expected]
        count+=len(answers)
        receipts[lang]={'cold_reload_answers':len(answers),'cold_reload_exact':True,'generation_seconds':seconds,
                        'model_manifest_sha256':q.digest((folder/'native-run/model/model.json').read_bytes())}
        del predictor
        q.cleanup()
        checkpoint=folder/'native-run/checkpoints/step-00000200/model.arti.st'
        def branch(name,source,steps):
            target=root/lang/name
            model,tokenizer,recipe,params=q.model_for('arti',base,work,lang,17,target)
            recipe.session.load_checkpoint(source)
            run=ARTITrainingRunConfig(output_dir=target/'run',checkpoint_steps=steps,checkpoint_seconds=86400,log_steps=steps)
            fitted=recipe.session.fit(recipe.train_data,steps=steps,run=run,max_grad_norm=1.,objective_weight=recipe.training_weight)
            state=[p.detach().cpu().clone() for p in params]
            cursor,step,saved=recipe.train_data.state_dict(),recipe.session.global_step,fitted.checkpoint_path
            del model,recipe,params
            q.cleanup()
            return state,cursor,step,saved
        a,ca,sa,_=branch('continuous',checkpoint,3)
        _,_,_,saved=branch('interrupted',checkpoint,1)
        b,cb,sb,_=branch('reloaded',saved,2)
        maximum=max((x-y).abs().max().item() for x,y in zip(a,b))
        assert maximum==0 and ca==cb and sa==sb==203
        receipts[lang].update(continuation_steps=3,checkpoint_step=200,final_step=203,
                              parameter_max_abs=maximum,data_cursor_identical=True)
        print(json.dumps({'lifecycle_verified':lang,**receipts[lang]}),flush=True)
    result={'cold_reload':True,'resume':True,'reloaded_ids':count,'resumed_steps':3,
            'scope':'both languages, seed 17; separate process; continuation diagnostics excluded from effect scores',
            'receipts':receipts}
    q.write(work/'lifecycle.json',result)
    return result


if __name__=='__main__':
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--work',type=Path,default=Path('reading-qa-work'))
    verify(parser.parse_args().work.resolve())