import math
import random
import heapq

random.seed(42)

def output(prefix, order, k):
    n = len(order)
    topk = heapq.nlargest(k, range(n), key=lambda x: order[x])

    with open(f'../data/secret/{prefix}.in', 'w') as ifile:
        ifile.write(f'{n} {k}\n')
        ifile.write(' '.join(map(str, order)))
        ifile.write('\n')

    with open(f'../data/secret/{prefix}.ans', 'w') as ofile:
        ofile.write(' '.join(map(str, topk)))
        ofile.write('\n')


def generate(prefix, count, n, k):
    for i in range(count):
        order = list(range(n))
        random.shuffle(order)

        output(f'{prefix}-{i:02d}', order, k)

generate('01-small-random', 20, 10, 3)

max_size = 15_000
max_k = 30
generate('02-large-random', 10, max_size, max_k)

output('03-edge-00', list(range(max_size)), max_k)
output('03-edge-01', list(range(max_size))[::-1], max_k)
output('03-edge-02', [0], 1)
output('03-edge-03', list(range(10)), 10)

# generate heap worst-case
low = 0
high = 10**6
ans = [-1]*max_k

for i in range(max_k-1,-1,-1):
    if ans[i] == -1 and i != max_k-1:
        ans[i] = high
        high += 1
    else:
        ans[i] = low
        low += 1
        p = (i-1) // 2
        if ans[p] == -1: ans[p] = -2

ans += range(low, max_size - max_k + low)

output('03-edge-04', ans, max_k)
output('03-edge-05', ans[::-1], max_k)