"""
Dominion setup

Create files describing trees, dominions, and corresponding polymorphisms
"""
import random, json
import output
from graphs import random_tree, load_graph_from_file
from dominion import random_dominion
from relations import random_relation, random_adjacent_relation, Relation


def grow_forest(filename, num_of_trees, num_of_vertices):
    """
    Add a specified number of trees to a given file.

    Arguments:
        filename (str): The name of the output file.
        num_of_trees (int): The number of trees to be created.
        num_of_vertices (int): How many vertices each of these trees should
            have.
    """

    for _ in range(num_of_trees):
        T = random_tree(range(num_of_vertices))
        T.write_to_file(filename)


def build_dominions(tree_filename, dominion_filename, num_of_dominions,
    dominion_size):
    """
    Use the trees stored in a given file as constraint graphs for creating
    dominions. These dominions are then stored in their own file, along with a
    note about which tree was used to create them.

    Arguments:
        tree_filename (str): The name of the file where trees are stored.
        dominion_filename (str): The name of the output file.
        num_of_dominions (int): The number of dominions to be created.
        dominion_size (int): The number of rows (and columns) of the dominions.
    """

    with open(output.path + '//{}.json'.format(tree_filename), 'r') \
    as read_file:
        num_of_trees = sum(1 for _ in read_file)
    for _ in range(num_of_dominions):
        tree_number = random.randrange(num_of_trees)
        T = load_graph_from_file(tree_filename, tree_number)
        D = random_dominion(dominion_size, tuple(T.vertices), T)
        with open(output.path + '//{}.json'.format(dominion_filename), 'a') \
        as write_file:
            json.dump((tree_number, D.array), write_file)
            write_file.write('\n')


def find_homomorphisms(tree_filename, homomorphism_filename, universe_size):
    """
    Produce a file detailing homomorphisms from a given family of trees to a
    given Hamming graph.

    Arguments:
        tree_filename (str): The name of the file where trees are stored.
        homomorphism_filename (str): The name of the output file.
        universe_size (int): The number of elements in the universe of the
            relations to be produced.
    """

    with open(output.path + '//{}.json'.format(tree_filename), 'r') \
    as read_file:
        num_of_trees = sum(1 for _ in read_file)
    for tree_number in range(num_of_trees):
        T = load_graph_from_file(tree_filename, tree_number)
        # Choose a root of the tree and build a list of (parent, child) node
        # pairs.
        unexplored_vertices = list(T.vertices)
        next_vertices_to_check = [unexplored_vertices.pop()]
        explored_vertices = set()
        pairs = []
        while unexplored_vertices:
            next_vertex = next_vertices_to_check.pop()
            new_neighbors = frozenset(
            T.neighbors(next_vertex)).difference(explored_vertices)
            for neighbor in new_neighbors:
                pairs.append((next_vertex, neighbor))
                unexplored_vertices.remove(neighbor)
                next_vertices_to_check.append(neighbor)
            explored_vertices.add(next_vertex)
        # Create a list whose entries will become the images of each label
        # under the homomorphism. Initialize every spot to 0.
        homomorphism_values = len(T.vertices)*[0]
        homomorphism_values[pairs[0][0]] = random_relation(universe_size)
        # Starting all homomorphisms at empty relation for an experiment.
        # homomorphism_values[pairs[0][0]] = Relation([], 28, 2)
        for (parent, child) in pairs:
            homomorphism_values[child] = \
            random_adjacent_relation(homomorphism_values[parent])
        homomorphism_values = tuple(tuple(rel.tuples)
        for rel in homomorphism_values)
        with open(output.path + '//{}.json'.format(homomorphism_filename),
        'a') as write_file:
            json.dump((tree_number, homomorphism_values), write_file)
            write_file.write('\n')