'''
Created on Oct 11, 2016

@author: yunfanrao
'''
#from xml.dom.minidom import parse
import xml.dom.minidom
import re
from datetime import datetime

class SightingsList(object):

    def __init__(self, sightingFile = None):
        if sightingFile == None:
            raise ValueError("SightingsLish.__inint__:  must f.xml as a paramount")
        self.DOMTree = xml.dom.minidom.parse(sightingFile)
        self.Data = self.DOMTree.documentElement
        self.fix = None
        self.sightings = None
        self.body = []
        self.date = []
        self.time = []
        self.observation = []
        self.height = []
        self.temperature = []
        self.pressure = []
        self.horizon = []
        self.count = 0

    def get_fix(self):
        self.fix = self.Data.getElementsByTagName("fix")

    def get_sightings(self):
        self.sightings = self.Data.getElementsByTagName("sighting")
        for sighting in self.sightings:
            self.count += 1

    def get_body(self):
        for sighting in self.sightings:
            try:
                self.body.append(str(sighting.getElementsByTagName("body")[0].childNodes[0].data))
            except:
                self.body.append("")
        return self.body

    def get_date(self):
        for sighting in self.sightings:
            try:
                self.date.append(str(sighting.getElementsByTagName("date")[0].childNodes[0].data))
            except:
                self.date.append("")
        return self.date

    def get_time(self):
        for sighting in self.sightings:
            try:
                self.time.append(str(sighting.getElementsByTagName("time")[0].childNodes[0].data))
            except:
                self.time.append("")
        return self.time

    def get_observation(self):
        for sighting in self.sightings:
            try:
                self.observation.append(str(sighting.getElementsByTagName("observation")[0].childNodes[0].data))
            except:
                self.observation.append("")
        return self.observation

    def get_height(self):
        for sighting in self.sightings:
            try:
                self.height.append(str(sighting.getElementsByTagName("height")[0].childNodes[0].data))
            except:
                self.height.append("0")
        return self.height

    def get_temperature(self):
        for sighting in self.sightings:
            try:
                self.temperature.append(str(sighting.getElementsByTagName("temperature")[0].childNodes[0].data))
            except:
                self.temperature.append("72")
        return self.temperature

    def get_pressure(self):
        for sighting in self.sightings:
            try:
                self.pressure.append(str(sighting.getElementsByTagName("pressure")[0].childNodes[0].data))
            except:
                self.pressure.append("1010")
        return self.pressure

    def get_horizon(self):
        for sighting in self.sightings:
            try:
                self.horizon.append(str(sighting.getElementsByTagName("horizon")[0].childNodes[0].data))
            except:
                self.horizon.append("natural")
        return self.horizon

    def get_count(self):
        return self.count

    def checkError(self):
        if self.fix == None or self.sightings == None:
            return True
        errors = [0 for i in range(self.count)]
        for i in range(self.count):
            try:
                for c in self.body[i]:
                    if not (c.isalpha() or c.isspace() or c == "."):
                        errors[i] +=1
                datetime.strptime(self.date[i],"%Y-%m-%d")
                datetime.strptime(self.time[i],"%H:%M:%S")
            except:
                errors[i] +=1
            isObservationWrong = self.isObservationError(self.observation[i])
            if isObservationWrong:
                errors[i] +=1
            try:
                float(self.height[i])
                int(self.temperature[i])
                int(self.pressure[i])
                if float(self.height[i]) < 0:
                    errors[i] +=1
                if (int(self.temperature[i]) < -2) or (int(self.temperature[i]) > 120):
                    errors[i] +=1
                if (int(self.pressure[i]) < 100) or (int(self.pressure[i]) > 1100):
                    errors[i] +=1
            except:
                errors[i] +=1

            natural = "natural"
            artificial = "artificial"
            if not ((self.horizon[i].lower() == natural)
                                   or (self.horizon[i].lower() == artificial)):
                errors[i] +=1

        return errors

    def isObservationError(self,observation):
        matchAngle = re.match( r'^(\-?\d+) d ((\d+)$ | (\d+\.{1}\d{1})$)',observation,re.X)
        if matchAngle:
            degree = int(matchAngle.group(1))
            if '.' in matchAngle.group(2):
                minute = float(matchAngle.group(2))
                minute = round(minute,1)
            else:
                minute = int(matchAngle.group(2))
            if (degree < 0) or (degree >= 90):
                return True
            if (minute < 0 ) or (minute >= 60):
                return True
            return False


import tkinter
import mover
import config

class BigPlane(mover.BaseMover):
    '''
        移动的敌机 - 大飞机
    '''
    def __init__(self, root, canvas, position, x, y, tags):
        super().__init__(root, canvas, position, x, y, tags,
                         config.image_bigplane_width, config.image_bigplane_height, True)

        # 移动者的移动步长
        self.steps = [config.step_length_bigplane_x, config.step_length_bigplane_y]
        # 移动方向 - 向下
        self.move_direction = [0, 1]
        # 移动者加载背景图像
        self.bg_image_fullname = config.images_path + config.filename_bigplane + config.filename_suffix
        self.bg_image = tkinter.PhotoImage(file=self.bg_image_fullname)
        self.bg_image_tags = tags

    def exec_move(self):
        if self.nw[1] < config.window_boundary_row:
            # Y轴边界之内正常移动
            x = self.steps[0] * self.move_direction[0]
            y = self.steps[1] * self.move_direction[1]
            self.base_move(self.bg_image_tags, x, y)
        else:
            # Y轴边界之外错误处理
            self.base_move(self.bg_image_tags, 0, -config.window_boundary_row)

    # 获取死亡图片
    def get_dead_images(self):
        img = []
        if self.do_dead_play:
            for i in self.dead_image_index:
                image_fullname = config.images_path + config.filename_bigplane + str(i) + config.filename_suffix
                image = tkinter.PhotoImage(file=self.bg_image_fullname)
                img.append(image)
            return img


from cryptography.hazmat.primitives.ciphers import Cipher
from cryptography.hazmat.primitives.ciphers.modes import CTR
from cryptography.hazmat.primitives.ciphers.algorithms import AES
from cryptography.hazmat.primitives.hmac import HMAC
from cryptography.hazmat.primitives.hashes import SHA256, SHA1, Hash
from cryptography.hazmat.primitives.kdf.hkdf import HKDFExpand
from cryptography.hazmat.backends import default_backend
import curve25519

import struct
import base64

bend = default_backend()

def b64decode(data):
    return base64.b64decode(data + '=' * (len(data) % 4))

def sha1(msg):
    sha = Hash(SHA1(), backend=bend)
    sha.update(msg)
    return sha.finalize()

def hmac(key, msg):
    hmac = HMAC(key, algorithm=SHA256(), backend=bend)
    hmac.update(msg)
    return hmac.finalize()

def hkdf(key, length=16, info=''):
    hkdf = HKDFExpand(algorithm=SHA256(), length=length, info=info, backend=bend)
    return hkdf.derive(key)

def hash_func(shared):
    return shared

class NtorError(Exception):
    pass

class ntor(object):
    def __init__(self, node):
        # 5.1.4. The "ntor" handshake

        # This handshake uses a set of DH handshakes to compute a set of
        # shared keys which the client knows are shared only with a particular
        # server, and the server knows are shared with whomever sent the
        # original handshake (or with nobody at all).  Here we use the
        # "curve25519" group and representation as specified in "Curve25519:
        # new Diffie-Hellman speed records" by D. J. Bernstein.

        # [The ntor handshake was added in Tor 0.2.4.8-alpha.]

        self.node = node

        # In this section, define:
        #   H(x,t) as HMAC_SHA256 with message x and key t.
        #   H_LENGTH  = 32.
        #   ID_LENGTH = 20.
        #   G_LENGTH  = 32
        #   PROTOID   = "ntor-curve25519-sha256-1"
        #   t_mac     = PROTOID | ":mac"
        #   t_key     = PROTOID | ":key_extract"
        #   t_verify  = PROTOID | ":verify"
        #   MULT(a,b) = the multiplication of the curve25519 point 'a' by the
        #               scalar 'b'.
        #   G         = The preferred base point for curve25519 ([9])
        #   KEYGEN()  = The curve25519 key generation algorithm, returning
        #               a private/public keypair.
        #   m_expand  = PROTOID | ":key_expand"

        # H is defined as hmac()
        # MULT is included in the curve25519 library as get_shared_key()
        # KEYGEN() is curve25519.Private()
        self.protoid = 'ntor-curve25519-sha256-1'
        self.t_mac = self.protoid + ':mac'
        self.t_key = self.protoid + ':key_extract'
        self.t_verify = self.protoid + ':verify'
        self.m_expand = self.protoid + ':key_expand'

        # To perform the handshake, the client needs to know an identity key
        # digest for the server, and an ntor onion key (a curve25519 public
        # key) for that server. Call the ntor onion key "B".  The client
        # generates a temporary keypair:
        #     x,X = KEYGEN()
        self.x = curve25519.Private()
        self.X = self.x.get_public()

        self.B = curve25519.Public(b64decode(self.node['ntor-onion-key']))

        # and generates a client-side handshake with contents:
        #   NODEID      Server identity digest  [ID_LENGTH bytes]
        #   KEYID       KEYID(B)                [H_LENGTH bytes]
        #   CLIENT_PK   X                       [G_LENGTH bytes]
        self.handshake  = b64decode(self.node['identity'])
        self.handshake += self.B.serialize()
        self.handshake += self.X.serialize()

    def complete_handshake(self, Y, auth):
        # The server's handshake reply is:
        # SERVER_PK   Y                       [G_LENGTH bytes]
        # AUTH        H(auth_input, t_mac)    [H_LENGTH bytes]

        # The client then checks Y is in G^* [see NOTE below], and computes

        # secret_input = EXP(Y,x) | EXP(B,x) | ID | B | X | Y | PROTOID
        si  = self.x.get_shared_key(curve25519.Public(Y), hash_func)
        si += self.x.get_shared_key(self.B, hash_func)
        si += b64decode(self.node['identity'])
        si += self.B.serialize()
        si += self.X.serialize()
        si += Y
        si += 'ntor-curve25519-sha256-1'

        # KEY_SEED = H(secret_input, t_key)
        # verify = H(secret_input, t_verify)
        key_seed = hmac(self.t_key, si)
        verify = hmac(self.t_verify, si)

        # auth_input = verify | ID | B | Y | X | PROTOID | "Server"
        ai = verify
        ai += b64decode(self.node['identity'])
        ai += self.B.serialize()
        ai += Y
        ai += self.X.serialize()
        ai += self.protoid
        ai += 'Server'

        # The client verifies that AUTH == H(auth_input, t_mac).
        if auth != hmac(self.t_mac, ai):
            raise NtorError('auth input does not match.')

        # Both parties check that none of the EXP() operations produced the
        # point at infinity. [NOTE: This is an adequate replacement for
        # checking Y for group membership, if the group is curve25519.]

        # Both parties now have a shared value for KEY_SEED.  They expand this
        # into the keys needed for the Tor relay protocol, using the KDF
        # described in 5.2.2 and the tag m_expand.

        # 5.2.2. KDF-RFC5869

        # For newer KDF needs, Tor uses the key derivation function HKDF from
        # RFC5869, instantiated with SHA256.  (This is due to a construction
        # from Krawczyk.)  The generated key material is:

        #     K = K_1 | K_2 | K_3 | ...

        #     Where H(x,t) is HMAC_SHA256 with value x and key t
        #       and K_1     = H(m_expand | INT8(1) , KEY_SEED )
        #       and K_(i+1) = H(K_i | m_expand | INT8(i+1) , KEY_SEED )
        #       and m_expand is an arbitrarily chosen value,
        #       and INT8(i) is a octet with the value "i".

        # In RFC5869's vocabulary, this is HKDF-SHA256 with info == m_expand,
        # salt == t_key, and IKM == secret_input.
        keys = hkdf(key_seed, length=72, info=self.m_expand)

        # When used in the ntor handshake, the first HASH_LEN bytes form the
        # forward digest Df; the next HASH_LEN form the backward digest Db; the
        # next KEY_LEN form Kf, the next KEY_LEN form Kb, and the final
        # DIGEST_LEN bytes are taken as a nonce to use in the place of KH in the
        # hidden service protocol.  Excess bytes from K are discarded.
        Df, Db, Kf, Kb = struct.unpack('>20s20s16s16s', keys)

        # we do what we can with what we've got.
        del self.X
        del self.x
        del self.B
        del key_seed
        del keys
        del verify
        del ai
        del auth
        del si
        del Y

        self.send_digest = Hash(SHA1(), backend=bend)
        self.send_digest.update(Df)
        self.recv_digest = Hash(SHA1(), backend=bend)
        self.recv_digest.update(Db)

        self.encrypt = Cipher(AES(Kf), CTR('\x00' * 16), backend=bend).encryptor()
        self.decrypt = Cipher(AES(Kb), CTR('\x00' * 16), backend=bend).decryptor()

    def get_handshake(self):
        return self.handshake


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import random
import tempfile

