import tensorflow_probability as tfp

from tensorflow.python.eager.context import graph_mode
from tensorflow.python.framework.ops import disable_tensor_equality

from symbolic_pymc.tensorflow.printing import tf_dprint


disable_tensor_equality()

with graph_mode(), tf.Graph().as_default() as test_graph:
    mu_tf = tf.compat.v1.placeholder(tf.float32, name='mu',
                                     shape=tf.TensorShape([None]))
    tau_tf = tf.compat.v1.placeholder(tf.float32, name='tau',
                                      shape=tf.TensorShape([None]))

    normal_tfp = tfp.distributions.normal.Normal(mu_tf, tau_tf)

    value_tf = tf.compat.v1.placeholder(tf.float32, name='value',
                                        shape=tf.TensorShape([None]))

    normal_log_lik = normal_tfp.log_prob(value_tf)

tf_dprint(normal_log_lik)

square_diff_outs = [o.outputs[0] for o in logpfn_fg.get_operations()
                    if o.type == 'SquaredDifference' or o.type.startswith('Gather')]

for t in square_diff_outs:
    tf_dprint(t)

from pprint import pprint

tfp_names_to_pymc = {i.name: k for i, k in zip(logpfn_cf.structured_input_signature[0], init.keys())}

pprint(tfp_names_to_pymc)

{'values_0': 'hierarchical_model/__log_sigma_alpha',
 'values_1': 'hierarchical_model/mu_alpha',
 'values_2': 'hierarchical_model/beta',
 'values_3': 'hierarchical_model/alpha',
 'values_4': 'hierarchical_model/mu_beta',
 'values_5': 'hierarchical_model/__log_sigma_beta',
 'values_6': 'hierarchical_model/__log_eps'}

from tensorflow.core.protobuf import config_pb2

from tensorflow.python.framework import ops
from tensorflow.python.framework import importer
from tensorflow.python.framework import meta_graph

from tensorflow.python.grappler import cluster
from tensorflow.python.grappler import tf_optimizer


try:
    gcluster = cluster.Cluster()
except tf.errors.UnavailableError:
    pass

config = config_pb2.ConfigProto()


def normalize_tf_graph(graph_output, graph_inputs=[]):
    """Use grappler to normalize a graph.

    Arguments
    =========
    graph_output: Tensor
      A tensor we want to consider as "output" of a FuncGraph.
    graph_inputs: list of Tensor (optional)
      Any tensors that correspond to inputs for the given output node.

    Returns
    =======
    The simplified graph.
    """
    train_op = graph_output.graph.get_collection_ref(ops.GraphKeys.TRAIN_OP)
    train_op.clear()
    train_op.extend([graph_output] + graph_inputs)

    # if graph_inputs is not None:
    #     # ops.GraphKeys.MODEL_VARIABLES?
    #     train_vars = graph_output.graph.get_collection_ref(ops.GraphKeys.TRAINABLE_VARIABLES),
    #     train_vars.clear()
    #     train_vars.extend(graph_inputs)

    metagraph = meta_graph.create_meta_graph_def(graph=graph_output.graph)

    optimized_graphdef = tf_optimizer.OptimizeGraph(
        config, metagraph, verbose=True, cluster=gcluster)

    optimized_graph = ops.Graph()
    with optimized_graph.as_default():
        importer.import_graph_def(optimized_graphdef, name="")

    opt_graph_output = optimized_graph.get_tensor_by_name(graph_output.name)

    return opt_graph_output

normal_log_lik_opt = normalize_tf_graph(normal_log_lik)

res_unopt = normal_log_lik.eval({'mu:0': np.r_[3], 'tau:0': np.r_[1], 'value:0': np.r_[1]},
                                 session=tf.compat.v1.Session(graph=normal_log_lik.graph))

res_opt = normal_log_lik_opt.eval({'mu:0': np.r_[3], 'tau:0': np.r_[1], 'value:0': np.r_[1]},
                                  session=tf.compat.v1.Session(graph=normal_log_lik_opt.graph))

# They should be equal, naturally
assert np.array_equal(res_unopt, res_opt)

_ = [res_unopt, res_opt]

[array([-2.9189386], dtype=float32), array([-2.9189386], dtype=float32)]

tf_dprint(normal_log_lik_opt)

from itertools import chain
from functools import partial

from unification import var, reify, unify

from kanren import run, eq, lall, conde
from kanren.goals import not_equalo
from kanren.core import goaleval

from symbolic_pymc.tensorflow.meta import mt
from symbolic_pymc.relations import buildo
from symbolic_pymc.relations.graph import graph_applyo, reduceo
from symbolic_pymc.etuple import ExpressionTuple, etuple


def onceo(goal):
    """A non-relational operator that yields only the first result from a relation."""
    def onceo_goal(s):
        nonlocal goal
        g = reify(goal, s)
        g_stream = goaleval(g)(s)
        s = next(g_stream)
        yield s

    return onceo_goal


