"""A search router against the fan-out, on the public panel.

A router sends each question to ONE engine. Three ways of choosing it are scored against the two
fan-out tiers, on hit@10 (an answer page in the first ten results), phrasing 1, trials 1-3.

  fixed      always the engine that is best overall
  by_set     the best engine for the question's set, chosen on other questions (cross-validated)
  learned    a text classifier per engine predicts whether it will hit; the highest wins (cross-validated)
"""
import pickle, sys, os, json, statistics, random, collections
HERE = os.path.dirname(os.path.abspath(__file__))
D = os.environ.get('JEVLIN_PANEL_DIR', os.path.join(HERE, 'panel'))
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']
T = (1, 2, 3)
def hit(a, i, t):
    r = calls[a].get((i, 1, t))
    if r is None: return None
    return 1.0 if score.match_rank(score.dedupe([u for u in r['urls'] if u])[:10], items[i], 'exact') else 0.0
H = {a: {i: [hit(a, i, t) for t in T] for i in items} for a in E + ['fast', 'balanced']}
ids = [i for i in items if all(None not in H[a][i] for a in H)]
m = lambda a, i: statistics.mean(H[a][i])
def rate(pick):  # pick: id -> engine
    return {i: m(pick[i], i) for i in pick}
pct = lambda per: round(100 * statistics.mean(per.values()), 1)
def ci(per):
    lo, hi = score.boot_ci(list(per.values())) if hasattr(score, 'boot_ci') else (None, None)
    return [lo, hi]
out = {'questions': len(ids), 'engines': {a: pct({i: m(a, i) for i in ids}) for a in E}}
best = max(E, key=lambda a: out['engines'][a])
per = {'fixed': {i: m(best, i) for i in ids}}
out['fixed_engine'] = best
# by_set, repeated 2-fold CV
rng = random.Random(7); acc = collections.defaultdict(list); picks = collections.Counter()
for rep in range(200):
    sh = ids[:]; rng.shuffle(sh); half = len(sh) // 2
    for train, test in ((sh[:half], sh[half:]), (sh[half:], sh[:half])):
        by = {}
        for s in {items[i]['set'] for i in ids}:
            tr = [i for i in train if items[i]['set'] == s]
            by[s] = max(E, key=lambda a: statistics.mean(m(a, i) for i in tr)) if tr else best
            picks[(s, by[s])] += 1
        for i in test: acc[i].append(m(by[items[i]['set']], i))
per['by_set'] = {i: statistics.mean(v) for i, v in acc.items()}
out['by_set_picks'] = {s: collections.Counter({e: c for (s2, e), c in picks.items() if s2 == s}).most_common(1)[0][0] for s in {items[i]['set'] for i in ids}}
out['by_set_in_sample_best'] = {s: {a: round(100 * statistics.mean(m(a, i) for i in ids if items[i]['set'] == s), 1) for a in E} for s in sorted({items[i]['set'] for i in ids})}
# learned router: TF-IDF + logistic regression per engine, 5-fold CV repeated
try:
    from sklearn.feature_extraction.text import TfidfVectorizer
    from sklearn.linear_model import LogisticRegression
    import numpy as np
    acc = collections.defaultdict(list); chosen = collections.Counter()
    for rep in range(10):
        sh = ids[:]; random.Random(rep).shuffle(sh)
        for f in range(5):
            test = sh[f::5]; ts = set(test); train = [i for i in sh if i not in ts]
            vec = TfidfVectorizer(ngram_range=(1, 2), min_df=2, sublinear_tf=True)
            Xtr = vec.fit_transform([items[i]['query'] for i in train]); Xte = vec.transform([items[i]['query'] for i in test])
            setcol = lambda L: np.array([[items[i]['set'] == s for s in ('obcode', 'devdex', 'seal')] for i in L], dtype=float)
            from scipy.sparse import hstack, csr_matrix
            Xtr = hstack([Xtr, csr_matrix(setcol(train))]).tocsr(); Xte = hstack([Xte, csr_matrix(setcol(test))]).tocsr()
            P = []
            for a in E:
                y = np.array([1 if m(a, i) >= 0.5 else 0 for i in train])
                P.append(LogisticRegression(C=1.0, max_iter=2000).fit(Xtr, y).predict_proba(Xte)[:, 1] if 0 < y.sum() < len(y) else np.full(len(test), y.mean()))
            P = np.array(P)
            for j, i in enumerate(test):
                a = E[int(P[:, j].argmax())]; chosen[a] += 1; acc[i].append(m(a, i))
    per['learned'] = {i: statistics.mean(v) for i, v in acc.items()}
    tot = sum(chosen.values()); out['learned_picks_pct'] = {a: round(100 * c / tot, 1) for a, c in chosen.most_common()}
except ImportError as e:
    out['learned_error'] = str(e)
per['fast'] = {i: m('fast', i) for i in ids}; per['balanced'] = {i: m('balanced', i) for i in ids}
out['hit_at_10'] = {k: pct(v) for k, v in per.items()}
out['vs_balanced'] = {}
for k in ('fixed', 'by_set', 'learned'):
    if k in per:
        d = score.boot_diff(per['balanced'], per[k]); out['vs_balanced'][k] = dict(diff=d['mean'], ci=d['ci'])
out['vs_fast'] = {}
for k in ('fixed', 'by_set', 'learned'):
    if k in per:
        d = score.boot_diff(per['fast'], per[k]); out['vs_fast'][k] = dict(diff=d['mean'], ci=d['ci'])
# when the fixed router's engine misses (majority of trials)
miss = [i for i in ids if m(best, i) < 0.5]
out['when_fixed_misses'] = dict(questions=len(miss), share_pct=round(100 * len(miss) / len(ids), 1),
    another_engine_had_it_pct=round(100 * statistics.mean(1 if max(m(a, i) for a in E if a != best) >= 0.5 else 0 for i in miss), 1))
# how often the single best engine for a question is each engine (ties split)
w = collections.Counter(); n = 0
for i in ids:
    top = max(m(a, i) for a in E)
    if top == 0: continue
    n += 1; t = [a for a in E if m(a, i) == top]
    for a in t: w[a] += 1 / len(t)
out['best_engine_share_pct'] = {a: round(100 * c / n, 1) for a, c in w.most_common()}
out['only_one_engine_pct'] = round(100 * statistics.mean(1 if sum(1 for a in E if m(a, i) >= 0.5) == 1 else 0 for i in ids), 1)
json.dump(out, open(os.path.join(D, 'router.json'), 'w'), indent=1); print(json.dumps(out, indent=1))