"""同场景配对比较手写策略、旧PPO及修订模型；真值仅在评测端使用。"""
import argparse
from dataclasses import asdict
import hashlib
import json
import math
from pathlib import Path
from statistics import mean, median
from time import perf_counter

import torch

from mcm26b.config import ProblemConfig
from mcm26b.evaluation import recompute_virtual_time
from mcm26b.harness import LocalSimTransport, mixed_scene, random_omni_scene, referee_spec
from mcm26b.planner import DirectionalExtension
from mcm26b.planner.bayes_loop import BayesLoop
from mcm26b.planner.rl_loop import RLBayesLoop, RLDirectionalLoop
from mcm26b.protocol.client import SimulatorClient
from referee import LocalSimulator, Source


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--split', choices=('development','validation','stress'),
                        default='development')
    parser.add_argument('--variants', default='baseline,legacy,revised')
    parser.add_argument('--problems', default='3,4')
    parser.add_argument('--model-dir', type=Path, default=Path('output/rl/review-20260911'))
    parser.add_argument('--label', required=True)
    args = parser.parse_args()
    torch.set_num_threads(1)
    config = ProblemConfig()
    output = args.model_dir / 'evaluation' / args.split / args.label
    output.mkdir(parents=True, exist_ok=False)
    variants = args.variants.split(',')
    rows = []
    def save():
        aggregates = {}
        for problem in (3,4):
            for variant in variants:
                part = [r for r in rows if r['problem']==problem and r['variant']==variant]
                if part:
                    aggregates[f'q{problem}-{variant}'] = dict(
                        runs=len(part), valid=sum(r['valid'] for r in part),
                        mean_virtual_s=mean(r['virtual_s'] for r in part),
                        median_virtual_s=median(r['virtual_s'] for r in part),
                        max_virtual_s=max(r['virtual_s'] for r in part),
                        failed_clears=sum(r['failed_clears'] for r in part),
                        repeated_measures=sum(r['repeated_measures'] for r in part))
        result = dict(source='local_referee_only', config=asdict(config), split=args.split,
                      rows=rows, aggregates=aggregates)
        (output/'results.json').write_text(
            json.dumps(result,ensure_ascii=False,indent=2),encoding='utf-8')
        return aggregates
    for problem in map(int, args.problems.split(',')):
        scenes = []
        if args.split == 'stress':
            for count in (10,16):
                scenes.append((f'boundary-{count}', count, 99000+count, [
                    Source(c+1,1800*math.cos(2*math.pi*c/count),1800*math.sin(2*math.pi*c/count),
                           1000,directional=(problem==4 and c%2==0),
                           heading_deg=360*c/count if problem==4 and c%2==0 else 0)
                    for c in range(count)]))
        else:
            seeds = range(700,703) if args.split=='development' else range(91000,91006)
            counts = (10,16) if args.split=='development' else (10,13,16)
            for count in counts:
                for seed in seeds:
                    directional = (1, count//2, count-1)[seed%3]
                    sources = (random_omni_scene(seed, count, config) if problem==3 else
                               mixed_scene(seed, count, directional, config))
                    scenes.append((f'{count}-{seed}',count,seed,sources))
        for name,count,seed,sources in scenes:
            scene_data = dict(sources=[asdict(s) for s in sources], error_seed=seed+2000000)
            scene_path = output/f'q{problem}-{name}-scene.json'
            scene_path.write_text(json.dumps(scene_data),encoding='utf-8')
            for variant in variants:
                model = None
                cls = BayesLoop if problem==3 else DirectionalExtension
                kwargs = {}
                if variant != 'baseline':
                    cls = RLBayesLoop if problem==3 else RLDirectionalLoop
                    model = (Path(f'output/rl/q{problem}_ppo.zip') if variant=='legacy'
                             else args.model_dir/f'q{problem}_masked.zip')
                    kwargs = dict(model_path=model, device='cpu')
                sim = LocalSimulator(sources,robot_id='RL-AUDIT',spec=referee_spec(config),
                                     error_seed=seed+2000000)
                client = SimulatorClient('RL-AUDIT',config=config,transport=LocalSimTransport(sim))
                start = perf_counter()
                summary = cls(client,config,**kwargs).run()
                events = client.log.action_events()
                truth = {s.channel for s in sources}
                successful = {e['request']['channel'] for e in events
                              if e['response'].get('clear_result')=='success'}
                seen = set()
                repeats = 0
                for e in events:
                    if e['path']=='/measure':
                        p=e['request']['position']
                        tag=(e['request']['channel'],round(p['x'],6),round(p['y'],6))
                        repeats += tag in seen
                        seen.add(tag)
                mismatches = recompute_virtual_time(events)['mismatches']
                valid = (successful==truth and summary['completed'] and not summary['anomalies']
                         and not set(summary['absent_channels']) & truth
                         and summary['exit_reason']=='user_exit' and not mismatches)
                row = dict(case=name,problem=problem,variant=variant,valid=valid,
                           virtual_s=client.state.virtual_time_s,all_cleared=successful==truth,
                           measures=summary['measures'],
                           failed_clears=summary['clears']-len(successful),
                           repeated_measures=repeats,wall_s=perf_counter()-start,
                           model_sha256=(hashlib.sha256(model.read_bytes()).hexdigest()
                                         if model else None))
                rows.append(row)
                (output/f'q{problem}-{name}-{variant}.json').write_text(json.dumps(
                    dict(row=row,summary=summary,events=events,timing_mismatches=mismatches),
                    ensure_ascii=False),encoding='utf-8')
                save()
                print(json.dumps(row),flush=True)
    print(json.dumps(save(),indent=2),flush=True)


if __name__=='__main__':
    main()