"""More held-out rows no training source contains: HellaSwag (validation) and dair-ai/emotion (test).

    python heldout2.py /content/out/heldout2_rows.jsonl 400
"""

import json
import random
import sys

from datasets import load_dataset

EMOTIONS = {
    "sadness": "The writer feels sad, hurt or lonely.",
    "joy": "The writer feels happy, pleased or content.",
    "love": "The writer feels love, affection or tenderness.",
    "anger": "The writer feels angry, annoyed or resentful.",
    "fear": "The writer feels afraid, anxious or nervous.",
    "surprise": "The writer feels surprised or amazed.",
}


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


def hellaswag(n, rng):
    rows = []
    for i, row in enumerate(load_dataset("Rowan/hellaswag", split = "validation")):
        if not str(row.get("label", "")).strip().isdigit():
            continue
        keys = ["A", "B", "C", "D"]
        label = keys[int(row["label"])]
        rows.append(
            {
                "id": f"hellaswag_val_{i}",
                "workflow": "hellaswag",
                "state": {"activity": row["activity_label"], "context": row["ctx"]},
                "questions": {
                    "ending": {
                        "type": "choice",
                        "instructions": "Which ending most plausibly continues the context?",
                        "criteria": dict(zip(keys, row["endings"])),
                    }
                },
                "gold": {"ending": one_hot(keys, label)},
            }
        )
    rng.shuffle(rows)
    return rows[:n]


def emotion(n, rng):
    dataset = load_dataset("dair-ai/emotion", "split", split = "test")
    names = dataset.features["label"].names
    rows = []
    for i, row in enumerate(dataset):
        label = names[int(row["label"])]
        rows.append(
            {
                "id": f"emotion_test_{i}",
                "workflow": "emotion",
                "state": row["text"],
                "questions": {
                    "emotion": {
                        "type": "choice",
                        "instructions": "Which emotion does the writer express?",
                        "criteria": {name: EMOTIONS[name] for name in names},
                    }
                },
                "gold": {"emotion": one_hot(names, label)},
            }
        )
    rng.shuffle(rows)
    return rows[:n]


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