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

def main():
    parser = argparse.ArgumentParser(description="Generator for an undirected graph. Adapted for the purposes of placingbombs.")
    parser.add_argument('-n', metavar='N', type=int, required=True, help='number of nodes (at least 1)')
    parser.add_argument('-m', metavar='M', type=int, default=None, help='number of edges (default: n(n-1)//2)')
    parser.add_argument('-z', metavar='Z', type=str, default="True", help='use zero-indexing (default: True). If false, use 1-indexing.')
    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')

    parser.add_argument('-k', metavar='K', type=int, default=1, help='(custom) number of bombs')
    args = parser.parse_args()
    random.seed(args.s)
    
    ## Adaption to placingbombs
    n = args.n
    m = args.m if args.m is not None else n*(n+1)//2
    k = args.k

    use_zero_index = args.z.lower() == "true"
    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.w
    use_edge_weights = max_edge_weights > 0
    options = args.o.split('-')

    edjs = undirected_graph_edge_list_generator(n + 1, m, allow_parallel, allow_loops, force_connected, max_edge_weights, use_edge_weights, do_shuffle, options)
    if do_shuffle:
        print("Doing shuffle", args.u, file=stderr)
        edjs = shuffle_vertices(n+1, edjs)
        random.shuffle(edjs)
    if not use_zero_index:
        edjs = add_one_to_vertices(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)
    

    bomb_locations = [random.randrange(1, n+1) for _ in range(k)]
    
    print(n, len(edjs), k)
    print(" ".join(map(str, bomb_locations)))
    print("\n".join(" ".join(map(str, edj)) for edj in edjs))