from absl.testing import parameterized
import numpy as np
import six
import tensorflow as tf
from tensorflow.core.example import example_pb2
from tensorflow.core.example import feature_pb2
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow_estimator.python.estimator import model_fn
from tensorflow_estimator.python.estimator.util import tf_keras
from tensorflow_estimator.python.estimator.canned import metric_keys
from tensorflow_estimator.python.estimator.canned import parsing_utils
from tensorflow_estimator.python.estimator.canned import prediction_keys
from tensorflow_estimator.python.estimator.canned import rnn
from tensorflow_estimator.python.estimator.export import export
from tensorflow_estimator.python.estimator.head import multi_class_head as multi_head_lib
from tensorflow_estimator.python.estimator.head import sequential_head as seq_head_lib
from tensorflow_estimator.python.estimator.inputs import numpy_io

# Names of variables created by BasicRNNCell model.
CELL_KERNEL_NAME = 'rnn_model/rnn/kernel'
CELL_RECURRENT_KERNEL_NAME = 'rnn_model/rnn/recurrent_kernel'
CELL_BIAS_NAME = 'rnn_model/rnn/bias'
LOGITS_WEIGHTS_NAME = 'rnn_model/logits/kernel'
LOGITS_BIAS_NAME = 'rnn_model/logits/bias'

def _assert_close(expected, actual, rtol=1e-04, name='assert_close'):
  with ops.name_scope(name, 'assert_close', (expected, actual, rtol)) as scope:
    expected = ops.convert_to_tensor(expected, name='expected')
    actual = ops.convert_to_tensor(actual, name='actual')
    rdiff = tf.math.abs(expected - actual, 'diff') / tf.math.abs(expected)
    rtol = ops.convert_to_tensor(rtol, name='rtol')
    return tf.compat.v1.debugging.assert_less(
        rdiff,
        rtol,
        data=('Condition expected =~ actual did not hold element-wise:'
              'expected = ', expected, 'actual = ', actual, 'rdiff = ', rdiff,
              'rtol = ', rtol,),
        name=scope)

def create_checkpoint(kernel, recurrent, bias, dense_kernel, dense_bias,
                      global_step, model_dir):
  """Create checkpoint file with provided model weights.

  Args:
    kernel: Iterable of values of input weights for the RNN cell.
    recurrent: Iterable of values of recurrent weights for the RNN cell.
    bias: Iterable of values of biases for the RNN cell.
    dense_kernel: Iterable of values for matrix connecting RNN output to logits.
    dense_bias: Iterable of values for logits bias term.
    global_step: Initial global step to save in checkpoint.
    model_dir: Directory into which checkpoint is saved.
  """
  model_weights = {}
  model_weights[CELL_KERNEL_NAME] = kernel
  model_weights[CELL_RECURRENT_KERNEL_NAME] = recurrent
  model_weights[CELL_BIAS_NAME] = bias
  model_weights[LOGITS_WEIGHTS_NAME] = dense_kernel
  model_weights[LOGITS_BIAS_NAME] = dense_bias

  with tf.Graph().as_default():
    # Create model variables.
    for k, v in six.iteritems(model_weights):
      tf.Variable(v, name=k, dtype=tf.dtypes.float32)

    # Create non-model variables.
    global_step_var = tf.compat.v1.train.create_global_step()
    assign_op = global_step_var.assign(global_step)

    # Initialize vars and save checkpoint.
    with tf.compat.v1.train.MonitoredTrainingSession(
        checkpoint_dir=model_dir) as sess:
      sess.run(assign_op)

def _make_rnn_layer(rnn_cell_fn=None,
                    units=None,
                    cell_type=rnn.USE_DEFAULT,
                    return_sequences=False):
  return rnn._make_rnn_layer(
      rnn_cell_fn=rnn_cell_fn,
      units=units,
      cell_type=cell_type,
      return_sequences=return_sequences)

@test_util.run_all_in_graph_and_eager_modes
class RNNLayerFnTest(tf.test.TestCase, parameterized.TestCase):
  """Tests for rnn layer function."""

  def testWrongClassProvided(self):
    """Tests that an error is raised if the class doesn't have a call method."""
    with self.assertRaisesRegexp(
        ValueError, 'RNN cell should have a `call` and `state_size` method.'):
      _make_rnn_layer(units=[10], cell_type=lambda units: object())

  def testWrongStringProvided(self):
    """Tests that an error is raised if cell type is unknown."""
    with self.assertRaisesRegexp(
        ValueError,
        'cell_type` should be a class producing a RNN cell, or a string .*.'):
      _make_rnn_layer(units=[10], cell_type='unknown-cell-name')

  @parameterized.parameters(['simple_rnn', rnn.USE_DEFAULT])
  def testDefaultCellProvided(self, cell_type):
    """Tests behavior when the default cell type is provided."""
    layer = _make_rnn_layer(cell_type=cell_type, units=[1])
    self.assertIsInstance(layer, tf_keras.layers.RNN)
    self.assertIsInstance(layer.cell, tf_keras.layers.SimpleRNNCell)

  @parameterized.parameters([('gru', tf_keras.layers.GRU),
                             ('lstm', tf_keras.layers.LSTM),
                             ('simple_rnn', tf_keras.layers.SimpleRNN)])
  def testSpecificLayerTypeProvided(self, cell_type, layer_type):
    """Tests specific layer type for GRU and LSTM."""
    layer = _make_rnn_layer(cell_type=cell_type, units=1)
    self.assertIsInstance(layer, layer_type)

  def testSpecificLayerTypeArguments(self):
    """Tests arguments for specific layer types (GRU and LSTM)."""
    mock_layer_type = tf.compat.v1.test.mock.Mock()
    with tf.compat.v1.test.mock.patch.object(rnn, '_CELL_TYPE_TO_LAYER_MAPPING',
                                             {'custom-type': mock_layer_type}):
      _make_rnn_layer(
          cell_type='custom-type',
          units=11,
          return_sequences='return-seq-value')
      mock_layer_type.assert_called_once_with(
          units=11, return_sequences='return-seq-value')

  @tf.compat.v1.test.mock.patch.object(tf_keras.layers, 'RNN')
  def testCustomCellProvided(self, mock_rnn_layer_type):
    """Tests behavior when a custom cell type is provided."""
    mock_custom_cell = tf.compat.v1.test.mock.Mock()
    _make_rnn_layer(
        units=[10],
        cell_type=lambda units: mock_custom_cell,
        return_sequences='return-seq-value')
    mock_rnn_layer_type.assert_called_once_with(
        cell=mock_custom_cell, return_sequences='return-seq-value')

  def testMultipleCellsProvided(self):
    """Tests behavior when multiple cells are provided."""
    layer = _make_rnn_layer(cell_type='simple_rnn', units=[1, 2])
    self.assertIsInstance(layer, tf_keras.layers.RNN)
    self.assertIsInstance(layer.cell, tf_keras.layers.StackedRNNCells)
    self.assertLen(layer.cell.cells, 2)
    self.assertIsInstance(layer.cell.cells[0], tf_keras.layers.SimpleRNNCell)

  @tf.compat.v1.test.mock.patch.object(tf_keras.layers, 'RNN')
  def testCustomCellFnProvided(self, mock_rnn_layer_type):
    """Tests behavior when a custom cell function is provided."""
    mock_cell_fn = tf.compat.v1.test.mock.Mock(return_value='custom-cell')
    _make_rnn_layer(
        rnn_cell_fn=mock_cell_fn, return_sequences='return-seq-value')
    mock_rnn_layer_type.assert_called_once_with(
        cell='custom-cell', return_sequences='return-seq-value')

def _mock_logits_layer(kernel, bias):
  """Sets initialization values to dense `logits` layers used in context."""

  class _MockDenseLayer(tf_keras.layers.Dense):

    def __init__(self, units, activation, name):
      kwargs = {}
      if name == 'logits':
        kwargs = {
            'kernel_initializer': tf.compat.v1.initializers.constant(kernel),
            'bias_initializer': tf.compat.v1.initializers.constant(bias)
        }

      super(_MockDenseLayer, self).__init__(
          units=units, name=name, activation=activation, **kwargs)

  return tf.compat.v1.test.mock.patch.object(tf_keras.layers, 'Dense',
                                             _MockDenseLayer)

def _default_features_fn():
  return {
      'price':
          tf.sparse.SparseTensor(
              values=[10., 5.], indices=[[0, 0], [0, 1]], dense_shape=[1, 2]),
  }

def _get_mock_head():
  mock_head = multi_head_lib.MultiClassHead(3)
  mock_head.create_estimator_spec = tf.compat.v1.test.mock.Mock(
      return_value=model_fn.EstimatorSpec(None))
  return mock_head

