"""Export a labeled, frozen comparison set from YOR training snapshots."""
import argparse
import copy
import hashlib
import json
from pathlib import Path
import shutil
import sys

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

import torch
import yaml
from omegaconf import OmegaConf
from utils.export_thor_ckpt import ACTION_COLS, measured_tick_dt
from models.encoder.dino import DinoV2Encoder


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--out', required=True)
    ap.add_argument('--runner-config', required=True)
    args = ap.parse_args()
    out = Path(args.out).resolve()
    out.mkdir(parents=True, exist_ok=True)
    torch.set_num_threads(4)
    shutil.copytree(ROOT / 'models', out / 'runtime/models', dirs_exist_ok=True,
                    ignore=shutil.ignore_patterns('__pycache__', '*.pyc'))
    base = yaml.safe_load(Path(args.runner_config).read_text())
    hz = round(1 / measured_tick_dt(Path('/home_shared/grail_charles/data/yor_episodes/df767cdb')), 3)
    encoders = {}
    manifest = []
    recipes = [('C-V3-M', 'cls_v3_norm'), ('C-V1-M', 'cls_v1_norm'),
               ('C-V3-R', 'cls_v3_raw'), ('C-V1-R', 'cls_v1_raw'),
               ('P-V3-M', 'patch_v3_norm'), ('P-V1-M', 'patch_v1_norm'),
               ('P-V3-R', 'patch_v3_raw')]
    for code, recipe in recipes:
        paths = list((ROOT / 'outputs/yor' / code).glob('2026-09-10/*/snapshot.pt'))
        assert len(paths) == 1, (code, paths)
        source = paths[0]
        # Training replaces snapshots atomically: this open pins one complete version.
        with source.open('rb') as f:
            snapshot = torch.load(f, map_location='cpu', weights_only=False)
        policy = snapshot['model']
        epoch = int(snapshot['epoch'])
        cfg = OmegaConf.load(source.parent / '.hydra/config.yaml')
        views = ['odin', 'wl', 'wr'][:int(cfg.dataset.views)]
        horizon = int(cfg.model.obs_window_size)
        chunk = int(cfg.model.act_window_size)
        feature = str(cfg.encoder.feature_key)
        patches = int(cfg.encoder.n_patches)
        assert policy.vqvae_is_fit
        assert policy._obs_window_size == horizon
        assert policy._gpt_model.config.n_patches == len(views) * patches
        if policy.normalize_actions:
            assert policy.action_norm is not None
            assert all(torch.isfinite(t).all() for t in policy.action_norm)
        label = f'{recipe}_e{epoch + 1:04d}'
        dest = out / label
        dest.mkdir(exist_ok=True)
        policy.to('cpu')
        policy.eval()
        torch.save(policy, dest / 'policy.pt')
        if feature not in encoders:
            enc = DinoV2Encoder(name=str(cfg.encoder.name), feature_key=feature,
                                n_patches=patches)
            enc.to('cpu')
            enc.eval()
            for p in enc.parameters():
                p.requires_grad = False
            encoders[feature] = enc
        enc = encoders[feature]
        torch.save(enc, dest / 'encoder.pt')
        card = dict(views=views, img_size=224, obs_horizon=horizon, chunk=chunk,
                    exec_n=4, action_hz=hz, act_fields=ACTION_COLS, proprio_fields=[])
        (dest / 'card.yaml').write_text(yaml.safe_dump(card, sort_keys=False))
        runner = copy.deepcopy(base)
        runner['policy'].update(ckpt_dir=f'~/ckpts/{out.name}/{label}', chunk=chunk, exec_n=4)
        runner['action'].update(hz=hz, preview_time=1/hz)
        runner['cameras']['wrist'] = len(views) > 1
        (dest / 'runner.yaml').write_text(yaml.safe_dump(runner, sort_keys=False))
        OmegaConf.save(cfg, dest / 'training_config.yaml', resolve=False)
        with torch.no_grad():
            tokens = enc(torch.zeros(len(views), 3, 224, 224))
            assert tokens.shape == (len(views), patches, 384)
            obs = tokens.reshape(1, 1, len(views)*patches, 384).repeat(1,horizon,1,1)
            prediction, _, _ = policy(obs, None, None)
            assert prediction.shape == (1,horizon,chunk,58), prediction.shape
            assert torch.isfinite(prediction).all()
        row = dict(label=label, recipe=recipe, source=str(source), epoch_index=epoch,
                   completed_epochs=epoch+1, views=views, obs_horizon=horizon,
                   feature_key=feature, normalized=bool(policy.normalize_actions),
                   action_hz=hz, stride=1, wandb_run_id=snapshot['wandb_run_id'])
        (dest / 'provenance.json').write_text(json.dumps(row, indent=2)+'\n')
        manifest.append(row)
        print(json.dumps(row), flush=True)
        del snapshot, policy
    (out / 'manifest.json').write_text(json.dumps(manifest, indent=2)+'\n')
    sums = []
    for path in sorted(out.rglob('*')):
        if path.is_file() and path.name != 'SHA256SUMS' and '__pycache__' not in path.parts:
            with path.open('rb') as f:
                h = hashlib.sha256()
                for block in iter(lambda: f.read(1024*1024), b''):
                    h.update(block)
            sums.append(f'{h.hexdigest()}  {path.relative_to(out)}')
    (out / 'SHA256SUMS').write_text('\n'.join(sums)+'\n')


if __name__ == '__main__':
    main()