def tf_graph_applyo(relation, a, b):
    """Construct a `graph_applyo` goal that evaluates a relation only at tensor nodes in a meta graph.

    Parameters
    ----------
    relation: function
      A binary relation/goal constructor function
    a: lvar, meta graph, or etuple
      The left-hand side of the relation.
    b: lvar, meta graph, or etuple
      The right-hand side of the relation
    """

    def _expand_some_nodes(node):
        if isinstance(node, mt.Tensor) and node.op is not None:
            return etuple(node.operator, *node.inputs, eval_obj=node)
        return None

    gapplyo = partial(graph_applyo, relation, preprocess_graph=_expand_some_nodes)
    return gapplyo(a, b)


def tfp_normal_log_prob(loc, scale):
    log_unnormalized = -0.5 * tf.math.squared_difference(
        x / scale, loc / scale)
    log_normalization = 0.5 * np.log(2. * np.pi) + tf.math.log(scale)
    return log_unnormalized - log_normalization

def shift_squared_subso(in_graph, out_subs):
    """Construct a goal that produces transforms for chains like (y + x)**2, (x + z)**2."""

    Y_lv, X_lv, mu_X_lv = var(), var(), var()
    scale_Y_lv = var()

    X_form_lv = mt.Placeholder(dtype=var(), shape=var(), name=var())
    # The actual base object's placeholder might have `_user_specified_name` as
    # an extra `op.node_def.attr`, so let's just make the entire NodeDef a
    # logic variable.
    X_form_lv.op.node_def = var()

    mu_Y_lv = mt.realdiv(X_lv, scale_Y_lv, name=var())

    # Y_T_reshaped_lv = mt.Transpose(mt.reshape(Y_lv, var(), name=var()), var())
    Y_reshaped_lv = mt.reshape(Y_lv, var(), name=var())

    sqr_diff_Y_lv = mt.SquaredDifference(
        mt.realdiv(Y_reshaped_lv,
                   scale_Y_lv,
                   name=var()),
        mu_Y_lv,
        name=var())

    def Y_sqrdiffo(in_g, out_g):
        return lall(eq(in_g, sqr_diff_Y_lv),
                    # This just makes sure that we're only considering X's
                    # that are Placeholders.
                    eq(X_lv, X_form_lv))

    scale_X_lv = var()
    sqr_diff_X_lv = mt.SquaredDifference(
        # Mul is only used because RealDiv with 1 is changed by grappler
        # mt.realdiv(X_lv, X_denom_lv, name=var()),
        mt.mul(scale_X_lv, X_lv, name=var()),
        mu_X_lv,
        name=var())

    def X_sqrdiffo(in_g, out_g):
        return eq(in_g, sqr_diff_X_lv)

    Y_new_mt = mt.addv2(X_lv, mt.mul(scale_Y_lv, Y_lv))
    Y_log_scale = mt.log(scale_Y_lv, name=var())

    res = lall(
        # The first (y - x/a)**2 (anywhere in the graph)
        tf_graph_applyo(Y_sqrdiffo, in_graph, in_graph),

        # The corresponding (x/b - z)**2 (also anywhere else in the graph)
        tf_graph_applyo(X_sqrdiffo, in_graph, in_graph),

        # Find the log-scale factor (at this point, we might as well match an
        # entire normal log-likelihood!)
        tf_graph_applyo(lambda x, y: eq(x, Y_log_scale), in_graph, in_graph),

        # Not sure if we need this, but we definitely don't want X == Y
        (not_equalo, [Y_lv, X_lv], True),

        # Create replacement rule pairs
        eq(out_subs, [[Y_lv, Y_new_mt],
                      [Y_log_scale, 0.0]]))

    return res

def shift_squared_terms(in_obj, graph_inputs=[]):
    """Re-center/scale SquaredDifference terms corresponding to hierarchical normals."""

    # Normalize and convert to a meta graph
    in_obj = mt(normalize_tf_graph(in_obj, graph_inputs=graph_inputs))

    # This run returns all the substitutions found in the graph
    subs_lv = var()
    subs_res = run(0, subs_lv, shift_squared_subso(in_obj, subs_lv))

    if not subs_res:
        print("Failed to find the required forms within the graph.")
        return

    # NOTE: We're only going to apply the first transformation pair for now.
    subs_res = [subs_res[0]]

    def subs_replaceo(in_g, out_g):
        """Create a goal that applies substitutions to a graph."""
        def _subs_replaceo(in_g, out_g):
            nonlocal subs_res
            # Each result is a pair of replacement pairs:
            #   the first pair is the re-center/scale transform,
            #   the second pair is the cancellation of the log differential scale term.
            subs_goals = [[eq(in_g, x), eq(out_g, y)]
                          for x, y in chain.from_iterable(subs_res)]
            x_g = conde(*subs_goals)
            return x_g

        g = onceo(tf_graph_applyo(_subs_replaceo, in_g, out_g))
        return g

    # Apply each substitution once
    out_graph_lv = var()
    res = run(1, out_graph_lv, reduceo(subs_replaceo, in_obj, out_graph_lv))

    if res:

        def reify_res(graph_res):
            """Reconstruct and/or reify meta object results."""
            from_etuple = graph_res.eval_obj if isinstance(graph_res, ExpressionTuple) else graph_res
            if hasattr(from_etuple, 'reify'):
                return from_etuple.reify()
            else:
                return from_etuple

        res = [reify_res(r) for r in res]

    if len(res) == 1 and isinstance(res[0], tf.Tensor):
        graph_res = res[0]
        return normalize_tf_graph(graph_res, graph_inputs=graph_inputs), subs_res