@test_util.run_all_in_graph_and_eager_modes
class RNNLogitFnTest(tf.test.TestCase, parameterized.TestCase):
  """Tests correctness of logits calculated from RNNModel."""

  def setUp(self):
    # Sets layers default weights for testing purpose.
    self.kernel = [[.1, -.2]]
    self.recurrent = [[.2, -.3], [.3, -.4]]
    self.bias = [.2, .5]
    self.dense_kernel = [[-1.], [1.]]
    self.dense_bias = [0.3]
    self.sequence_feature_columns = [
        tf.feature_column.sequence_numeric_column('price', shape=(1,))
    ]
    self.context_feature_columns = []
    super(RNNLogitFnTest, self).setUp()

  def _mock_logits_layer(self):
    return _mock_logits_layer(self.dense_kernel, bias=self.dense_bias)

  def _test_logits(self,
                   logits_dimension,
                   features_fn,
                   expected_logits,
                   expected_mask,
                   return_sequences=False,
                   training=False):
    """Tests that the expected logits are calculated."""
    rnn_layer = tf_keras.layers.SimpleRNN(
        2,
        return_sequences=return_sequences,
        kernel_initializer=tf.compat.v1.initializers.constant(self.kernel),
        recurrent_initializer=tf.compat.v1.initializers.constant(
            self.recurrent),
        bias_initializer=tf.compat.v1.initializers.constant(self.bias))
    with self._mock_logits_layer():
      logit_layer = rnn.RNNModel(
          rnn_layer=rnn_layer,
          units=logits_dimension,
          sequence_feature_columns=self.sequence_feature_columns,
          context_feature_columns=self.context_feature_columns,
          return_sequences=return_sequences)
    logits = logit_layer(features_fn(), training=training)
    if return_sequences:
      logits = (logits, logits._keras_mask)
      expected_logits = (expected_logits, expected_mask)
    self.evaluate(tf.compat.v1.initializers.global_variables())
    self.assertAllClose(expected_logits, self.evaluate(logits), atol=1e-4)

  @parameterized.named_parameters(
      {
          'testcase_name': 'Static',
          'return_sequences': False,
          'expected_logits': [[-0.6033]]
      }, {
          'testcase_name': 'Sequential',
          'return_sequences': True,
          'expected_logits': [[[-1.4388], [-0.6033]]]
      }, {
          'testcase_name': 'SequentialTrain',
          'return_sequences': True,
          'expected_logits': [[[-1.4388], [-0.6033]]],
          'training': True
      }, {
          'testcase_name': 'SequentialInfer',
          'return_sequences': True,
          'expected_logits': [[[-1.4388], [-0.6033]]],
          'training': False
      })
  def testOneDimLogits(self, return_sequences, expected_logits, training=False):
    """Tests one-dimensional logits.

    Intermediate values are rounded for ease in reading.
    input_layer = [[[10]], [[5]]]
    sequence_mask = [[1, 1]]
    initial_state = [0, 0]
    rnn_output_timestep_1 = [[tanh(.1*10 + .2*0 + .3*0 +.2),
                              tanh(-.2*10 - .3*0 - .4*0 +.5)]]
                          = [[0.83, -0.91]]
    rnn_output_timestep_2 = [[tanh(.1*5 + .2*.83 - .3*.91 +.2),
                              tanh(-.2*5 - .3*.83 + .4*.91 +.5)]]
                          = [[0.53, -0.37]]
    logits_timestep_1 = [[-1*0.83 - 1*0.91 + 0.3]] = [[-1.4388]]
    logits_timestep_2 = [[-1*0.53 - 1*0.37 + 0.3]] = [[-0.6033]]

    Args:
      return_sequences: A boolean indicating whether to return the last output
        in the output sequence, or the full sequence.
      expected_logits: An array with expected logits result.
      training: Specifies if this training or evaluation / prediction mode.
    """
    expected_mask = [[1, 1]]

    self._test_logits(
        logits_dimension=1,
        features_fn=_default_features_fn,
        expected_mask=expected_mask,
        expected_logits=expected_logits,
        return_sequences=return_sequences,
        training=training)

  @parameterized.named_parameters(
      {
          'testcase_name': 'Static',
          'return_sequences': False,
          'expected_logits': [[-0.6033, 0.7777, 0.5698]]
      }, {
          'testcase_name': 'Sequential',
          'return_sequences': True,
          'expected_logits': [[[-1.4388, 1.0884, 0.5762],
                               [-0.6033, 0.7777, 0.5698]]]
      })
  def testMultiDimLogits(self, return_sequences, expected_logits):
    """Tests multi-dimensional logits.

    Intermediate values are rounded for ease in reading.
    input_layer = [[[10]], [[5]]]
    sequence_mask = [[1, 1]]
    initial_state = [0, 0]
    rnn_output_timestep_1 = [[tanh(.1*10 + .2*0 + .3*0 +.2),
                              tanh(-.2*10 - .3*0 - .4*0 +.5)]]
                          = [[0.83, -0.91]]
    rnn_output_timestep_2 = [[tanh(.1*5 + .2*.83 - .3*.91 +.2),
                              tanh(-.2*5 - .3*.83 + .4*.91 +.5)]]
                          = [[0.53, -0.37]]
    logits_timestep_1 = [[-1*0.83 - 1*0.91 + 0.3],
                         [0.5*0.83 + 0.3*0.91 + 0.4],
                         [0.2*0.83 - 0.1*0.91 + 0.5]]
                      = [[-1.4388, 1.0884, 0.5762]]
    logits_timestep_2 = [[-1*0.53 - 1*0.37 + 0.3],
                         [0.5*0.53 + 0.3*0.37 + 0.4],
                         [0.2*0.53 - 0.1*0.37 + 0.5]]
                      = [[-0.6033, 0.7777, 0.5698]]

    Args:
      return_sequences: A boolean indicating whether to return the last output
        in the output sequence, or the full sequence.
      expected_logits: An array with expected logits result.
    """
    expected_mask = [[1, 1]]

    self.dense_kernel = [[-1., 0.5, 0.2], [1., -0.3, 0.1]]
    self.dense_bias = [0.3, 0.4, 0.5]
    self._test_logits(
        logits_dimension=3,
        features_fn=_default_features_fn,
        expected_mask=expected_mask,
        expected_logits=expected_logits,
        return_sequences=return_sequences)

  @parameterized.named_parameters(
      {
          'testcase_name': 'Static',
          'return_sequences': False,
          'expected_logits': [[-0.6033, 0.7777, 0.5698],
                              [-1.2473, 1.0170, 0.5745]]
      }, {
          'testcase_name': 'Sequential',
          'return_sequences': True,
          'expected_logits': [[
              [-1.4388, 1.0884, 0.5762], [-0.6033, 0.7777, 0.5698]
          ], [[0.0197, 0.5601, 0.5860], [-1.2473, 1.0170, 0.5745]]]
      })
  def testMultiExampleMultiDim(self, return_sequences, expected_logits):
    """Tests multiple examples and multi-dimensional logits.

    Intermediate values are rounded for ease in reading.
    input_layer = [[[10], [5]], [[2], [7]]]
    sequence_mask = [[1, 1], [1, 1]]
    initial_state = [[0, 0], [0, 0]]
    rnn_output_timestep_1 = [[tanh(.1*10 + .2*0 + .3*0 +.2),
                              tanh(-.2*10 - .3*0 - .4*0 +.5)],
                             [tanh(.1*2 + .2*0 + .3*0 +.2),
                              tanh(-.2*2 - .3*0 - .4*0 +.5)]]
                          = [[0.83, -0.91], [0.38, 0.10]]
    rnn_output_timestep_2 = [[tanh(.1*5 + .2*.83 - .3*.91 +.2),
                              tanh(-.2*5 - .3*.83 + .4*.91 +.5)],
                             [tanh(.1*7 + .2*.38 + .3*.10 +.2),
                              tanh(-.2*7 - .3*.38 - .4*.10 +.5)]]
                          = [[0.53, -0.37], [0.76, -0.78]
    logits_timestep_1 = [[-1*0.83 - 1*0.91 + 0.3,
                          0.5*0.83 + 0.3*0.91 + 0.4,
                          0.2*0.83 - 0.1*0.91 + 0.5],
                         [-1*0.38 + 1*0.10 + 0.3,
                          0.5*0.38 - 0.3*0.10 + 0.4,
                          0.2*0.38 + 0.1*0.10 + 0.5]]
                      = [[-1.4388, 1.0884, 0.5762], [0.0197, 0.5601, 0.5860]]
    logits_timestep_2 = [[-1*0.53 - 1*0.37 + 0.3,
                          0.5*0.53 + 0.3*0.37 + 0.4,
                          0.2*0.53 - 0.1*0.37 + 0.5],
                         [-1*0.76 - 1*0.78 + 0.3,
                          0.5*0.76 +0.3*0.78 + 0.4,
                          0.2*0.76 -0.1*0.78 + 0.5]]
                      = [[-0.6033, 0.7777, 0.5698], [-1.2473, 1.0170, 0.5745]]

    Args:
      return_sequences: A boolean indicating whether to return the last output
        in the output sequence, or the full sequence.
      expected_logits: An array with expected logits result.
    """
    expected_mask = [[1, 1], [1, 1]]

    def features_fn():
      return {
          'price':
              tf.sparse.SparseTensor(
                  values=[10., 5., 2., 7.],
                  indices=[[0, 0], [0, 1], [1, 0], [1, 1]],
                  dense_shape=[2, 2]),
      }

    self.dense_kernel = [[-1., 0.5, 0.2], [1., -0.3, 0.1]]
    self.dense_bias = [0.3, 0.4, 0.5]
    self._test_logits(
        logits_dimension=3,
        features_fn=features_fn,
        expected_mask=expected_mask,
        expected_logits=expected_logits,
        return_sequences=return_sequences)

  @parameterized.named_parameters(
      {
          'testcase_name': 'Static',
          'return_sequences': False,
          'expected_logits': [[-0.6033], [0.0197]]
      }, {
          'testcase_name': 'Sequential',
          'return_sequences': True,
          'expected_logits': [[[-1.4388], [-0.6033]], [[0.0197], [0.0197]]]
      })
  def testMultiExamplesDifferentLength(self, return_sequences, expected_logits):
    """Tests multiple examples with different lengths.

    Intermediate values are rounded for ease in reading.
    input_layer = [[[10], [5]], [[2], [0]]]
    sequence_mask = [[1, 1], [1, 0]]
    initial_state = [[0, 0], [0, 0]]
    rnn_output_timestep_1 = [[tanh(.1*10 + .2*0 + .3*0 +.2),
                              tanh(-.2*10 - .3*0 - .4*0 +.5)],
                             [tanh(.1*2 + .2*0 + .3*0 +.2),
                              tanh(-.2*2 - .3*0 - .4*0 +.5)]]
                          = [[0.83, -0.91], [0.38, 0.10]]
    rnn_output_timestep_2 = [[tanh(.1*5 + .2*.83 - .3*.91 +.2),
                              tanh(-.2*5 - .3*.83 + .4*.91 +.5)],
                             [_]]
                          = [[0.53, -0.37], [_, _]]
    logits_timestep_1 = [[-1*0.83 - 1*0.91 + 0.3],
                         [-1*0.38 + 1*0.10 + 0.3]]
                      = [[-0.4388], [0.0197]]
    logits_timestep_2 = [[-1*0.53 - 1*0.37 + 0.3],
                         [_]]
                      = [[-0.6033], [_]]

    Args:
      return_sequences: A boolean indicating whether to return the last output
        in the output sequence, or the full sequence.
      expected_logits: An array with expected logits result.
    """
    expected_mask = [[1, 1], [1, 0]]

    def features_fn():
      return {
          'price':
              tf.sparse.SparseTensor(
                  values=[10., 5., 2.],
                  indices=[[0, 0], [0, 1], [1, 0]],
                  dense_shape=[2, 2]),
      }

    self._test_logits(
        logits_dimension=1,
        features_fn=features_fn,
        expected_mask=expected_mask,
        expected_logits=expected_logits,
        return_sequences=return_sequences)

  def testMultiExamplesWithContext(self):
    """Tests multiple examples with context features.

    Intermediate values are rounded for ease in reading.
    input_layer = [[[10, -0.5], [5, -0.5]], [[2, 0.8], [0, 0]]]
    sequence_mask = [[1, 1], [1, 0]]
    initial_state = [[0, 0], [0, 0]]
    rnn_output_timestep_1 = [[tanh(.1*10 - 1*.5 + .2*0 + .3*0 +.2),
                              tanh(-.2*10 - 0.9*.5 - .3*0 - .4*0 +.5)],
                             [tanh(.1*2 + 1*.8 + .2*0 + .3*0 +.2),
                              tanh(-.2*2 + .9*.8 - .3*0 - .4*0 +.5)]]
                          = [[0.60, -0.96], [0.83, 0.68]]
    rnn_output_timestep_2 = [[tanh(.1*5 - 1*.5 + .2*.60 - .3*.96 +.2),
                              tanh(-.2*5 - .9*.5 - .3*.60 + .4*.96 +.5)],
                             [<ignored-padding>]]
                          = [[0.03, -0.63], [<ignored-padding>]]
    logits = [[-1*0.03 - 1*0.63 + 0.3],
              [-1*0.83 + 1*0.68 + 0.3]]
           = [[-0.3662], [0.1414]]
    """
    expected_mask = [[1, 1], [1, 0]]

    def features_fn():
      return {
          'price':
              tf.sparse.SparseTensor(
                  values=[10., 5., 2.],
                  indices=[[0, 0], [0, 1], [1, 0]],
                  dense_shape=[2, 2]),
          'context': [[-0.5], [0.8]],
      }

    self.context_feature_columns = [
        tf.feature_column.numeric_column('context', shape=(1,))
    ]

    self.kernel = [[.1, -.2], [1., 0.9]]
    self._test_logits(
        logits_dimension=1,
        features_fn=features_fn,
        expected_mask=expected_mask,
        expected_logits=[[-0.3662], [0.1414]])

  def testMultiExamplesMultiFeatures(self):
    """Tests examples with multiple sequential feature columns.

    Intermediate values are rounded for ease in reading.
    input_layer = [[[1, 0, 10], [0, 1, 5]], [[1, 0, 2], [0, 0, 0]]]
    sequence_mask = [[1, 1], [1, 0]]
    initial_state = [[0, 0], [0, 0]]
    rnn_output_timestep_1 = [[tanh(.5*1 + 1*0 + .1*10 + .2*0 + .3*0 +.2),
                              tanh(-.5*1 - 1*0 - .2*10 - .3*0 - .4*0 +.5)],
                             [tanh(.5*1 + 1*0 + .1*2 + .2*0 + .3*0 +.2),
                              tanh(-.5*1 - 1*0 - .2*2 - .3*0 - .4*0 +.5)]]
                          = [[0.94, -0.96], [0.72, -0.38]]
    rnn_output_timestep_2 = [[tanh(.5*0 + 1*1 + .1*5 + .2*.94 - .3*.96 +.2),
                              tanh(-.5*0 - 1*1 - .2*5 - .3*.94 + .4*.96 +.5)],
                             [<ignored-padding>]]
                          = [[0.92, -0.88], [<ignored-padding>]]
    logits = [[-1*0.92 - 1*0.88 + 0.3],
              [-1*0.72 - 1*0.38 + 0.3]]
           = [[-1.5056], [-0.7962]]
    """
    expected_mask = [[1, 1], [1, 0]]

    def features_fn():
      return {
          'price':
              tf.sparse.SparseTensor(
                  values=[10., 5., 2.],
                  indices=[[0, 0], [0, 1], [1, 0]],
                  dense_shape=[2, 2]),
          'on_sale':
              tf.sparse.SparseTensor(
                  values=[0, 1, 0],
                  indices=[[0, 0], [0, 1], [1, 0]],
                  dense_shape=[2, 2]),
      }

    price_column = tf.feature_column.sequence_numeric_column(
        'price', shape=(1,))
    on_sale_column = tf.feature_column.indicator_column(
        tf.feature_column.sequence_categorical_column_with_identity(
            'on_sale', num_buckets=2))
    self.sequence_feature_columns = [price_column, on_sale_column]

    self.kernel = [[.5, -.5], [1., -1.], [.1, -.2]]
    self._test_logits(
        logits_dimension=1,
        features_fn=features_fn,
        expected_mask=expected_mask,
        expected_logits=[[-1.5056], [-0.7962]])

  @parameterized.parameters([(model_fn.ModeKeys.TRAIN, True),
                             (model_fn.ModeKeys.EVAL, False),
                             (model_fn.ModeKeys.PREDICT, False)])
  def testTrainingMode(self, mode, expected_training_mode):
    """Tests that `training` argument is properly used."""

    class _MockRNNCell(tf_keras.layers.SimpleRNNCell):
      """Used to test that `training` argument is properly used."""

      def __init__(self, test_case):
        self._test_case = test_case
        super(_MockRNNCell, self).__init__(units=10)

      def call(self, inputs, states, training=None):
        self._test_case.assertEqual(training, expected_training_mode)
        return super(_MockRNNCell, self).call(
            inputs=inputs, states=states, training=training)

    estimator = rnn.RNNEstimator(
        head=_get_mock_head(),
        rnn_cell_fn=lambda: _MockRNNCell(self),
        sequence_feature_columns=self.sequence_feature_columns)
    features = {
        'price':
            tf.sparse.SparseTensor(
                values=[
                    10.,
                ], indices=[[0, 0]], dense_shape=[1, 1]),
    }
    estimator.model_fn(features=features, labels=None, mode=mode, config=None)

