import subprocess, os
from random import randint

def generate(n, m, k, m_per_stage=float('inf')):

    assert m_per_stage * (n - 1) > maxm

    output = ""
    output += f"{n} {m}\n"

    s = randint(1, k)
    output += f"{s} {k}\n"

    # randomly distribute m machines into n-1 buckets (stages)
    # absolutely not uniformly distributed, this will skew more towards the later stages
    # but that doesn't really matter

    machines = n - 1
    buckets = [1 for _ in range(n - 1)]

    while machines < m:
        target = randint(0, n - 2)
        while buckets[target] >= m_per_stage:
            target += 1
            target %= n - 1
        buckets[target] += 1
        machines += 1

    for stage, bucket_size in enumerate(buckets):
        for _ in range(bucket_size):
            output += f"{stage+1} {randint(1, k)} {randint(1, k)}\n"

    return output


if __name__ == '__main__':
    minn, maxn = 2, 30
    minm, maxm = 1, 500
    mink, maxk = 1, int(1e4)

    def random_case():
        n=randint(minn,maxn)
        m=randint(n-1, maxm)
        k=randint(mink, maxk)

        return n, m, k

    #
    # change these bits
    #

    random_case_count = 7

    # generate random cases with specific parameters:
    # [n, m, k, count] (s, m_i, m_o are randomised <k)
    cases = [
        [2, 1, 5, 1],
        [3, 3, 5, 2],
        [3, 10, 15, 2],
        [3, 50, 1000, 2],
        [10, 10, 100, 2],
        [25, 200, 5000, 2],
        [30, 500, 10000, 2],

        # fully random cases
        *[[*random_case(), 1] for _ in range(random_case_count)]
    ]

    #
    # ^^^
    #

    os.makedirs("../data/secret", exist_ok=True) # make data/secret directory if doesn't exist

    test_number = 1
    for i, values in enumerate(cases):
        n, m, k, count = values
        for j in range(count):
            path = f"../data/secret/r_{test_number:02d}"
            test_number += 1

            test = generate(n, m, k)
            with open(path + ".in", "w+", newline="\n") as f:
                f.write(test)

            # apply solution to find result
            answer = subprocess.run(
                ["py", "../submissions/accepted/dp.py"],
                input=test, text=True, capture_output=True)

            with open(path + ".ans", "w+", newline="\n") as f:
                f.write(answer.stdout)