#!/usr/bin/python3

from random import randint, seed, shuffle
from string import ascii_lowercase as abc
seed(47)

MINN = [ None, 5, 2, 2, 2, 1 ]
MAXN = [ None, 5, 10, 100, 3000, 100000 ]

class RandomSet:
    def __init__(self, N):
        self.inside = [ n for n in range(N) ]
        self.where = [ n for n in range(N) ]

    def random_element(self):
        return self.inside[ randint( 0, len(self.inside)-1 ) ]

    def remove(self, x):
        wx = self.where[x]
        if wx != len(self.inside) - 1:
            last = self.inside[-1]
            self.where[last] = wx
            self.inside[wx] = last
        self.inside.pop()

class DisjointSet:
    def __init__(self, N, K):
        self.degree = [ 0 for n in range(N) ]
        self.color = [ n for n in range(N) ]
        self.components = [ [n] for n in range(N) ]
        self.living_vertices = RandomSet(N)
        self.limit = K

    def union(self, x, y):
        cx, cy = self.color[x], self.color[y]
        if cx == cy: return
        nx, ny = len(self.components[cx]), len(self.components[cy])
        if nx < ny: cx, cy = cy, cx
        for ex in self.components[cy]:
            self.components[cx].append(ex)
            self.color[ex] = cx
        self.components[cy] = []

    def get_new_edge(self):
        fails = 0
        while True:
            x = self.living_vertices.random_element()
            y = self.living_vertices.random_element()
            if self.color[x] != self.color[y]:
                self.degree[x] += 1
                if self.degree[x] >= self.limit: self.living_vertices.remove(x)
                self.degree[y] += 1
                if self.degree[y] >= self.limit: self.living_vertices.remove(y)
                self.union(x, y)
                return x, y
            fails += 1
            if fails == 123456: return None

def relabel(N, edges):
    ''' nahodne premenuje vsetky vrcholy a nahodne otoci kazdu hranu '''
    perm = list(range(N))
    shuffle(perm)
    new_edges = [ (y, x) if randint(0,1) else (x, y) for x, y in edges ]
    return [ (perm[x], perm[y]) for x,y in new_edges ]

def save(fname, N, K, edges, do_relabel=True):
    print('saving', fname)
    if do_relabel:
        edges = relabel(N, edges)
    f = open(fname, 'w')
    print(N, K, len(edges), file=f)
    for x, y in edges: print(x, y, file=f)
    f.close()

def gen_random(N, K, P):
    print('gen_random', N, K, P)
    assert N*K >= 2*P
    DS = DisjointSet(N, K)
    edges = []
    for p in range(P):
        e = DS.get_new_edge()
        if e is None: break
        edges.append(e)
    return N, K, edges

def gen_snakes(N, count):
    print('gen_snakes', N, count)
    lengths = [ 1 for c in range(count) ]
    for i in range(N-count):
        lengths[ randint(0, count-1) ] += 1
    where = 0
    edges = []
    for l in lengths:
        for i in range(l-1):
            edges.append( (where+i, where+i+1) )
        where += l
    return N, 2, edges

def gen_stars(N, K):
    print('gen_stars', N, K)
    min_count = (N+K) // (K+1)
    max_count = min(N, int(min_count*1.1))
    count = randint(min_count, max_count)

    lengths = [ 1 for c in range(count) ]
    total = count
    while total < N:
        x = randint(0, count-1)
        if lengths[x] == K+1: continue
        lengths[x] += 1
        total += 1 
    where = 0
    edges = []
    for l in lengths:
        for i in range(l-1):
            edges.append( (where, where+i+1) )
        where += l
    return N, max(lengths)-1, edges

def gen_ktrees(N, K, count):
    print('gen_ktrees', N, K, count)
    assert K >= 3

    lengths = [ 1 for c in range(count) ]
    total = count
    while True:
        x = randint(0, count-1)
        delta = K if lengths[x] == 1 else K-1
        if total + delta > N: break
        total += delta
        lengths[x] += delta

    where = 0
    edges = []
    for l in lengths:
        big = l // (K-1)
        if K == 3: big -= 1
        big_degs = [ 0 for _ in range(big) ]
        for i in range(big-1):
            edges.append( (where+i, where+i+1) )
            big_degs[i] += 1
            big_degs[i+1] += 1
        step = 0
        for i in range(big):
            for j in range(K-big_degs[i]):
                edges.append( (where+i, where+big+step) )
                step += 1
        where += l
    return N, max(lengths)-1, edges

