from symbols.data.data import write_propositions, write_schemata
from symbols.data.partitioned_option import Rule
from symbols.data.unique_symbol_list import UniqueSymbolList
from symbols.experimental.generate_ppddl import PPDDLSymbolWrapper, build_ppddl
from symbols.experimental.generate_propositions import factorise, generate_starting_symbol, generate_symbol_set, \
    _is_similar
from symbols.file_utils import make_path
from symbols.pddl.description import Description
from symbols.pddl.schema import Schema, NOT_FAILED
from symbols.render.render import visualise_ppddl_symbols
from symbols.symbols.distribution_symbol import DistributionSymbol
import itertools
import numpy as np
import pickle


def compute_combinations(operators):
    effects = list()
    combs = []
    for operator in operators:
        for x in operator.list_effects:
            effects.append(x.get_symbols())
    for i in range(1, len(effects) + 1):
        els = [list(x) for x in itertools.combinations(effects, i)]
        combs.extend(els)
    ret = dict
    for x in combs:
        mask = list()
        eff = list()
        for m, e in x:
            mask.append(m)
            eff.append(e)
        ret[tuple(mask)] = eff
    return ret


def describe(option):
    return str(option)


def generate_symbol_vocabulary(domain, ppddl_symbol_dir, operators, n_objects, initial_states, verbose=True):
    (factors, options) = factorise(operators, n_objects, verbose=verbose)
    if verbose:
        print("Generating start state symbols...")

    symbol_list = generate_starting_symbol(domain, factors, n_objects, initial_states)

    if verbose:
        print("Start position generated " + str(len(symbol_list)) + " propositions")
        print("Generating propositions...")

    (option_symbols, symbol_list) = generate_symbol_set(operators, factors, symbol_list)
    write_propositions(ppddl_symbol_dir, symbol_list, verbose=verbose)
    return (factors, options), (option_symbols, symbol_list)


# TODO: reduce same objects to class
def build_forward_model(domain, output_directory, ppddl_symbol_dir, operators, factors, option_symbols, symbol_list,
                        verbose=True):

    if verbose:
        print("Generating PPDDL description...")
    schemata = build_ppddl(factors, operators, symbol_list, option_symbols, verbose=verbose)
    write_schemata(ppddl_symbol_dir, schemata, verbose=verbose)
    desc = Description(str(domain), symbol_list, describe)
    for schema in schemata:
        desc.add_schema(schema)
    filename = make_path(output_directory, 'test.ppddl')
    if verbose:
        print("Writing to file " + filename)

    with open(filename, 'w') as file:
        file.write(str(desc))

    return desc