with graph_mode(), tf.Graph().as_default() as demo_graph:
    X_tfp = tfp.distributions.normal.Normal(0.0, 1.0, name='X')

    x_tf = tf.compat.v1.placeholder(tf.float32, name='value_x',
                                    shape=tf.TensorShape([None]))

    tau_tf = tf.compat.v1.placeholder(tf.float32, name='tau',
                                      shape=tf.TensorShape([None]))

    Y_tfp = tfp.distributions.normal.Normal(x_tf, tau_tf, name='Y')

    y_tf = tf.compat.v1.placeholder(tf.float32, name='value_y',
                                    shape=tf.TensorShape([None]))

    y_T_reshaped = tf.transpose(tf.reshape(y_tf, []))

    hier_norm_lik = tf.math.log(y_tf) + Y_tfp.log_prob(y_T_reshaped) + X_tfp.log_prob(x_tf)
    hier_norm_lik = normalize_tf_graph(hier_norm_lik)

tf_dprint(hier_norm_lik)

with graph_mode(), demo_graph.as_default():
    test_output_res, test_remaps = shift_squared_terms(hier_norm_lik, graph_inputs=[x_tf, y_tf])

for rm in test_remaps:
    for r in rm:
      tf_dprint(r[0])
      print("->")
      tf_dprint(r[1])
      print("------")

tf_dprint(test_output_res)

def compute_point_diff():
    with graph_mode(), demo_graph.as_default():

        Y_trans_tfp = tfp.distributions.normal.Normal(0.0, 1.0, name='Y_trans')

        y_shifted_tf = x_tf + tau_tf * y_tf

        hier_norm_trans_lik = tf.math.log(y_shifted_tf) + Y_trans_tfp.log_prob(y_T_reshaped) + X_tfp.log_prob(x_tf)
        hier_norm_trans_lik = normalize_tf_graph(hier_norm_trans_lik)


    test_point = {x_tf.name: np.r_[1.0],
                  tau_tf.name: np.r_[1e-20],
                  y_tf.name: np.r_[1000.1]}

    with tf.compat.v1.Session(graph=test_output_res.graph).as_default():
        val = test_output_res.eval(test_point)

    with tf.compat.v1.Session(graph=hier_norm_trans_lik.graph).as_default():
        val_2 = hier_norm_trans_lik.eval(test_point)

    return val, val_2

_ = np.subtract(*compute_point_diff())

def recenter_sqrdiffo(in_g, out_g):
    """Create a goal that essentially reduces `(a / d - (a + d * c) / d)**2` to `d**2`"""
    a_sqd_lv, b_sqd_lv, d_sqd_lv = var(), var(), var()

    # Pattern: (a / d - b / d)**2
    target_sqrdiff_lv = mt.SquaredDifference(
        mt.realdiv(a_sqd_lv, d_sqd_lv, name=var()),
        mt.realdiv(b_sqd_lv, d_sqd_lv, name=var()),
        name=var()
    )

    # Pattern: d * c + a
    c_sqd_lv = var()
    b_part_lv = mt.addv2(mt.mul(d_sqd_lv, c_sqd_lv, name=var()), a_sqd_lv, name=var())

    # Replacement: c**2
    simplified_sqrdiff_lv = mt.SquaredDifference(
        c_sqd_lv,
        0.0
    )

    reshape_lv = var()
    simplified_sqrdiff_reshaped_lv = mt.SquaredDifference(
        mt.reshape(c_sqd_lv, reshape_lv),
        0.0
    )

    res = lall(
        # input == (a / d - b / d)**2 must be "true"
        eq(in_g, target_sqrdiff_lv),
        # "and"
        conde([
            # "if" b == d * c + a is "true"
            eq(b_sqd_lv, b_part_lv),
            # "then" output ==  (c - 0)**2 is also "true"
            eq(out_g, simplified_sqrdiff_lv)

            # "or"
        ], [
            # We have to use this to cover some variation also not
            # sufficiently/consistently "normalized" by `grappler`.

            # "if" b == reshape(d * c + a, ?) is "true"
            eq(b_sqd_lv, mt.reshape(b_part_lv, reshape_lv, name=var())),
            # "then" output == (reshape(c, ?) - 0)**2 is also "true"
            eq(out_g, simplified_sqrdiff_reshaped_lv)
        ]))
    return res

with graph_mode(), test_output_res.graph.as_default():

    res = run(1, var('q'),
              reduceo(lambda x, y: tf_graph_applyo(recenter_sqrdiffo, x, y),
                      test_output_res, var('q')))

    test_output_res = normalize_tf_graph(res[0].eval_obj.reify())

tf_dprint(test_output_res.graph.get_tensor_by_name('SquaredDifference:0'))

_ = np.subtract(*compute_point_diff())