# -*- coding: utf-8
"""
Functions for dealing with a logical language.
"""

import copy
import functools
import inspect
import itertools
import logging
import re

from nltk.sem import logic as l

L = logging.getLogger(__name__)


class TypeSystem(object):

  ANY_TYPE = l.ANY_TYPE
  EVENT_TYPE = l.EVENT_TYPE

  def __init__(self, primitive_types):
    assert "?" not in primitive_types, "Cannot override ANY_TYPE name"
    assert "v" not in primitive_types, "Cannot override EVENT_TYPE name"
    self._types = {primitive_type_name: l.BasicType(name=primitive_type_name,
                                                    parent=l.ENTITY_TYPE)
                   for primitive_type_name in primitive_types}
    self._types["?"] = self.ANY_TYPE
    self._types["v"] = self.EVENT_TYPE

  def __getitem__(self, type_expr):
    if isinstance(type_expr, l.Type):
      return type_expr
    if isinstance(type_expr, str):
      return self._types[type_expr]
    return self.make_function_type(type_expr)

  def __iter__(self):
    return iter(self._types.values())

  def make_function_type(self, type_expr):
    ret = self[type_expr[-1]]
    for type_expr_i in type_expr[:-1][::-1]:
      ret = l.ComplexType(self[type_expr_i], ret)
    return ret

  def new_function(self, name, type, defn, **kwargs):
    type = self[type]
    return Function(name, type, defn, **kwargs)

  def new_constant(self, name, type, **kwargs):
    type = self[type]
    return l.Variable(name, type)


# Wrapper for a typed function.
class Function(object):
  """
  Wrapper for a typed function.
  """

  def __init__(self, name, type, defn, weight=0.0):
    self.name = name
    self.type = type
    self.defn = defn
    self.weight = weight

  @property
  def arity(self):
    return len(self.type.flat) - 1

  @property
  def arg_types(self):
    return self.type.flat[:-1]

  @property
  def return_type(self):
    return self.type.flat[-1]

  def __hash__(self):
    return hash((self.name, self.type, self.defn))

  def __eq__(self, other):
    return hash(self) == hash(other)

  def __str__(self):
    return "function %s : %s" % (self.name, self.type)

  __repr__ = __str__


def make_application(pred, args):
  pred = l.ConstantExpression(l.Variable(pred)) if isinstance(pred, str) else pred
  expr = l.ApplicationExpression(pred, args[0])
  return functools.reduce(lambda x, y: l.ApplicationExpression(x, y),
                          args[1:], expr)


def next_bound_var(bound_vars, type):
  """
  Helper function: generate the next bound variable in a context where there
  are currently the bound variables `bound_vars` (assumed to be generated by
  this function).
  """
  name_length = 1 + len(bound_vars) // 26
  name_id = len(bound_vars) % 26
  name = chr(97 + name_id)
  return l.Variable(name * name_length, type)


def listify(fn=None, wrapper=list):
  """
  A decorator which wraps a function's return value in ``list(...)``.

  Useful when an algorithm can be expressed more cleanly as a generator but
  the function should return an list.
  """
  def listify_return(fn):
    @functools.wraps(fn)
    def listify_helper(*args, **kw):
      return wrapper(fn(*args, **kw))
    return listify_helper
  if fn is None:
    return listify_return
  return listify_return(fn)


def extract_lambda(expr):
  """
  Extract `LambdaExpression` arguments to the top of a semantic form.
  This makes them compatible with the CCG parsing setup, which needs top-level
  lambdas in order to perform function application during parsing.
  """
  variables = []

  def process_lambda(lambda_expr):
    # Create a new unique variable and substitute.
    unique = l.unique_variable()
    unique.type = lambda_expr.variable.type
    new_expr = lambda_expr.term.replace(lambda_expr.variable, l.IndividualVariableExpression(unique))
    return unique, new_expr

  # Traverse the LF and replace lambda expressions wherever necessary.
  def inner(node):
    if isinstance(node, l.ApplicationExpression):
      new_args = []

      for arg in node.args:
        if isinstance(arg, l.LambdaExpression):
          new_var, new_arg = process_lambda(arg)

          variables.append(new_var)
          new_args.append(new_arg)
        else:
          new_args.append(inner(arg))

      return make_application(node.pred.variable.name, new_args)
    else:
      return node

  expr = inner(expr)
  wrappings = []

  for variable_ordering in itertools.permutations(variables):
    wrapping = expr
    for variable in variable_ordering:
      wrapping = l.LambdaExpression(variable, wrapping)

    wrappings.append(wrapping.normalize())

  return wrappings


