"""Average compatible checkpoints from the same randomly initialized training lineage."""
import argparse
import json
from pathlib import Path
import torch
from safetensors.torch import load_file,save_file
from tinyquery.prepare import file_sha256


def main():
    p=argparse.ArgumentParser();p.add_argument('--checkpoints',nargs='+',required=True);p.add_argument('--weights',nargs='+',type=float)
    p.add_argument('--out',required=True);args=p.parse_args();weights=args.weights or [1]*len(args.checkpoints)
    assert len(weights)==len(args.checkpoints) and all(w>=0 for w in weights) and sum(weights)>0
    weights=[w/sum(weights) for w in weights];total={};config=None;sources=[]
    for name,weight in zip(args.checkpoints,weights):
        path=Path(name);c=json.loads((path.parent/'config.json').read_text());assert config is None or config==c;config=c
        info=json.loads((path.parent/'checkpoint-info.json').read_text());assert info['random_initialization']
        state=load_file(str(path));assert not total or total.keys()==state.keys()
        for key,value in state.items():
            if key not in total:total[key]=value.float()*weight
            else:total[key].add_(value.float(),alpha=weight)
        sources.append({'checkpoint':str(path),'sha256':file_sha256(path),'weight':weight,'info':info});del state
    dest=Path(args.out);dest.mkdir(parents=True,exist_ok=True)
    save_file({k:v.to(torch.bfloat16).contiguous() for k,v in total.items()},str(dest/'model.safetensors'),
              metadata={'method':'weighted_parameter_average','random_initialization':'true'})
    (dest/'config.json').write_text(json.dumps(config,indent=2))
    info={'method':'weighted_parameter_average','sources':sources,'random_initialization':True,
          'step':max(s['info']['step'] for s in sources),'step_interpretation':'Latest source step; these are averaged weights, not that raw checkpoint.'}
    for key in ['processed_tokens','response_tokens','training_seconds']:info[key]=max(s['info'][key] for s in sources)
    (dest/'checkpoint-info.json').write_text(json.dumps(info,indent=2));print(json.dumps({'out':str(dest),'sources':[(s['info']['step'],s['weight']) for s in sources]}))


if __name__=='__main__':main()