def undirected_graph_edge_list_generator(n, m, allow_parallel, allow_loops, force_connected, max_edge_weights, use_edge_weights, do_shuffle, options):
    """ Generates an edge list
    """
    edges = {}
    edge_budget = m

    def containsEdge(u, v, w=None):
        nonlocal n, edges
        if w is None:
            return (u, v) in edges or (v, u) in edges
        else:
            return (u, v, w) in edges or (v, u, w) in edges

    def addEdge(u, v, w=None, howmany=1, ignore_budget=False):
        nonlocal n, edges, edge_budget, allow_parallel, allow_loops
        if not allow_parallel and containsEdge(u, v, w):
            return False
        if not allow_loops and u == v:
            return False
        if not ignore_budget and howmany > edge_budget:
            return False
        assert(howmany > 0)
        assert(howmany == 1 or allow_parallel)
        assert(0 <= u < n)
        assert(0 <= v < n)
        if w is not None and not use_edge_weights:
            raise Exception("Can't add a weighted edge in an unweigthed graph")
        
        e = (u, v) if w is None else (u, v, w)
        edges[e] = howmany if not e in edges else edges[e] + howmany
        edge_budget -= howmany
        return True

    def addEdges(new_edges, shuffle_order, item_name, allow_partial=False):
        if (len(new_edges) > edge_budget):
            print(f"Not enough budget ({edge_budget}) to create a {item_name} with {len(new_edges)} edges.", end=" ", file=stderr)
            if not allow_partial:
                print("Aborting.", file=stderr)
                return False
            else:
                print(f"Adding {edge_budget}.", file=stderr)
        random.shuffle(new_edges) if shuffle_order else None
        for u, v in new_edges[:edge_budget]:
            addEdge(u, v)
        return True

    def get_order(size, shuffle_order):
        nonlocal n
        assert(0 <= size <= n)
        order = random.sample(range(n), size)
        order.sort() if not shuffle_order else None
        return order

    def get_prospective_edge(i, j, order, shuffle_order):
        e = [order[i], order[j]]
        random.shuffle(e) if shuffle_order else None
        return e[0], e[1] 

    def add_prospective_edge(i, j, order, shuffle_order, new_edges, add_parallel):
        u, v = get_prospective_edge(i, j, order, shuffle_order)
        if add_parallel or not containsEdge(u, v):
            new_edges.append((u,v))

    def get_random_nonedge():
        nonlocal n, edges, allow_loops
        u, v = (random.randrange(n), random.randrange(n)) if allow_loops else random.sample(range(n), 2)
        while containsEdge(u, v):
            u, v = (random.randrange(n), random.randrange(n)) if allow_loops else random.sample(range(n), 2)
        return u, v

    def addAPath(size=n, shuffle_order=False, add_parallel=False):
        nonlocal n
        order = get_order(size, shuffle_order)
        new_edges = []
        for i in range(1, size):
            add_prospective_edge(i-1, i, order, shuffle_order, new_edges, add_parallel)
        return addEdges(new_edges, shuffle_order, "path")

    def addACycle(size=n, shuffle_order=False, add_parallel=False):
        nonlocal n
        order = get_order(size, shuffle_order)
        new_edges = []
        for i in range(0, size):
            add_prospective_edge(i-1, i, order, shuffle_order, new_edges, add_parallel)
        return addEdges(new_edges, shuffle_order, "cycle")

    def addAStar(size=n, shuffle_order=False, add_parallel=False):
        nonlocal n
        order = get_order(size, shuffle_order)
        new_edges = []
        for i in range(1, size):
            add_prospective_edge(0, i, order, shuffle_order, new_edges, add_parallel)
        return addEdges(new_edges, shuffle_order, "star")

    def addASausage(size=n, width=3, edge_count=2*n, shuffle_order=False, add_parallel=False):
        print(f"Adding sausage of size {size} width {width} edges {edge_count}", file=stderr)
        nonlocal n
        assert(1 < width <= n)
        if edge_count < n - 1:
            print(f"Not enough edges to make a sausage", file=stderr)
            return False
        order = get_order(size, shuffle_order)

        new_edges = []
        # Step 1: Create a star in first box
        for i in range(1, width):
            add_prospective_edge(0, i, order, shuffle_order, new_edges, add_parallel)

        # Step 2: Create the sausage connection
        for i in range(width, size):
            prev_i = (i // width - 1) * width + random.randrange(width)
            add_prospective_edge(prev_i, i, order, shuffle_order, new_edges, add_parallel)

        # Step 3: Add additional edges (note: if parallel is disabled, no guarantee that all
        # edges will be used)
        for _ in range(edge_count - n - 1):
            base_i = random.randrange(width, size)
            prev_i = (base_i // width - 1) * width + random.randrange(width)
            add_prospective_edge(prev_i, base_i, order, shuffle_order, new_edges, add_parallel)

        return addEdges(new_edges, shuffle_order, "sausage")

    def addABinaryTree(size=n, shuffle_order=False, add_parallel=False):
        nonlocal n
        order = get_order(size, shuffle_order)
        new_edges = []
        for i in range(1, size):
            add_prospective_edge((i + 1)//2 - 1, i, order, shuffle_order, new_edges, add_parallel)
        return addEdges(new_edges, shuffle_order, "binary tree")

    def addAStarOfStars(size=n, shuffle_order=False, add_parallel=False):
        nonlocal n
        sqrt = int(size**0.5)
        order = get_order(size, shuffle_order)
        new_edges = []
        for i in range(1, size):
            add_prospective_edge((i + sqrt - 1)//sqrt - 1, i, order, shuffle_order, new_edges, add_parallel)
        return addEdges(new_edges, shuffle_order, "tree of trees")

    def addAClique(size=n, shuffle_order=False, add_parallel=False):
        nonlocal n
        order = get_order(size, shuffle_order)
        new_edges = []
        for i in range(size):
            for j in range(i+1, size):
                add_prospective_edge(i, j, order, shuffle_order, new_edges, add_parallel)
        return addEdges(new_edges, shuffle_order, "clique", allow_partial=True)

    def addRandom(edges_to_add=edge_budget, add_parallel=False):
        """Add random edges"""
        nonlocal n, edge_budget, edges, allow_loops
        order = get_order(n, True)

        if not add_parallel:
            edges_so_far = len(edges)
            if edges_so_far + edges_to_add >= 0.8*(n*(n-1)//2):
                ## We want a graph which is almost complete; generate
                ## all non-added edges and pick randomly among them
                all_possible_edges = []
                for i in range(n):
                    for j in range(i if allow_loops else i + 1, n):
                        add_prospective_edge(i, j, order, True, all_possible_edges, False)
                
                random.shuffle(all_possible_edges)
                addEdges(all_possible_edges[:edges_to_add], True, "random")
            else:
                # Draw random edges not yet added
                for _ in range(edges_to_add):
                    u, v = get_random_nonedge()
                    addEdge(u, v)

        else:
            for _ in range(edges_to_add):
                u, v = (random.randrange(n), random.randrange(n)) if allow_loops else random.sample(range(n), 2)
                addEdge(u, v)


    # 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")

    for option in options:
        understood_option = False
        for f, description in [(addAPath, "path"),
                                (addACycle, "hamcycle"),
                                (addABinaryTree, "binarytree"),
                                (addAStar, "star"),
                                (addAStarOfStars, "sqrstar"),
                                (addAClique, "clique")]:
            if option.startswith(description):
                size = n if len(option) == len(description) else int(option[len(description):])
                f(size=size, shuffle_order=do_shuffle, add_parallel=allow_parallel)
                understood_option = True
        
        if option.startswith("sausage"):
            size = n
            width = int(n**0.5)
            sausage_edges = min(3*n, edge_budget)

            option_subparts = option[len("sausage"):].split("_") if len(option) > len("sausage") else []
            print(option_subparts, option, file=stderr)
            if len(option_subparts) == 1:
                width = int(option_subparts[0])
            elif len(option_subparts) == 2:
                width = int(option_subparts[0])
                sausage_edges = int(option_subparts[1])
            elif len(option_subparts) >= 3:
                size = int(option_subparts[0])
                width = int(option_subparts[1])
                sausage_edges = int(option_subparts[2])
            addASausage(size=n, width=width, edge_count=sausage_edges, shuffle_order=do_shuffle, add_parallel=allow_parallel)
            understood_option = True

        elif option.startswith("random"):
            count = edge_budget if len(option) == len("random") else int(option[len("random"):])
            addRandom(edges_to_add=count, add_parallel=allow_parallel)
            understood_option = True

    if force_connected:
        # Forcing connectivity after the fact --- can increase m.
        # To avoid this, use an option which creates a connected graph
        # first.
        uf = list(range(n))
        islands = set(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 (pi == pj): return
            if random.choice([True, False]):
                uf[pi] = pj
                islands.remove(pi)
            else:
                uf[pj] = pi
                islands.remove(pj)

        for u, v in edges:
            union(u, v)
        
        queue = deque(islands)
        while len(queue) >= 2:
            a, b = queue.popleft(), queue.popleft()
            union(a, b)
            addEdge(a, b)
            queue.append(find(a))


    # 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

def add_one_to_vertices(edjs):
    res = []
    for e in edjs:
        if len(e) == 2:
            res.append((e[0] + 1, e[1] + 1))
        elif len(e) == 3:
            res.append((e[0] + 1, e[1] + 1, e[2]))
        else:
            raise Exception("Unexpected edge: {}".format(e))
    return res

def shuffle_vertices(n, edjs):
    """ Renames the vertices randomly, and updates the provided edge list accordingly
    """
    rename = list(range(n))
    random.shuffle(rename)
    res = []
    for e in edjs:
        if len(e) == 2:
            res.append((rename[e[0]], rename[e[1]]))
        elif len(e) == 3:
            res.append((rename[e[0]], rename[e[1]], e[2]))
        else:
            raise Exception("Unexpexted edge: {}".format(e))
    return res

if __name__ == "__main__":
    main()
