import pickle, sys, os, json, itertools, collections, statistics
HERE = os.path.dirname(os.path.abspath(__file__))
D = os.environ.get('JEVLIN_PANEL_DIR', os.path.join(HERE, 'panel'))  # a checkout of the supersearch-tests dataset
sys.path.insert(0, D)
import score
items, calls = pickle.load(open(os.path.join(D, 'panel.pkl'), 'rb'))
E = ['exa', 'brave', 'parallel-search', 'perplexity-search']; K = 10
def top(arm, key):
    r = calls[arm].get(key)
    return score.dedupe([score.canon(u) for u in r['urls'] if u])[:K] if r else None
def shared(A, B): return len(set(A) & set(B))
out = {}
# self-agreement: same engine, same query, two fresh runs (trial 1 vs 2); and same engine, reworded query (phrasing 1 vs 2, trial 1)
selfrun, selfword, crossword = {}, {}, {}
for a in E:
    s1 = [shared(top(a, (i, p, 1)), top(a, (i, p, 2))) for i in items for p in (1, 2, 3) if top(a, (i, p, 1)) and top(a, (i, p, 2))]
    s2 = [shared(top(a, (i, p, 1)), top(a, (i, q, 1))) for i in items for p, q in ((1, 2), (1, 3), (2, 3)) if top(a, (i, p, 1)) and top(a, (i, q, 1))]
    selfrun[a] = dict(n=len(s1), mean_shared=round(statistics.mean(s1), 2)); selfword[a] = dict(n=len(s2), mean_shared=round(statistics.mean(s2), 2))
out['same_engine_two_runs'] = selfrun; out['same_engine_reworded'] = selfword
# gold-level agreement: of the questions engine A found, how many did engine B also find (phrasing 1, trial-avg over paired trials)
def found(a, i, t):
    r = calls[a].get((i, 1, t))
    return None if r is None else bool(score.match_rank(score.dedupe([u for u in r['urls'] if u]), items[i], 'exact'))
cond = {}
for a in E:
    for b in E:
        if a == b: continue
        num = den = 0
        for i in items:
            for t in (1, 2, 3):
                fa, fb = found(a, i, t), found(b, i, t)
                if fa is None or fb is None: continue
                if fa: den += 1; num += 1 if fb else 0
        cond[a + '>' + b] = round(100 * num / den, 1)
out['of_A_finds_also_found_by_B'] = cond
# Wilson intervals for agreement-as-evidence
ag = json.load(open(os.path.join(D, 'agree.json')))['agreement_as_evidence']
out['evidence_ci'] = {n: list(score.wilson(v['gold'] / v['pages'], v['pages'])) for n, v in ag.items()}  # already percent
# examples: one question per engine where only that engine had the gold page in all three trials; and two where all four agreed at rank 1
ex = collections.defaultdict(list); allfour = []
for i, it in items.items():
    f = {a: [found(a, i, t) for t in (1, 2, 3)] for a in E}
    if any(None in v for v in f.values()): continue
    always = [a for a in E if all(f[a])]; never = [a for a in E if not any(f[a])]
    if len(always) == 1 and len(never) == 3:
        r = calls[always[0]][(i, 1, 1)]; rk = score.match_rank(score.dedupe([u for u in r['urls'] if u]), it, 'exact')
        ex[always[0]].append(dict(id=i, set=it['set'], query=it['query'], gold=it['gold_urls'][0], rank=rk))
    if len(always) == 4: allfour.append(i)
out['only_one_engine_all_trials'] = {a: len(v) for a, v in ex.items()}
out['all_four_all_trials'] = len(allfour); out['none_any_trial'] = sum(1 for i in items if all(not any(found(a, i, t) for t in (1, 2, 3)) for a in E))
out['examples'] = {a: sorted(v, key=lambda x: (x['set'] != 'devdex', len(x['query'])))[:2] for a, v in ex.items()}
json.dump(out, open(os.path.join(D, 'extra.json'), 'w'), indent=1)
print(json.dumps({k: v for k, v in out.items() if k != 'examples'}, indent=1))
for a, v in out['examples'].items():
    for e in v: print(a, '|', e['set'], '|', e['query'][:110], '|', e['gold'][:90], '| rank', e['rank'])