"""Build an openjev dataset from your own retrieval logs, to fine-tune or re-calibrate Ragzilla.

One JSONL line per question your retriever served, with the passages in the order it returned them
and which of them answered (from clicks, citations, annotators, or a judge you already run):

    {"id": "q17", "question": "how long is a cdl permit good for in texas",
     "passages": ["...", {"title": "...", "text": "..."}, "..."], "relevant": [0, 3], "group": "optional"}

    python my_rows.py my_rows.jsonl --out my_rows            # 50 / 10 / 20 / 20 split by group
    openjev tod train my_rows --method encoder --model ragzilla --out my_ragzilla \
        --loss focal --focal-gamma 2 --focal-alpha 0.5 --brier 0.1 --learning-rate 1e-5 --max-steps 25
    openjev tod calibrate my_ragzilla my_rows
    openjev tod evaluate my_ragzilla my_rows

Rows are prepared the way Ragzilla's training data was: whitespace collapsed, a titled passage rendered
"<title>\\n<text>", passages over 4,000 characters cut at a word boundary, duplicate passages merged,
empty passages dropped, rows with fewer than 2 passages dropped (and counted). Rows with the same `group`
(a user, a document, the same question) share a split; a small workload needs a calibration split large
enough to choose an operating point on, hence 20% calibration and 20% test by default.

Needs only `pip install openjev-sdk==0.5.0`. The spec below is Ragzilla's own (`spec.json` in this
folder has the same fingerprint), so the rows load into the model as its training rows did.
"""
import argparse
import json
import re
from pathlib import Path

from openjev.tod import Episode, ListSpec, Subset, build_dataset

SPEC = ListSpec('ragzilla', [dict(name='relevant', question='Which passages are relevant to answering the question?',
                                  candidates='passages', version='1')])
MAX_CHARS = 4000


def clean(text):
    text = re.sub(r'\s+', ' ', text).strip()
    if len(text) > MAX_CHARS:
        text = text[:MAX_CHARS].rsplit(' ', 1)[0] + ' ...'
    return text


def passage_text(p):
    if isinstance(p, dict):
        title = re.sub(r'\s+', ' ', p.get('title') or '').strip()
        body = clean(p.get('text') or '')
        return f'{title}\n{body}' if title and body else body
    return clean(str(p))


def decision(row, drops):
    """One SDK decision from a row, or None (counted in `drops`)."""
    texts, keep_idx, seen = [], [], {}
    for i, p in enumerate(row['passages']):
        t = passage_text(p)
        if not t:
            drops['empty passages'] += 1
            continue
        if t in seen:                       # duplicate: the first copy stands for both
            if i in set(row.get('relevant', [])):
                keep_idx.append(seen[t])
            drops['duplicate passages merged'] += 1
            continue
        seen[t] = len(texts)
        if i in set(row.get('relevant', [])):
            keep_idx.append(len(texts))
        texts.append(t)
    if len(texts) < 2:
        drops['rows with fewer than 2 passages'] += 1
        return None
    state = {'question': re.sub(r'\s+', ' ', row['question']).strip(),
             'passages': [{'id': f'p{i}', 'text': t} for i, t in enumerate(texts)]}
    target = Subset(sorted({f'p{i}' for i in keep_idx}))
    try:
        return SPEC.decision(str(row['id']), 'relevant', state, target=target, group_id=str(row.get('group') or row['id']))
    except ValueError as exc:               # e.g. more than 255 passages
        drops[f'refused by the SDK: {exc}'[:80]] += 1
        return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('rows', type=Path)
    ap.add_argument('--out', type=Path, required=True)
    ap.add_argument('--fractions', default='0.5,0.1,0.2,0.2', help='train,validation,calibration,test')
    ap.add_argument('--seed', default='0')
    a = ap.parse_args()
    from collections import Counter
    drops = Counter()
    rows = [json.loads(l) for l in a.rows.open(encoding='utf-8') if l.strip()]
    decisions = [d for d in (decision(r, drops) for r in rows) if d is not None]
    f = [float(x) for x in a.fractions.split(',')]
    fractions = dict(train=f[0], validation=f[1], calibration=f[2], test=f[3])
    ds = build_dataset(SPEC, [Episode(d.id, (d,)) for d in decisions], fractions=fractions, seed=a.seed)
    ds.write(a.out)
    print(json.dumps(dict(out=str(a.out), rows_read=len(rows), decisions=len(decisions), dropped=dict(drops),
                          spec_fingerprint=SPEC.fingerprint, splits=ds.report['splits']), indent=2))


if __name__ == '__main__':
    main()