import pickle, sys, os, json, collections, statistics, re
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']; ALL = E + ['exa-lite', 'fast', 'balanced']
def hits(arm, i):  # per-trial ranks for phrasing 1
    out = []
    for t in (1, 2, 3):
        r = calls[arm].get((i, 1, t))
        if r is None: continue
        out.append(score.match_rank(score.dedupe([u for u in r['urls'] if u]), items[i], 'exact'))
    return out
def gold_kind(it):
    g = (it['gold_urls'] or [''])[0]; d = score.reg_domain(g); h = score.canon(g)
    if it['set'] == 'devdex':
        return 'GitHub file' if 'github.com' in h else 'Product documentation'
    if 'wikipedia.org' in d: return 'Wikipedia'
    if d == 'github.com': return 'GitHub file'
    if d in ('arxiv.org', 'nature.com', 'sciencedirect.com', 'springer.com', 'nih.gov', 'acm.org', 'ieee.org', 'doi.org', 'researchgate.net', 'semanticscholar.org', 'openreview.net', 'biorxiv.org', 'wiley.com', 'tandfonline.com', 'jstor.org', 'plos.org', 'mdpi.com') or d.endswith('.edu'): return 'Papers and academic'
    if d.endswith('.gov') or '.gov.' in d or d.endswith('.int') or d in ('europa.eu', 'un.org', 'worldbank.org', 'who.int', 'oecd.org'): return 'Government and institutions'
    if it['set'] == 'obcode': return 'Vendor API reference'
    return 'Other sites'
kind = {i: gold_kind(it) for i, it in items.items()}
print(collections.Counter(kind.values())); print(collections.Counter((it['set'], kind[i]) for i, it in items.items()))
dom = collections.Counter(score.reg_domain((it['gold_urls'] or [''])[0]) for it in items.values() if it['set'] == 'seal'); print(dom.most_common(25))
def table(group):
    res = {}
    for g in sorted(set(group.values())):
        ids = [i for i in items if group[i] == g]; row = {'n': len(ids)}
        for a in ALL:
            per = [statistics.mean([1.0 if r else 0.0 for r in hits(a, i)]) for i in ids if hits(a, i)]
            p = statistics.mean(per); lo, hi = score.wilson(p, len(per))
            row[a] = dict(hit=round(100 * p, 1), lo=lo, hi=hi)  # score.wilson already returns percent
        best = max(E, key=lambda a: row[a]['hit']); second = sorted(E, key=lambda a: -row[a]['hit'])[1]
        da = {i: statistics.mean([1.0 if r else 0.0 for r in hits(best, i)]) for i in ids}; db = {i: statistics.mean([1.0 if r else 0.0 for r in hits(second, i)]) for i in ids}
        d = score.boot_diff(da, db)
        row['best_single'] = best; row['second'] = second; row['best_minus_second'] = d
        res[g] = row
    return res
out = {'by_gold_kind': table(kind), 'by_set': table({i: it['set'] for i, it in items.items()})}
# rank position of the gold page per arm (phrasing 1, all trials)
pos = {}
for a in ALL:
    c = collections.Counter(); n = 0
    for i in items:
        for r in hits(a, i):
            n += 1; c['1' if r == 1 else '2-3' if r and r <= 3 else '4-10' if r else 'miss'] += 1
    pos[a] = {k: round(100 * v / n, 1) for k, v in c.items()}
out['gold_position'] = pos
# query length (phrasing 1 text) vs hit: short <=6 words, medium 7-10, long >10 -- from the request query
ph = {}
for l in open(os.path.join(D, 'phrasings.jsonl')):
    r = json.loads(l); ph[r['id']] = r
print({k: (v if not isinstance(v, list) else v[:1]) for k, v in list(ph.values())[0].items()})
json.dump(out, open(os.path.join(D, 'kinds.json'), 'w'), indent=1, default=str)
for g, row in out['by_gold_kind'].items(): print(f"{g:32} n={row['n']:3}", ' '.join(f"{a[:5]}={row[a]['hit']}" for a in ALL), '| best', row['best_single'], row['best_minus_second'])
for g, row in out['by_set'].items(): print(f"{g:32} n={row['n']:3}", ' '.join(f"{a[:5]}={row[a]['hit']}" for a in ALL), '| best', row['best_single'], row['best_minus_second'])
print(pos)