"""Export the settings-ablation showcase: one frozen policy and start, four physics settings.

Usage: python replay_export.py --simcompare EXPERIMENT_DIR --results RESULTS_DIR --output DIR
Seed 4 (baseline) on held-out start 54: MuJoCo default contact (10 ms) and 2 ms contact from the
A2 screen; PhysX TGS 100/1 (default) and PGS from the held-out confirmation run.
"""
import argparse
import hashlib
import json
from pathlib import Path
import sys

SEED4 = '9540ebdca9d33494f588dc60c8ac820b711d5a36ff1c3a84ef1333400aa0af47'
START = 54
PANELS = [  # (run root, condition, label, display colour)
    ('ablation-a2-01', 'mujoco-nominal', 'MuJoCo · contact 10 ms (default)', 0x548344),
    ('ablation-a2-01', 'mujoco-timeconst-2ms', 'MuJoCo · stiffer contact, 2 ms', 0x9bb58a),
    ('ablation-a2-confirm-01', 'physx-nominal', 'PhysX · TGS 100/1 (default)', 0x315d9e),
    ('ablation-a2-confirm-01', 'physx-solver-pgs', 'PhysX · PGS solver', 0x7f9fd1),
]


def sha(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def main(simcompare, results, output):
    sys.path.insert(0, str(simcompare))
    from rl_replay import read_episode, viewer_trace, write_gzip
    traces, sources = [], []
    for root, condition, label, color in PANELS:
        rows = [json.loads(line) for f in sorted((results/root/'runs'/condition).glob('shard-*/episodes.jsonl'))
                for line in f.read_text().splitlines()]
        row, = [r for r in rows if r['scene_index'] == START and r['policy_sha256'] == SEED4]
        shard = next(p for p in (results/root/'runs'/condition).glob('shard-*')
                     if (p/row['replay']['file']).exists() and sha(p/row['replay']['file']) == row['replay']['sha256'])
        recording = shard/row['replay']['file']
        trace = viewer_trace(*read_episode(recording))
        trace.update(label=label, color=color, condition=condition)
        traces.append(trace)
        sources.append(dict(condition=condition, run=root, recording=str(recording.relative_to(results)),
                            recording_sha256=row['replay']['sha256'], outcome=row['outcome'],
                            peak_force_N=row['peak_force_N'], stop_time_s=row['physics_steps']*.001))
    if len({t['scene_sha256'] for t in traces}) != 1 or len({t['reference_sha256'] for t in traces}) != 1:
        raise ValueError('Panels must share one frozen scene and one policy')
    output.mkdir(parents=True, exist_ok=True)
    bundle = output/'settings-seed4-test054.json.gz'
    write_gzip(bundle, dict(schema='simcompare.bundle.v1', traces=traces,
        notice='One frozen policy (seed 4) on one held-out start under four physics settings. Native recorded poses every 10 ms plus exact terminal poses; stopped runs hold their last pose.'))
    (output/'settings-replay-provenance.json').write_text(json.dumps(dict(
        schema='simcompare.blog.replay-provenance.v1', policy_sha256=SEED4, scene_index=START,
        selection=('Start 54 shows the pattern of 18 of 25 held-out starts (MuJoCo default succeeds, 2 ms overloads, '
                   'PhysX default overloads, PGS succeeds); among those 18 its PhysX peak (91 N) is closest to the median (90 N). '
                   'Selected after inspecting outcomes.'),
        bundle=bundle.name, bundle_sha256=sha(bundle), sources=sources), indent=2)+'\n')


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--simcompare', type=Path, required=True)
    p.add_argument('--results', type=Path, required=True)
    p.add_argument('--output', type=Path, required=True)
    a = p.parse_args(); main(a.simcompare, a.results, a.output)