"""Rebuilds data/extraction.json from the extraction coders' work under data/coding/.

For each eligible trial (in data/trials.json's order): the third reader's extraction where the two coders
differed in any fact, otherwise coder A's, which is then identical in every fact to coder B's. A later round
(data/coding/extraction2) extracted again the trials whose full text arrived after the first round, and
replaces the first round's extraction of them. The author's annotations (data/annotations.json) mark counts
of improvement (responders, remitters) and say which arm each between-arm difference subtracts from which.

    python3 -I code/assemble_extraction.py [bundle dir] [output path]

With no output path it writes data/extraction.json; analyze.R reads that file, so this script is not part
of code/run.
"""

import json
import pathlib
import sys

root = pathlib.Path(sys.argv[1] if len(sys.argv) > 1 else ".")
out_path = pathlib.Path(sys.argv[2]) if len(sys.argv) > 2 else root / "data" / "extraction.json"
coding = root / "data" / "coding"


def read_all(folder, pattern):
    return [x for f in sorted(folder.glob(pattern)) for x in json.loads(f.read_text())]


coder_a, third = {}, {}
for round_dir in (coding / "extraction", coding / "extraction2"):
    for x in read_all(round_dir / "coderA", "coderA_batch*.json"):
        coder_a[x["trial"]] = x
        third.pop(x["trial"], None)
    for x in read_all(round_dir / "third", "decisions*.json"):
        third[x["trial"]] = x

notes = json.loads((root / "data" / "annotations.json").read_text())
improvement = {(c["trial"], c["measure"], c["day"], c["form"], c["definition"]) for c in notes["counts"]}
contrast = {(c["trial"], c["form"]): c["contrast"] for c in notes["contrasts"]}

out = []
for t in json.loads((root / "data" / "trials.json").read_text()):
    if not t["design"].get("eligible"):
        continue
    tid = t["trial"]
    x = json.loads(json.dumps(third.get(tid) or coder_a[tid]))
    x["decided_by"] = "third reader" if tid in third else "both coders"
    x.pop("reason", None)
    for e in x.get("depression") or []:
        if (tid, e.get("measure"), e.get("day"), e.get("form"), e.get("definition")) in improvement:
            e["counts"] = "improvement"
        if (tid, e.get("form")) in contrast:
            e["contrast_sign"] = contrast[(tid, e.get("form"))]
    out.append(x)
out_path.write_text(json.dumps(out, indent=1, ensure_ascii=False) + "\n")
print(len(out), "trials;", sum(1 for x in out if x["decided_by"] == "third reader"), "decided by the third reader")