##################### samples

save('00.sample.a.in', 5, 2, [ (0,3), (1,4) ], False)
save('00.sample.b.in', 6, 4, [ (0,1), (1,2), (3,4), (4,5) ], False)
save('00.sample.c.in', 9, 1, [ (4,7) ], False)

##################### batch 1

for p, letter in enumerate('abcde'):
    save(f'01.{letter}.in', *gen_random(5, 4, p) )

##################### batch 2

strategies = [ gen_snakes, gen_snakes, gen_stars, gen_random, gen_random, gen_random, gen_random ]
args_list = [ (2,), (6,), (3,), (2,0), (2,4), (3,3), (4,4) ]

for letter, strategy, args in zip( abc[:len(strategies)], strategies, args_list):
    save(f'02.{letter}.in', *strategy(MAXN[2], *args) )

##################### batch 3

strategies = [ gen_snakes, gen_snakes, \
               gen_stars, gen_stars, \
               gen_ktrees, gen_ktrees, gen_ktrees, \
               gen_random, gen_random, gen_random, gen_random, \
               gen_random, gen_random, gen_random, gen_random ]
args_list = [ (4,), (17,), \
              (3,), (7,), \
              (3, 4), (3, 10), (7, 7), \
              (3, 77), (4, 63), (3, 23), (7, 15), \
              (1, 1, 2), (1, 0, 2), (1, 1, 3), (1, 2, 4) ]

for letter, strategy, args in zip( abc[:len(strategies)], strategies, args_list):
    N = randint( int(0.9*MAXN[3]), MAXN[3] )
    if strategy == gen_random and len(args) == 3:
        N, args = args[2], args[:2]
    save(f'03.{letter}.in', *strategy(N, *args) )

##################### batch 4

strategies = [ gen_snakes, gen_snakes, gen_snakes, \
               gen_stars, gen_stars, gen_stars, \
               gen_ktrees, gen_ktrees, gen_ktrees, \
               gen_random, gen_random, gen_random, gen_random ]
args_list = [ (4,), (17,), (MAXN[4]//5,), \
              (3,), (7,), (MAXN[4]//5,), \
              (3, 7), (3, 27), (7, 11), \
              (3, int(0.7*MAXN[4])), (4, int(0.5*MAXN[4])), (3, int(0.1*MAXN[4])), (100, 100) ]

for letter, strategy, args in zip( abc[:len(strategies)], strategies, args_list):
    N = randint( int(0.9*MAXN[4]), MAXN[4] )
    save(f'04.{letter}.in', *strategy(N, *args) )

##################### batch 5

strategies = [ gen_snakes, gen_snakes, gen_snakes, \
               gen_stars, gen_stars, gen_stars, \
               gen_ktrees, gen_ktrees, gen_ktrees, gen_ktrees, \
               gen_random, gen_random, gen_random, gen_random, gen_random, \
               gen_random, gen_random, gen_random, gen_random, gen_random, gen_random, gen_random, gen_random ]
args_list = [ (9,), (47,), (MAXN[5]//5,), \
              (3,), (7,), (MAXN[5]//5,), \
              (3, 11), (3, 111), (4, 23), (7, 123), \
              (3, int(0.7*MAXN[5])), (4, int(0.5*MAXN[5])), (3, int(0.1*MAXN[5])), (100, 9876), (4, MAXN[5]-1, MAXN[5]), \
              (1, 1, 2), (1, 0, 2), (1, 1, 3), (1, 2, 4), (0, 0, 1), (0, 0, 7), (1, 0, 1), (2, 0, MAXN[5]) ]

for letter, strategy, args in zip( abc[:len(strategies)], strategies, args_list):
    N = randint( int(0.9*MAXN[5]), MAXN[5] )
    if strategy == gen_random and len(args) == 3:
        N, args = args[2], args[:2]
    save(f'05.{letter}.in', *strategy(N, *args) )

