"""Paired, role-balanced tests; bootstrap units are complete seeded deals."""
import argparse
import json
import statistics
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if __package__ in (None, ''):
    sys.path.insert(0, str(ROOT))

from qojbot.history import write_json
from qojbot.model_benchmark import run, _bootstrap
from QOJDoudizhuModel.policy import QOJPolicy
from QOJDoudizhuModel.storage import digest, source_hashes, teacher_hashes


def evaluate(policy, deals, seed, modes=('fixed', 'bidding'), opponent=None):
    opponent = opponent or policy.teacher
    result = run({'lstm': opponent, 'resnet': policy}, deals=deals, seed=seed, modes=modes)
    # Reuse the tested paired evaluator without presenting the new network as ResNet.
    result['labels'] = {'lstm': type(opponent).__name__, 'resnet': 'QOJNet'}
    result['selection_score'] = statistics.mean(
        metrics['resnet']['mean_score'] for metrics in result['summary'].values())
    clusters = {}
    for record in result['records']:
        clusters.setdefault(record['seed'], []).extend(
            record['deltas'][seat] for seat, name in enumerate(record['lineup']) if name == 'resnet')
    result['selection_score_95pct_cluster_bootstrap'] = _bootstrap(
        [statistics.mean(values) for values in clusters.values()])
    result['positive_score_supported'] = (
        deals >= 100 and result['selection_score_95pct_cluster_bootstrap'][0] > 0 and
        all(metrics['resnet']['mean_score'] > 0 for metrics in result['summary'].values()))
    result['teacher_sha256'] = teacher_hashes(policy.teacher)
    if hasattr(opponent, 'paths'):
        result['opponent'] = {'model_type': opponent.model_type,
                              'paths': {key: str(path) for key, path in opponent.paths.items()},
                              'sha256': teacher_hashes(opponent)}
    result['source_sha256'] = source_hashes()
    result['protocol'] = {'lineups_per_deal_per_mode': 6,
                          'selection': 'validation only; test seed reserved',
                          'score': 'QOJ bid * (1+bombs+spring), landlord x2',
                          'confidence': 'bootstrap by original seed; few deals yield unreliable intervals',
                          'information': 'own hand and public information only',
                          'limitation': 'not online rating or six-game match placement'}
    return result


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--checkpoint', type=Path, required=True)
    parser.add_argument('--device', choices=('cpu', 'auto', 'mps', 'cuda'), default='mps')
    parser.add_argument('--deals', type=int, default=100)
    parser.add_argument('--seed', type=int, default=40261002)
    parser.add_argument('--modes', nargs='+', choices=('fixed', 'bidding'), default=['fixed', 'bidding'])
    parser.add_argument('--opponent', choices=('douzero', 'search'), default='douzero')
    parser.add_argument('--opponent-dir', type=Path,
                        help='explicit DouZero benchmark weights, independent of the frozen prior')
    parser.add_argument('--teacher-dir', type=Path,
                        help='frozen prior used by the checkpoint')
    parser.add_argument('--output', type=Path, default=Path('QOJDoudizhuModel/results/test.json'))
    args = parser.parse_args()
    if args.opponent == 'search' and args.opponent_dir:
        parser.error('--opponent-dir requires --opponent douzero')
    policy = QOJPolicy(args.checkpoint, device=args.device, teacher_dir=args.teacher_dir)
    opponent = None
    if args.opponent == 'search':
        from QOJDoudizhuModel.learning import SearchOpponent
        opponent = SearchOpponent(.1)
    elif args.opponent_dir:
        from qojbot.douzero import DouZeroPolicy
        opponent = DouZeroPolicy(args.opponent_dir, device=str(policy.device), batch_size=256)
    report = evaluate(policy, args.deals, args.seed, tuple(args.modes), opponent)
    path = args.checkpoint / 'best.pt' if args.checkpoint.is_dir() else args.checkpoint
    report['checkpoint_sha256'] = digest(path)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    write_json(args.output, report)
    print(json.dumps({key: report[key] for key in ('games', 'selection_score', 'summary')}, indent=2))


if __name__ == '__main__':
    main()