"""Round-2 candidate list from the pool scores (pure Python). Run as ``python -m classifier.candidates``.

Spec section 7, "Round-2 candidates": rank the local-context pool posts by the predicted probability of the
thin classes and by per-field entropy, with a Korean quota and at most 50 posts per author. The output is a
ranked list only; context lookups and Astra labelling need the user's go.

Rules (fixed here, before the list is read):
- **Probabilities:** calibrated with the run's dev temperatures, not gated (``calibrate.probs``). The gate is a
  decision rule that renormalises a field after the hard argmax of an earlier one: with valence ``mixed`` it
  sets fine emotion ``mixed`` to 1, which made rare fine emotions certain for 35 of the first gated top 100. The
  ungated heads keep the model's own uncertainty, which is what a ranking for labelling needs.
- **Thin events:** an event's probability is the sum of its options' probabilities.
  - direction bearish; emotion valence negative;
  - any behaviour; any P&L;
  - any of the five specific commercial intents (not ``unclear``);
  - the rare fine emotions, the options other than the absence that hold under 1% of the train labels.
- **thin_score** is the largest event probability; **entropy** is the mean over all fields of H(p) / ln K.
- **Order:** thin_score, then entropy (both descending), then id.
- **Author cap:** an author's posts after the first 50 in that order are left out.
- **Korean quota:** position n is filled from the Korean queue whenever fewer than ceil(0.30 n) of the first
  n - 1 are Korean (while Korean posts remain); otherwise the better-ranked head of the two queues is taken.
  30% sits above the train (26.8%) and pool (23.6%) Korean shares, because Korean bearish labels are the thinnest.

Pool rows already exclude labelled ids and text keys, P4-window context and unresolved context (the export's
``pool`` command); this module refuses any pool row flagged for P4 context or overlap anyway. Posts whose
context post is unavailable (deleted) stay in, as in training, and carry ``has_unavailable_context``.
"""
from __future__ import annotations

import argparse
from collections import Counter
import hashlib
import json
import math
from pathlib import Path

from .calibrate import probs as calibrated
from .data import load_jsonl, write_jsonl
from .taxonomy import Taxonomy

AUTHOR_CAP = 50
KOREAN_SHARE = 0.30
RARE_SHARE = 0.01
COMMERCIAL = ('sponsored_disclosed', 'own_paid_service', 'referral_or_affiliate', 'giveaway_or_bait', 'asset_shill')
VERSION = 'round2-candidates-v2'


def rare_fine(train_rows, tax: Taxonomy, share: float = RARE_SHARE) -> list[str]:
    counts = Counter(r['targets'].get('emotion_fine') for r in train_rows if r.get('targets'))
    counts.pop(None, None)
    total = sum(counts.values())
    absence = tax.absence['emotion_fine']
    return [o for o in tax.fields['emotion_fine'] if o != absence and counts.get(o, 0) < share * total]


def events(tax: Taxonomy, rare: list[str]) -> dict[str, tuple[str, list[str]]]:
    def present(f):
        return [o for o in tax.fields[f] if o != tax.absence[f]]
    return {'bearish': ('direction_primary', ['bearish']), 'negative': ('emotion_valence', ['negative']),
            'behavior': ('behavior', present('behavior')), 'pnl': ('pnl', present('pnl')),
            'commercial': ('commercial_intent', list(COMMERCIAL)), 'rare_fine': ('emotion_fine', list(rare))}


def score(pred: dict, tax: Taxonomy, temps: dict, ev: dict) -> dict:
    p = calibrated(pred, tax, temps)
    probs = {name: sum(p[f][o] for o in opts) for name, (f, opts) in ev.items()}
    top = max(probs, key=lambda k: (probs[k], k))
    ent = []
    for f in tax.names:
        k = len(tax.fields[f])
        ent.append(-sum(x * math.log(x) for x in p[f].values() if x > 0) / math.log(k))
    return {'thin_score': probs[top], 'top_event': top, 'events': probs, 'entropy': sum(ent) / len(ent)}


def rank(items: list[dict], cap: int = AUTHOR_CAP, korean: float = KOREAN_SHARE) -> tuple[list[dict], list[dict]]:
    """items carry id, author_id, lang, thin_score, entropy. Returns (ranked, left out by the author cap)."""
    order = sorted(items, key=lambda x: (-x['thin_score'], -x['entropy'], x['id']))
    seen, kept, capped = Counter(), [], []
    for x in order:
        seen[x['author_id']] += 1
        (kept if seen[x['author_id']] <= cap else capped).append(x)
    ko = [x for x in kept if x['lang'] == 'ko']
    other = [x for x in kept if x['lang'] != 'ko']
    position = {x['id']: i for i, x in enumerate(kept)}
    out, i, j, n_ko = [], 0, 0, 0
    while i < len(ko) or j < len(other):
        n = len(out) + 1
        need_ko = n_ko < math.ceil(korean * n)
        if i < len(ko) and (j >= len(other) or need_ko or position[ko[i]['id']] < position[other[j]['id']]):
            out.append(ko[i]); i += 1; n_ko += 1
        else:
            out.append(other[j]); j += 1
    return out, capped


