#!/usr/bin/env python3
import random

random.seed(0)

TNAMES = [];OA=ord('A')
for i in range(26):
    for j in range(26):
        TNAMES.append( f'{chr(OA+i)}{chr(OA+j)}' )

MAX_BALLS=50
MAX_QUERIES=200

def gen_random_case(ind,nteams,nqueries):
    random.shuffle(TNAMES)
    TEAMS = TNAMES[:nteams]
    #print(TEAMS)
    freqs = [ random.randint(1,MAX_BALLS) for _ in range(nteams)]
    freqs.sort(reverse=True)
    #print(freqs)
    with open(f'../data/secret/{ind:02}-random.in','w') as out:
        print(f'{nteams}',file=out)
        for i in range(nteams):
            print( f'{freqs[i]} {TEAMS[i]}', file=out)

        print(f'{nqueries}',file=out)
        used = set()
        for q in range(nqueries):
            while True:
                t = TEAMS[random.randint(0,nteams-1)]
                k = random.randint(1,nteams)
                if (t,k) in used:continue

                print(t,k,file=out)
                used.add( (t,k) )
                break

def gen_corner(ind,nteams,nqueries):
    TEAMS = ['A'*8 + chr(OA+i) + chr(OA+j) for i in range(26) for j in range(26) ][:nteams]
    freqs = [MAX_BALLS] * 40 + [1] * (nteams - 40)

    queries = []
    for i in range(nteams):
        for k in range(3):
            queries.append((TEAMS[i], k+1))
        for k in range(5):
            if 1 <= i + (k - 1) + 1 <= nteams:
                queries.append((TEAMS[i], i + (k - 1) + 1))

    while len(queries) < nqueries:
        t = TEAMS[random.randint(0,nteams-1)]
        k = random.randint(30, nteams)
        if (t,k) not in queries:
          queries.append( (t,k) )

    if len(queries)>nqueries:
        random.shuffle(queries)
        queries = queries[:nqueries]

    with open(f'../data/secret/{ind:02}-corner.in','w') as out:
        print(f'{nteams}',file=out)
        for i in range(nteams):
            print( f'{freqs[i]} {TEAMS[i]}', file=out)

        print(f'{nqueries}',file=out)
        for t,k in queries:
            print(t,k,file=out)

def gen_low_probability(ind):
    nteams = 50
    TEAMS = [chr(OA+i) + chr(OA+j) for i in range(26) for j in range(26) ][:nteams]
    freqs = [MAX_BALLS] * (nteams - 3) + [1] * 3

    queries = []
    for i in range(nteams):
        for k in range(3):
            queries.append((TEAMS[i], k+1))
        for k in range(5):
            if 1 <= i + (k - 1) + 1 <= nteams:
                queries.append((TEAMS[i], i + (k - 1) + 1))

    if len(queries)>MAX_QUERIES:
        random.shuffle(queries)
        queries = queries[:MAX_QUERIES]

    with open(f'../data/secret/{ind:02}-lowprob.in','w') as out:
        print(f'{nteams}',file=out)
        for i in range(nteams):
            print( f'{freqs[i]} {TEAMS[i]}', file=out)

        print(f'{len(queries)}',file=out)
        for t,k in queries:
            print(t,k,file=out)


gen_random_case( 1, 1,1)
gen_random_case( 2, 2,4)
gen_random_case( 3, 3,9)
gen_random_case( 4, 4,16)
gen_random_case( 5, 5,25)
gen_random_case( 6, 10,50)
gen_random_case( 7, 10,100)
gen_random_case( 8, 20,200)
gen_random_case( 9, 20,200)
gen_random_case(10, 50,200)
gen_random_case(11, 50,200)
gen_random_case(12, 50,200)
gen_corner(13, 50, 200)
gen_low_probability(14)