class RNNModelTest(tf.test.TestCase, parameterized.TestCase):
  """Tests for RNNModel."""

  def setUp(self):
    super(RNNModelTest, self).setUp()
    self.kernel = [[.1, -.2]]
    self.recurrent = [[.2, -.3], [.3, -.4]]
    self.bias = [.2, .5]
    self.dense_kernel = [[-1.], [1.]]
    self.dense_bias = [0.3]
    self.sequence_feature_columns = [
        tf.feature_column.sequence_numeric_column('price', shape=(1,))
    ]
    self.x = {
        'price':
            tf.sparse.SparseTensor(
                values=[10., 5., 2.],
                indices=[[0, 0], [0, 1], [1, 0]],
                dense_shape=[2, 2]),
    }
    self.y = ops.convert_to_tensor([[[0], [1]], [[0], [1]]])

  def _get_compiled_model(self,
                          return_sequences=False,
                          optimizer='Adam',
                          **kwargs):
    """Initializes and compiles a RNN model with specific weights."""
    rnn_layer = tf_keras.layers.SimpleRNN(
        2,
        return_sequences=return_sequences,
        kernel_initializer=tf.compat.v1.initializers.constant(self.kernel),
        recurrent_initializer=tf.compat.v1.initializers.constant(
            self.recurrent),
        bias_initializer=tf.compat.v1.initializers.constant(self.bias))
    with _mock_logits_layer(self.dense_kernel, bias=self.dense_bias):
      model = rnn.RNNModel(
          units=1,
          rnn_layer=rnn_layer,
          sequence_feature_columns=self.sequence_feature_columns,
          activation=tf_keras.activations.sigmoid,
          return_sequences=return_sequences,
          **kwargs)
      model.compile(
          optimizer=optimizer,
          loss=tf_keras.losses.BinaryCrossentropy(reduction='sum'),
          metrics=['accuracy'])
    return model

  def testModelWeights(self):
    """Tests that the layers weights are properly added to the model weights."""
    col = tf.feature_column.categorical_column_with_hash_bucket(
        'tokens', hash_bucket_size=1)
    context_feature_columns = [
        tf.feature_column.embedding_column(col, dimension=1)
    ]
    seq_col = tf.feature_column.sequence_categorical_column_with_hash_bucket(
        'seq-tokens', hash_bucket_size=1)
    sequence_feature_columns = [
        tf.feature_column.embedding_column(seq_col, dimension=1)
    ]
    model = rnn.RNNModel(
        units=1,
        rnn_layer=tf_keras.layers.SimpleRNN(2),
        sequence_feature_columns=sequence_feature_columns,
        activation=tf_keras.activations.sigmoid,
        context_feature_columns=context_feature_columns)
    model.compile(
        optimizer='Adam',
        loss=tf_keras.losses.BinaryCrossentropy(reduction='sum'),
        metrics=['accuracy'])

    model.predict(
        x={
            'tokens': ops.convert_to_tensor([['a']]),
            'seq-tokens': ops.convert_to_tensor([[['a']]])
        },
        steps=1)
    # Weights included are:
    # - recurrent, kernel and bias from RNN layer
    # - kernel and bias from logits layer
    # - sequential feature column embedding
    # - context feature column embedding.
    self.assertLen(model.get_weights(), 7)

  def _testModelConfig(self, **kwargs):
    """Tests the parameters of a RNNModel stored to and restored from config.

    Args:
      **kwargs: Additional keyword arguments to initialize the RNNModel before
        calling `get_config`.

    Returns:
      A dictionary with RNNModel initialization arguments from the `from_config`
      call.
    """
    seq_col = tf.feature_column.sequence_categorical_column_with_hash_bucket(
        'seq-tokens', hash_bucket_size=1)
    sequence_feature_columns = [
        tf.feature_column.embedding_column(
            seq_col, dimension=1, initializer=tf.compat.v1.initializers.zeros())
    ]
    model = rnn.RNNModel(
        units=11,
        rnn_layer=tf_keras.layers.SimpleRNN(3),
        sequence_feature_columns=sequence_feature_columns,
        return_sequences=True,
        name='rnn-model',
        **kwargs)

    with tf.compat.v1.test.mock.patch.object(
        rnn.RNNModel, '__init__', return_value=None) as init:
      rnn.RNNModel.from_config(
          model.get_config(),
          custom_objects={'Zeros': tf.compat.v1.initializers.zeros})
      return list(init.call_args_list[0])[1]

  def testModelConfig(self):
    """Tests that a RNNModel can be stored to and restored from config."""
    init_kwargs = self._testModelConfig()
    self.assertEqual(init_kwargs['name'], 'rnn-model')
    self.assertEqual(init_kwargs['units'], 11)
    self.assertEqual(init_kwargs['return_sequences'], True)
    self.assertEqual(
        init_kwargs['sequence_feature_columns'][0].categorical_column.name,
        'seq-tokens')
    self.assertEqual(init_kwargs['context_feature_columns'], None)
    self.assertEqual(init_kwargs['activation'].__name__, 'linear')
    self.assertEqual(init_kwargs['rnn_layer'].cell.units, 3)

  def testModelConfigWithActivation(self):
    """Tests store / restore from config with logits activation."""
    init_kwargs = self._testModelConfig(activation=tf_keras.activations.sigmoid)
    self.assertEqual(init_kwargs['activation'].__name__, 'sigmoid')

  def testModelConfigWithContextFeatures(self):
    """Tests store / restore from config with context features."""
    init_kwargs = self._testModelConfig(context_feature_columns=[
        tf.feature_column.numeric_column('context', shape=(1,))
    ])
    self.assertEqual(init_kwargs['context_feature_columns'][0].name, 'context')

  def DISABLED_testSaveModelWeights(self):  # See b/129842600.
    """Tests that model weights can be saved and restored."""
    model = self._get_compiled_model(return_sequences=True)
    model.fit(x=self.x, y=self.y, batch_size=1, steps_per_epoch=1, epochs=1)
    y1 = model.predict(x=self.x, steps=1)
    model.save_weights(self.get_temp_dir() + 'model')

    model = self._get_compiled_model(return_sequences=True, name='model-2')
    model.load_weights(self.get_temp_dir() + 'model')
    y2 = model.predict(x=self.x, steps=1)
    self.assertAllClose(y1, y2)

  def DISABLED_testEvaluationMetrics(self):  # See b/129842600.
    """Tests evaluation metrics computation in non-sequential case."""
    model = self._get_compiled_model()
    metrics = model.evaluate(
        x=self.x, y=ops.convert_to_tensor([[0], [1]]), steps=1)
    # See `RNNClassifierEvaluationTest` for details on computation.
    self.assertAllClose(metrics, (1.1196611, 1.), atol=1e-4)

  def DISABLED_testEvaluationSequential(self):  # See b/129842600.
    """Tests that the sequence mask is properly used to aggregate loss."""
    model = self._get_compiled_model(return_sequences=True)
    metrics = model.evaluate(x=self.x, y=self.y, steps=1)
    # See `RNNClassifierEvaluationTest` for details on computation.
    self.assertAllClose(metrics, (1.9556, 1. / 3.), atol=1e-4)

  def DISABLED_testPredictions(self):  # See b/129842600.
    """Tests predictions with RNN model."""
    model = self._get_compiled_model()
    # See `RNNClassifierPredictionTest` for details on computation.
    self.assertAllClose(
        model.predict(x=self.x, steps=1), [[0.353593], [0.5049296]], atol=1e-4)

  def DISABLED_testPredictionsSequential(self):  # See b/129842600.
    """Tests sequential predictions with RNN model."""
    model = self._get_compiled_model(return_sequences=True)
    # See `RNNClassifierPredictionTest` for details on computation.
    self.assertAllClose(
        model.predict(x=self.x, steps=1),
        [[[0.191731], [0.353593]], [[0.5049296], [0.5049296]]],
        atol=1e-4)

  @parameterized.named_parameters(
      ('StringOptimizer', 'Adam'),
      ('OptimizerInstance', tf_keras.optimizers.Adam()))
  def DISABLED_testTraining(self, optimizer):  # See b/129842600.
    """Tests the loss computed in training step."""
    model = self._get_compiled_model(optimizer=optimizer)
    history = model.fit(
        x=self.x,
        y=ops.convert_to_tensor([[0], [1]]),
        batch_size=1,
        steps_per_epoch=1)
    # See `RNNClassifierTrainingTest` for details on computation.
    self.assertAllClose(history.history['loss'], [1.1196611], atol=1e-4)

  def DISABLED_testTrainingSequential(self):  # See b/129842600.
    """Tests the loss computed in training step in sequential case."""
    model = self._get_compiled_model(return_sequences=True)
    history = model.fit(x=self.x, y=self.y, batch_size=1, steps_per_epoch=1)
    # See `RNNClassifierTrainingTest` for details on computation.
    self.assertAllClose(history.history['loss'], [1.9556], atol=1e-4)

