"""Aggregate Study C (ANYmal-D flat) cross-evaluations into summary.json.

Usage: python export.py RESULTS_DIR OUT.json CELL...
RESULTS_DIR holds <cell>-eval-01/*.json from eval_policy.py, <cell>/logs/<backend>-seed<n>.log training logs,
and engine readbacks in phase0-03/ or phase0-hc2/.
Values are per training seed (each the mean over 64 envs x 10 s after 2 s warmup);
falls are counted over all 5 seeds x 64 envs = 320 robots per cell.
"""
import glob
import json
import re
import statistics
import sys
from pathlib import Path

METRICS = ('err_xy_m_s', 'err_yaw_rad_s', 'torque_rms_Nm', 'cost_of_transport', 'stance_slip_m_s',
           'p99_foot_force_N', 'duty_factor', 'gait_frequency_hz', 'attitude_rms_rad')


def readback(path, backend):
    d = json.loads(Path(path).read_text())
    if backend == 'newton_mjwarp':
        n = d['newton']
        return dict(substep_s=n['substep_s'], cone='elliptic' if n['opt_cone'] == 1 else 'pyramidal',
                    impratio=n['opt_impratio']['unique'], pair_friction=n['pair_sliding_friction_max_rule']['unique'],
                    contact_timeconst_s=n['effective_timeconst_s']['unique'])
    p = d['physx']
    return dict(steps_per_second=p['scene']['steps_per_second'], solver=p['scene']['solver'],
                iterations=[p['articulation'][0]['position'], p['articulation'][0]['velocity']],
                foot_static=p['runtime_robot_material']['static']['unique'],
                foot_dynamic=p['runtime_robot_material']['dynamic']['unique'],
                ground=list(p['usd_materials'].values())[0])


def main(results, out, cells):
    results = Path(results)
    summary = dict(schema='locomotion.study-c.summary.v1', task='Isaac-Velocity-Flat-AnymalD', cells={})
    for cell in cells:
        rows = {}
        for f in sorted(glob.glob(str(results/f'{cell}-eval-01'/'*.json'))):
            d = json.loads(Path(f).read_text())
            name = Path(f).stem
            train, rest = name.split('-seed')
            seed, evaluated = rest.split('-on-')
            for command, r in d['results'].items():
                rows.setdefault(f'{train}>{evaluated}', {}).setdefault(command, []).append(
                    dict(seed=int(seed), fall_rate=r['fall_rate'], **{k: r[k]['mean'] for k in METRICS}))
        table = {}
        for pair, commands in rows.items():
            table[pair] = {}
            for command, seeds in commands.items():
                table[pair][command] = dict(
                    seeds=len(seeds), falls=round(sum(s['fall_rate'] for s in seeds)*64), robots=64*len(seeds),
                    **{k: statistics.mean(s[k] for s in seeds) for k in METRICS}, per_seed=sorted(seeds, key=lambda s: s['seed']))
        rb, training = {}, {}
        for backend in ('newton_mjwarp', 'isaacsim_physx'):
            for folder in ('phase0-03', 'phase0-hc2'):
                path = results/folder/f'{cell}-{backend}.json'
                if path.exists():
                    rb[backend] = readback(path, backend)
            # Final logged mean episode reward per training seed (source competence).
            finals = []
            for log in sorted((results/cell/'logs').glob(f'{backend}-seed*.log')):
                rewards = re.findall(r'Mean reward: ([-\d.]+)', log.read_text())
                finals.append(float(rewards[-1]) if rewards else None)
            if finals:
                training[backend] = dict(final_mean_reward_per_seed=finals)
        summary['cells'][cell] = dict(transfer=table, readback=rb, training=training)
    Path(out).write_text(json.dumps(summary, indent=1))


if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2], sys.argv[3:])