# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
r"""Arbitrary operations from prensors to prensors in an expression.

This is useful if a single op generates an entire structure. In general, it is
better to use the existing expressions framework or design a custom expression
than use this op. So long as any of the output is required, all of the input
is required.

For example, suppose you have an op my_op, that takes a prensor of the form:

```
  event
   / \
 foo   bar
```

and produces a prensor of the form my_result_schema:

```
   event
    / \
 foo2 bar2
```

```
my_result_schema = create_schema(
    is_repeated=True,
    children={"foo2":{is_repeated:True, dtype:tf.int64},
              "bar2":{is_repeated:False, dtype:tf.int64}})
```

If you give it an expression original with the schema:

```
 session
    |
  event
  /  \
foo   bar

result = map_prensor_to_prensor(
  original,
  path.Path(["session","event"]),
  my_op,
  my_result_schema)
```

Result will have the schema:

```
 session
    |
  event--------
  /  \    \    \
foo   bar foo2 bar2
```

"""

from typing import Any, Callable, Dict, FrozenSet, Optional, Sequence, Union

from struct2tensor import calculate_options
from struct2tensor import expression
from struct2tensor import expression_add
from struct2tensor import path
from struct2tensor import prensor
import tensorflow as tf

from tensorflow_metadata.proto.v0 import schema_pb2


class Schema(object):
  """A finite schema for a prensor.

  Effectively, this stores everything for the prensor but the tensors
  themselves.

  Notice that this is slightly different than schema_pb2.Schema, although
  similar in nature. At present, there is no clear way to extract is_repeated
  and dtype from schema_pb2.Schema.

  See create_schema below for constructing a schema.

  Note that for LeafNodeTensor, dtype is not None.
  Also, for ChildNodeTensor and RootNodeTensor, dtype is None. However,
  a ChildNodeTensor or RootNodeTensor could be childless.
  """

  def __init__(self,
               is_repeated: bool = True,
               dtype: Optional[tf.DType] = None,
               schema_feature: Optional[schema_pb2.Feature] = None,
               children: Optional[Dict[path.Step, "Schema"]] = None):
    """Create a new Schema object.

    Args:
      is_repeated: is the root repeated?
      dtype: tf.dtype of the root if the root is a leaf, otherwise None.
      schema_feature: schema_pb2.Feature of the root (no struct_domain
        necessary)
      children: child schemas.
    """
    self._is_repeated = is_repeated
    self._type = dtype
    self._schema_feature = schema_feature
    self._children = children if children is not None else {}
    # Cannot have a type and children.
    assert (self._type is None or not self._children)

  @property
  def is_repeated(self) -> bool:
    return self._is_repeated

  @property
  def type(self) -> Optional[tf.DType]:
    return self._type

  @property
  def schema_feature(self) -> Optional[schema_pb2.Feature]:
    return self._schema_feature

  def get_child(self, key: path.Step):
    return self._children[key]

  def known_field_names(self) -> FrozenSet[path.Step]:
    return frozenset(self._children.keys())

  def __str__(self) -> str:
    return ("Schema(is_repeated:{is_repeated} type:{type}"
            " "
            "schema_feature:({schema_feature})"
            " "
            "children:{children})").format(
                is_repeated=self._is_repeated,
                type=self._type,
                schema_feature=self._schema_feature,
                children=self._children)


def create_schema(is_repeated: bool = True,
                  dtype: Optional[tf.DType] = None,
                  schema_feature: Optional[schema_pb2.Feature] = None,
                  children: Optional[Dict[path.Step, Any]] = None) -> Schema:
  """Create a schema recursively.

  Example:
  my_result_schema = create_schema(
    is_repeated=True,
    children={"foo2":{is_repeated=True, dtype=tf.int64},
              "bar2":{is_repeated=False, dtype=tf.int64}})

  Args:
    is_repeated: whether the root is repeated.
    dtype: the dtype of a leaf (None for non-leaves).
    schema_feature: the schema_pb2.Feature describing this expression. name and
      struct_domain need not be specified.
    children: the child schemas. Note that the value type of children is either
      a Schema or a dictionary of arguments to create_schema.

  Returns:
    a new Schema represented by the inputs.
  """
  children_dict = children or {}
  child_schemas = {
      k: _create_schema_helper(v) for k, v in children_dict.items()
  }
  return Schema(
      is_repeated=is_repeated,
      dtype=dtype,
      schema_feature=schema_feature,
      children=child_schemas)


def _create_schema_helper(my_input) -> Schema:
  """Helper for create_schema.

  If my_input is a Schema, it is just returned.

  Otherwise, my_input should be a dictionary of arguments to create_schema.

  Args:
    my_input: either a Schema or a dictionary which optionally has the keys
      is_repeated, dtype, and children.

  Returns:
    a Schema.
  """
  if isinstance(my_input, Schema):
    return my_input
  return create_schema(**my_input)


def map_prensor_to_prensor(
    root_expr: expression.Expression, source: path.Path,
    paths_needed: Sequence[path.Path],
    prensor_op: Callable[[prensor.Prensor], prensor.Prensor],
    output_schema: Schema) -> expression.Expression:
  r"""Maps an expression to a prensor, and merges that prensor.

  For example, suppose you have an op my_op, that takes a prensor of the form:

    event
     / \