"""Choose the two dim sets on train2017 and write rules.json.

    python choose.py

The decision is sum(feat[pos]) > sum(feat[neg]), evaluated at zero. A threshold
would absorb the offset between the two sums; choosing the sets so the offset is
already zero removes it instead, which is why the rule carries no free parameter.

Dims are added greedily, alternating sides, each step taking whichever remaining
dim most improves F1 with the boundary pinned at zero. Selection reads only
train2017. Nothing here touches val2017.
"""
import argparse
from pathlib import Path

import torch

from common import COCO_ROOT, prf1, write_artifact
from common.cached import load_pooled
from common.pools import TRAIN2017

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


def greedy(X, y, k, candidates):
    """Grow pos and neg sets to k each, scoring only at threshold zero."""
    pos, neg = [], []
    cur = torch.zeros(X.shape[0])
    for step in range(2 * k):
        side, sign = (pos, 1.0) if step % 2 == 0 else (neg, -1.0)
        used = set(pos) | set(neg)
        best_f1, best_d = -1.0, None
        for d in candidates:
            if d in used:
                continue
            f1 = prf1(cur + sign * X[:, d] > 0, y).f1
            if f1 > best_f1:
                best_f1, best_d = f1, d
        side.append(int(best_d))
        cur = cur + sign * X[:, best_d]
    return pos, neg, prf1(cur > 0, y).f1


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--cache', type=Path, default=None)
    ap.add_argument('--sizes', type=int, nargs='+',
                    default=[1, 2, 3, 4, 6, 8, 10, 20])
    ap.add_argument('--candidates', type=int, default=192)
    ap.add_argument('--out', type=Path, default=HERE / 'rules.json')
    args = ap.parse_args()

    cache = args.cache or COCO_ROOT / 'pooled_train2017'
    X, y = load_pooled(cache, 'train2017')
    print(f'[train] {X.shape[0]} images, person rate {y.float().mean():.3f}',
          flush=True)

    # Restrict the search to dims with real class separation, for tractability.
    sep = (X[y].mean(0) - X[~y].mean(0)).abs()
    candidates = torch.topk(sep, args.candidates).indices.tolist()

    rules = {}
    for k in args.sizes:
        pos, neg, f1 = greedy(X, y, k, candidates)
        rules[f'd{2 * k}'] = {'pos_dims': pos, 'neg_dims': neg,
                              'n_dims': 2 * k, 'free_parameters': 0,
                              'F1_train': round(f1, 4)}
        print(f'  d{2 * k:<3} F1 {f1:.4f}  {pos} > {neg}', flush=True)

    write_artifact(args.out, {
        'candidates': args.candidates,
        'sizes': [2 * k for k in args.sizes],
        'rules': rules,
    }, generator='choose.py',
        pool_info={'pool': TRAIN2017.name, 'split': TRAIN2017.split,
                   'n_images': int(X.shape[0]),
                   'positive_rate': round(y.float().mean().item(), 4),
                   'selection': TRAIN2017.selection},
        decision='sum(feat[pos_dims]) > sum(feat[neg_dims])',
        search='greedy, alternating sides, boundary pinned at zero')
    print(f'[done] wrote {args.out}', flush=True)


if __name__ == '__main__':
    main()