"""Package frozen recent YOR snapshots for Thor, retaining temporal semantics."""
import argparse
import hashlib
import json
import os
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


def sha256_stream(stream):
    digest = hashlib.sha256()
    for block in iter(lambda: stream.read(8*1024*1024), b''):
        digest.update(block)
    return digest.hexdigest()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--selection', required=True)
    parser.add_argument('--out', required=True)
    parser.add_argument('--encoder-bundle', required=True)
    args = parser.parse_args()
    out = Path(args.out).resolve()
    out.mkdir(parents=True, exist_ok=False)
    base = Path(args.encoder_bundle).resolve()
    torch.set_num_threads(4)
    for repo in sorted(Path(torch.hub.get_dir()).glob('facebookresearch_dinov2*')):
        if (repo/'dinov2').is_dir():
            sys.path.insert(0, str(repo))
            break
    encoder = torch.load(base/'encoder.pt', map_location='cpu', weights_only=False)
    encoder.eval()
    base_cfg = OmegaConf.load(base/'training_config.yaml')
    base_card = yaml.safe_load((base/'card.yaml').read_text())
    source_tick_hz = float(base_card['action_hz'])
    (out/'common').mkdir()
    shutil.copy2(base/'encoder.pt', out/'common/encoder.pt')
    shutil.copytree(ROOT/'models', out/'runtime/models',
                    ignore=shutil.ignore_patterns('__pycache__', '*.pyc'))
    manifest = []
    with torch.no_grad():
        tokens = encoder(torch.zeros(3, 3, 224, 224))
    assert tokens.shape == (3, 1, 384)
    for item in json.loads(Path(args.selection).read_text()):
        run = Path(item['run']).resolve()
        cfg = OmegaConf.load(run/'.hydra/config.yaml')
        for key in ['name', 'feature_key', 'n_patches', 'output_dim']:
            assert cfg.encoder[key] == base_cfg.encoder[key], key
        assert cfg.dataset.views == 3
        source = run/'snapshot.pt'
        # Pin the inode while training atomically replaces later snapshots.
        with source.open('rb') as stream:
            source_stat = os.fstat(stream.fileno())
            source_hash = sha256_stream(stream)
            stream.seek(0)
            snapshot = torch.load(stream, map_location='cpu', weights_only=False)
        completed = int(snapshot['epoch'])+1
        policy = snapshot['model']
        policy.to('cpu')
        policy.eval()
        history, chunk = int(cfg.model.obs_window_size), int(cfg.model.act_window_size)
        assert policy.vqvae_is_fit and policy.normalize_actions
        assert (policy._obs_window_size, policy._act_window_size, policy._act_dim) == (history, chunk, 58)
        assert policy.action_norm is not None and all(torch.isfinite(v).all() for v in policy.action_norm)
        assert (policy.action_norm[1] > 0).all()
        temporal = OmegaConf.to_container(cfg.data.get('temporal', {}), resolve=True)
        skip = int(temporal.get('frame_skip', 1))
        online_filter = temporal.get('motion_filter', {}).get('enabled', False)
        prefiltered = bool(cfg.data.get('source_prefiltered',
                           cfg.dataset['_target_'] == 'datasets.yor_jittered.YorJitteredDataset'))
        variable = bool(online_filter or prefiltered)
        offset = int(cfg.data.get('action_offset', 0))
        label = f"{item['label']}_e{completed:04d}"
        bundle = out/label
        bundle.mkdir()
        torch.save(policy, bundle/'policy.pt')
        os.link(out/'common/encoder.pt', bundle/'encoder.pt')
        frame_hz = float(item.get('source_frame_hz', source_tick_hz))
        tick_hz = float(item.get('source_tick_hz', source_tick_hz))
        card = dict(views=['odin', 'wl', 'wr'], img_size=224, obs_horizon=history,
                    chunk=chunk, exec_n=4, act_fields=ACTION_COLS, proprio_fields=[],
                    action_offset=offset, frame_skip=skip, normalize_actions=True,
                    prediction_units='raw_command_units', source_tick_hz=tick_hz,
                    source_frame_hz=frame_hz,
                    timing_mode='variable_waypoints' if variable else 'fixed_rate',
                    training_action_hz=None if variable else frame_hz/skip,
                    deployment_status='offline_test_bundle' if variable else 'smoke_tested_bundle')
        if not variable:
            card['action_hz'] = frame_hz/skip
        (bundle/'card.yaml').write_text(yaml.safe_dump(card, sort_keys=False))
        OmegaConf.save(cfg, bundle/'training_config.yaml', resolve=False)
        obs = tokens.reshape(1, 1, 3, 384).repeat(1, history, 1, 1)
        with torch.no_grad():
            torch.manual_seed(7)
            expected, _, _ = policy(obs, None, None)
            restored = torch.load(bundle/'policy.pt', map_location='cpu', weights_only=False)
            restored.eval()
            torch.manual_seed(7)
            actual, _, _ = restored(obs, None, None)
        assert actual.shape == (1, history, chunk, 58) and torch.isfinite(actual).all()
        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
        for before, after in zip(policy.action_norm, restored.action_norm):
            torch.testing.assert_close(before, after, rtol=0, atol=0)
        row = dict(label=label, recipe_name=item['recipe'], source=str(source),
                   source_snapshot_sha256=source_hash, source_bytes=source_stat.st_size,
                   source_mtime_ns=source_stat.st_mtime_ns, completed_epochs=completed,
                   planned_epochs=int(cfg.epochs), training_complete=completed >= int(cfg.epochs),
                   wandb_run_id=snapshot.get('wandb_run_id'), batch_size=int(cfg.batch_size),
                   observation_history=history, action_chunk=chunk, action_offset=offset,
                   frame_skip=skip, motion_filter=temporal.get('motion_filter'),
                   dataset_prefiltered=prefiltered, timing_mode=card['timing_mode'],
                   normalized=True, outputs_raw_units=True, smoke_shape=list(actual.shape))
        (bundle/'provenance.json').write_text(json.dumps(row, indent=2)+'\n')
        timing = ('Targets are variable-time waypoints. There is no fixed training action Hz. '
                  'The existing fixed-rate Thor runner needs an explicit waypoint execution/history strategy '
                  'before live use; do not assign a 30 Hz training rate to this model.' if variable else
                  f'Targets/history are sampled every {skip} source frames, approximately {frame_hz/skip:.3f} Hz.')
        (bundle/'README.md').write_text(
            f'# {label}\n\nCompleted {completed}/{cfg.epochs} epochs. Batch {cfg.batch_size}. '
            f'V3 CLS; normalized internally; 58D outputs are already in raw command units.\n\n'
            f'History {history}; chunk {chunk}; target offset {offset}; frame skip {skip}. '
            f'obs[t] predicts action rows starting at t+{offset} in the sampled sequence.\n\n'
            f'{timing}\n\nRun the package `smoke_all.py --only {label}` on Thor to test this policy '
            'without cameras or robot RPC. This export does not change the robot runner.\n')
        manifest.append(row)
        print(json.dumps(row), flush=True)
        del snapshot, policy, restored
    (out/'manifest.json').write_text(json.dumps(manifest, indent=2)+'\n')
    shutil.copy2(ROOT/'tools/cls_debug/smoke_thor_bundle.py', out/'smoke_all.py')
    lines = [f'# Recent YOR CLS policies\n',
             'Each directory contains policy.pt, encoder.pt, card.yaml, training_config.yaml and provenance.json. '
             'Encoder files share storage through hard links. Policy normalization is internal; outputs are raw.\n',
             '| Bundle | Epochs | Batch | Chunk | Skip | Offset | Timing |',
             '|---|---:|---:|---:|---:|---:|---|']
    for row in manifest:
        lines.append(f"| {row['label']} | {row['completed_epochs']}/{row['planned_epochs']} | {row['batch_size']} | {row['action_chunk']} | {row['frame_skip']} | {row['action_offset']} | {row['timing_mode']} |")
    lines.extend(['', 'Offline smoke on Thor (no cameras or robot RPC):', '', '```bash',
                  f'cd ~/ckpts/{out.name}',
                  '/home/thor2/miniconda3/envs/charles/bin/python smoke_all.py --device cuda', '```', '',
                  'Use `--only <bundle-name>` to test one policy. Variable-waypoint cards intentionally omit '
                  '`action_hz`: the existing fixed-rate live runner cannot infer their physical timing from a chunk. '
                  'These models need explicit waypoint timing/history handling before live execution.', '',
                  'The unfiltered chunk-20 baseline is already installed at '
                  '`~/ckpts/yor_cls_chunk20_20260912/cls_v3_norm_chunk20_unfiltered_e1000`.'])
    (out/'README.md').write_text('\n'.join(lines)+'\n')
    sums = []
    for path in sorted(out.rglob('*')):
        if path.is_file():
            with path.open('rb') as stream:
                sums.append(sha256_stream(stream)+'  '+str(path.relative_to(out)))
    (out/'SHA256SUMS').write_text('\n'.join(sums)+'\n')


if __name__ == '__main__':
    main()