def get_arity(expr):
  """
  Get the arity of a lambda-extracted expression.
  """
  if isinstance(expr, l.LambdaExpression):
    return 1 + get_arity(expr.term)
  else:
    return 0


def read_ec_sexpr(sexpr):
  """
  Parse an EC-style S-expression into an untyped NLTK representation.
  """
  tokens = re.split(r"([()\s])", sexpr)

  bound_vars = set()
  bound_var_stack = []

  is_call = False
  stack = [(None, None, [])]
  for token in tokens:
    token = token.strip()
    if not token:
      continue

    if token == "(":
      if is_call:
        # Second consecutive left-paren -- this means we have a complex
        # function expression.
        stack.append((l.ApplicationExpression, None, []))
      is_call = True
    elif token == "lambda":
      is_call = False
      variable = next_bound_var(bound_vars, l.ANY_TYPE)
      bound_vars.add(variable)
      bound_var_stack.append(variable)

      stack.append((l.LambdaExpression, None, []))
    elif is_call:
      head = token
      if head.startswith("$"):
        bruijn_index = int(head[1:])
        # Bound variable is the head of an application expression.
        # First replace with a function-looking variable, then update parser
        # state.
        var_idx = -1 - bruijn_index
        var = bound_var_stack[var_idx]
        bound_vars.remove(var)

        new_var = l.Variable(var.name.upper())
        bound_var_stack[var_idx] = new_var
        bound_vars.add(new_var)

        head = l.FunctionVariableExpression(new_var)

      stack.append((l.ApplicationExpression, head, []))
      is_call = False
    elif token == ")":
      stack_top = stack.pop()
      if stack_top[0] == l.ApplicationExpression:
        _, pred, args = stack_top
        result = make_application(pred, args)
      elif stack_top[0] == l.LambdaExpression:
        _, _, term = stack_top
        variable = bound_var_stack.pop()
        result = l.LambdaExpression(variable, term[0])
      else:
        raise RuntimeError("unknown element on stack", stack_top)

      stack_parent = stack[-1]
      if stack_parent[0] == l.ApplicationExpression and stack_parent[1] is None:
        # We have just finished reading the head of an application expression.
        expr, _, args = stack_parent
        stack[-1] = (expr, result, args)
      else:
        # Add to children of parent node.
        stack_parent[2].append(result)
    elif token.startswith("$"):
      bruijn_index = int(token[1:])
      stack[-1][2].append(l.IndividualVariableExpression(bound_var_stack[-1 - bruijn_index]))
    else:
      stack[-1][2].append(l.ConstantExpression(l.Variable(token)))

  assert len(stack) == 1
  assert len(stack[0][2]) == 1
  return stack[0][2][0], bound_vars


