import pickle, json, sys, numpy as np, collections
import os; S = sys.argv[1]; L = {f['index']: f['label'] for f in json.load(open(f'{S}/' + os.environ.get('LABELS','../../../../eval/labels/image6.json')))['facings']}
RANK = {
 'inl': lambda e: e['inliers'],
 'inl*col': lambda e: e['inliers'] * (0.25 + e['color']),
}
def conf_feats(ranked):
    (k1, e1), (k2, e2) = ranked[0], ranked[1]
    i1, i2 = e1['inliers'], e2['inliers']
    return {
      'inl': i1,
      'margin': i1 - i2,
      'ratio': 1 - i2 / max(i1, 1),
      'inl*ratio': i1 * (1 - i2 / max(i1, 1)),
      'cover': e1['cover'],
      'inl*ratio*qcov': i1 * (1 - i2 / max(i1, 1)) * np.sqrt(e1.get('qcover', 0)),
      'inl*ratio*col': i1 * (1 - i2 / max(i1, 1)) * e1['color'],
    }
def evaluate(path, rank='inl'):
    r = pickle.load(open(path, 'rb')); E = set(r['enrolled'])
    rows = [x for x in r['rows'] if L[x['index']] != 'exclude']
    n_in = sum(L[x['index']] != 'unknown' for x in rows); n_en = sum(L[x['index']] in E for x in rows)
    recs = []; top1 = 0
    for x in rows:
        t = L[x['index']]
        if not x['ev']: recs.append((t, None, {k: 0 for k in conf_feats([('a', dict(inliers=0,cover=0,color=0)),('b',dict(inliers=0,cover=0,color=0))])})); continue
        ranked = sorted(x['ev'].items(), key=lambda kv: -RANK[rank](kv[1]))
        b = ranked[0][0] if ranked[0][1]['inliers'] > 0 else None
        if t in E and b == t: top1 += 1
        recs.append((t, b, conf_feats(ranked)))
    out = dict(top1=f'{top1}/{n_en}', argmax_e2e=f'{top1}/{n_in}')
    for f in ['inl', 'margin', 'ratio', 'inl*ratio', 'inl*ratio*qcov', 'inl*ratio*col']:
        pts = sorted(((c[f], t == b and b is not None) for t, b, c in recs if b), key=lambda p: -p[0])
        best = {0.9: 0, 0.95: 0, 1.0: 0}; c = a = 0; i = 0
        while i < len(pts):
            j = i
            while j < len(pts) and pts[j][0] == pts[i][0]: c += pts[j][1]; a += 1; j += 1
            for p in best:
                if c / a >= p: best[p] = max(best[p], c)
            i = j
        out[f] = '/'.join(str(best[p]) for p in (1.0, 0.95, 0.9))
    return out
if __name__ == '__main__':
    print('feature cells = correct accepted at precision 100%/95%/90%  (enrolled in-catalog n=42, all in-catalog n=61, unknown n=48)')
    for p in sys.argv[2:]:
        for rk in ['inl', 'inl*col']:
            print(p.split('/')[-1].replace('.pkl',''), rk, json.dumps(evaluate(f'{S}/{p}', rk)))