import argparse
import json
import math
import numpy as np
from scipy.stats import binomtest
from common import rows, save_json, inside


def compare(original, edited):
    def index(records):
        result = {(r['id'], r['sample']): r for r in records}
        if len(result) != len(records):
            raise ValueError('Повторные результаты')
        return result
    a, b = index(original), index(edited)
    if a.keys() != b.keys() or not a:
        raise ValueError('Нужны одинаковые непустые наборы задач и образцов')
    groups = {}
    rescues = losses = 0
    for key in sorted(a):
        if a[key]['seed'] != b[key]['seed']:
            raise ValueError('Seed не совпадает')
        delta = int(b[key]['correct']) - int(a[key]['correct'])
        groups.setdefault(key[0], []).append(delta)
        rescues += delta == 1
        losses += delta == -1
    per_problem = np.array([np.mean(x) for x in groups.values()])
    gain = float(per_problem.mean())
    se = float(per_problem.std(ddof=1) / math.sqrt(len(per_problem))) if len(per_problem) > 1 else None
    rng = np.random.default_rng(1701)
    boot = np.array([rng.choice(per_problem, len(per_problem), replace=True).mean() for _ in range(10000)])
    return {'problems': len(groups), 'samples': len(a), 'gain': gain, 'cluster_paired_se': se,
            'problem_bootstrap_ci95': np.quantile(boot, [.025, .975]).tolist(),
            'rescues': rescues, 'regressions': losses,
            'mcnemar_exact_p': binomtest(rescues, rescues + losses).pvalue if rescues + losses and all(len(v) == 1 for v in groups.values()) else None,
            'warning': 'Повторные образцы одной задачи объединены перед оценкой неопределённости. Подтверждение допустимо только на отложенной выборке.'}


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('original')
    p.add_argument('edited')
    p.add_argument('--out', required=True)
    a = p.parse_args()
    ja = json.loads((inside(a.original).parent / 'job.json').read_text())
    jb = json.loads((inside(a.edited).parent / 'job.json').read_text())
    for key in ('data_sha256', 'protocol', 'samples', 'max_tokens', 'runtime_versions', 'limit', 'set', 'seed_scheme'):
        if ja[key] != jb[key]:
            raise ValueError(f'Несопоставимые запуски: {key}')
    save_json(a.out, compare(rows(a.original), rows(a.edited)))