"""Verify the fixed-budget, one-seed PhysX direction pilot and export its results."""
import argparse
import gzip
import hashlib
import json
from pathlib import Path


def export(root, previous, output):
    inputs = {}
    def digest(path):
        return hashlib.sha256(path.read_bytes()).hexdigest()
    def read(path):
        inputs[str(path)] = digest(path)
        return json.loads(path.read_text())
    def checkpoint(path):
        manifest = read(path/'checkpoint-manifest.json')
        assert manifest['environment_steps'] == 131072
        for name, expected in manifest['files'].items():
            assert digest(path/name) == expected
        return dict(files=manifest['files'], interface=read(path/'interface.json'))
    def summary(rows):
        return dict(episodes=len(rows), successes=sum(r['success'] for r in rows),
                    overloads=sum(r['failure'] == 'overload_stop' for r in rows),
                    other_unsuccessful=sum(not r['success'] and r['failure'] != 'overload_stop' for r in rows))
    def recordings(directory, rows):
        for row in rows:
            p = directory/row['replay']['file']
            assert digest(p) == row['replay']['sha256']
            data = json.loads(gzip.decompress(p.read_bytes()))
            assert digest(directory/data['scene_file']) == data['scene_file_sha256']
            assert data['terminal_pose_recorded']

    original = root/'results/physx-cpu-direction-pilot-01'
    recovered = root/'results/physx-cpu-direction-recovery-01'
    protocol = read(original/'protocol.json')
    assert protocol['seed'] == 0 and protocol['sample_budget'] == 131072
    assert protocol['training_engine'] == 'physx' and protocol['force_penalties'] == [0, 10]
    assert read(recovered/'status.json')['status'] == 'complete'
    policies, validations, comparisons, episodes = {}, {}, {}, []
    for weight in (0, 10):
        physx = (original/'runs/physx-force0-seed0/checkpoints/000131072' if weight == 0 else
                 recovered/'runs/physx-force10-seed0')
        mujoco = previous/f'results/review-resumed-01/runs/insert-force{weight}-seed0/checkpoints/000131072'
        policies[str(weight)] = dict(physx_trained=checkpoint(physx), mujoco_trained=checkpoint(mujoco))
        assert policies[str(weight)]['mujoco_trained']['files']['policy.zip'] == protocol['comparisons'][str(weight)]['sha256']
        for value in policies[str(weight)].values():
            assert value['interface']['force_penalty'] == weight
        source = recovered/'validation'/f'force{weight}'
        assert read(source/'status.json')['status'] == 'complete'
        rows = read(source/'outcomes.json')
        assert len(rows) == 24 and sorted(r['worker'] for r in rows) == list(range(24))
        assert {r['seed'] for r in rows} == {1000}
        recordings(source, rows)
        validations[str(weight)] = dict(physx=summary(rows),
            mujoco=summary(read(mujoco/'validation-outcomes.json')))
        for row in rows:
            episodes.append(dict(partition='source_validation', force_weight=weight,
                                 evaluation_engine='physx', training_seed=0, **row))
        comparisons[str(weight)] = {}
        for engine in ('physx', 'physx_gpu', 'mujoco_cpu'):
            directory = recovered/'evaluations'/f'force{weight}-{engine}'
            status = read(directory/'status.json')
            assert status['status'] == 'complete' and status['episodes'] == 24 and not status['errors']
            rows = read(directory/'outcomes.json')
            assert len(rows) == 24 and len({(r['scene_index'], r['policy'], r['seed']) for r in rows}) == 24
            recordings(directory, rows)
            stored = read(recovered/f'comparison-force{weight}-{engine}.json')
            assert stored['source_outcomes_sha256'] == digest(directory/'outcomes.json')
            groups = {}
            for name in ('physx_trained', 'mujoco_trained', 'scripted'):
                identity = None if name == 'scripted' else policies[str(weight)][name]['files']['policy.zip']
                selected = [r for r in rows if r['policy_sha256'] == identity]
                assert len(selected) == 8 and sorted(r['scene_index'] for r in selected) == list(range(8))
                assert {r['seed'] for r in selected} == {2000}
                groups[name] = summary(selected)
                for key in ('episodes', 'successes', 'overloads'):
                    assert groups[name][key] == stored['policies'][name][key]
                for row in selected:
                    episodes.append(dict(partition='prior_eight',force_weight=weight,
                        evaluation_engine=engine,policy_group=name,training_seed=None if name == 'scripted' else 0,**row))
            comparisons[str(weight)][engine] = groups
    assert len(episodes) == 192
    result = dict(schema='simcompare.blog.physx-direction-pilot.v1', status='complete',
        training_seed=0,sample_budget_per_policy=131072,training_engine='PhysX CPU',
        additional_trainings=2,episode_records=192,protocol=protocol,
        selection='Fixed final checkpoint declared before outcomes; no checkpoint selection on these tests.',
        scope='One seed per reward, 24 source-validation starts and eight previously used test starts. This is not the five-seed, 50-new-start cohort. MuJoCo comparators use the same 131072-transition budget.',
        interpretation='The PhysX CPU baseline is less successful than the matched MuJoCo policy in every tested target. The PhysX force-penalty policy has zero successful insertions across all targets. Lower overload counts without task completion are not robustness. Source learning is not solved at this budget; successful behavior or engine-wide robustness differences are not established.',
        recovery='Baseline completed all PPO updates and saved the final checkpoint before an intermittent validation XML import failure. Unchanged scene passed an isolated check; the complete validation was repeated in a fresh journaled pass. No baseline retraining, physics change, or outcome selection. Original failure evidence retained.',
        replay='All 192 original native pose/action recordings and geometry hashes verified in the compact workspace backup; camera-independent 100 Hz samples plus finite terminal state, with 1 kHz peak diagnostics. Raw archive is not published.',
        validations=validations,comparisons=comparisons,checkpoints=policies,
        source_files_sha256=inputs,exporter_sha256=digest(Path(__file__)))
    output.mkdir(parents=True,exist_ok=True)
    for name, data in [('summary.json', result), ('episodes.json', episodes)]:
        (output/name).write_text(json.dumps(data,indent=2,allow_nan=False)+'\n')
    manifest = dict(schema='simcompare.blog.assets.v1',status='complete',files={
        p.name:dict(sha256=digest(p),bytes=p.stat().st_size)
        for p in sorted(output.iterdir()) if p.is_file() and p.name != 'manifest.json'})
    (output/'manifest.json').write_text(json.dumps(manifest,indent=2)+'\n')
    print(json.dumps(dict(verified_episodes=len(episodes),validations=validations,comparisons=comparisons),indent=2))


if __name__ == '__main__':
    parser=argparse.ArgumentParser()
    for name in ('root','previous','output'):parser.add_argument('--'+name,type=Path,required=True)
    args=parser.parse_args()
    export(args.root,args.previous,args.output)