#!/usr/bin/env python3
import sys, argparse, random
from collections import deque
from sys import stderr
from types import new_class

def main():
    parser = argparse.ArgumentParser(description="Generator for categories and items in categories.")
    parser.add_argument('-n', metavar='N', type=int, required=True, help='number of nodes')
    parser.add_argument('-m', metavar='M', type=int, required=True, help='number of edges')
    parser.add_argument('-p', metavar='P', type=str, default="False", help='allow parallel edges (default: False)' )
    parser.add_argument('-l', metavar='L', type=str, default="False", help='allow self-loops (default: False)' )
    parser.add_argument('-c', metavar='C', type=str, default="True", help='force graph to be connected (default: True)')
    parser.add_argument('-u', metavar='U', type=str, default="True", help='shuffle vertices after generation (default: True)')
    parser.add_argument('-w', metavar='W', type=int, default=0, help='maximum edge weight. If 0, then no weights')
    parser.add_argument('-o', metavar='O', type=str, default="random", help='options (dash-separated)')
    parser.add_argument('-s', metavar='S', type=int, default=None, help='seed')

    #Custom
    parser.add_argument('-a', metavar='A', type=int, required=True, help='number of water depots')
    parser.add_argument('-f', metavar='F', type=int, required=True, help='number of food depots')
    parser.add_argument('-t', metavar='T', type=int, required=True, help='max edge weight')
    args = parser.parse_args()
    random.seed(args.s)

    #result = water_food_depot_generator(args.n, args.m, args.w, args.f, args.s)
    # prints first line consisting of n and k

    ## Adaption to doomsday
    n = args.n
    m = args.m
    w = args.a
    f = args.f
    if w > n:
        w = n
    if f > n:
        f = n
    
    allow_parallel = args.p.lower() == "true"
    allow_loops = args.l.lower() == "true"
    force_connected = args.c.lower() == "true"
    do_shuffle = args.u.lower() == "true"
    max_edge_weights = args.t
    use_edge_weights = max_edge_weights > 0
    options = args.o.split('-')

    edjs = undirected_graph_edge_list_generator(n, m, allow_parallel, allow_loops, force_connected, max_edge_weights, use_edge_weights, options)
    #if do_shuffle:
    #    print("Doing shuffle", args.u, file=stderr)
    #    edjs = shuffle_vertices(n+1, edjs)
    #    random.shuffle(edjs)

    for option in options:
        if option.startswith("extraZeroEdges"):
            for _ in range(int(option[len("extraZeroEdges"):])):
                extraEdge = [0, random.randrange(1, n+1)]
                random.shuffle(extraEdge)
                edjs.append(tuple(extraEdge))
            random.shuffle(edjs)
    
    water_locations = random.sample(range(n), min(w, n))
    food_locations = random.sample(range(n), min(f, n))

    print(str(n) + " " + str(len(edjs)) + " " + str(w) + " " + str(f))
    print(" ".join(map(str, water_locations)))
    print(" ".join(map(str, food_locations)))
    print("\n".join(" ".join(map(str, edj)) for edj in edjs))

    # print all the generated lines
    #for line in result:
    #    print(" ".join(line))

def water_food_depot_generator(n, m, w, f, seed):
    random.seed(seed)

    # The result should be stored like this
    # index 0 -> list of water depots
    # index 1 -> list of food depots
    # index 2..m+2 -> weighted edges
    result = []

    # Firstly, set the first w locations as water depots
    # Seconds, set the last f locations as food depots
    water_depots = list(range(1, w))
    food_depots = list(range(n - f, n))

    # TODO
    # Next, create edges that guarantees paths from
    # the source to at least 1 water depot and 1 food depot

    return result

def undirected_graph_edge_list_generator(n, m, allow_parallel, allow_loops, force_connected, max_edge_weights, use_edge_weights, options):
    """ Generates an edge list
    """
    # Sanity checks
    if m < n-1 and force_connected: raise Exception("Not enough edges to enforce connectivity")
    if m > n * (n - 1) // 2 + (n if allow_loops else 0) and not allow_parallel: raise Exception("Can't have that many edges without parallel edges")
    vertices = list(range(n))

    edges = {}

    if ("path" in options):
        print("Generating path", n, m, file=stderr)
        if m < n - 1: raise Exception("Not enough edges to make a path")
        for i in range(n-1):
            edges[(i, i+1)] = 1
        m -= n - 1
    elif ("hamcircle" in options):
        if m < n - 1: raise Exception("Not enough edges to make a hamiltonian circle")
        for i in range(n):
            edges[(i, (i+1)%n)] = 1
        m -= n
    elif ("complete" in options):
        print("Warning: not enough edges to make a complete graph. Edges will be removed.", file=stderr)
        for a in range(n):
            for b in range(a + 1, n):
                edges[(a, b)] = 1
        m -= n * (n - 1) // 2
    elif ("star" in options):
        if m < n - 1: raise Exception("Not enough edges to make a star")
        for i in range(1, n):
            edges[(0, i)] = 1
        m -= n - 1
    elif ("starofstars" in options):
        if m < n - 1: raise Exception("Not enough edges to make a star of stars")
        sqrtN = int(n**0.5)
        for i in range(1, n):
            edges[(i // sqrtN, i)] = 1
        m -= n - 1
    elif ("twostarstogether" in options):
        if m < n - 1: raise Exception("Not enough edges to make two stars that are connected.")
        if n < 2: raise Exception("Not enough vertices to make two stars.")
        edges[(0, 1)] = 1
        for i in range(2, n):
            edges[(i % 2, i)] = 1
        m -= n - 1
    elif force_connected:
        # First we create a connected component. We connect two adjacent components by joining their top guys in UF 
        # structure. join is random
        uf = list(range(n))
        def find(i):
            if uf[i] == i:
                return i
            uf[i] = find(uf[i])
            return uf[i]

        def union(i, j):
            pi, pj = find(i), find(j)
            if random.choice([True, False]):
                uf[pi] = pj
            else:
                uf[pj] = pi
        
        queue = deque(range(n))
        while len(queue) >= 2:
            a, b = queue.popleft(), queue.popleft()
            union(a, b)
            edges[(a, b)] = 1
            queue.append(find(a))

        m -= n-1


    if ("random" in options):
        # Note: can be very slow if m approaches n*(n-1)//2 and parallel edges are not allowed
        print("Generating {} random edges".format(m), file=stderr)
        for _ in range(m):
            e = (random.randrange(n), random.randrange(n)) if allow_loops else tuple(random.sample(vertices, 2))
            while (e in edges.keys() and not allow_parallel):
                e = (random.randrange(n), random.randrange(n)) if allow_loops else tuple(random.sample(vertices, 2))
            edges[e] = 1 if e not in edges else edges[e] + 1

    # Convert edge dict to list
    edjs = []
    for a, b in edges.keys():
        for _ in range(edges[(a, b)]):
            if use_edge_weights:
                min_edge_weights = 1 if "nozeroweight" in options else 0
                edjs.append((a, b, random.randrange(min_edge_weights, max_edge_weights + 1)))
            else:
                edjs.append((a, b))

    return edjs

if __name__ == "__main__":
    main()
