import pickle, sys, os, json, itertools, collections, statistics, random
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)
    if not r: return None
    return score.dedupe([score.canon(u) for u in r['urls'] if u])[:K]
def raw(arm, key):
    r = calls[arm].get(key); return r['urls'][:K] if r else None
out = {}
# 0. reproduce headline: hit@10, phrasing 1, trials 1-3, mean over trials then items
def hit(arm, key, item):
    u = raw(arm, key)
    if u is None: return None
    mode = 'exact'
    return 1 if score.match_rank(u, item, mode) else 0
repro = {}
for a in E + ['exa-lite', 'fast', 'balanced']:
    per = []
    for i, it in items.items():
        hs = [hit(a, (i, 1, t), it) for t in (1, 2, 3)]; hs = [h for h in hs if h is not None]
        if hs: per.append(sum(hs) / len(hs))
    repro[a] = round(100 * sum(per) / len(per), 1)
out['repro_hit10'] = repro
# 1. pairwise overlap of top-10 URL sets, every query instance (485 items x 3 phrasings x 3 trials)
keys = [(i, p, t) for i in items for p in (1, 2, 3) for t in (1, 2, 3)]
pair = {}
for a, b in itertools.combinations(E, 2):
    shared, jac, dom = [], [], []
    for k in keys:
        A, B = top(a, k), top(b, k)
        if not A or not B: continue
        sa, sb = set(A), set(B)
        shared.append(len(sa & sb)); jac.append(len(sa & sb) / len(sa | sb))
        da, db = {score.reg_domain(u) for u in A}, {score.reg_domain(u) for u in B}
        dom.append(len(da & db) / len(da | db))
    pair[a + '|' + b] = dict(n=len(shared), mean_shared_urls=round(statistics.mean(shared), 2), jaccard=round(100 * statistics.mean(jac), 1),
                             zero_overlap_pct=round(100 * sum(1 for s in shared if s == 0) / len(shared), 1), domain_jaccard=round(100 * statistics.mean(dom), 1))
out['pairwise'] = pair
# 2. per query: how many engines returned each URL in the pooled top-10s
votes_hist = collections.Counter(); per_engine_excl = collections.defaultdict(list); pool_sizes = []
for k in keys:
    tops = {a: top(a, k) for a in E}
    if any(v is None for v in tops.values()): continue
    cnt = collections.Counter(u for a in E for u in set(tops[a]))
    pool_sizes.append(len(cnt))
    for u, c in cnt.items(): votes_hist[c] += 1
    for a in E:
        if tops[a]: per_engine_excl[a].append(sum(1 for u in set(tops[a]) if cnt[u] == 1) / len(set(tops[a])))
tot = sum(votes_hist.values())
out['pool'] = dict(queries=len(pool_sizes), mean_distinct_pages=round(statistics.mean(pool_sizes), 1), max_possible=40,
                   by_engines_returning={str(c): round(100 * n / tot, 1) for c, n in sorted(votes_hist.items())},
                   exclusive_share_of_own_results={a: round(100 * statistics.mean(v), 1) for a, v in per_engine_excl.items()})
# 3. agreement as evidence: P(page is the gold page | n engines returned it), phrasing 1, trials 1-3
prec = collections.defaultdict(lambda: [0, 0]); gold_found_by = collections.Counter(); excl_gold = collections.Counter(); n_items = 0
for i, it in items.items():
    mode = 'exact'
    for t in (1, 2, 3):
        k = (i, 1, t); raws = {a: raw(a, k) for a in E}
        if any(v is None for v in raws.values()): continue
        n_items += 1
        # map canonical url -> engines, and whether it is gold
        eng = collections.defaultdict(set); isgold = {}
        for a in E:
            for u in raws[a]:
                c = score.canon(u)
                if not c: continue
                eng[c].add(a); isgold[c] = isgold.get(c, False) or bool(score.match_rank([u], it, mode))
        for c, s in eng.items():
            prec[len(s)][1] += 1; prec[len(s)][0] += 1 if isgold[c] else 0
        finders = set().union(*[s for c, s in eng.items() if isgold[c]]) if any(isgold.values()) else set()
        gold_found_by[len(finders)] += 1
        if len(finders) == 1: excl_gold[next(iter(finders))] += 1
out['agreement_as_evidence'] = {str(n): dict(pages=v[1], gold=v[0], pct_gold=round(100 * v[0] / v[1], 2)) for n, v in sorted(prec.items())}
out['gold_found_by_n_engines'] = {str(n): round(100 * c / n_items, 1) for n, c in sorted(gold_found_by.items())}
out['gold_found_by_exactly_one'] = {a: round(100 * c / n_items, 1) for a, c in excl_gold.most_common()}
out['n_query_trials_for_gold'] = n_items
# 4. source concentration per engine: share of results from its 10 most-returned domains; distinct domains
conc = {}
for a in E:
    dc = collections.Counter()
    for k in keys:
        for u in (top(a, k) or []): dc[score.reg_domain(u)] += 1
    n = sum(dc.values())
    conc[a] = dict(results=n, distinct_domains=len(dc), top10_share=round(100 * sum(c for _, c in dc.most_common(10)) / n, 1),
                   top5=[(d, round(100 * c / n, 1)) for d, c in dc.most_common(5)])
out['sources'] = conc
json.dump(out, open(os.path.join(D, 'agree.json'), 'w'), indent=1)
print(json.dumps(out, indent=1))