@test_util.run_all_in_graph_and_eager_modes
class RNNEstimatorInitTest(tf.test.TestCase):

  def setUp(self):
    col = tf.feature_column.sequence_categorical_column_with_hash_bucket(
        'tokens', hash_bucket_size=10)
    self.feature_columns = [
        tf.feature_column.embedding_column(col, dimension=2)
    ]
    self.cell_units = [4, 2]
    super(RNNEstimatorInitTest, self).setUp()

  def testConflictingRNNCellFn(self):
    with self.assertRaisesRegexp(
        ValueError,
        'units and cell_type must not be specified when using rnn_cell_fn'):
      rnn.RNNClassifier(
          sequence_feature_columns=self.feature_columns,
          rnn_cell_fn=lambda: 'mock-cell',
          units=self.cell_units)

    with self.assertRaisesRegexp(
        ValueError,
        'units and cell_type must not be specified when using rnn_cell_fn'):
      rnn.RNNClassifier(
          sequence_feature_columns=self.feature_columns,
          rnn_cell_fn=lambda: 'mock-cell',
          cell_type='lstm')

  def testNonSequentialHeadProvided(self):
    with self.assertRaisesRegexp(
        ValueError, 'Provided head must be a `_SequentialHead` object when '
        '`return_sequences` is set to True.'):
      rnn.RNNEstimator(
          head=multi_head_lib.MultiClassHead(n_classes=3),
          sequence_feature_columns=self.feature_columns,
          return_sequences=True)

  def testWrongOptimizerTypeProvided(self):
    classifier = rnn.RNNClassifier(
        self.feature_columns, units=[1], optimizer=object())
    with self.assertRaisesRegexp(
        ValueError,
        'The given object is not a tf_keras.optimizers.Optimizer instance.'):
      classifier.model_fn(
          features=None, labels=None, mode=model_fn.ModeKeys.TRAIN, config=None)

@test_util.run_all_in_graph_and_eager_modes
class RNNClassifierTrainingTest(tf.test.TestCase):

  def setUp(self):
    self.kernel = [[.1, -.2]]
    self.recurrent = [[.2, -.3], [.3, -.4]]
    self.bias = [.2, .5]
    self.dense_kernel = [[-1.], [1.]]
    self.dense_bias = [0.3]
    self.sequence_feature_columns = [
        tf.feature_column.sequence_numeric_column('price', shape=(1,))
    ]
    super(RNNClassifierTrainingTest, self).setUp()

  def _assert_checkpoint(self, n_classes, input_units, cell_units,
                         expected_global_step):

    shapes = {
        name: shape
        for (name, shape) in tf.train.list_variables(self.get_temp_dir())
    }

    self.assertEqual([], shapes[tf.compat.v1.GraphKeys.GLOBAL_STEP])
    self.assertEqual(
        expected_global_step,
        tf.train.load_variable(self.get_temp_dir(),
                               tf.compat.v1.GraphKeys.GLOBAL_STEP))

    # RNN Cell variables.
    for i, cell_unit in enumerate(cell_units):
      name_suffix = '_%d' % i if i else ''
      self.assertEqual([input_units, cell_unit],
                       shapes[CELL_KERNEL_NAME + name_suffix])
      self.assertEqual([cell_unit, cell_unit],
                       shapes[CELL_RECURRENT_KERNEL_NAME + name_suffix])
      self.assertEqual([cell_unit], shapes[CELL_BIAS_NAME + name_suffix])
      input_units = cell_unit

    # Logits variables.
    logits_dimension = n_classes if n_classes > 2 else 1
    self.assertEqual([cell_units[-1], logits_dimension],
                     shapes[LOGITS_WEIGHTS_NAME])
    self.assertEqual([logits_dimension], shapes[LOGITS_BIAS_NAME])

  def _mock_optimizer(self, expected_loss=None):
    var_names = (CELL_BIAS_NAME, CELL_KERNEL_NAME, CELL_RECURRENT_KERNEL_NAME,
                 LOGITS_BIAS_NAME, LOGITS_WEIGHTS_NAME)
    expected_var_names = ['%s:0' % name for name in var_names]

    class _Optimizer(tf_keras.optimizers.Optimizer):
      """Mock optimizer checking that loss has the proper value."""

      def __init__(self, test_case):
        super(_Optimizer, self).__init__(name='my-optimizer')
        self.call_count = 0
        self._test_case = test_case

      def get_updates(self, loss, params):
        self.call_count += 1
        trainable_vars = tf.compat.v1.get_collection(
            tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES)
        self._test_case.assertItemsEqual(expected_var_names,
                                         [var.name for var in trainable_vars])

        # Verify loss. We can't check the value directly so we add an assert op.
        self._test_case.assertEquals(0, loss.shape.ndims)
        if expected_loss is None:
          return [self.iterations.assign_add(1).op]
        assert_loss = _assert_close(
            tf.cast(expected_loss, name='expected', dtype=tf.dtypes.float32),
            loss,
            name='assert_loss')
        with tf.control_dependencies((assert_loss,)):
          return [self.iterations.assign_add(1).op]

      def get_config(self):
        pass

    return _Optimizer(test_case=self)

  def _testFromScratchWithDefaultOptimizer(self, n_classes):

    def train_input_fn():
      return {
          'tokens':
              tf.sparse.SparseTensor(
                  values=['the', 'cat', 'sat'],
                  indices=[[0, 0], [0, 1], [0, 2]],
                  dense_shape=[1, 3]),
      }, [[1]]

    col = tf.feature_column.sequence_categorical_column_with_hash_bucket(
        'tokens', hash_bucket_size=10)
    embed = tf.feature_column.embedding_column(col, dimension=2)
    input_units = 2

    cell_units = [4, 2]
    est = rnn.RNNClassifier(
        sequence_feature_columns=[embed],
        units=cell_units,
        n_classes=n_classes,
        model_dir=self.get_temp_dir())

    # Train for a few steps, and validate final checkpoint.
    num_steps = 10
    est.train(input_fn=train_input_fn, steps=num_steps)
    self._assert_checkpoint(n_classes, input_units, cell_units, num_steps)

  def testBinaryClassFromScratchWithDefaultOptimizer(self):
    self._testFromScratchWithDefaultOptimizer(n_classes=2)

  def testMultiClassFromScratchWithDefaultOptimizer(self):
    self._testFromScratchWithDefaultOptimizer(n_classes=4)

  def testFromScratchWithCustomRNNCellFn(self):

    def train_input_fn():
      return {
          'tokens':
              tf.sparse.SparseTensor(
                  values=['the', 'cat', 'sat'],
                  indices=[[0, 0], [0, 1], [0, 2]],
                  dense_shape=[1, 3]),
      }, [[1]]

    col = tf.feature_column.sequence_categorical_column_with_hash_bucket(
        'tokens', hash_bucket_size=10)
    embed = tf.feature_column.embedding_column(col, dimension=2)
    input_units = 2
    cell_units = [4, 2]
    n_classes = 2

    def rnn_cell_fn():
      cells = [tf_keras.layers.SimpleRNNCell(units=n) for n in cell_units]
      return tf_keras.layers.StackedRNNCells(cells)

    est = rnn.RNNClassifier(
        sequence_feature_columns=[embed],
        rnn_cell_fn=rnn_cell_fn,
        n_classes=n_classes,
        model_dir=self.get_temp_dir())

    # Train for a few steps, and validate final checkpoint.
    num_steps = 10
    est.train(input_fn=train_input_fn, steps=num_steps)
    self._assert_checkpoint(n_classes, input_units, cell_units, num_steps)

  def _testExampleWeight(self, n_classes):

    def train_input_fn():
      return {
          'tokens':
              tf.sparse.SparseTensor(
                  values=['the', 'cat', 'sat', 'dog', 'barked'],
                  indices=[[0, 0], [0, 1], [0, 2], [1, 0], [1, 1]],
                  dense_shape=[2, 3]),
          'w': [[1], [2]],
      }, [[1], [0]]

    col = tf.feature_column.sequence_categorical_column_with_hash_bucket(
        'tokens', hash_bucket_size=10)
    embed = tf.feature_column.embedding_column(col, dimension=2)
    input_units = 2

    cell_units = [4, 2]
    est = rnn.RNNClassifier(
        units=cell_units,
        sequence_feature_columns=[embed],
        n_classes=n_classes,
        weight_column='w',
        model_dir=self.get_temp_dir())

    # Train for a few steps, and validate final checkpoint.
    num_steps = 10
    est.train(input_fn=train_input_fn, steps=num_steps)
    self._assert_checkpoint(n_classes, input_units, cell_units, num_steps)

  def testBinaryClassWithExampleWeight(self):
    self._testExampleWeight(n_classes=2)

  def testMultiClassWithExampleWeight(self):
    self._testExampleWeight(n_classes=4)

  def _testFromCheckpoint(self, input_fn, expected_loss, **kwargs):
    """Loads classifier from checkpoint, runs training and checks loss."""
    create_checkpoint(
        kernel=self.kernel,
        recurrent=self.recurrent,
        bias=self.bias,
        dense_kernel=self.dense_kernel,
        dense_bias=self.dense_bias,
        global_step=100,
        model_dir=self.get_temp_dir())

    mock_optimizer = self._mock_optimizer(expected_loss=expected_loss)

    est = rnn.RNNClassifier(
        units=[2],
        sequence_feature_columns=self.sequence_feature_columns,
        optimizer=mock_optimizer,
        model_dir=self.get_temp_dir(),
        **kwargs)
    self.assertEqual(0, mock_optimizer.call_count)
    est.train(input_fn=input_fn, steps=10)
    self.assertEqual(1, mock_optimizer.call_count)

  def testBinaryClassFromCheckpoint(self):

    def train_input_fn():
      return {
          'price':
              tf.sparse.SparseTensor(
                  values=[10., 5., 2.],
                  indices=[[0, 0], [0, 1], [1, 0]],
                  dense_shape=[2, 2]),
      }, [[0], [1]]

    # Uses same checkpoint and examples as testBinaryClassEvaluationMetrics.
    # See that test for loss calculation.
    self._testFromCheckpoint(train_input_fn, expected_loss=0.559831)

  def testMultiClassFromCheckpoint(self):

    def train_input_fn():
      return {
          'price':
              tf.sparse.SparseTensor(
                  values=[10., 5., 2., 7.],
                  indices=[[0, 0], [0, 1], [1, 0], [1, 1]],
                  dense_shape=[2, 2]),
      }, [[0], [1]]

    # Uses same checkpoint and examples as testMultiClassEvaluationMetrics.
    # See that test for loss calculation.
    self.dense_kernel = [[-1., 0.5, 0.2], [1., -0.3, 0.1]]
    self.dense_bias = [0.3, 0.4, 0.5]
    self._testFromCheckpoint(
        train_input_fn, expected_loss=1.331465, n_classes=3)

  def testBinaryClassFromCheckpointSequential(self):

    def train_input_fn():
      return {
          'price':
              tf.sparse.SparseTensor(
                  values=[10., 5., 2.],
                  indices=[[0, 0], [0, 1], [1, 0]],
                  dense_shape=[2, 2]),
      }, tf.sparse.SparseTensor(
          values=[0, 1, 0],
          indices=[[0, 0], [0, 1], [1, 0]],
          dense_shape=[2, 2])

    # Same example as testBinaryClassEvaluationMetricsSequential.
    # logits = [[[-1.4388], [-0.6033]],
    #            [[0.0197], [_]]]
    # probability = np.exp(logits) / (1 + np.exp(logits))
    #             = [[0.1917, 0.3536],
    #                [0.5049, _]]
    # loss = -label * ln(p) - (1 - label) * ln(1 - p)
    # loss = [[0.2129,  1.0396],
    #         [0.7031, _]]
    # aggregated_loss = sum(loss) / 3
    # aggregated_loss = 0.6518
    self._testFromCheckpoint(
        train_input_fn, expected_loss=0.651841, return_sequences=True)

  def testBinaryClassFromCheckpointSequentialWithWeights(self):

    def train_input_fn():
      return {
          'price':
              tf.sparse.SparseTensor(
                  values=[10., 5., 2.],
                  indices=[[0, 0], [0, 1], [1, 0]],
                  dense_shape=[2, 2]),
          'weights':
              tf.sparse.SparseTensor(
                  values=[0., 0.5, 0.5],
                  indices=[[0, 0], [0, 1], [1, 0]],
                  dense_shape=[2, 2])
      }, tf.sparse.SparseTensor(
          values=[0, 0, 1],
          indices=[[0, 0], [0, 1], [1, 0]],
          dense_shape=[2, 2])

    # Checkpoint and input are the same as testBinaryClassEvaluationMetrics, and
    # expected loss is the same as we use non-zero weights only for the last
    # step of each sequence.
    # loss = [[_,  0.436326],
    #         [0.6833351, _]]
    # weights = [[0, 0.5], [0.5, 0]]
    # aggregated_loss = (0.436326 + 0.6833351) / 2.
    #                 = 0.559831
    self._testFromCheckpoint(
        train_input_fn,
        expected_loss=0.559831,
        return_sequences=True,
        weight_column='weights',
        loss_reduction=tf_keras.losses.Reduction.SUM)

  def testDefaultGradientClipping(self):
    """Tests that optimizer applies default gradient clipping value."""

    def train_input_fn():
      return {
          'price':
              tf.sparse.SparseTensor(
                  values=[
                      1.,
                  ], indices=[[0, 0]], dense_shape=[1, 1]),
      }, [[1]]

    def _wrap_create_estimator_spec(create_estimator_spec):
      """Wraps function and asserts that the optimizer applies clipping."""

      def _wrapped_create_estimator_spec(obj,
                                         features,
                                         mode,
                                         logits,
                                         labels=None,
                                         optimizer=None,
                                         trainable_variables=None,
                                         train_op_fn=None,
                                         update_ops=None,
                                         regularization_losses=None):
        var = tf.Variable([1.0])
        mock_loss = 10 * var
        gradients = optimizer.get_gradients(mock_loss, [var])
        self.assertLen(gradients, 1)
        # Initial gradient value is 10 and expected to be clipped to 5 (default
        # clipping value).
        with tf.control_dependencies(
            (tf.compat.v1.debugging.assert_equal(gradients[0], 5.0),)):
          return create_estimator_spec(obj, features, mode, logits, labels,
                                       optimizer, trainable_variables,
                                       train_op_fn, update_ops,
                                       regularization_losses)

      return _wrapped_create_estimator_spec

    with tf.compat.v1.test.mock.patch.object(
        multi_head_lib.MultiClassHead, 'create_estimator_spec',
        _wrap_create_estimator_spec(
            multi_head_lib.MultiClassHead.create_estimator_spec)):
      est = rnn.RNNClassifier(
          n_classes=3,
          sequence_feature_columns=[
              tf.feature_column.sequence_numeric_column('price')
          ],
          units=[2],
          model_dir=self.get_temp_dir())
      est.train(input_fn=train_input_fn, steps=1)