class Ontology(object):
  """
  Defines an ontology for expressing and evaluating logical forms.
  """

  def __init__(self, types, functions, constants, variable_weight=0.1):
    """
    Arguments:
      types: TypeSystem
      functions: List of `k` `Function` instances
      constants: List of constants as (optionally typed) `Variable` instances
      variable_weight: log-probability of observing any variable
    """
    self.types = types

    self.functions = []
    self.functions_dict = {}
    self.constants, self.constants_dict = [], {}
    self.variable_weight = variable_weight

    self.add_functions(functions)
    self.add_constants(constants)

    self._prepare()

  EXPR_TYPES = [l.ApplicationExpression, l.ConstantExpression,
                l.IndividualVariableExpression, l.LambdaExpression,
                l.FunctionVariableExpression]

  def add_functions(self, functions):
    # Ignore functions which already exist.
    new_functions = []
    for function in functions:
      if function.name in self.functions_dict:
        existing_function = self.functions_dict[function.name]
        assert existing_function == function, \
            "Function name clash: existing %r, inserting %r" % (existing_function, function)
      else:
        new_functions.append(function)

    self.functions.extend(functions)
    self.functions_dict.update({fn.name: fn for fn in functions})

    for function in functions:
      # We can't statically verify the type of the definition, but we can at
      # least verify the arity.
      L.debug("verifying arity: %s stated %i actual %i (%s)",
              function.name, function.arity, self.get_expr_arity(function.defn),
              function.defn)
      assert function.arity == self.get_expr_arity(function.defn), function.name

  def add_constants(self, constants):
    self.constants = constants
    self.constants_dict = {c.name: c for c in constants}

  def _prepare(self):
    self._nltk_type_signature = self._make_nltk_type_signature()

  def iter_expressions(self, max_depth=3, **kwargs):
    ret = self._iter_expressions_inner(max_depth, bound_vars=(), **kwargs)

    # Extract lambda arguments to the top level.
    # NB, this breaks the type record. Should be fine.
    #
    # TODO: Currently the program enumeration is generating *both* programs of
    # the structure `\x y z.(..x..y..z..)` and `..(\x.(..),\y.(..),\z.(..))` --
    # that is, with broad and narrow scope. The `extract_lambda` operation
    # reduces a lot of these expressions to be overlapping. It'd be more
    # efficient to pick a single enumeration strategy and stick with it.
    ret = itertools.chain.from_iterable(
        extract_lambda(expr) for expr in ret)

    return ret

  @listify
  def _iter_expressions_inner(self, max_depth, bound_vars,
                              type_request=None, function_weights=None):
    """
    Enumerate all legal expressions.

    Arguments:
      max_depth: Maximum tree depth to traverse.
      bound_vars: Bound variables (and their types) in the parent context. The
        returned expressions may reference these variables. List of `(name,
        type)` tuples.
      type_request: Optional requested type of the expression. This helps
        greatly restrict the space of enumerations when the type system is
        strong.
      function_weights: Override for function weights to determine the order in
        which we consider proposing function application expressions.
    """
    if max_depth == 0:
      return
    elif max_depth == 1 and not bound_vars:
      # require some bound variables to generate a valid lexical entry
      # semantics
      return

    for expr_type in self.EXPR_TYPES:
      if expr_type == l.ApplicationExpression:
        # Loop over functions according to their weights.
        # from pprint import pprint
        # pprint(sorted([(fn.weight, fn.name) for fn in self.functions], key=lambda x: x[0]))
        fn_weight_key = (lambda fn: function_weights[fn.name]) if function_weights is not None \
                        else (lambda fn: fn.weight)
        fns_sorted = sorted(self.functions_dict.values(), key=fn_weight_key,
                            reverse=True)

        if max_depth > 1:
          for fn in fns_sorted:
            # If there is a present type request, only consider functions with
            # the correct return type.
            # print("\t" * (6 - max_depth), fn.name, fn.return_type, " // request: ", type_request, bound_vars)
            if type_request is not None and fn.return_type != type_request:
              continue

            # Special case: yield fast event queries without recursion.
            if fn.arity == 1 and fn.arg_types[0] == self.types.EVENT_TYPE:
              yield make_application(fn.name, (l.ConstantExpression(l.Variable("e")),))
            elif fn.arity == 0:
              # 0-arity functions are represented in the logic as
              # `ConstantExpression`s.
              # print("\t" * (6 - max_depth + 1), "yielding const ", fn.name)
              yield l.ConstantExpression(l.Variable(fn.name))
            else:
              # print("\t" * (6 - max_depth), fn, fn.arg_types)
              sub_args = []
              for i, arg_type_request in enumerate(fn.arg_types):
                # print("\t" * (6 - max_depth + 1), "ARGUMENT %i %s (max_depth %i)" % (i, arg_type_request, max_depth - 1))
                sub_args.append(
                    self._iter_expressions_inner(max_depth=max_depth - 1,
                                                bound_vars=bound_vars,
                                                type_request=arg_type_request,
                                                function_weights=function_weights))

              for arg_combs in itertools.product(*sub_args):
                candidate = make_application(fn.name, arg_combs)
                valid = self._valid_application_expr(candidate)
                # print("\t" * (6 - max_depth + 1), "valid %s? %s" % (candidate, valid))
                if valid:
                  yield candidate
      elif expr_type == l.LambdaExpression and max_depth > 1:
        for bound_var_type in self.observed_argument_types:
          bound_var = next_bound_var(bound_vars, bound_var_type)
          subexpr_bound_vars = bound_vars + (bound_var,)

          subexpr_type_requests = []
          if type_request is None:
            subexpr_type_requests = [None]
          else:
            # Build new type requests using the flat structure.
            type_request_flat = type_request.flat

            subexpr_type_requests.append(type_request_flat)

            # Basic case: the variable is used as one of the existing
            # arguments of a target subexpression.
            subexpr_type_requests.extend([type_request_flat[:i] + type_request_flat[i + 1:]
                                          for i, type_i in enumerate(type_request_flat)
                                          if type_i == bound_var_type])

            # # The subexpression might take this variable as an additional
            # # argument in any position. We have to enumerate all possibilities
            # # -- yikes!
            # for insertion_point in range(len(type_request)):
            #   subexpr_type_requests.append(type_request[:insertion_point] + (bound_var_type,)
            #       + type_request[insertion_point:])
          # print("\t" * (6-max_depth), "λ %s :: %s" % (bound_var, bound_var_type), type_request, subexpr_type_requests)
          # print("\t" * (6-max_depth), "Now recursing with max_depth=%i" % (max_depth - 1))

          for subexpr_type_request in subexpr_type_requests:
            if isinstance(subexpr_type_request, tuple):
              if not subexpr_type_request:
                continue
              else:
                subexpr_type_request = self.types.make_function_type(subexpr_type_request)

            results = self._iter_expressions_inner(max_depth=max_depth - 1,
                                                    bound_vars=subexpr_bound_vars,
                                                    type_request=subexpr_type_request,
                                                    function_weights=function_weights)
            for expr in results:
              candidate = l.LambdaExpression(bound_var, expr)
              valid = self._valid_lambda_expr(candidate, bound_vars)
              # print("\t" * (6 - max_depth), "valid lambda %s? %s" % (candidate, valid))
              if self._valid_lambda_expr(candidate, bound_vars):
                # Assign variable types before returning.
                extra_types = {bound_var.name: bound_var.type
                               for bound_var in subexpr_bound_vars}

                try:
                  # TODO make sure variable names are unique before this happens
                  self.typecheck(candidate, extra_types)
                except l.InconsistentTypeHierarchyException:
                  pass
                else:
                  yield candidate
      elif expr_type == l.IndividualVariableExpression:
        for bound_var in bound_vars:
          if type_request and not bound_var.type.matches(type_request):
            continue

          # print("\t" * (6-max_depth), "var %s" % bound_var)

          yield l.IndividualVariableExpression(bound_var)
      elif expr_type == l.ConstantExpression:
        for constant in self.constants:
          if type_request is not None and not constant.type.matches(type_request):
            continue

          yield l.ConstantExpression(constant)
      elif expr_type == l.FunctionVariableExpression:
        # NB we don't support enumerating bound variables with function types
        # right now -- the following only considers yielding fixed functions
        # from the ontology.
        for function in self.functions:
          # Be a little strict here to avoid excessive enumeration -- only
          # consider emitting functions when the type request specifically
          # demands a function, not e.g. AnyType
          if type_request is None or type_request == self.types.ANY_TYPE \
              or not function.type.matches(type_request):
            continue

          yield l.FunctionVariableExpression(l.Variable(function.name, function.type))

  def typecheck(self, expr, extra_type_signature=None):
    type_signature = self._nltk_type_signature
    if extra_type_signature is not None:
      type_signature = copy.copy(type_signature)
      type_signature.update(extra_type_signature)

    expr.typecheck(signature=type_signature)

  def infer_type(self, expr, variable_name, extra_types=None):
    """
    Infer the type of a bound variable with name `variable_name` used in `expr`.

    Args:
      expr:
      variable_name:
      extra_types: Optional dictionary of provisional function types, mapping
        from function name to a type expression. Useful for doing type
        inference on elements of new functions before adding them to an
        ontology instance.
    """
    apparent_types = set()
    extra_types = extra_types or {}

    def visitor(node):
      if isinstance(node, l.ApplicationExpression):
        fn_name = node.pred.variable.name

        try:
          function_type = self.functions_dict[fn_name].type
        except KeyError:
          try:
            function_type = extra_types[fn_name]
          except KeyError:
            # No function information available. Ditch.
            return self.types.ANY_TYPE

        for i, arg in enumerate(node.args):
          visitor(arg)

          if isinstance(arg, l.IndividualVariableExpression) and arg.variable.name == variable_name:
            # We've found a use of the value as a function argument -- extract
            # the apparent type.
            apparent_types.add(function_type.flat[i])
          elif isinstance(arg, l.ApplicationExpression) and isinstance(arg.pred, l.FunctionVariableExpression) \
              and arg.pred.variable.name == variable_name:
            arg_types = ((self.types.ANY_TYPE,) * len(arg.args))
            apparent_types.add(arg_types + (function_type.flat[i],))
      elif isinstance(node, l.LambdaExpression):
        visitor(node.term)

    visitor(expr)
    if len(apparent_types) > 1:
      if len(apparent_types) == 2 and self.types.ANY_TYPE in apparent_types:
        # Good, just remove the AnyType.
        apparent_types.remove(self.types.ANY_TYPE)
      else:
        # TODO check type compatibility
        raise NotImplementedError("Multiple apparent types: %s" % apparent_types)

    type_ret = next(iter(apparent_types))
    return self.types[type_ret]

  @property
  def observed_argument_types(self):
    """
    Collection of types (primitive or functional) which have been observed as
    arguments in this ontology's function.
    """
    obs_types = set(itertools.chain.from_iterable(
      fn.arg_types for fn in self.functions))
    return obs_types - set([self.types.ANY_TYPE])

  def get_expr_arity(self, expr):
    """
    Get the arity (number of bound variables) of a function definition.
    """
    if isinstance(expr, l.LambdaExpression):
      return 1 + self.get_expr_arity(expr.term)
    elif isinstance(expr, l.ApplicationExpression):
      function = self.functions_dict[expr.pred.variable.name]
      return function.arity - len(expr.args)
    elif isinstance(expr, (l.FunctionVariableExpression, l.ConstantExpression)) \
        and expr.variable.name in self.functions_dict:
      return self.functions_dict[expr.variable.name].arity
    elif isinstance(expr, l.IndividualVariableExpression):
      return 0
    elif callable(expr):
      return len(inspect.signature(expr).parameters)
    else:
      raise ValueError("non-callable object: %r" % expr)

  def _valid_application_expr(self, application_expr):
    """
    Check whether this `ApplicationExpression` should be considered when
    enumerating programs.
    """
    # TODO check type consistency
    return True

  def _valid_lambda_expr(self, lambda_expr, ctx_bound_vars):
    """
    Check whether this `LambdaExpression` should be considered when enumerating
    programs.

    Arguments:
      lambda_expr: `LambdaExpression`
      ctx_bound_vars: Bound variables from the containing context
    """

    # Collect bound arguments and the body expression.
    bound_args = []
    expr = lambda_expr
    while isinstance(expr, l.LambdaExpression):
      bound_args.append(expr.variable)
      expr = expr.term
    body = expr

    # Exclude exprs which do not use all of their bound arguments.
    available_vars = set(bound_args) | set(ctx_bound_vars)
    if available_vars != set(body.variables()):
      return False

    # # Exclude exprs with simplistic bodies.
    # if isinstance(body, l.IndividualVariableExpression):
    #   return False

    return True

  def _make_nltk_type_expr(self, type_expr):
    if isinstance(type_expr, tuple) and len(type_expr) == 1:
      type_expr = type_expr[0]

    if type_expr in self.nltk_types:
      return self.nltk_types[type_expr]
    elif len(type_expr) > 1:
      return l.ComplexType(self._make_nltk_type_expr(type_expr[0]),
                           self._make_nltk_type_expr(type_expr[1:]))
    else:
      raise RuntimeError("unknown basic type %s" % (type_expr,))

  def _make_nltk_type_signature(self):
    return {fn.name: fn.type for fn in self.functions}

  def as_ec_sexpr(self, expr):
    """
    Convert an `nltk.sem.logic` `Expression` to an S-expr string.
    """

    # Expressions which might contain a function reference
    func_exprs = (l.ConstantExpression, l.EventVariableExpression,
                  l.FunctionVariableExpression)

    def inner(expr, var_stack):
      if isinstance(expr, l.LambdaExpression):
        # Add lambda variable to var map.
        return "(lambda %s)" % inner(expr.term, var_stack + [expr.variable.name])
      elif isinstance(expr, l.ApplicationExpression):
        args = [inner(arg, var_stack) for arg in expr.args]
        return "(%s %s)" % (expr.pred.variable.name, " ".join(args))
      # elif isinstance(expr, l.AndExpression):
      #   return "(and %s %s)" % (inner(expr.first), inner(expr.second))
      elif isinstance(expr, func_exprs) and expr.variable.name in self.functions_dict:
        # EC requires S-expressions in normal form -- i.e. functions need to
        # appear in their applied form. We'll need a valid function type here
        # to get anything done.
        arity = self.functions_dict[expr.variable.name].arity

        if arity > 0:
          return ("(lambda " * arity) + \
              ("(%s %s)" % (expr.variable.name,
                            " ".join("$%i" % (idx - 1) for idx in range(arity, 0, -1)))) + \
              (")" * arity)
        else:
          return expr.variable.name
      elif isinstance(expr, l.IndividualVariableExpression):
        bruijn_index = len(var_stack) - var_stack.index(expr.variable.name) - 1
        return "$%i" % bruijn_index
      elif isinstance(expr, l.FunctionVariableExpression):
        raise ValueError("unknown function %s" % expr)
      elif isinstance(expr, l.ConstantExpression):
        return expr.variable.name
      else:
        raise ValueError("un-handled expression component %r" % expr)

    return inner(expr, [])


def is_negation(expr):
  return isinstance(expr, l.ApplicationExpression) and isinstance(expr.pred, l.NegatedExpression)
