"""Held-out evaluation rows that neither training stage sees: ANLI (test r1-r3) and WinoGrande (dev).

Written once to /content/out/heldout_rows.jsonl in the typed-decisions layout
(id, workflow, state, questions, gold), so our model and Jev answer identical requests.
"""

import json
import random
import sys

from datasets import load_dataset

NLI = ("entailment", "neutral", "contradiction")
NLI_CRITERIA = {
    "entailment": "The premise shows the hypothesis is true.",
    "neutral": "The premise neither confirms nor rules out the hypothesis.",
    "contradiction": "The premise shows the hypothesis is false.",
}


def one_hot(keys, label):
    return {"label": label, "probabilities": {k: 1.0 if k == label else 0.0 for k in keys}}


def anli(n, rng):
    rows = []
    for split in ("test_r1", "test_r2", "test_r3"):
        for i, row in enumerate(load_dataset("facebook/anli", split = split)):
            if int(row["label"]) < 0:
                continue
            label = NLI[int(row["label"])]
            rows.append(
                {
                    "id": f"anli_{split}_{i}",
                    "workflow": "anli",
                    "state": {"premise": row["premise"], "hypothesis": row["hypothesis"]},
                    "questions": {
                        "relation": {
                            "type": "choice",
                            "instructions": "How does the premise relate to the hypothesis?",
                            "criteria": dict(NLI_CRITERIA),
                        }
                    },
                    "gold": {"relation": one_hot(NLI, label)},
                }
            )
    rng.shuffle(rows)
    return rows[:n]


def winogrande(n, rng):
    rows = []
    for i, row in enumerate(load_dataset("allenai/winogrande", "winogrande_xl", split = "validation")):
        label = "option_1" if str(row["answer"]) == "1" else "option_2"
        rows.append(
            {
                "id": f"winogrande_dev_{i}",
                "workflow": "winogrande",
                "state": {"sentence": row["sentence"]},
                "questions": {
                    "blank": {
                        "type": "choice",
                        "instructions": "Which option correctly fills the blank (_) in the sentence?",
                        "criteria": {"option_1": row["option1"], "option_2": row["option2"]},
                    }
                },
                "gold": {"blank": one_hot(("option_1", "option_2"), label)},
            }
        )
    rng.shuffle(rows)
    return rows[:n]


if __name__ == "__main__":
    out = sys.argv[1] if len(sys.argv) > 1 else "/content/out/heldout_rows.jsonl"
    n = int(sys.argv[2]) if len(sys.argv) > 2 else 400
    rng = random.Random(0)
    rows = anli(n, rng) + winogrande(n, rng)
    with open(out, "w", encoding = "utf-8") as f:
        for row in rows:
            f.write(json.dumps(row) + "\n")
    print("HELDOUT_DONE", out, len(rows), flush = True)