def sorted_key_dict(unsorted_dict):
  return {k: unsorted_dict[k] for k in sorted(unsorted_dict)}

@test_util.run_all_in_graph_and_eager_modes
class RNNClassifierEvaluationTest(tf.test.TestCase):

  def setUp(self):
    self.kernel = [[.1, -.2]]
    self.recurrent = [[.2, -.3], [.3, -.4]]
    self.bias = [.2, .5]
    self.dense_kernel = [[-1.], [1.]]
    self.dense_bias = [0.3]
    self.global_step = 100
    self.sequence_feature_columns = [
        tf.feature_column.sequence_numeric_column('price', shape=(1,))
    ]
    super(RNNClassifierEvaluationTest, self).setUp()

  def _testFromCheckpoint(self, input_fn, **kwargs):
    create_checkpoint(
        kernel=self.kernel,
        recurrent=self.recurrent,
        bias=self.bias,
        dense_kernel=self.dense_kernel,
        dense_bias=self.dense_bias,
        global_step=self.global_step,
        model_dir=self.get_temp_dir())

    est = rnn.RNNClassifier(
        units=[2],
        sequence_feature_columns=self.sequence_feature_columns,
        model_dir=self.get_temp_dir(),
        **kwargs)
    return est.evaluate(input_fn, steps=1)

  def testBinaryClassEvaluationMetrics(self):

    def eval_input_fn():
      return {
          'price':
              tf.sparse.SparseTensor(
                  values=[10., 5., 2.],
                  indices=[[0, 0], [0, 1], [1, 0]],
                  dense_shape=[2, 2]),
      }, [[0], [1]]

    eval_metrics = self._testFromCheckpoint(eval_input_fn)

    # Uses identical numbers to testMultiExamplesWithDifferentLength.
    # See that test for logits calculation.
    # logits = [[-0.603282], [0.019719]]
    # probability = exp(logits) / (1 + exp(logits)) = [[0.353593], [0.504930]]
    # loss = -label * ln(p) - (1 - label) * ln(1 - p)
    #      = [[0.436326], [0.683335]]
    # sum_over_batch_size = (0.436326 + 0.683335)/2
    expected_metrics = {
        tf.compat.v1.GraphKeys.GLOBAL_STEP: self.global_step,
        metric_keys.MetricKeys.LOSS: 0.559831,
        metric_keys.MetricKeys.LOSS_MEAN: 0.559831,
        metric_keys.MetricKeys.ACCURACY: 1.0,
        metric_keys.MetricKeys.PREDICTION_MEAN: 0.429262,
        metric_keys.MetricKeys.LABEL_MEAN: 0.5,
        metric_keys.MetricKeys.ACCURACY_BASELINE: 0.5,
        # With default threshold of 0.5, the model is a perfect classifier.
        metric_keys.MetricKeys.RECALL: 1.0,
        metric_keys.MetricKeys.PRECISION: 1.0,
        # Positive example is scored above negative, so AUC = 1.0.
        metric_keys.MetricKeys.AUC: 1.0,
        metric_keys.MetricKeys.AUC_PR: 1.0,
    }
    self.assertAllClose(
        sorted_key_dict(expected_metrics), sorted_key_dict(eval_metrics))

  def testBinaryClassEvaluationMetricsSequential(self):

    def eval_input_fn():
      return {
          'price':
              tf.sparse.SparseTensor(
                  values=[10., 5., 2.],
                  indices=[[0, 0], [0, 1], [1, 0]],
                  dense_shape=[2, 2]),
      }, tf.sparse.SparseTensor(
          values=[0, 1, 0],
          indices=[[0, 0], [0, 1], [1, 0]],
          dense_shape=[2, 2])

    eval_metrics = self._testFromCheckpoint(
        eval_input_fn, return_sequences=True)

    # logits = [[[-1.4388], [-0.6033]],
    #            [[0.0197], [_]]]
    # probability = np.exp(logits) / (1 + np.exp(logits))
    #             = [[0.1917, 0.3536],
    #                [0.5049, _]]
    # labels = [[0, 1],
    #           [0, _]]
    # loss = -label * ln(p) - (1 - label) * ln(1 - p)
    # loss = [[0.2129,  1.0396],
    #         [0.7031, _]]
    # aggregated_loss = sum(loss) / 3
    # aggregated_loss = 0.6518
    # accuracy = 1/3
    # prediction_mean = mean(probability) = 0.3501
    expected_metrics = {
        tf.compat.v1.GraphKeys.GLOBAL_STEP: self.global_step,
        metric_keys.MetricKeys.LOSS: 0.651841,
        metric_keys.MetricKeys.LOSS_MEAN: 0.651841,
        metric_keys.MetricKeys.ACCURACY: 1.0 / 3,
        metric_keys.MetricKeys.PREDICTION_MEAN: 0.350085,
        metric_keys.MetricKeys.LABEL_MEAN: 1.0 / 3,
        metric_keys.MetricKeys.ACCURACY_BASELINE: 2.0 / 3,
        metric_keys.MetricKeys.RECALL: 0.0,
        metric_keys.MetricKeys.PRECISION: 0.0,
        metric_keys.MetricKeys.AUC: 0.5,
        metric_keys.MetricKeys.AUC_PR: 0.30685282,
    }
    self.assertAllClose(
        sorted_key_dict(expected_metrics), sorted_key_dict(eval_metrics))

  def testMultiClassEvaluationMetrics(self):

    def eval_input_fn():
      return {
          'price':
              tf.sparse.SparseTensor(
                  values=[10., 5., 2., 7.],
                  indices=[[0, 0], [0, 1], [1, 0], [1, 1]],
                  dense_shape=[2, 2]),
      }, [[0], [1]]

    self.dense_kernel = [[-1., 0.5, 0.2], [1., -0.3, 0.1]]
    self.dense_bias = [0.3, 0.4, 0.5]
    # Uses identical numbers to testMultiExampleMultiDim.
    # See that test for logits calculation.
    # logits = [[-0.603282, 0.777708, 0.569756],
    #           [-1.247356, 1.017018, 0.574481]]
    # logits_exp = exp(logits) / (1 + exp(logits))
    #            = [[0.547013, 2.176468, 1.767836],
    #               [0.287263, 2.764937, 1.776208]]
    # softmax_probabilities = logits_exp / logits_exp.sum()
    #                       = [[0.121793, 0.484596, 0.393611],
    #                          [0.059494, 0.572639, 0.367866]]
    # loss = -1. * log(softmax[label])
    #      = [[2.105432], [0.557500]]
    # sum_over_batch_size = (2.105432 + 0.557500)/2
    eval_metrics = self._testFromCheckpoint(eval_input_fn, n_classes=3)

    expected_metrics = {
        tf.compat.v1.GraphKeys.GLOBAL_STEP: self.global_step,
        metric_keys.MetricKeys.LOSS: 1.331465,
        metric_keys.MetricKeys.LOSS_MEAN: 1.331466,
        metric_keys.MetricKeys.ACCURACY: 0.5,
    }

    self.assertAllClose(
        sorted_key_dict(expected_metrics), sorted_key_dict(eval_metrics))

@test_util.run_all_in_graph_and_eager_modes
class RNNClassifierPredictionTest(tf.test.TestCase):

  def setUp(self):
    self.kernel = [[.1, -.2]]
    self.recurrent = [[.2, -.3], [.3, -.4]]
    self.bias = [.2, .5]
    self.dense_kernel = [[-1.], [1.]]
    self.dense_bias = [0.3]
    self.sequence_feature_columns = [
        tf.feature_column.sequence_numeric_column('price', shape=(1,))
    ]
    super(RNNClassifierPredictionTest, self).setUp()

  def _testFromCheckpoint(self, input_fn, **kwargs):
    create_checkpoint(
        kernel=self.kernel,
        recurrent=self.recurrent,
        bias=self.bias,
        dense_kernel=self.dense_kernel,
        dense_bias=self.dense_bias,
        global_step=100,
        model_dir=self.get_temp_dir())

    n_classes = 2
    if 'n_classes' in kwargs:
      n_classes = kwargs['n_classes']
      assert n_classes >= 2
    label_vocabulary = [
        'class_{}'.format(class_idx) for class_idx in range(n_classes)
    ]

    est = rnn.RNNClassifier(
        units=[2],
        sequence_feature_columns=self.sequence_feature_columns,
        label_vocabulary=label_vocabulary,
        model_dir=self.get_temp_dir(),
        **kwargs)
    return next(est.predict(input_fn))

  def testBinaryClassPredictions(self):
    # Uses identical numbers to testOneDimLogits.
    # See that test for logits calculation.
    # logits = [-0.603282]
    # logistic = exp(-0.6033) / (1 + exp(-0.6033)) = [0.353593]
    # probabilities = [0.646407, 0.353593]
    # class_ids = argmax(probabilities) = [0]
    predictions = self._testFromCheckpoint(_default_features_fn)
    self.assertAllClose([-0.603282],
                        predictions[prediction_keys.PredictionKeys.LOGITS])
    self.assertAllClose([0.353593],
                        predictions[prediction_keys.PredictionKeys.LOGISTIC])
    self.assertAllClose(
        [0.646407, 0.353593],
        predictions[prediction_keys.PredictionKeys.PROBABILITIES])
    self.assertAllClose([0],
                        predictions[prediction_keys.PredictionKeys.CLASS_IDS])
    self.assertEqual([b'class_0'],
                     predictions[prediction_keys.PredictionKeys.CLASSES])

  def testMultiClassPredictions(self):
    self.dense_kernel = [[-1., 0.5, 0.2], [1., -0.3, 0.1]]
    self.dense_bias = [0.3, 0.4, 0.5]
    # Uses identical numbers to testMultiDimLogits.
    # See that test for logits calculation.
    # logits = [-0.603282, 0.777708, 0.569756]
    # logits_exp = exp(logits) = [0.547013, 2.176468, 1.767836]
    # softmax_probabilities = logits_exp / logits_exp.sum()
    #                       = [0.121793, 0.484596, 0.393611]
    # class_ids = argmax(probabilities) = [1]
    predictions = self._testFromCheckpoint(_default_features_fn, n_classes=3)
    self.assertAllClose([-0.603282, 0.777708, 0.569756],
                        predictions[prediction_keys.PredictionKeys.LOGITS])
    self.assertAllClose(
        [0.121793, 0.484596, 0.393611],
        predictions[prediction_keys.PredictionKeys.PROBABILITIES])
    self.assertAllClose([1],
                        predictions[prediction_keys.PredictionKeys.CLASS_IDS])
    self.assertEqual([b'class_1'],
                     predictions[prediction_keys.PredictionKeys.CLASSES])

  def testBinaryClassPredictionsSequential(self):

    def predict_input_fn():
      return {
          'price':
              tf.sparse.SparseTensor(
                  values=[10., 5.],
                  indices=[[0, 0], [0, 1]],
                  dense_shape=[1, 2]),
      }

    # Same as first record of testBinaryClassEvaluationMetricsSequential.
    # Last step values are carried over.
    # logits = [[-1.4388], [-0.6033], [_]]
    # probabilities = np.exp(logits) / (1 + np.exp(logits))
    #               = [[0.8083, 0.1917], [0.6464, 0.3536], [_, _]]
    # class_ids = [[0], [0], [_]]
    # classes = [['class_0'], ['class_0'], [_]]
    predictions = self._testFromCheckpoint(
        predict_input_fn, return_sequences=True, sequence_mask='my-mask')
    self.assertAllEqual([1, 1], predictions['my-mask'])
    self.assertAllClose([[-1.438803], [-0.603282]],
                        predictions[prediction_keys.PredictionKeys.LOGITS])
    self.assertAllClose([[0.191731], [0.353593]],
                        predictions[prediction_keys.PredictionKeys.LOGISTIC])
    self.assertAllClose(
        [[0.808269, 0.191731], [0.646407, 0.353593]],
        predictions[prediction_keys.PredictionKeys.PROBABILITIES])
    self.assertAllClose([[0], [0]],
                        predictions[prediction_keys.PredictionKeys.CLASS_IDS])
    self.assertAllEqual([[b'class_0'], [b'class_0']],
                        predictions[prediction_keys.PredictionKeys.CLASSES])

