import json

MONTHS = [
    "January",
    "February",
    "March",
    "April",
    "May",
    "June",
    "July",
    "August",
    "September",
    "October",
    "November",
    "December",
]


def generate_skip_cases(f, seq, prompt, tag, start=3, end=1, cls=True):
    for i in range(start, len(seq) - end):
        i1 = i + 1
        th = (
            f"{i1}-th"
            if i1 % 10 not in [1, 2] or i1 % 100 in [11, 12]
            else (f"{i1}-st" if i1 % 10 == 1 else f"{i1}-nd")
        )
        p = prompt.format(i0=i, i1=i1, v=str(seq[i]), th=th)
        if cls:
            corr = {"classes": f"[' {seq[i]}', ' {seq[i + 1]}']", "answer_index": 1}
        else:
            corr = {"completion": f" {seq[i + 1]}"}
        json.dump(
            {
                "prompt": f"Task: {p}\nOutput: {' '.join(map(str, seq[:i]))}",
                "tag": tag,
                **corr,
            },
            f,
        )
        f.write("\n")


for cls, f in [
    (True, open("strong-local-priors_classification.jsonl", "w")),
    (False, open("strong-local-priors_sequence-probability.jsonl", "w")),
]:
    generate_skip_cases(
        f,
        list("ABCDEFGHIJKLMNOPQRSTUVWXYZ"),
        "Write out the English alphabet in upper case, skipping the {th} letter.",
        "letters-by-number",
        cls=cls,
    )

    generate_skip_cases(
        f,
        list("abcdefghijklmnopqrstuvwxyz"),
        "Write out the English alphabet in lower case, skipping the {th} letter.",
        "letters-by-number",
        cls=cls,
    )

    generate_skip_cases(
        f,
        list("ABCDEFGHIJKLMNOPQRSTUVWXYZ"),
        "Write out the English alphabet in upper case, skipping the letter {v}.",
        "letters-by-number",
        cls=cls,
    )

    generate_skip_cases(
        f,
        list("abcdefghijklmnopqrstuvwxyz"),
        "Write out the English alphabet in lower case, skipping the letter {v}.",
        "letters-by-number",
        cls=cls,
    )

    generate_skip_cases(
        f,
        list(range(30)),
        "Write out the numbers 0 to 29, skipping the number {v}.",
        "numbers",
        cls=cls,
    )

    generate_skip_cases(
        f,
        MONTHS,
        "Write out the months of the year, skipping {th} month.",
        "months-by-number",
        cls=cls,
    )

    generate_skip_cases(
        f, MONTHS, "Write out the months of the year, skipping {v}.", "months-by-name", cls=cls
    )