def sha256(path) -> str:
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def main(argv=None):
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument('--data', required=True, help='Export directory (pool/, train.jsonl, taxonomy.json)')
    p.add_argument('--run', required=True, help='Run directory with temperatures.json and predictions/pool-*.jsonl')
    p.add_argument('--out', required=True, help='Output directory (refuses to overwrite)')
    args = p.parse_args(argv)
    data, run, out = Path(args.data), Path(args.run), Path(args.out)
    if (out / 'candidates.jsonl').exists() or (out / 'summary.json').exists():
        raise SystemExit('Refusing to overwrite ' + str(out))
    tax = Taxonomy.load(data / 'taxonomy.json')
    if Taxonomy.load(run / 'taxonomy.json').to_json() != tax.to_json():
        raise SystemExit('The run taxonomy differs from the export taxonomy')
    temps = json.loads((run / 'temperatures.json').read_text())
    rare = rare_fine(load_jsonl(data / 'train.jsonl'), tax)
    ev = events(tax, rare)
    pool_files = sorted((data / 'pool').glob('*.jsonl'))
    pred_files = [run / 'predictions' / ('pool-' + f.name) for f in pool_files]
    rows = [r for f in pool_files for r in load_jsonl(f)]
    preds = {}
    for f in pred_files:
        for r in load_jsonl(f):
            if r['id'] in preds:
                raise SystemExit('Duplicate prediction id: ' + r['id'])
            preds[r['id']] = r
    ids = [r['id'] for r in rows]
    if len(set(ids)) != len(ids) or set(ids) != set(preds):
        raise SystemExit('Pool rows and pool predictions do not match one to one')
    flagged = [r['id'] for r in rows if r.get('p4_context') or any((r.get('overlap') or {}).values())]
    if flagged:
        raise SystemExit('Pool rows carry exclusion flags: %s' % flagged[:5])
    items = []
    for r in rows:
        s = score(preds[r['id']], tax, temps, ev)
        items.append({'id': r['id'], 'tweet_id': r['tweet_id'], 'author_id': r['author_id'], 'lang': r['lang'],
                      'day_utc': r['day_utc'], 'has_unavailable_context': bool(r.get('has_unavailable_context')), **s})
    ranked, capped = rank(items)
    out.mkdir(parents=True, exist_ok=True)
    rounded = []
    for n, x in enumerate(ranked, 1):
        rounded.append({'rank': n, **{k: x[k] for k in ('id', 'tweet_id', 'author_id', 'lang', 'day_utc', 'has_unavailable_context',
                                                 'top_event')},
                        'thin_score': round(x['thin_score'], 6), 'entropy': round(x['entropy'], 6),
                        'events': {k: round(v, 6) for k, v in x['events'].items()}})
    write_jsonl(out / 'candidates.jsonl', rounded)

    def block(xs):
        return {'posts': len(xs), 'korean': sum(1 for x in xs if x['lang'] == 'ko'),
                'top_event': dict(Counter(x['top_event'] for x in xs).most_common()),
                'expected_events': {k: round(sum(x['events'][k] for x in xs), 1) for k in ev},
                'authors': len({x['author_id'] for x in xs})}
    summary = {
        'version': VERSION, 'rules': {'author_cap': AUTHOR_CAP, 'korean_share': KOREAN_SHARE, 'rare_share': RARE_SHARE,
                                      'events': {k: {'field': f, 'options': o} for k, (f, o) in ev.items()},
                                      'order': 'thin_score desc, entropy desc, id', 'probabilities': 'calibrated, not gated'},
        'pool_posts': len(rows), 'ranked': len(ranked), 'left_out_by_author_cap': len(capped),
        'authors_capped': len({x['author_id'] for x in capped}),
        'top_100': block(ranked[:100]), 'top_1000': block(ranked[:1000]), 'top_4000': block(ranked[:4000]),
        'all': block(ranked),
        'inputs': {str(f): sha256(f) for f in [data / 'taxonomy.json', data / 'train.jsonl', run / 'temperatures.json',
                                               *pool_files, *pred_files]},
        'candidates_sha256': sha256(out / 'candidates.jsonl'),
        'note': 'A ranked list only: context lookups and Astra labelling need the user\'s go.'}
    (out / 'summary.json').write_text(json.dumps(summary, indent=1, ensure_ascii=False) + '\n')
    print(json.dumps({k: summary[k] for k in ('pool_posts', 'ranked', 'left_out_by_author_cap', 'authors_capped')}))
    return 0


if __name__ == '__main__':
    raise SystemExit(main())