class BaseRNNClassificationIntegrationTest(object):

  def setUp(self):
    col = tf.feature_column.sequence_categorical_column_with_hash_bucket(
        'tokens', hash_bucket_size=10)
    embed = tf.feature_column.embedding_column(col, dimension=2)
    self.feature_columns = [embed]
    super(BaseRNNClassificationIntegrationTest, self).setUp()

  def __init__(self, _create_estimator_fn):
    self._create_estimator_fn = _create_estimator_fn

  def _test_complete_flow(self,
                          train_input_fn,
                          eval_input_fn,
                          predict_input_fn,
                          n_classes,
                          batch_size,
                          optimizer='Adam'):
    cell_units = [4, 2]
    est = self._create_estimator_fn(
        self.feature_columns,
        n_classes,
        cell_units,
        self.get_temp_dir(),
        optimizer=optimizer)

    # TRAIN
    num_steps = 10
    est.train(train_input_fn, steps=num_steps)

    # EVALUATE
    scores = est.evaluate(eval_input_fn)
    self.assertEqual(num_steps, scores[tf.compat.v1.GraphKeys.GLOBAL_STEP])
    self.assertIn('loss', six.iterkeys(scores))

    # PREDICT
    predicted_proba = np.array([
        x[prediction_keys.PredictionKeys.PROBABILITIES]
        for x in est.predict(predict_input_fn)
    ])
    self.assertAllEqual((batch_size, n_classes), predicted_proba.shape)

    # EXPORT
    feature_spec = parsing_utils.classifier_parse_example_spec(
        self.feature_columns, label_key='label', label_dtype=tf.dtypes.int64)
    serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
        feature_spec)
    export_dir = est.export_savedmodel(tempfile.mkdtemp(),
                                       serving_input_receiver_fn)
    self.assertTrue(tf.compat.v1.gfile.Exists(export_dir))

  def _testNumpyInputFn(self, optimizer):
    """Tests complete flow with numpy_input_fn."""
    n_classes = 3
    batch_size = 10
    words = ['dog', 'cat', 'bird', 'the', 'a', 'sat', 'flew', 'slept']
    # Numpy only supports dense input, so all examples will have same length.
    # TODO(b/73160931): Update test when support for prepadded data exists.
    sequence_length = 3

    features = []
    for _ in range(batch_size):
      sentence = random.sample(words, sequence_length)
      features.append(sentence)

    x_data = np.array(features)
    y_data = np.random.randint(n_classes, size=batch_size)

    train_input_fn = numpy_io.numpy_input_fn(
        x={'tokens': x_data},
        y=y_data,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    eval_input_fn = numpy_io.numpy_input_fn(
        x={'tokens': x_data}, y=y_data, batch_size=batch_size, shuffle=False)
    predict_input_fn = numpy_io.numpy_input_fn(
        x={'tokens': x_data}, batch_size=batch_size, shuffle=False)

    self._test_complete_flow(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        n_classes=n_classes,
        batch_size=batch_size,
        optimizer=optimizer)

  def testNumpyInputFnStringOptimizer(self):
    self._testNumpyInputFn(optimizer='Adam')

  def testNumpyInputFnOptimizerInstance(self):
    self._testNumpyInputFn(optimizer=tf_keras.optimizers.Adam())

  def testParseExampleInputFn(self):
    """Tests complete flow with input_fn constructed from parse_example."""
    n_classes = 3
    batch_size = 10
    words = [b'dog', b'cat', b'bird', b'the', b'a', b'sat', b'flew', b'slept']

    _, examples_file = tempfile.mkstemp()
    writer = tf.io.TFRecordWriter(examples_file)
    for _ in range(batch_size):
      sequence_length = random.randint(1, len(words))
      sentence = random.sample(words, sequence_length)
      label = random.randint(0, n_classes - 1)
      example = example_pb2.Example(
          features=feature_pb2.Features(
              feature={
                  'tokens':
                      feature_pb2.Feature(
                          bytes_list=feature_pb2.BytesList(value=sentence)),
                  'label':
                      feature_pb2.Feature(
                          int64_list=feature_pb2.Int64List(value=[label])),
              }))
      writer.write(example.SerializeToString())
    writer.close()

    feature_spec = parsing_utils.classifier_parse_example_spec(
        self.feature_columns, label_key='label', label_dtype=tf.dtypes.int64)

    def _train_input_fn():
      dataset = tf.compat.v1.data.experimental.make_batched_features_dataset(
          examples_file, batch_size, feature_spec)
      return dataset.map(lambda features: (features, features.pop('label')))

    def _eval_input_fn():
      dataset = tf.compat.v1.data.experimental.make_batched_features_dataset(
          examples_file, batch_size, feature_spec, num_epochs=1)
      return dataset.map(lambda features: (features, features.pop('label')))

    def _predict_input_fn():
      dataset = tf.compat.v1.data.experimental.make_batched_features_dataset(
          examples_file, batch_size, feature_spec, num_epochs=1)

      def features_fn(features):
        features.pop('label')
        return features

      return dataset.map(features_fn)

    self._test_complete_flow(
        train_input_fn=_train_input_fn,
        eval_input_fn=_eval_input_fn,
        predict_input_fn=_predict_input_fn,
        n_classes=n_classes,
        batch_size=batch_size)

def _rnn_classifier_fn(feature_columns, n_classes, cell_units, model_dir,
                       optimizer):
  return rnn.RNNClassifier(
      units=cell_units,
      sequence_feature_columns=feature_columns,
      n_classes=n_classes,
      optimizer=optimizer,
      model_dir=model_dir)

@test_util.run_all_in_graph_and_eager_modes
class RNNClassifierIntegrationTest(BaseRNNClassificationIntegrationTest,
                                   tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    BaseRNNClassificationIntegrationTest.__init__(self, _rnn_classifier_fn)

def _rnn_classifier_dropout_fn(feature_columns, n_classes, cell_units,
                               model_dir, optimizer):

  def _rnn_cell_fn():
    cells = []
    for units in cell_units:
      cells.append(tf_keras.layers.SimpleRNNCell(units, dropout=0.5))
    return tf_keras.layers.StackedRNNCells(cells)

  return rnn.RNNClassifier(
      rnn_cell_fn=_rnn_cell_fn,
      sequence_feature_columns=feature_columns,
      n_classes=n_classes,
      optimizer=optimizer,
      model_dir=model_dir)

@test_util.run_all_in_graph_and_eager_modes
class RNNClassifierDropoutIntegrationTest(BaseRNNClassificationIntegrationTest,
                                          tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    BaseRNNClassificationIntegrationTest.__init__(self,
                                                  _rnn_classifier_dropout_fn)

def _rnn_estimator_fn(feature_columns, n_classes, cell_units, model_dir,
                      optimizer):
  return rnn.RNNEstimator(
      head=multi_head_lib.MultiClassHead(n_classes=n_classes),
      units=cell_units,
      sequence_feature_columns=feature_columns,
      optimizer=optimizer,
      model_dir=model_dir)

@test_util.run_all_in_graph_and_eager_modes
class RNNEstimatorIntegrationTest(BaseRNNClassificationIntegrationTest,
                                  tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    BaseRNNClassificationIntegrationTest.__init__(self, _rnn_estimator_fn)

@test_util.run_all_in_graph_and_eager_modes
class ModelFnTest(tf.test.TestCase):
  """Tests correctness of RNNEstimator's model function."""

  def _test_sequential_mask_in_head(self, mask=None):
    features = {
        'price':
            tf.sparse.SparseTensor(
                values=[10., 5., 4.],
                indices=[[0, 0], [0, 1], [1, 0]],
                dense_shape=[2, 2])
    }
    if mask:
      features['sequence_mask'] = ops.convert_to_tensor(mask)
    expected_mask = mask or [[1, 1], [1, 0]]

    sequence_feature_columns = [
        tf.feature_column.sequence_numeric_column('price', shape=(1,))
    ]

    mock_head = _get_mock_head()
    seq_head = seq_head_lib.SequentialHeadWrapper(
        mock_head, sequence_length_mask='sequence_mask')
    estimator = rnn.RNNEstimator(
        head=seq_head,
        units=[10],
        sequence_feature_columns=sequence_feature_columns,
        return_sequences=True)
    estimator.model_fn(
        features=features,
        labels=None,
        mode=model_fn.ModeKeys.PREDICT,
        config=None)
    passed_features = list(
        mock_head.create_estimator_spec.call_args)[1]['features']
    self.assertIn('sequence_mask', passed_features)
    sequence_mask = self.evaluate(passed_features['sequence_mask'])
    self.assertAllEqual(sequence_mask, expected_mask)

  def testSequentialMaskInHead(self):
    self._test_sequential_mask_in_head()

  def testSequentialMaskInHeadWithMasks(self):
    self._test_sequential_mask_in_head([[1, 1], [1, 1]])

if __name__ == '__main__':
  tf.test.main()


"""
Utilities to interact with XML ReleaseNotes.
"""
__author__ = "Marco Clemencic <marco.clemencic@cern.ch>"
__version__ = "$Id: ReleaseNotes.py,v 1.10 2009-09-14 15:56:58 ishapova Exp $"

# exported symbols
__all__ = [ "ReleaseNotes" ]

import xml.etree.ElementTree as ET
# define needed namespaces
ET._namespace_map["http://lhcb.cern.ch"] = "lhcb"
ET._namespace_map["http://www.w3.org/2001/XMLSchema-instance"] = "xsi"

import time

from _internals import log

def indent(elem, level=0, indentation = "  "):
    """
    Indent an ElementTree instance to allow pretty print.
    Code taken from http://effbot.org/zone/element-lib.htm and adapted to our
    needs.
    """
    i = "\n" + level*indentation
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + indentation
        for elem in elem:
            indent(elem, level+1, indentation)
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i

# helper function
_xel = lambda t: ET.QName("http://lhcb.cern.ch",t)

class ReleaseNotes(object):
    def __init__(self, filename):
        self.filename = filename
        # load XML
        self.tree = ET.parse(self.filename)
        import logging
        self.log = logging.getLogger("CondDB.Admin.ReleaseNotes")

    def _setDescription(self, node, description, patch = None):
        """
        Internal function to add the description to a node ("note" or "global_tag").
        """
        desc_el = None
        if description:
            desc_el = ET.SubElement(node,_xel("description"))
            for d in description:
                ET.SubElement(desc_el, _xel("di")).text = d
            if patch:
                ET.SubElement(desc_el, _xel("newpatch")).text = str(patch)
        return desc_el

    def _makeEntry(self, type_name, contributor, date = None):
        """
        Internal function to prepare a base entry ("note" or "global_tag").
        """
        if date is None:
            date = time.strftime("%Y-%m-%d")

        entry = ET.Element(_xel(type_name))
        ET.SubElement(entry,_xel("contributor")).text = contributor
        ET.SubElement(entry,_xel("date")).text = date

        return entry

    def _prependEntry(self, entry):
        """
        Internal function to add the entry at the beginning of the release notes
        (after the list of maintainers.)
        """
        # this should be safe enough: the first element is always "maintainer"
        pos = 0
        for i in self.tree.getroot().getchildren():
            if "maintainer" not in str(i.tag):
                break
            pos += 1

        self.tree.getroot().insert(pos, entry)

    def _findHomogeneous(self, contributor, date, tag, datatypes, patch = None, description = None):
        """
        This internal function does look-up in the release_note.xml file for possible previously added homogeneous
        tags entries, checking whether coincidence takes place for "tag" name, date, contributor and data types in case of
        a global tag processing, and date, contributor, patch and description for a local tag entry processing.
        It returns the object of found entry for further modifications. The homogeneous entry is looked only in
        the area of release_notes.xml which is after the last one sqldddb release.

        contributor: person providing the changes
        date: date in the format "YYYY-MM-DD", the current date is used if omitted
        tag: the name of the global tag
        patch: numeric id of the patch (on JIRA)
        description: list of comments, ["comment1","comment2"]
        """

        if date is None:
            date = time.strftime("%Y-%m-%d")

        # Accessing the root element of release notes DOM
        rootelement = self.tree.getroot()

        # Starting of iteration procedure for analyzing each element in the tree for duplications
        counter = 0 # counter for number of current entries found with proposed global tag name per one partition
        partitionList = [] # list of partitions found already containing proposed global tag name
        if not description:
            description = []

        for element in rootelement:
            if counter > 0: break
            if element.tag != _xel("sqldddb_tag"):
                bingoCommon = 0 # Number of common for global and local tags matchings found in a target entry
                bingoLT = 0 # Number of local tag matchings (for description text and patch number) found in a target entry
                found_datatypes = []
                if element.tag == _xel("global_tag") or element.tag == _xel("note"):
                    for subelement in element:
                        if subelement.text == contributor or subelement.text == date or subelement.text == tag:
                            bingoCommon += 1
                            continue
                        if subelement.tag == _xel("type"):
                            found_datatypes.append(subelement.text)
                        if bingoCommon >= 2 and subelement.tag == _xel("partition"):
                            for subsubelement in subelement:
                                if subsubelement.tag == _xel("name"): partitionList.append(subsubelement.text)
                        elif bingoCommon == 2 and subelement.tag == _xel("description"):
                            for subsubelement in subelement:
                                for line in description:
                                    if subsubelement.text == line:
                                        bingoLT += 1
                                if subsubelement.text == str(patch):
                                    bingoLT += 1
                    if (bingoCommon == 3 and set(found_datatypes) == set(datatypes)) or bingoLT == len(description)+1:
                        counter += 1
                        homogenEntry = element
                        self.log.info('Homogeneous entry with tag name "%s" is found to be used in the \
                        CondDB for %d partition(s): %s' % (tag,counter,partitionList))
                        self.log.info('Merging will be done in release notes for homogeneous entries \
                        for patch "%s" ...' % patch)
                    else: homogenEntry = None
                    if bingoCommon == 3 and not set(found_datatypes) == set(datatypes):
                        log.warning("Check the data types! Previous global tag is homogeneous with the \
new one except the data types. Previous data type set:%s. New one:%s" %(found_datatypes, datatypes))
                        log.info("Now we are going to create the global tag entry in release_notes.xml without merging.")
            else:
                homogenEntry = None
                break

        return homogenEntry

    def addNote(self, contributor, partitions, description, datatypes, date = None, patch = None, forceNewLT = False):
        """
        Add a basic entry to the release notes.

        contributor: person providing the changes
        partitions: dictionary in the format {"PARTITION":("tag",{'modified':["file1","file2"],'added':["file3","file4"]})}
        description: list of comments, ["comment1","comment2"]
        date: date in the format "YYYY-MM-DD", the current date is used if omitted
        patch: numeric id of the patch (on JIRA)
        forceNewLT: boolean flag to control unconditional creation of local tag entry in release_notes.xml
        """
        local_tag = partitions[partitions.keys()[0]][0]

        homogenEntry = self._findHomogeneous(contributor, date, local_tag, datatypes, patch, description)
        if forceNewLT or homogenEntry is None:
            note = self._makeEntry("note", contributor, date)
            # Adding new local tag entry element to the root element of release_notes.xml tree
            self._prependEntry(note)
        elif homogenEntry is not None:
            note = homogenEntry
            descr_ChildElem = note.find(str(_xel("description")))
            if descr_ChildElem is not None:
                note.remove(descr_ChildElem)

        # Attaching sub-sub-elements to new sub-element or element of a new local tag entry
        for part in partitions:
            tag, files = partitions[part]
            part_el = ET.SubElement(note,_xel("partition"))
            ET.SubElement(part_el,_xel("name")).text = part
            ET.SubElement(part_el,_xel("tag")).text = tag
            modified_files = list(files['modified'])
            modified_files.sort()
            added_files = list(files['added'])
            added_files.sort()
            for f in modified_files:
                ET.SubElement(part_el,_xel("file")).text = f
            for f in added_files:
                felem = ET.SubElement(part_el,_xel("file"))
                felem.text = f
                felem.set('status','new')

            if len(datatypes) != 0:
                for type in datatypes:
                    ET.SubElement(part_el,_xel("type")).text = type

        self._setDescription(note, description, patch)

    def addGlobalTag(self, contributor, tag, partitions, datatypes, description = None, date = None, patch = None, forceNewGT = False):
        """
        And an entry for a global tag to the release notes.

        contributor: person providing the changes
        tag: the name of the global tag
        partitions: dictionary in the format {"PARTITION":("base_tag",["tag1","tag2"])}
        description: list of comments, ["comment1","comment2"]
        date: date in the format "YYYY-MM-DD", the current date is used if omitted
        patch: numeric id of the patch (on JIRA)
        forceNewGT: flag for forcing creation of a new global tag entry in a release_notes.xml file
                    as a stand-alone entry in case any previous homogeneous entries are found.
        """

        homogenEntry = self._findHomogeneous(contributor,date,tag,datatypes)
        if forceNewGT or not homogenEntry:
            global_tag = self._makeEntry("global_tag", contributor, date)
            ET.SubElement(global_tag,_xel("tag")).text = tag
            if len(datatypes) != 0:
                for type in datatypes:
                    ET.SubElement(global_tag,_xel("type")).text = type
            # Adding new global tag entry element to the root element of release_notes.xml tree
            self._prependEntry(global_tag)
        elif homogenEntry:
            global_tag = homogenEntry
            descr_ChildElem = global_tag.find(str(_xel("description")))
            if descr_ChildElem:
                global_tag.remove(descr_ChildElem)

        # Attaching sub-sub-elements to new sub-element or element of a new global tag entry
        for part in partitions:
            base_tag, tags = partitions[part]
            part_el = ET.SubElement(global_tag,_xel("partition"))
            ET.SubElement(part_el,_xel("name")).text = part
            ET.SubElement(part_el,_xel("base")).text = base_tag
            tags = list(tags)
            for t in tags:
                ET.SubElement(part_el,_xel("tag")).text = t

        self._setDescription(global_tag, description, patch)

    def addHashToGlobalTag(self, hashSumObj, partition, gtag):
        """Function to add hash sum element to the global tag entry in the release notes DOM."""

        hashSumAdded = False
        # Accessing the root element of release notes DOM
        rootelement = self.tree.getroot()
        # Create new hash element
        hashelement = ET.Element(_xel("hash"),{'alg':hashSumObj.name})
        hashelement.text = hashSumObj.hexdigest()
        # Find needed global tag element
        for element in rootelement:
            if element.tag == _xel("global_tag"):
                for subelement1 in element:
                    if subelement1.tag == _xel("tag") and subelement1.text == gtag:
                        # Find subelement with requested partition.
                        for subelement2 in element:
                            if subelement2.tag == _xel("partition"):
                                for part_subelement2 in subelement2:
                                    if part_subelement2.tag == _xel("name") and part_subelement2.text == partition:
                                        # Check if the hash sum element with the
                                        # same as requested hash type already
                                        # exists and remove it
                                        oldHashSumElem = subelement2.find(str(_xel("hash")))
                                        if oldHashSumElem is not None and oldHashSumElem.get('alg') == hashSumObj.name:
                                            subelement2.remove(oldHashSumElem)
                                        # Add hash sum element to release notes DOM
                                        subelement2.append(hashelement)
                                        hashSumAdded = True
                                        return
        if hashSumAdded:
            self.log.info("Hash sum added to release notes for tag '%s' and  partition '%s'." %(gtag, partition))
        else:
            self.log.info("Hash sum wasn't added to release notes since no entry "
             "was found for tag '%s' and  partition '%s'." %(gtag, partition))

    def getGlobalTagElement(self, partition, GT_name):
        """Finds and returns the xml element of global tag.

         It looks for a global tag element with the name 'GT_name'
         in 'partition' partition.
         """
        rootelement = self.tree.getroot()
        for element in rootelement:
            if element.tag == _xel("global_tag"):
                for subelement1 in element:
                    if subelement1.tag == _xel("tag") and subelement1.text == GT_name:
                        # Checking partition
                        for subelement2 in element:
                            if subelement2.tag == _xel("partition"):
                                for part_subelement2 in subelement2:
                                    if (part_subelement2.tag == _xel("name") and
                                        part_subelement2.text == partition):
                                        return element

    def getGlobalTags(self, partition = None, datatype = None):
        """The function finds all global tags names for the requested partition and data type."""

        # Accessing the root element of release notes DOM
        rootelement = self.tree.getroot()
        gts = [] # global tags to return
        for element in rootelement:
            if element.tag == _xel("global_tag"):
                # Checking data type
                if datatype:
                    datatypes = []
                    for subelement in element:
                        if subelement.tag == _xel("type"):
                            datatypes.append(subelement.text)
                    if datatype not in datatypes:
                        continue

                # Checking partition
                for subelement2 in element:
                    if subelement2.tag == _xel("partition"):
                        for part_subelement2 in subelement2:
                            if part_subelement2.tag == _xel("name") and ( not partition or
                                                      part_subelement2.text == partition):
                                for subelement1 in element:
                                    if subelement1.tag == _xel("tag"):
                                        gts.append(subelement1.text)
        return gts

    def getDataType(self, partition, global_tag_name):
        """Get data types which correspond to a global tag in a specific partition"""

        datatypes = []
        for element in [e for e in self.getGlobalTagElement(partition, global_tag_name) if e.tag == _xel("type")]:
            datatype = element.text.lstrip('#')
            if datatype not in datatypes: datatypes.append(datatype)

        return datatypes

    def collectAllDataTypes(self, partitions = None):
        """Get all known data types for a partition or a list of them."""

        if partitions and type(partitions).__name__ != 'list': partitions = [partitions]

        rootelement = self.tree.getroot()

        datatypes = []
        for element in [e for e in rootelement if e.tag == _xel("global_tag")]:
            for subelement2 in [e2 for e2 in element if e2.tag == _xel("partition")]:
                for part_subelement2 in [ parte2 for parte2 in subelement2
                                         if parte2.tag == _xel("name") and
                                         (not partitions or parte2.text in partitions)]:
                    for subelement in [sube for sube in element if sube.tag == _xel("type")]:
                        datatype = subelement.text.lstrip('#')
                        if datatype not in datatypes: datatypes.append(datatype)
        return datatypes

    def write(self, filename = None , encoding = "utf-8"):
        if filename is None:
            filename = self.filename

        from StringIO import StringIO
        f = StringIO()

        # FIXME: (MCl) work-around for a limitation of E.T. in Python 2.5
        xml_hdr = '<?xml version="1.0" encoding="UTF-8"?>\n<?xml-stylesheet type="text/xsl" href="release_notes.xsl"?>\n'
        f.write(xml_hdr)
        indent(self.tree.getroot(), indentation = "\t")
        self.tree.write(f,encoding)

        open(filename,"w").write(f.getvalue()+"\n")