"""Score every rule in rules.json on val2017 and write eval.json.

    python verify.py

Dimensions were chosen on train2017 and are not refit here.
"""
import argparse
from pathlib import Path

from common import COCO_ROOT, prf1, read_artifact, write_artifact
from common.cached import load_pooled
from common.pools import VAL5000

HERE = Path(__file__).resolve().parent


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--cache', type=Path, default=None)
    ap.add_argument('--rules', type=Path, default=HERE / 'rules.json')
    ap.add_argument('--out', type=Path, default=HERE / 'eval.json')
    args = ap.parse_args()

    cache = args.cache or COCO_ROOT / 'pooled_val2017'
    X, y = load_pooled(cache, 'val2017')
    print(f'[val] {X.shape[0]} images, person rate {y.float().mean():.3f}\n',
          flush=True)

    doc = read_artifact(args.rules)
    out = {}
    print(f"{'rule':>6}{'dims':>6}{'F1 train':>10}{'F1 val':>9}{'P':>9}{'R':>9}")
    for name, r in doc['rules'].items():
        s = X[:, r['pos_dims']].sum(1) - X[:, r['neg_dims']].sum(1)
        m = prf1(s > 0, y)
        out[name] = {'n_dims': r['n_dims'], 'free_parameters': 0,
                     'pos_dims': r['pos_dims'], 'neg_dims': r['neg_dims'],
                     'F1_train': r['F1_train'],
                     'F1': round(m.f1, 4), 'precision': round(m.precision, 4),
                     'recall': round(m.recall, 4)}
        print(f'{name:>6}{r["n_dims"]:>6}{r["F1_train"]:>10.4f}{m.f1:>9.4f}'
              f'{m.precision:>9.4f}{m.recall:>9.4f}', flush=True)

    write_artifact(args.out, {'rules': out},
                   generator='verify.py',
                   pool_info={'pool': VAL5000.name, 'split': VAL5000.split,
                              'n_images': int(X.shape[0]),
                              'positive_rate': round(y.float().mean().item(), 4),
                              'selection': VAL5000.selection},
                   task='image-level person presence (binary)',
                   protocol='dims selected on train2017, not refit here')
    print(f'\n[done] wrote {args.out}', flush=True)


if __name__ == '__main__':
    main()