"""Convert rsac checkpoints (.pt) into the Hugging Face release layout (safetensors + JSON).

For each NAME=CKPT it writes OUT/NAME/actor.safetensors (policy only, for deployment),
OUT/NAME/full.safetensors (actor.* and critics.* plus log_alpha, for fine-tuning) and
OUT/NAME/training.json (iteration, samples, run config, last validation row, reward scales),
then reloads both files with modeling_aio and checks the actor gives the same outputs as the
checkpoint. Copy a checkpoint before converting it if a run may still overwrite it.

  python release/huggingface/export.py --out DIR apc-v3-finetune=PATH/best.pt mixed-it6000=PATH/best.pt
"""
import argparse, json, sys
from pathlib import Path
import torch
from safetensors.torch import save_file

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import modeling_aio as M  # noqa: E402


def tensors(state, prefix=''):
    return {prefix+k: v.detach().cpu().clone().contiguous() for k, v in state.items()}


def load_pt(ckpt):
    """torch.load that also works off Windows: run configs pickle pathlib.WindowsPath objects."""
    import pathlib
    saved = pathlib.WindowsPath
    if sys.platform != 'win32':
        pathlib.WindowsPath = pathlib.PureWindowsPath
    try:
        return torch.load(ckpt, map_location='cpu', weights_only=False)
    finally:
        pathlib.WindowsPath = saved


def export(name, ckpt, out):
    ck = load_pt(ckpt)
    d = out/name; d.mkdir(parents=True, exist_ok=True)
    meta = {'format': 'pt', 'source': Path(ckpt).name, 'iteration': str(ck.get('iteration')),
            'samples': str(ck.get('samples'))}
    save_file(tensors(ck['model']['actor']), d/'actor.safetensors', metadata=meta)
    full = tensors(ck['model']['actor'], 'actor.')
    if 'critics' in ck['model']:
        full.update(tensors(ck['model']['critics'], 'critics.'))
    if 'log_alpha' in ck:
        full['log_alpha'] = torch.tensor([float(ck['log_alpha'])])
    save_file(full, d/'full.safetensors', metadata=meta)
    info = {k: ck.get(k) for k in ('iteration', 'samples', 'log_alpha', 'best', 'stale', 'reward_scale', 'val', 'config')}
    (d/'training.json').write_text(json.dumps(info, indent=1, default=str))

    ref = M.Actor(); ref.load_state_dict(ck['model']['actor']); ref.eval()
    g = torch.Generator().manual_seed(0)
    obs = torch.randn(3, 7, M.SIZE, generator=g); prev = torch.rand(3, 7, M.NMV, generator=g)*2-1
    with torch.no_grad():
        a = ref(obs, prev)[0]
        for f in ('actor.safetensors', 'full.safetensors'):
            b = M.load_actor(d/f)(obs, prev)[0]
            assert torch.equal(a, b), f'{name}/{f}: actor output differs'
    if 'critics' in ck['model']:
        M.load_full(d/'full.safetensors')
    n = M.count(ref)
    print(json.dumps(dict(name=name, iteration=ck.get('iteration'), samples=ck.get('samples'), actor_params=n,
                          files={p.name: p.stat().st_size for p in sorted(d.iterdir())})))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--out', type=Path, required=True)
    ap.add_argument('ckpts', nargs='+', help='NAME=CKPT')
    a = ap.parse_args()
    for item in a.ckpts:
        name, path = item.split('=', 1)
        export(name, Path(path), a.out)


if __name__ == '__main__':
    main()