import copy
from absl.testing import parameterized
import numpy as np

from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_spec
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops.ragged import row_partition
from tensorflow.python.ops.ragged.row_partition import RowPartition
from tensorflow.python.ops.ragged.row_partition import RowPartitionSpec
from tensorflow.python.platform import googletest

def _get_specified_row_partition():
  """Needed for merge_with_spec tests. Normally, nvals isn't set."""
  return RowPartition(
      row_splits=constant_op.constant([0, 3, 8], dtype=dtypes.int64),
      nrows=constant_op.constant(2, dtype=dtypes.int64),
      nvals=constant_op.constant(8, dtype=dtypes.int64),
      internal=row_partition._row_partition_factory_key)

@test_util.run_all_in_graph_and_eager_modes
class RowPartitionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
  #=============================================================================
  # RowPartition class docstring examples
  #=============================================================================

  def testClassDocStringExamples(self):
    # From section: "Component Tensors"
    rp = RowPartition.from_row_splits(row_splits=[0, 4, 4, 7, 8, 8])
    self.assertAllEqual(rp.row_splits(), [0, 4, 4, 7, 8, 8])
    del rp

    # From section: "Alternative Row-Partitioning Schemes"
    rt1 = RowPartition.from_row_splits(row_splits=[0, 4, 4, 7, 8, 8])
    rt2 = RowPartition.from_row_lengths(row_lengths=[4, 0, 3, 1, 0])
    rt3 = RowPartition.from_value_rowids(
        value_rowids=[0, 0, 0, 0, 2, 2, 2, 3], nrows=5)
    rt4 = RowPartition.from_row_starts(row_starts=[0, 4, 4, 7, 8], nvals=8)
    rt5 = RowPartition.from_row_limits(row_limits=[4, 4, 7, 8, 8])
    for rp in (rt1, rt2, rt3, rt4, rt5):
      self.assertAllEqual(rp.row_splits(), [0, 4, 4, 7, 8, 8])
    del rt1, rt2, rt3, rt4, rt5

    # From section: "Multiple Ragged Dimensions"
    inner_rt = RowPartition.from_row_splits(row_splits=[0, 4, 4, 7, 8, 8])
    outer_rt = RowPartition.from_row_splits(row_splits=[0, 3, 3, 5])
    del inner_rt, outer_rt

  #=============================================================================
  # RowPartition Constructor (private)
  #=============================================================================

  def testRowPartitionConstruction(self):
    row_splits = constant_op.constant([0, 2, 2, 5, 6, 7], dtypes.int64)
    rp = RowPartition(
        row_splits=row_splits,
        internal=row_partition._row_partition_factory_key)
    self.assertAllEqual(rp.row_splits(), [0, 2, 2, 5, 6, 7])

  def testRowPartitionConstructionErrors(self):
    row_splits = constant_op.constant([0, 2, 2, 5, 6, 7], dtypes.int64)

    with self.assertRaisesRegex(ValueError,
                                'RowPartition constructor is private'):
      RowPartition(row_splits=row_splits)

    with self.assertRaisesRegex(TypeError,
                                'Row-partitioning argument must be a Tensor'):
      RowPartition(
          row_splits=[0, 2, 2, 5, 6, 7],
          internal=row_partition._row_partition_factory_key)

    with self.assertRaisesRegex(ValueError, r'Shape \(6, 1\) must have rank 1'):
      RowPartition(
          row_splits=array_ops.expand_dims(row_splits, 1),
          internal=row_partition._row_partition_factory_key)

    with self.assertRaisesRegex(TypeError,
                                'Cached value must be a Tensor or None.'):
      RowPartition(
          row_splits=row_splits,
          row_lengths=[2, 3, 4],
          internal=row_partition._row_partition_factory_key)

    with self.assertRaisesRegex(ValueError, 'Inconsistent dtype'):
      RowPartition(
          row_splits=constant_op.constant([0, 3], dtypes.int64),
          nrows=constant_op.constant(1, dtypes.int32),
          internal=row_partition._row_partition_factory_key)

  #=============================================================================
  # RowPartition Factory Ops
  #=============================================================================

  def testFromValueRowIdsWithDerivedNRows(self):
    # nrows is known at graph creation time.
    value_rowids = constant_op.constant([0, 0, 2, 2, 2, 3, 4], dtypes.int64)
    # TODO(martinz): add nrows
    rp = RowPartition.from_value_rowids(value_rowids, validate=False)
    self.assertEqual(rp.dtype, dtypes.int64)

    rp_row_splits = rp.row_splits()
    rp_value_rowids = rp.value_rowids()
    rp_nrows = rp.nrows()

    self.assertIs(rp_value_rowids, value_rowids)  # value_rowids
    self.assertAllEqual(rp_value_rowids, value_rowids)
    self.assertAllEqual(rp_nrows, 5)
    self.assertAllEqual(rp_row_splits, [0, 2, 2, 5, 6, 7])

  def testFromValueRowIdsWithDerivedNRowsDynamic(self):
    # nrows is not known at graph creation time.
    value_rowids = constant_op.constant([0, 0, 2, 2, 2, 3, 4], dtypes.int64)
    value_rowids = array_ops.placeholder_with_default(value_rowids, shape=None)

    rp = RowPartition.from_value_rowids(value_rowids, validate=False)

    rp_value_rowids = rp.value_rowids()
    rp_nrows = rp.nrows()

    self.assertIs(rp_value_rowids, value_rowids)  # value_rowids
    self.assertAllEqual(rp_value_rowids, value_rowids)
    self.assertAllEqual(rp_nrows, 5)

  def testFromValueRowIdsWithExplicitNRows(self):
    value_rowids = constant_op.constant([0, 0, 2, 2, 2, 3, 4], dtypes.int64)
    nrows = constant_op.constant(7, dtypes.int64)

    rp = RowPartition.from_value_rowids(value_rowids, nrows, validate=False)

    rp_value_rowids = rp.value_rowids()
    rp_nrows = rp.nrows()
    rp_row_splits = rp.row_splits()

    self.assertIs(rp_value_rowids, value_rowids)  # value_rowids
    self.assertIs(rp_nrows, nrows)  # nrows
    self.assertAllEqual(rp_row_splits, [0, 2, 2, 5, 6, 7, 7, 7])

  def testFromValueRowIdsWithExplicitNRowsEqualToDefault(self):
    value_rowids = constant_op.constant([0, 0, 2, 2, 2, 3, 4], dtypes.int64)
    nrows = constant_op.constant(5, dtypes.int64)

    rp = RowPartition.from_value_rowids(value_rowids, nrows, validate=False)

    rp_value_rowids = rp.value_rowids()
    rp_nrows = rp.nrows()
    rp_row_splits = rp.row_splits()

    self.assertIs(rp_value_rowids, value_rowids)  # value_rowids
    self.assertIs(rp_nrows, nrows)  # nrows
    self.assertAllEqual(rp_value_rowids, value_rowids)
    self.assertAllEqual(rp_nrows, nrows)
    self.assertAllEqual(rp_row_splits, [0, 2, 2, 5, 6, 7])

  def testFromValueRowIdsWithEmptyValues(self):
    rp = RowPartition.from_value_rowids([])
    rp_nrows = rp.nrows()
    self.assertEqual(rp.dtype, dtypes.int64)
    self.assertEqual(rp.value_rowids().shape.as_list(), [0])
    self.assertAllEqual(rp_nrows, 0)

  def testFromRowSplits(self):
    row_splits = constant_op.constant([0, 2, 2, 5, 6, 7], dtypes.int64)

    rp = RowPartition.from_row_splits(row_splits, validate=False)
    self.assertEqual(rp.dtype, dtypes.int64)

    rp_row_splits = rp.row_splits()
    rp_nrows = rp.nrows()

    self.assertIs(rp_row_splits, row_splits)
    self.assertAllEqual(rp_nrows, 5)

  def testFromRowSplitsWithDifferentSplitTypes(self):
    splits1 = [0, 2, 2, 5, 6, 7]
    splits2 = np.array([0, 2, 2, 5, 6, 7], np.int64)
    splits3 = np.array([0, 2, 2, 5, 6, 7], np.int32)
    splits4 = constant_op.constant([0, 2, 2, 5, 6, 7], dtypes.int64)
    splits5 = constant_op.constant([0, 2, 2, 5, 6, 7], dtypes.int32)
    rt1 = RowPartition.from_row_splits(splits1)
    rt2 = RowPartition.from_row_splits(splits2)
    rt3 = RowPartition.from_row_splits(splits3)
    rt4 = RowPartition.from_row_splits(splits4)
    rt5 = RowPartition.from_row_splits(splits5)
    self.assertEqual(rt1.row_splits().dtype, dtypes.int64)
    self.assertEqual(rt2.row_splits().dtype, dtypes.int64)
    self.assertEqual(rt3.row_splits().dtype, dtypes.int32)
    self.assertEqual(rt4.row_splits().dtype, dtypes.int64)
    self.assertEqual(rt5.row_splits().dtype, dtypes.int32)

  def testFromRowSplitsWithEmptySplits(self):
    err_msg = 'row_splits tensor may not be empty'
    with self.assertRaisesRegex(ValueError, err_msg):
      RowPartition.from_row_splits([])

  def testFromRowStarts(self):
    nvals = constant_op.constant(7)
    row_starts = constant_op.constant([0, 2, 2, 5, 6], dtypes.int64)

    rp = RowPartition.from_row_starts(row_starts, nvals, validate=False)
    self.assertEqual(rp.dtype, dtypes.int64)

    rp_row_starts = rp.row_starts()
    rp_row_splits = rp.row_splits()
    rp_nrows = rp.nrows()

    self.assertAllEqual(rp_nrows, 5)
    self.assertAllEqual(rp_row_starts, row_starts)
    self.assertAllEqual(rp_row_splits, [0, 2, 2, 5, 6, 7])

  def testFromRowLimits(self):
    row_limits = constant_op.constant([2, 2, 5, 6, 7], dtypes.int64)

    rp = RowPartition.from_row_limits(row_limits, validate=False)
    self.assertEqual(rp.dtype, dtypes.int64)

    rp_row_limits = rp.row_limits()
    rp_row_splits = rp.row_splits()
    rp_nrows = rp.nrows()

    self.assertAllEqual(rp_nrows, 5)
    self.assertAllEqual(rp_row_limits, row_limits)
    self.assertAllEqual(rp_row_splits, [0, 2, 2, 5, 6, 7])

  def testFromRowLengths(self):
    row_lengths = constant_op.constant([2, 0, 3, 1, 1], dtypes.int64)

    rp = RowPartition.from_row_lengths(row_lengths, validate=False)
    self.assertEqual(rp.dtype, dtypes.int64)

    rp_row_lengths = rp.row_lengths()
    rp_nrows = rp.nrows()

    self.assertIs(rp_row_lengths, row_lengths)  # nrows
    self.assertAllEqual(rp_nrows, 5)
    self.assertAllEqual(rp_row_lengths, row_lengths)

  def testFromUniformRowLength(self):
    nvals = 16
    a1 = RowPartition.from_uniform_row_length(
        nvals=nvals, uniform_row_length=2)
    self.assertAllEqual(a1.uniform_row_length(), 2)
    self.assertAllEqual(a1.nrows(), 8)

  def testFromUniformRowLengthWithEmptyValues(self):
    a = RowPartition.from_uniform_row_length(
        nvals=0, uniform_row_length=0, nrows=10)
    self.assertEqual(self.evaluate(a.nvals()), 0)
    self.assertEqual(self.evaluate(a.nrows()), 10)

  def testFromUniformRowLengthWithPlaceholders1(self):
    nvals = array_ops.placeholder_with_default(
        constant_op.constant(6, dtype=dtypes.int64), None)
    rt1 = RowPartition.from_uniform_row_length(
        nvals=nvals, uniform_row_length=3)
    const_nvals1 = self.evaluate(rt1.nvals())
    self.assertEqual(const_nvals1, 6)

  def testFromUniformRowLengthWithPlaceholders2(self):
    nvals = array_ops.placeholder_with_default(6, None)
    ph_rowlen = array_ops.placeholder_with_default(3, None)
    rt2 = RowPartition.from_uniform_row_length(
        nvals=nvals, uniform_row_length=ph_rowlen)
    const_nvals2 = self.evaluate(rt2.nvals())
    self.assertEqual(const_nvals2, 6)

  def testFromValueRowIdsWithBadNRows(self):
    value_rowids = constant_op.constant([0, 0, 2, 2, 2, 3, 4], dtypes.int64)
    nrows = constant_op.constant(5, dtypes.int64)

    with self.assertRaisesRegex(ValueError, r'Expected nrows >= 0; got -2'):
      RowPartition.from_value_rowids(
          value_rowids=array_ops.placeholder_with_default(value_rowids, None),
          nrows=-2)

    with self.assertRaisesRegex(
        ValueError, r'Expected nrows >= value_rowids\[-1\] \+ 1; got nrows=2, '
        r'value_rowids\[-1\]=4'):
      RowPartition.from_value_rowids(value_rowids=value_rowids, nrows=2)

    with self.assertRaisesRegex(
        ValueError, r'Expected nrows >= value_rowids\[-1\] \+ 1; got nrows=4, '
        r'value_rowids\[-1\]=4'):
      RowPartition.from_value_rowids(value_rowids=value_rowids, nrows=4)

    with self.assertRaisesRegex(ValueError, r'Shape \(7, 1\) must have rank 1'):
      RowPartition.from_value_rowids(
          value_rowids=array_ops.expand_dims(value_rowids, 1), nrows=nrows)

    with self.assertRaisesRegex(ValueError, r'Shape \(1,\) must have rank 0'):
      RowPartition.from_value_rowids(
          value_rowids=value_rowids, nrows=array_ops.expand_dims(nrows, 0))

  #=============================================================================
  # RowPartition.__str__
  #=============================================================================
  def testRowPartitionStr(self):
    row_splits = [0, 2, 5, 6, 6, 7]
    rp = RowPartition.from_row_splits(row_splits, validate=False)
    if context.executing_eagerly():
      expected_repr = 'tf.RowPartition(row_splits=[0 2 5 6 6 7])'
    else:
      expected_repr = ('tf.RowPartition(row_splits='
                       'Tensor("RowPartitionFromRowSplits/row_splits:0", '
                       'shape=(6,), dtype=int64))')
    self.assertEqual(repr(rp), expected_repr)
    self.assertEqual(str(rp), expected_repr)

  def testRowPartitionStrUniformRowLength(self):
    rp = RowPartition.from_uniform_row_length(5, nvals=10, nrows=2)
    if context.executing_eagerly():
      expected_repr = ('tf.RowPartition(nrows=2, uniform_row_length=5)')
    else:
      expected_repr = (
          'tf.RowPartition(nrows='
          'Tensor("RowPartitionFromUniformRowLength/'
          'nrows:0", shape=(), dtype=int64), '
          'uniform_row_length=Tensor("RowPartitionFromUniformRowLength/'
          'uniform_row_length:0", shape=(), dtype=int64))')
    self.assertEqual(repr(rp), expected_repr)
    self.assertEqual(str(rp), expected_repr)

  @parameterized.parameters([
      # from_value_rowids
      {
          'descr': 'bad rank for value_rowids',
          'factory': RowPartition.from_value_rowids,
          'value_rowids': [[1, 2], [3, 4]],
          'nrows': 10
      },
      {
          'descr': 'bad rank for nrows',
          'factory': RowPartition.from_value_rowids,
          'value_rowids': [1, 2, 3, 4],
          'nrows': [10]
      },
      {
          'descr': 'negative value_rowid',
          'factory': RowPartition.from_value_rowids,
          'value_rowids': [-5, 2, 3, 4],
          'nrows': 10
      },
      {
          'descr': 'non-monotonic-increasing value_rowid',
          'factory': RowPartition.from_value_rowids,
          'value_rowids': [4, 3, 2, 1],
          'nrows': 10
      },
      {
          'descr': 'value_rowid > nrows',
          'factory': RowPartition.from_value_rowids,
          'value_rowids': [1, 2, 3, 4],
          'nrows': 2
      },

      # from_row_splits
      {
          'descr': 'bad rank for row_splits',
          'factory': RowPartition.from_row_splits,
          'row_splits': [[1, 2], [3, 4]]
      },
      {
          'descr': 'row_splits[0] != 0',
          'factory': RowPartition.from_row_splits,
          'row_splits': [2, 3, 4]
      },
      {
          'descr': 'non-monotonic-increasing row_splits',
          'factory': RowPartition.from_row_splits,
          'row_splits': [0, 3, 2, 4]
      },

      # from_row_lengths
      {
          'descr': 'bad rank for row_lengths',
          'factory': RowPartition.from_row_lengths,
          'row_lengths': [[1, 2], [1, 0]]
      },
      {
          'descr': 'negatve row_lengths',
          'factory': RowPartition.from_row_lengths,
          'row_lengths': [3, -1, 2]
      },

      # from_row_starts
      {
          'descr': 'bad rank for row_starts',
          'factory': RowPartition.from_row_starts,
          'nvals': 2,
          'row_starts': [[1, 2], [3, 4]]
      },
      {
          'descr': 'row_starts[0] != 0',
          'factory': RowPartition.from_row_starts,
          'nvals': 5,
          'row_starts': [2, 3, 4]
      },
      {
          'descr': 'non-monotonic-increasing row_starts',
          'factory': RowPartition.from_row_starts,
          'nvals': 4,
          'row_starts': [0, 3, 2, 4]
      },
      {
          'descr': 'row_starts[0] > nvals',
          'factory': RowPartition.from_row_starts,
          'nvals': 4,
          'row_starts': [0, 2, 3, 5]
      },

      # from_row_limits
      {
          'descr': 'bad rank for row_limits',
          'factory': RowPartition.from_row_limits,
          'row_limits': [[1, 2], [3, 4]]
      },
      {
          'descr': 'row_limits[0] < 0',
          'factory': RowPartition.from_row_limits,
          'row_limits': [-1, 3, 4]
      },
      {
          'descr': 'non-monotonic-increasing row_limits',
          'factory': RowPartition.from_row_limits,
          'row_limits': [0, 3, 2, 4]
      },

      # from_uniform_row_length
      {
          'descr': 'rowlen * nrows != nvals (1)',
          'factory': RowPartition.from_uniform_row_length,
          'nvals': 5,
          'uniform_row_length': 3
      },
      {
          'descr': 'rowlen * nrows != nvals (2)',
          'factory': RowPartition.from_uniform_row_length,
          'nvals': 5,
          'uniform_row_length': 6
      },
      {
          'descr': 'rowlen * nrows != nvals (3)',
          'factory': RowPartition.from_uniform_row_length,
          'nvals': 6,
          'uniform_row_length': 3,
          'nrows': 3
      },
      {
          'descr': 'rowlen must be a scalar',
          'factory': RowPartition.from_uniform_row_length,
          'nvals': 4,
          'uniform_row_length': [2]
      },
      {
          'descr': 'rowlen must be nonnegative',
          'factory': RowPartition.from_uniform_row_length,
          'nvals': 4,
          'uniform_row_length': -1
      },
  ])
  def testFactoryValidation(self, descr, factory, **kwargs):
    # When input tensors have shape information, some of these errors will be
    # detected statically.
    with self.assertRaises((errors.InvalidArgumentError, ValueError)):
      partition = factory(**kwargs)
      self.evaluate(partition.row_splits())

    # Remove shape information (by wrapping tensors in placeholders), and check
    # that we detect the errors when the graph is run.
    if not context.executing_eagerly():

      def wrap_arg(v):
        return array_ops.placeholder_with_default(
            constant_op.constant(v, dtype=dtypes.int64),
            tensor_shape.TensorShape(None))

      kwargs = dict((k, wrap_arg(v)) for (k, v) in kwargs.items())

      with self.assertRaises(errors.InvalidArgumentError):
        partition = factory(**kwargs)
        self.evaluate(partition.row_splits())

  @parameterized.named_parameters([
      ('FromRowSplits', lambda: RowPartition.from_row_splits([0, 2, 8]),
       ['row_splits']),
      ('FromRowLengths', lambda: RowPartition.from_row_lengths([3, 0, 8]),
       ['row_splits', 'row_lengths']),
      ('FromValueRowIds',
       lambda: RowPartition.from_value_rowids([0, 0, 3, 4, 4, 4]),
       ['row_splits', 'value_rowids', 'row_lengths', 'nrows']),
      ('FromRowStarts',
       lambda: RowPartition.from_row_starts([0, 3, 7], nvals=10),
       ['row_splits']),
      ('FromRowLimits', lambda: RowPartition.from_row_limits([3, 7, 10]),
       ['row_splits']),
  ])
  def testPrecomputedSplits(self, rp_factory, expected_encodings):
    rp = rp_factory()
    self.assertEqual(rp._has_precomputed_row_splits(),
                     'row_splits' in expected_encodings)
    self.assertEqual(rp._has_precomputed_row_lengths(),
                     'row_lengths' in expected_encodings)
    self.assertEqual(rp._has_precomputed_value_rowids(),
                     'value_rowids' in expected_encodings)
    self.assertEqual(rp._has_precomputed_nrows(), 'nrows' in expected_encodings)

  def testWithPrecomputedSplits(self):
    rp = RowPartition.from_row_splits([0, 2, 8])

    rp_with_row_splits = rp._with_precomputed_row_splits()
    self.assertTrue(rp_with_row_splits._has_precomputed_row_splits())

    self.assertFalse(rp._has_precomputed_row_lengths())
    rp_with_row_lengths = rp._with_precomputed_row_lengths()
    self.assertTrue(rp_with_row_lengths._has_precomputed_row_lengths())

    self.assertFalse(rp._has_precomputed_value_rowids())
    rp_with_value_rowids = rp._with_precomputed_value_rowids()
    self.assertTrue(rp_with_value_rowids._has_precomputed_value_rowids())

    self.assertFalse(rp._has_precomputed_nrows())
    rp_with_nrows = rp._with_precomputed_nrows()
    self.assertTrue(rp_with_nrows._has_precomputed_nrows())

    self.assertFalse(rp._has_precomputed_nvals())
    rp_with_nvals = rp._with_precomputed_nvals()
    self.assertTrue(rp_with_nvals._has_precomputed_nvals())

  @parameterized.named_parameters([
      dict(
          testcase_name='FromRowSplitsAndRowSplits',
          x=lambda: RowPartition.from_row_splits([0, 3, 8]),
          y=lambda: RowPartition.from_row_splits([0, 3, 8]),
          expected_encodings=['row_splits']),
      dict(
          testcase_name='FromRowSplitsAndUniformRowLength',
          x=lambda: RowPartition.from_row_splits([0, 3, 6]),
          y=lambda: RowPartition.from_uniform_row_length(3, nvals=6),
          expected_encodings=['row_splits', 'uniform_row_length', 'nrows']),
      dict(
          testcase_name='FromRowSplitsAndRowLengths',
          x=lambda: RowPartition.from_row_splits([0, 3, 8]),
          y=lambda: RowPartition.from_row_lengths([3, 5]),
          expected_encodings=['row_splits', 'row_lengths']),
      dict(
          testcase_name='FromRowSplitsAndValueRowIds',
          x=lambda: RowPartition.from_row_splits([0, 3, 8]),
          y=lambda: RowPartition.from_value_rowids([0, 0, 0, 1, 1, 1, 1, 1]),
          expected_encodings=[
              'row_splits', 'row_lengths', 'value_rowids', 'nrows'
          ]),
      dict(
          testcase_name='FromRowSplitsAndRowSplitsPlusNRows',
          x=lambda: RowPartition.from_row_splits([0, 3, 8]),
          y=lambda: RowPartition.from_row_splits([0, 3, 8]).
          _with_precomputed_nrows(),
          expected_encodings=['row_splits', 'nrows']),
  ])
  def testMergePrecomputedEncodings(self, x, y, expected_encodings):
    x = x()
    y = y()
    for validate in (True, False):
      result = x._merge_precomputed_encodings(y, validate)
      self.assertEqual(result._has_precomputed_row_splits(),
                       'row_splits' in expected_encodings)
      self.assertEqual(result._has_precomputed_row_lengths(),
                       'row_lengths' in expected_encodings)
      self.assertEqual(result._has_precomputed_value_rowids(),
                       'value_rowids' in expected_encodings)
      self.assertEqual(result._has_precomputed_nrows(),
                       'nrows' in expected_encodings)
      self.assertEqual(result.uniform_row_length() is not None,
                       'uniform_row_length' in expected_encodings)
      for r in (x, y):
        if (r._has_precomputed_row_splits() and
            result._has_precomputed_row_splits()):
          self.assertAllEqual(r.row_splits(), result.row_splits())
        if (r._has_precomputed_row_lengths() and
            result._has_precomputed_row_lengths()):
          self.assertAllEqual(r.row_lengths(), result.row_lengths())
        if (r._has_precomputed_value_rowids() and
            result._has_precomputed_value_rowids()):
          self.assertAllEqual(r.value_rowids(), result.value_rowids())
        if r._has_precomputed_nrows() and result._has_precomputed_nrows():
          self.assertAllEqual(r.nrows(), result.nrows())
        if (r.uniform_row_length() is not None and
            result.uniform_row_length() is not None):
          self.assertAllEqual(r.uniform_row_length(),
                              result.uniform_row_length())

  def testMergePrecomputedEncodingsFastPaths(self):
    # Same object: x gets returned as-is.
    x = RowPartition.from_row_splits([0, 3, 8, 8])
    self.assertIs(x._merge_precomputed_encodings(x), x)

    # Same encoding tensor objects: x gets returned as-is.
    y = RowPartition.from_row_splits(x.row_splits(), validate=False)
    self.assertIs(x._merge_precomputed_encodings(y), x)

  def testMergePrecomputedEncodingsWithMatchingTensors(self):
    # The encoding tensors for `a` are a superset of the encoding tensors
    # for `b`, and where they overlap, they the same tensor objects.
    a = RowPartition.from_value_rowids([0, 0, 3, 4, 4, 4])
    b = RowPartition.from_row_splits(a.row_splits(), validate=False)
    self.assertIs(a._merge_precomputed_encodings(b), a)
    self.assertIs(b._merge_precomputed_encodings(a), a)
    self.assertIsNot(a, b)

  @parameterized.named_parameters([
      dict(
          testcase_name='RowSplitMismatch',
          x=lambda: RowPartition.from_row_splits([0, 3, 8]),
          y=lambda: RowPartition.from_row_splits([0, 3, 8, 9]),
          message='incompatible row_splits'),
      dict(
          testcase_name='RowLengthMismatch',
          x=lambda: RowPartition.from_row_lengths([2, 0, 2]),
          y=lambda: RowPartition.from_row_lengths([2, 0, 2, 1]),
          message='incompatible row_splits'),  # row_splits is checked first
      dict(
          testcase_name='ValueRowIdMismatch',
          x=lambda: RowPartition.from_value_rowids([0, 3, 3, 4]),
          y=lambda: RowPartition.from_value_rowids([0, 3, 4]),
          message='incompatible value_rowids'),
  ])
  def testMergePrecomputedEncodingStaticErrors(self, x, y, message):
    if context.executing_eagerly():
      return
    # Errors that are caught by static shape checks.
    x = x()
    y = y()
    with self.assertRaisesRegex(ValueError, message):
      x._merge_precomputed_encodings(y).row_splits()
    with self.assertRaisesRegex(ValueError, message):
      y._merge_precomputed_encodings(x).row_splits()

  @parameterized.named_parameters([
      dict(
          testcase_name='NRowsMismatchAlt',
          x=lambda: RowPartition.from_uniform_row_length(5, nrows=4, nvals=20),
          y=lambda: RowPartition.from_uniform_row_length(5, nrows=3, nvals=15),
          message='incompatible nrows'),
      dict(
          testcase_name='UniformRowLengthMismatch',
          x=lambda: RowPartition.from_uniform_row_length(5, nvals=20),
          y=lambda: RowPartition.from_uniform_row_length(2, nvals=8),
          message='incompatible (nvals|uniform_row_length)'),
      dict(
          testcase_name='RowSplitMismatch',
          x=lambda: RowPartition.from_row_splits([0, 3, 8]),
          y=lambda: RowPartition.from_row_splits([0, 5, 8]),
          message='incompatible row_splits'),
      dict(
          testcase_name='RowLengthMismatch',
          x=lambda: RowPartition.from_row_lengths([2, 0, 2]),
          y=lambda: RowPartition.from_row_lengths([0, 0, 2]),
          message='incompatible (row_splits|nvals)'),
      dict(
          testcase_name='ValueRowIdMismatch',
          x=lambda: RowPartition.from_value_rowids([0, 3, 3]),
          y=lambda: RowPartition.from_value_rowids([0, 0, 3]),
          message='incompatible row_splits'),  # row_splits is checked first
  ])
  def testMergePrecomputedEncodingRuntimeErrors(self, x, y, message):
    # Errors that are caught by runtime value checks.
    x = x()
    y = y()
    with self.assertRaisesRegex(errors.InvalidArgumentError, message):
      self.evaluate(x._merge_precomputed_encodings(y).row_splits())
    with self.assertRaisesRegex(errors.InvalidArgumentError, message):
      self.evaluate(y._merge_precomputed_encodings(x).row_splits())

  @parameterized.named_parameters([
      # It throws the right error, but it still complains.
      dict(
          testcase_name='NRowsMismatch',
          x=lambda: RowPartition.from_uniform_row_length(5, nvals=20),
          y=lambda: RowPartition.from_uniform_row_length(5, nvals=15),
          message='incompatible nvals',
          emessage='incompatible nrows'),
  ])
  def testMergePrecomputedEncodingStaticErrors2(self, x, y,
                                                message, emessage):
    # Message error and type varies depending upon eager execution.
    x = x()
    y = y()

    error_type = errors_impl.InvalidArgumentError
    expected_message = emessage if context.executing_eagerly() else message
    with self.assertRaisesRegex(error_type, expected_message):
      self.evaluate(x._merge_precomputed_encodings(y).row_splits())
    with self.assertRaisesRegex(error_type, expected_message):
      self.evaluate(y._merge_precomputed_encodings(x).row_splits())

  @parameterized.named_parameters([
      dict(
          testcase_name='NoneSpecified',
          rp=(lambda: RowPartition.from_row_splits([0, 3, 8])),
          spec=RowPartitionSpec(nrows=None, nvals=None, dtype=dtypes.int64)),
      dict(
          testcase_name='NRowsSpecified',
          rp=(lambda: RowPartition.from_row_splits([0, 3, 8])),
          spec=RowPartitionSpec(nrows=2, nvals=None, dtype=dtypes.int64)),
      dict(
          testcase_name='NValsSpecified',
          rp=_get_specified_row_partition,
          spec=RowPartitionSpec(nrows=None, nvals=8, dtype=dtypes.int64))
  ])
  def testMergeWithSpecNoop(self, rp, spec):
    rp = rp()
    actual = rp._merge_with_spec(spec)
    self.assertAllEqual(actual.row_splits(), rp.row_splits())
    self.assertAllEqual(actual.static_nrows, rp.static_nrows)
    self.assertAllEqual(actual.static_nvals, rp.static_nvals)

  @parameterized.named_parameters([
      dict(
          testcase_name='NRowsNValsUpdated',
          rp=(lambda: RowPartition.from_row_splits([0, 3, 8])),
          spec=RowPartitionSpec(nrows=2, nvals=8, dtype=dtypes.int64),
          expected=_get_specified_row_partition),
      dict(
          testcase_name='NValsUpdated',
          rp=(lambda: RowPartition.from_row_splits([0, 3, 8])),
          spec=RowPartitionSpec(nrows=None, nvals=8, dtype=dtypes.int64),
          expected=_get_specified_row_partition)])
  def testMergeWithSpecUpdate(self, rp, spec, expected):
    rp = rp()
    expected = expected()
    actual = rp._merge_with_spec(spec)
    self.assertAllEqual(actual.row_splits(), expected.row_splits())
    self.assertAllEqual(actual.static_nrows, expected.static_nrows)
    self.assertAllEqual(actual.static_nvals, expected.static_nvals)

  @parameterized.named_parameters([
      dict(
          testcase_name='from_uniform_row_length',
          x=lambda: RowPartition.from_uniform_row_length(5, nvals=20),
          expected=True),
      dict(
          testcase_name='from_row_splits',
          x=lambda: RowPartition.from_row_splits([0, 3, 8]),
          expected=False),
      dict(
          testcase_name='from_row_lengths',
          x=lambda: RowPartition.from_row_lengths([2, 0, 2]),
          expected=False),
      dict(
          testcase_name='from_row_lengths_uniform',
          x=lambda: RowPartition.from_row_lengths([3, 3, 3]),
          expected=False),
  ])
  def testIsUniform(self, x, expected):
    x = x()
    self.assertEqual(expected, x.is_uniform())

  @parameterized.named_parameters([
      dict(
          testcase_name='doc_example',
          x=lambda: RowPartition.from_row_lengths([3, 2, 0, 2]),
          expected=[0, 1, 2, 0, 1, 0, 1]),
      dict(
          testcase_name='from_uniform_row_length',
          x=lambda: RowPartition.from_uniform_row_length(4, nvals=12),
          expected=[0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]),
      dict(
          testcase_name='from_row_splits',
          x=lambda: RowPartition.from_row_splits([0, 3, 8]),
          expected=[0, 1, 2, 0, 1, 2, 3, 4]),
  ])
  def testOffsetsInRows(self, x, expected):
    x = x()
    actual = x.offsets_in_rows()
    self.assertAllEqual(expected, actual)

  def testFromUniformRowLengthBugConvertToTensor(self):
    # This originally failed to run because nrows was dtypes.int32. I think
    # we may need to consider the semantics of the type of a RowPartition
    # if preferred_dtype is unspecified. Also, looking at convert_to_tensor:
    # dtype specifies the type of the output.
    # preferred_dtype/dtype_hint is a suggestion, and dtype_hint is the new
    # name.
    nrows = constant_op.constant(3, dtype=dtypes.int32)
    nvals = constant_op.constant(12, dtype=dtypes.int64)
    row_length = constant_op.constant(4, dtype=dtypes.int64)
    rp = RowPartition.from_uniform_row_length(row_length, nvals=nvals,
                                              nrows=nrows, dtype=dtypes.int64)
    self.assertEqual(rp.nrows().dtype, dtypes.int64)

  def testFromUniformRowLengthNvalDynamic(self):
    # A key question is whether if nrows and uniform_row_length are known,
    # and nvals is given but not known statically, should we determine nvals?
    # TODO(martinz): Uncomment after nvals is fixed.
    # @def_function.function(
    #     input_signature=[tensor_spec.TensorSpec(None, dtypes.int32)])
    # def foo(nvals):
    #   rp = RowPartition.from_uniform_row_length(12, nvals=nvals, nrows=3)
    #   nval_output = tensor_util.constant_value(rp.nvals())
    #   self.assertEqual(nval_output, 36)
    # foo(constant_op.constant(36, dtype=dtypes.int32))
    pass

  def testFromUniformRowLengthNvalDynamicNoValidate(self):
    # A key question is whether if nrows and uniform_row_length are known,
    # and nvals is given but not known statically, should we determine nvals?
    # TODO(martinz): Uncomment after nvals is fixed.
    # @def_function.function(
    #     input_signature=[tensor_spec.TensorSpec(None, dtypes.int32)])
    # def foo(nvals):
    #   rp = RowPartition.from_uniform_row_length(12, nvals=nvals, nrows=3,
    #                                             validate=False)
    #   nval_output = tensor_util.constant_value(rp.nvals())
    #   self.assertEqual(nval_output, 36)
    # foo(constant_op.constant(36, dtype=dtypes.int32))
    pass

  def testFromUniformRowLengthNvalDynamicWrong(self):
    # A key question is whether if nrows and uniform_row_length are known,
    # and nvals is given but not known statically and WRONG,
    # what should we do? We add a check, but checks are only checked for
    # row_splits.
    @def_function.function(
        input_signature=[tensor_spec.TensorSpec(None, dtypes.int32)])
    def foo(nvals):
      rp = RowPartition.from_uniform_row_length(12, nvals=nvals, nrows=3)
      return rp.nvals()

    with self.assertRaises(errors.InvalidArgumentError):
      nvals = foo(constant_op.constant(7, dtype=dtypes.int32))
      self.evaluate(nvals)

  def testFromUniformRowLengthNvalDynamicWrongRowSplits(self):
    # A key question is whether if nrows and uniform_row_length are known,
    # and nvals is given but not known statically and WRONG,
    # what should we do?
    # A key question is whether if nrows and uniform_row_length are known,
    # and nvals is given but not known statically and WRONG,
    # what should we do? We add a check, but checks are only checked for
    # row_splits.
    @def_function.function(
        input_signature=[tensor_spec.TensorSpec(None, dtypes.int32)])
    def foo(nvals):
      rp = RowPartition.from_uniform_row_length(12, nvals=nvals, nrows=3)
      return rp.row_splits()

    with self.assertRaises(errors.InvalidArgumentError):
      rs = foo(constant_op.constant(7, dtype=dtypes.int32))
      self.evaluate(rs)

  def testFromUniformRowPartitionNrows(self):
    rp = RowPartition.from_uniform_row_length(3, nrows=4)
    self.assertAllEqual(4, rp.nrows())
    self.assertAllEqual(3, rp.uniform_row_length())
    self.assertAllEqual(12, rp.static_nvals)

  def testFromUniformRowPartitionNvalsStatic(self):
    rp = RowPartition.from_uniform_row_length(3, nvals=12)
    self.assertAllEqual(4, rp.static_nrows)
    self.assertAllEqual(3, rp.static_uniform_row_length)
    self.assertAllEqual(12, rp.static_nvals)

  def testFromUniformRowPartitionNvalsStaticNoValidate(self):
    rp = RowPartition.from_uniform_row_length(3, nrows=4, nvals=12,
                                              validate=False)
    self.assertAllEqual(4, rp.static_nrows)
    self.assertAllEqual(3, rp.static_uniform_row_length)
    self.assertAllEqual(12, rp.static_nvals)

  def testFromUniformRowPartitionNvalsIs(self):
    # TODO(martinz): Uncomment after nvals is fixed.
    # nvals = constant_op.constant(12)
    # rp = RowPartition.from_uniform_row_length(3, nvals=nvals)
    # self.assertIs(rp.nvals(), nvals)
    pass

  def testFromUniformRowPartitionRowStartsStatic(self):
    rp = RowPartition.from_row_starts([0, 3, 6], nvals=12)
    self.assertAllEqual(12, rp.static_nvals)

  def testStaticNrows(self):
    rp = RowPartition.from_row_splits([0, 3, 4, 5])
    static_nrows = rp.static_nrows
    self.assertIsInstance(static_nrows, int)
    self.assertAllEqual(3, static_nrows)

  def testStaticNrowsUnknown(self):
    @def_function.function(
        input_signature=[tensor_spec.TensorSpec(None, dtypes.int32)])
    def foo(rs):
      rp = RowPartition.from_row_splits(rs)
      static_nrows = rp.static_nrows
      self.assertIsNone(static_nrows)
    foo(array_ops.constant([0, 3, 4, 5], dtype=dtypes.int32))

@test_util.run_all_in_graph_and_eager_modes
class RowPartitionSpecTest(test_util.TensorFlowTestCase,
                           parameterized.TestCase):

  def testDefaultConstruction(self):
    spec = RowPartitionSpec()
    self.assertIsNone(spec.nrows)
    self.assertIsNone(spec.nvals)
    self.assertIsNone(spec.uniform_row_length)
    self.assertEqual(spec.dtype, dtypes.int64)

  @parameterized.parameters([
      (None, None, None, dtypes.int64, None, None, None, dtypes.int64),
      (5, None, None, dtypes.int32, 5, None, None, dtypes.int32),
      (None, 20, None, dtypes.int64, None, 20, None, dtypes.int64),
      (None, None, 8, dtypes.int64, None, None, 8, dtypes.int64),
      (5, None, 8, dtypes.int64, 5, 40, 8, dtypes.int64),  # nvals inferred
      (None, 20, 4, dtypes.int32, 5, 20, 4, dtypes.int32),  # nrows inferred
      (0, None, None, dtypes.int32, 0, 0, None, dtypes.int32),  # nvals inferred
      (None, None, 0, dtypes.int32, None, 0, 0, dtypes.int32),  # nvals inferred
  ])  # pyformat: disable
  def testConstruction(self, nrows, nvals, uniform_row_length, dtype,
                       expected_nrows, expected_nvals,
                       expected_uniform_row_length, expected_dtype):
    spec = RowPartitionSpec(nrows, nvals, uniform_row_length, dtype)
    self.assertEqual(spec.nrows, expected_nrows)
    self.assertEqual(spec.nvals, expected_nvals)
    self.assertEqual(spec.uniform_row_length, expected_uniform_row_length)
    self.assertEqual(spec.dtype, expected_dtype)

  @parameterized.parameters([
      dict(dtype=dtypes.float32, error='dtype must be tf.int32 or tf.int64'),
      dict(nrows=0, nvals=5, error='.* not compatible .*'),
      dict(uniform_row_length=0, nvals=5, error='.* not compatible .*'),
      dict(nvals=11, uniform_row_length=5, error='.* not compatible .*'),
      dict(
          nrows=8, nvals=10, uniform_row_length=5,
          error='.* not compatible .*'),
  ])
  def testConstructionError(self,
                            nrows=None,
                            nvals=None,
                            uniform_row_length=None,
                            dtype=dtypes.int64,
                            error=None):
    with self.assertRaisesRegex(ValueError, error):
      RowPartitionSpec(nrows, nvals, uniform_row_length, dtype)

  @parameterized.parameters([
      RowPartitionSpec(),
      RowPartitionSpec(dtype=dtypes.int64),
      RowPartitionSpec(uniform_row_length=3)
  ])  # pyformat: disable

  def testDeepcopy(self, spec):
    spec = RowPartitionSpec()
    spec_b = copy.deepcopy(spec)
    self.assertEqual(repr(spec), repr(spec_b))

  def testValueType(self):
    spec = RowPartitionSpec()
    self.assertEqual(spec.value_type, RowPartition)

  @parameterized.parameters([
      dict(
          spec=RowPartitionSpec(),
          expected=(tensor_shape.TensorShape([None]),
                    tensor_shape.TensorShape([None]),
                    tensor_shape.TensorShape([None]), dtypes.int64)),
      dict(
          spec=RowPartitionSpec(dtype=dtypes.int32),
          expected=(tensor_shape.TensorShape([None]),
                    tensor_shape.TensorShape([None]),
                    tensor_shape.TensorShape([None]), dtypes.int32)),
      dict(
          spec=RowPartitionSpec(nrows=8, nvals=13),
          expected=(tensor_shape.TensorShape([8]),
                    tensor_shape.TensorShape([13]),
                    tensor_shape.TensorShape([None]), dtypes.int64)),
      dict(
          spec=RowPartitionSpec(nrows=8, uniform_row_length=2),
          expected=(
              tensor_shape.TensorShape([8]),
              tensor_shape.TensorShape([16]),  # inferred
              tensor_shape.TensorShape([2]),
              dtypes.int64)),
  ])
  def testSerialize(self, spec, expected):
    serialization = spec._serialize()
    # TensorShape has an unconventional definition of equality, so we can't use
    # assertEqual directly here.  But repr() is deterministic and lossless for
    # the expected values, so we can use that instead.
    self.assertEqual(repr(serialization), repr(expected))

  @parameterized.parameters([
      dict(
          spec=RowPartitionSpec(),
          expected=tensor_spec.TensorSpec([None], dtypes.int64)),
      dict(
          spec=RowPartitionSpec(dtype=dtypes.int32),
          expected=tensor_spec.TensorSpec([None], dtypes.int32)),
      dict(
          spec=RowPartitionSpec(nrows=17, dtype=dtypes.int32),
          expected=tensor_spec.TensorSpec([18], dtypes.int32)),
      dict(
          spec=RowPartitionSpec(nvals=10, uniform_row_length=2),
          expected=tensor_spec.TensorSpec([6], dtypes.int64)),  # inferred nrow
  ])
  def testComponentSpecs(self, spec, expected):
    self.assertEqual(spec._component_specs, expected)

  @parameterized.parameters([
      dict(
          rp_factory=lambda: RowPartition.from_row_splits([0, 3, 7]),
          components=[0, 3, 7]),
  ])
  def testToFromComponents(self, rp_factory, components):
    rp = rp_factory()
    spec = rp._type_spec
    actual_components = spec._to_components(rp)
    self.assertAllEqual(actual_components, components)
    rp_reconstructed = spec._from_components(actual_components)
    _assert_row_partition_equal(self, rp, rp_reconstructed)

  @parameterized.parameters([
      (RowPartitionSpec(), RowPartitionSpec()),
      (RowPartitionSpec(nrows=8), RowPartitionSpec(nrows=8)),
      (RowPartitionSpec(nrows=8), RowPartitionSpec(nrows=None)),
      (RowPartitionSpec(nvals=8), RowPartitionSpec(nvals=8)),
      (RowPartitionSpec(nvals=8), RowPartitionSpec(nvals=None)),
      (RowPartitionSpec(uniform_row_length=8),
       RowPartitionSpec(uniform_row_length=8)),
      (RowPartitionSpec(uniform_row_length=8),
       RowPartitionSpec(uniform_row_length=None)),
      (RowPartitionSpec(nvals=12), RowPartitionSpec(uniform_row_length=3)),
      (RowPartitionSpec(nrows=12), RowPartitionSpec(uniform_row_length=72)),
      (RowPartitionSpec(nrows=5), RowPartitionSpec(nvals=15)),
      (RowPartitionSpec(nvals=0), RowPartitionSpec(nrows=0)),
      (RowPartitionSpec(nvals=0), RowPartitionSpec(uniform_row_length=0)),
  ])
  def testIsCompatibleWith(self, spec1, spec2):
    self.assertTrue(spec1.is_compatible_with(spec2))

  @parameterized.parameters([
      (RowPartitionSpec(), RowPartitionSpec(dtype=dtypes.int32)),
      (RowPartitionSpec(nvals=5), RowPartitionSpec(uniform_row_length=3)),
      (RowPartitionSpec(nrows=7,
                        nvals=12), RowPartitionSpec(uniform_row_length=3)),
      (RowPartitionSpec(nvals=5), RowPartitionSpec(nrows=0)),
      (RowPartitionSpec(nvals=5), RowPartitionSpec(uniform_row_length=0)),
  ])
  def testIsNotCompatibleWith(self, spec1, spec2):
    self.assertFalse(spec1.is_compatible_with(spec2))

  @parameterized.parameters([
      dict(
          spec1=RowPartitionSpec(nrows=8, nvals=3, dtype=dtypes.int32),
          spec2=RowPartitionSpec(nrows=8, nvals=3, dtype=dtypes.int32),
          expected=RowPartitionSpec(nrows=8, nvals=3, dtype=dtypes.int32)),
      dict(
          spec1=RowPartitionSpec(nrows=8, nvals=None),
          spec2=RowPartitionSpec(nrows=None, nvals=8),
          expected=RowPartitionSpec(nrows=None, nvals=None)),
      dict(
          spec1=RowPartitionSpec(nrows=8, nvals=33),
          spec2=RowPartitionSpec(nrows=3, nvals=13),
          expected=RowPartitionSpec(nrows=None, nvals=None)),
      dict(
          spec1=RowPartitionSpec(nrows=12, uniform_row_length=3),
          spec2=RowPartitionSpec(nrows=3, uniform_row_length=3),
          expected=RowPartitionSpec(nrows=None, uniform_row_length=3)),
      dict(
          spec1=RowPartitionSpec(5, 35, 7),
          spec2=RowPartitionSpec(8, 80, 10),
          expected=RowPartitionSpec(None, None, None)),
  ])
  def testMostSpecificCompatibleType(self, spec1, spec2, expected):
    actual = spec1.most_specific_compatible_type(spec2)
    self.assertEqual(actual, expected)

  @parameterized.parameters([
      (RowPartitionSpec(), RowPartitionSpec(dtype=dtypes.int32)),
  ])
  def testMostSpecificCompatibleTypeError(self, spec1, spec2):
    with self.assertRaisesRegex(ValueError, 'No TypeSpec is compatible'):
      spec1.most_specific_compatible_type(spec2)

  def testNumRowsInt64(self):
    row_splits = array_ops.constant([0, 2, 3, 5], dtype=dtypes.int64)
    values = RowPartition.from_row_splits(row_splits)
    nrows = values.nrows()
    self.assertEqual(dtypes.int64, nrows.dtype)

  def testFromValue(self):
    self.assertEqual(
        RowPartitionSpec.from_value(RowPartition.from_row_splits([0, 2, 8, 8])),
        RowPartitionSpec(nrows=3))
    self.assertEqual(
        RowPartitionSpec.from_value(
            RowPartition.from_row_lengths([5, 3, 0, 2])),
        RowPartitionSpec(nrows=4))
    self.assertEqual(
        RowPartitionSpec.from_value(
            RowPartition.from_value_rowids([0, 2, 2, 8])),
        RowPartitionSpec(nrows=9, nvals=4))
    self.assertEqual(
        RowPartitionSpec.from_value(
            RowPartition.from_uniform_row_length(
                nvals=12, uniform_row_length=3)),
        RowPartitionSpec(nvals=12, uniform_row_length=3))

  @parameterized.parameters([
      dict(original=RowPartitionSpec(),
           dtype=dtypes.int32,
           expected=RowPartitionSpec(dtype=dtypes.int32)),
      dict(original=RowPartitionSpec(dtype=dtypes.int32),
           dtype=dtypes.int64,
           expected=RowPartitionSpec()),
      dict(original=RowPartitionSpec(nvals=20, nrows=4, uniform_row_length=5),
           dtype=dtypes.int32,
           expected=RowPartitionSpec(nvals=20, nrows=4, uniform_row_length=5,
                                     dtype=dtypes.int32)),
  ])
  def testWithDType(self, original, dtype, expected):
    actual = original.with_dtype(dtype)
    self.assertEqual(actual, expected)

  @parameterized.parameters([
      dict(a=RowPartitionSpec(),
           b=RowPartitionSpec(nrows=3, uniform_row_length=5),
           expected=RowPartitionSpec(nrows=3, uniform_row_length=5)),
      dict(a=RowPartitionSpec(nrows=3),
           b=RowPartitionSpec(uniform_row_length=5),
           expected=RowPartitionSpec(nrows=3, uniform_row_length=5)),
      dict(a=RowPartitionSpec(nvals=20),
           b=RowPartitionSpec(nrows=3),
           expected=RowPartitionSpec(nvals=20, nrows=3)),
      dict(a=RowPartitionSpec(nvals=20, dtype=dtypes.int32),
           b=RowPartitionSpec(nrows=3, dtype=dtypes.int32),
           expected=RowPartitionSpec(nvals=20, nrows=3, dtype=dtypes.int32)),
  ])
  def testMergeWith(self, a, b, expected):
    actual = a._merge_with(b)
    actual_rev = b._merge_with(a)
    self.assertEqual(actual, expected)
    self.assertEqual(actual_rev, expected)

  @parameterized.parameters([
      dict(a=RowPartitionSpec(nrows=3, nvals=10),
           b=RowPartitionSpec(uniform_row_length=5),
           error_type=ValueError,
           error_regex='Merging incompatible RowPartitionSpecs'),
      dict(a=RowPartitionSpec(uniform_row_length=5, dtype=dtypes.int32),
           b=RowPartitionSpec(uniform_row_length=5, dtype=dtypes.int64),
           error_type=ValueError,
           error_regex='Merging RowPartitionSpecs with incompatible dtypes'),
  ])
  def testMergeWithRaises(self, a, b, error_type, error_regex):
    with self.assertRaisesRegex(error_type, error_regex):
      a._merge_with(b)

def _assert_row_partition_equal(test_class, actual, expected):
  assert isinstance(test_class, test_util.TensorFlowTestCase)
  assert isinstance(actual, RowPartition)
  assert isinstance(expected, RowPartition)

  test_class.assertEqual(actual._has_precomputed_row_splits(),
                         expected._has_precomputed_row_splits())
  test_class.assertEqual(actual._has_precomputed_row_lengths(),
                         expected._has_precomputed_row_lengths())
  test_class.assertEqual(actual._has_precomputed_value_rowids(),
                         expected._has_precomputed_value_rowids())
  test_class.assertEqual(actual._has_precomputed_nrows(),
                         expected._has_precomputed_nrows())
  test_class.assertEqual(actual.uniform_row_length() is None,
                         expected.uniform_row_length() is None)

  if expected._has_precomputed_row_splits():
    test_class.assertAllEqual(actual.row_splits(), expected.row_splits())
  if expected._has_precomputed_row_lengths():
    test_class.assertAllEqual(actual.row_lengths(), expected.row_lengths())
  if expected._has_precomputed_value_rowids():
    test_class.assertAllEqual(actual.value_rowids(), expected.value_rowids())
  if expected._has_precomputed_nrows():
    test_class.assertAllEqual(actual.nrows(), expected.nrows())
  if expected.uniform_row_length() is not None:
    test_class.assertAllEqual(actual.uniform_row_length(),
                              expected.uniform_row_length())

if __name__ == '__main__':
  googletest.main()


import psycopg2
import sys
from pathlib import Path

file = Path(__file__).resolve()
parent = file.parents[0]
sys.path.append(str(parent))

from authentication import verify_access_token

def get_all_pets(
    connection: psycopg2.extensions.connection,
    owner_id: int,
    access_token: str,
):
    """
    Returns all pets in the database

    Arguments:
        connection: postgres connection
        owner_id: id of the pet owner
        access_token: access token of the pet owner
    Returns:
        List of pets
    """
    # Verify access token
    if not verify_access_token(connection, owner_id, access_token):
        return False

    cur = connection.cursor()

    query = """SELECT * FROM pets WHERE owner_id = %s ORDER BY name;"""

    try:
        # Execute query
        cur.execute(query, (owner_id,))
        pets = cur.fetchall()
        print(f"Query executed successfully: {query}")
    except Exception as e:
        print(f"Error while executing query: {e}")
        return []

    cur.close()

    pets = [{
        "id": pet[0],
        "name": pet[1],
        "animal": pet[2],
        "breed": pet[3],
        "description": pet[4],
        "image_url": pet[5],
        "is_missing": pet[6],
        "owner_id": pet[7],
    } for pet in pets]

    return pets

def get_pet(
        connection: psycopg2.extensions.connection,
        pet_id: int,
):
    """
    Returns a pet by the given id

    Arguments:
        connection: postgres connection
        pet_id: id of pet
    Returns:
        A pet
    """

    cur = connection.cursor()

    query = """SELECT * FROM pets WHERE id = %s LIMIT 1;"""

    try:
        # Execute query
        cur.execute(query, (pet_id,))
        pet = cur.fetchall()
        print(f"Query executed successfully: {query}")
    except Exception as e:
        print(f"Error while executing query: {e}")
        return False

    cur.close()

    res = pet[0]
    pet = {
        "id": res[0],
        "name": res[1],
        "animal": res[2],
        "breed": res[3],
        "description": res[4],
        "image_url": res[5],
        "owner_id": res[6]
    }

    return pet

def add_pet(
    connection: psycopg2.extensions.connection,
    name: str,
    animal: str,
    breed: str,
    description: str,
    image_url: str,
    owner_id: int,
    access_token: str,
):
    """
    Adds a pet to the database

    Arguments:
        connection: postgres connection
        name: name of the pet
        animal: type of animal
        breed: breed of the pet
        description: description of the pet
        image_url: url of the image of the pet
        owner_id: id of the pet owner
        access_token: access token of the pet owner
    Returns:
        True if successful, False otherwise
    """
    # Verify access token
    if not verify_access_token(connection, owner_id, access_token):
        return False

    cur = connection.cursor()

    query = """INSERT INTO pets (name, animal, breed, description, image_url, owner_id) VALUES (%s, %s, %s, %s, %s, %s);"""

    try:
        # Execute query
        cur.execute(query, (name, animal, breed, description, image_url, owner_id,))
        print(f"Query executed successfully: {query}")
    except Exception as e:
        print(f"Error while executing query: {e}")
        return False

    cur.close()
    connection.commit()

    return True

def edit_pet(
    connection: psycopg2.extensions.connection,
    id: int,
    name: str,
    animal: str,
    breed: str,
    description: str,
    image_url: str,
    owner_id: int,
    access_token: str,
):
    """
    Edits an existing pet in the database

    Arguments:
        connection: postgres connection
        id: id of the pet
        name: name of the pet
        animal: type of animal
        breed: breed of the pet
        description: description of the pet
        image_url: url of the image of the pet
        owner_id: id of the pet owner
        access_token: access token of the pet owner
    Returns:
        True if successful, False otherwise
    """
    # Verify access token
    if not verify_access_token(connection, owner_id, access_token):
        return False

    cur = connection.cursor()

    query = """UPDATE pets SET name = %s, animal = %s, breed = %s, description = %s, image_url = %s, owner_id = %s WHERE id = %s;"""

    result = False
    try:
        # Execute query
        cur.execute(query, (name, animal, breed, description, image_url, owner_id, id))
        result = True
        print(f"Query executed successfully: {query}")
    except Exception as e:
        print(f"Error while executing query: {e}")

    cur.close()
    connection.commit()

    return result

def update_pet_missing_status(connection: psycopg2.extensions.connection, pet_id: int, new_status: bool) -> bool:

    result = False
    try:
        cur = connection.cursor()

        # Update the missing status in the database
        cur.execute("UPDATE pets SET isMissing = %s WHERE id = %s;", (new_status, pet_id))

        connection.commit()
        cur.close()
        result = True

    except Exception as e:
        print(f"Error updating status: {e}")
        connection.rollback()

    return result

def delete_all_pet_data_from_database(connection: psycopg2.extensions.connection, to_delete_id: int, user_id: int, access_token: str):
    """
    This function deletes the pet with the provided id

    """

    # Verify access token
    if not verify_access_token(connection, user_id, access_token):
        print("Access token verification failed.")
        return False

    if not unlink_pet_sightings_from_user_reports(connection=connection, user_id=user_id, pet_id=to_delete_id):
        print("Failed to unlink sightings from user reports.")
        return False

    # Delete pet's missing reports
    if not delete_pet_reports(connection=connection, user_id=user_id, pet_id=to_delete_id):
        print("Failed to delete pet's missing reports.")
        return False

    # Delete pet
    if not delete_pet(connection=connection, user_id=user_id, pet_id=to_delete_id):
        print("Failed to delete pet.")
        return False

    return True

def unlink_pet_sightings_from_user_reports(connection: psycopg2.extensions.connection, user_id: int, pet_id: int):
    """
    This function unlinks sightings on this user's pet's missing reports.
    """

    cur = connection.cursor()

    query = """
        UPDATE sightings SET missing_report_id = NULL WHERE missing_report_id IN (SELECT id FROM missing_reports WHERE author_id = %s AND pet_id = %s);
    """

    try:
        cur.execute(query, (user_id, pet_id))

        # Commit the change
        connection.commit()
        print(f"Sightings successfully unlinked")
        return True

    except Exception as e:
        print(f"Error with deleting user: {e}")

    cur.close()
    return False

def delete_pet_reports(connection: psycopg2.extensions.connection, user_id: int, pet_id: int):
    """
    This function deletes all missing reports created by the user for a pet.
    """

    cur = connection.cursor()

    query = """
        DELETE FROM missing_reports WHERE author_id = %s AND pet_id = %s;
    """

    try:
        cur.execute(query, (user_id, pet_id))

        # Commit the change
        connection.commit()
        print(f"Missing reports successfully deleted")
        return True

    except Exception as e:
        print(f"Error with deleting user: {e}")

    cur.close()
    return False

def delete_pet(connection, user_id, pet_id):
    """
    This function deletes the pet that is selected.
    """

    cur = connection.cursor()

    query = """
        DELETE FROM pets WHERE owner_id = %s AND id = %s;
    """

    try:
        cur.execute(query, (user_id, pet_id))
        connection.commit()
        print(f"Pet successfully deleted")
        return True

    except Exception as e:
        print(f"Error while deleting pet: {e}")

    cur.close()
    return False


import warnings
import matplotlib.pyplot as plt
import seaborn as sns
import itertools
import numpy as np
import pandas as pd
from joblib import Parallel, delayed
from sklearn.decomposition import PCA
from sklearn.model_selection import KFold, train_test_split
from sklearn.preprocessing import StandardScaler

TEST_SIZE = 57
K_FOLD_NUMBER = 10

def pca(x_data, num_features):
    scalar = StandardScaler()
    X = scalar.fit_transform(x_data.T)
    pca = PCA(n_components=num_features)
    principalComponents = pca.fit_transform(X)
    return principalComponents

def sigmoid(x):
    warnings.filterwarnings('ignore')
    return 1 / (1 + np.exp(-x))

def feed_forward(weights, inputs, keep_rate):
    activation_layers = [sigmoid(np.dot(inputs, weights[0]))]
    dropouts = [
        np.random.rand(activation_layers[0].shape[0], activation_layers[0].shape[1]) < keep_rate]
    activation_layers[0] = (activation_layers[0] * dropouts[0]) / keep_rate
    for i in range(1, len(weights)):
        activation_layers[i - 1] = np.hstack(
            (activation_layers[i - 1], np.ones((activation_layers[i - 1].shape[0], 1))))
        activation_layer = sigmoid(np.dot(activation_layers[i - 1], weights[i]))
        dropout = np.random.rand(activation_layer.shape[0], activation_layer.shape[1]) < keep_rate
        if i == len(weights) - 1:
            potential = (np.dot(activation_layers[i - 1], weights[i]) * dropout) / keep_rate
        activation_layer = (activation_layer * dropout) / keep_rate
        activation_layers.append(activation_layer)
        dropouts.append(dropout)
    return activation_layers, dropouts, potential

def derivative_loss_array(y_pred, potential, y_true):
    ret_arr = np.zeros(y_pred.shape)
    mask = (y_pred == 0) | (y_pred == 1)
    ret_arr[mask] = 1 * (y_pred[mask] - y_true[mask])
    ret_arr[~mask] = (1 - y_true[~mask] - sigmoid(-potential[~mask])) * (
            1 / (y_pred[~mask] * (1 - y_pred[~mask])))
    return ret_arr

def backpropagation(weights, activation_layers, potential, y_true, dropouts, keep_rate):
    diff = derivative_loss_array(activation_layers[-1], potential, y_true)
    # diff = (activation_layers[-1] - y_true)
    dot_one = activation_layers[-1] * diff
    delta_layers = [dot_one * (1 - activation_layers[-1])]
    delta_layers[0] = (delta_layers[0] * dropouts[-1]) / keep_rate

    for i in range(len(weights) - 1, 0, -1):
        a = np.dot(delta_layers[-1], weights[i][:-1].T)
        b = activation_layers[i - 1] * (1 - activation_layers[i - 1])
        delta_layers.append(a * (dropouts[-1] / keep_rate) * b[:, :-1])
    delta_layers.reverse()
    return delta_layers

# TODO : Real Loss Function
def loss(y_pred, potential, y_true):
    ret_arr = np.zeros(y_pred.shape)
    curr = np.zeros(y_pred.shape)
    mask = (y_pred == 0) | (y_pred == 1)
    ret_arr[mask] = 1000000 * (y_pred[mask] - y_true[mask])
    curr[~mask] = 1 - 2 * y_true[~mask]
    for i in range(curr.shape[0]):
    	if curr[i] == 255:
    		curr[i] = -1
    ret_arr[~mask] = np.log(1 + np.exp((curr[~mask])*potential[~mask]))
    ret_arr = np.nan_to_num(ret_arr,nan= 1000000)
    for i in range(ret_arr.shape[0]):
    	if ret_arr[i] > 1000000:
    		ret_arr[i] = 1000000
    # I do the list, flatten() and tolist() because there was an error TypeError: unhashable type: 'numpy.ndarray' when trying to plot
    return np.mean(ret_arr)

def loss_MSE(y_pred, potential, y_true):
	return np.mean((y_true - y_pred)**2)

def train_model(weights, keep_rate, learning_rate, epochs, x_train, y_train, x_test, y_test):
    N = y_train.size
    train_errors = []
    test_errors = []
    for i in range(epochs):
        train_outs, dropouts_train, potential_train = feed_forward(weights, x_train, keep_rate)

        # Backpropagation
        delta_layers = backpropagation(weights, train_outs, potential_train, y_train, dropouts_train,
                                       keep_rate)

        # Update weights
        weights[0] -= learning_rate * np.dot(x_train.T, delta_layers[0]) / N
        for i in range(1, len(weights)):
            weights[i] -= learning_rate * np.dot(train_outs[i - 1].T, delta_layers[i]) / N

        # added

        train_errors.append(loss(train_outs[-1], potential_train, y_train))
        test_outs, dropouts_test, potential_test = feed_forward(weights, x_test, keep_rate)
        test_errors.append(loss(test_outs[-1], potential_test, y_test))

    return weights, train_errors, test_errors

def train_transfer(weights, keep_rate, learning_rate, epochs, x_train, y_train):
    N = y_train.size
    for i in range(epochs):
        train_outs, dropouts, potential = feed_forward(weights, x_train, keep_rate)

        # Backpropagation
        delta_layers = backpropagation(weights, train_outs, potential, y_train, dropouts,
                                       keep_rate)

        # Update weights
        weights[len(weights) - 1] -= learning_rate * np.dot(train_outs[len(weights) - 2].T, delta_layers[len(weights) - 1]) / N
    return weights

def accuracy(weights, x_test, y_test, err_margin, train_errors, test_errors, epochs):
    activation_layers, dropout, potential = feed_forward(weights, x_test, 1)
    mask = y_test == 1
    correct = (activation_layers[-1][mask] >= 1 - err_margin).sum() + (
            activation_layers[-1][~mask] <= err_margin).sum()

    # for plotting
    print("Plotting")
    acc = correct / y_test.size
    plot(epochs, train_errors, test_errors, acc)

    return acc

def process(test_case, x_train, x_test, y_train, y_test, num_features):
    hidden_size_1 = test_case[0]
    hidden_size_2 = test_case[1]
    keep_rate = test_case[2]
    learning_rate = test_case[3]
    epochs = test_case[4]
    if hidden_size_2 == 0:
        W1 = np.random.normal(scale=0.1, size=(num_features + 1, hidden_size_1))
        W2 = np.random.normal(scale=0.1, size=(hidden_size_1 + 1, 1))
        weights = [W1, W2]
    else:
        W1 = np.random.normal(scale=0.1, size=(num_features + 1, hidden_size_1))
        W2 = np.random.normal(scale=0.1, size=(hidden_size_1 + 1, hidden_size_2))
        W3 = np.random.normal(scale=0.1, size=(hidden_size_2 + 1, 1))
        weights = [W1, W2, W3]
    final_weights, train_errors, test_errors = train_model(weights, keep_rate, learning_rate, epochs, x_train, y_train, x_test, y_test)
    final_accuracy = accuracy(final_weights, x_test, y_test, 0.1, train_errors, test_errors, epochs)
    return final_accuracy, test_case, num_features

def process_K_fold(test_case, x_train, x_test, y_train, y_test, num_features):
    hidden_size_1 = test_case[0]
    hidden_size_2 = test_case[1]
    keep_rate = test_case[2]
    learning_rate = test_case[3]
    epochs = test_case[4]
    if hidden_size_2 == 0:
        W1 = np.random.normal(scale=0.1, size=(num_features + 1, hidden_size_1))
        W2 = np.random.normal(scale=0.1, size=(hidden_size_1 + 1, 1))
        weights = [W1, W2]
    else:
        W1 = np.random.normal(scale=0.1, size=(num_features + 1, hidden_size_1))
        W2 = np.random.normal(scale=0.1, size=(hidden_size_1 + 1, hidden_size_2))
        W3 = np.random.normal(scale=0.1, size=(hidden_size_2 + 1, 1))
        weights = [W1, W2, W3]
    final_weights = train_model(weights, keep_rate, learning_rate, epochs, x_train, y_train)
    final_accuracy_10 = accuracy(final_weights, x_test, y_test, 0.1)
    final_accuracy_30 = accuracy(final_weights, x_test, y_test, 0.3)
    final_accuracy_50 = accuracy(final_weights, x_test, y_test, 0.5)
    return np.array([final_accuracy_10, final_accuracy_30, final_accuracy_50])

def transfer_test(x, y, test_case, num_features):

    x_train, x_test, y_train, y_test  = train_test_split(x, y, test_size=TEST_SIZE, random_state=33)

    x_first, x_second, y_first, y_second = train_test_split(x_train, y_train, test_size=TEST_SIZE*2, random_state=33)

    # initialize model
    hidden_size_1 = test_case[0]
    hidden_size_2 = test_case[1]
    keep_rate = test_case[2]
    learning_rate = test_case[3]
    epochs = test_case[4]
    if hidden_size_2 == 0:
        W1 = np.random.normal(scale=0.1, size=(num_features + 1, hidden_size_1))
        W2 = np.random.normal(scale=0.1, size=(hidden_size_1 + 1, 1))
        weights = [W1, W2]
    else:
        W1 = np.random.normal(scale=0.1, size=(num_features + 1, hidden_size_1))
        W2 = np.random.normal(scale=0.1, size=(hidden_size_1 + 1, hidden_size_2))
        W3 = np.random.normal(scale=0.1, size=(hidden_size_2 + 1, 1))
        weights = [W1, W2, W3]

    # train model
    final_weights = train_model(weights, keep_rate, learning_rate, epochs, x_first, y_first)

    # reinitialize final layer
    if hidden_size_2 == 0:
        W4 = np.random.normal(scale=0.1, size=(hidden_size_1 + 1, 1))
        weights_transfer = [final_weights[0], W4]
    else:
        W4 = np.random.normal(scale=0.1, size=(hidden_size_2 + 1, 1))
        weights_transfer = [final_weights[0], final_weights[1], W4]

    # train transfer
    final_weights_transfer = train_transfer(weights_transfer, keep_rate, learning_rate, epochs, x_second, y_second)
    final_accuracy_10 = accuracy(final_weights, x_test, y_test, 0.1)
    final_accuracy_30 = accuracy(final_weights, x_test, y_test, 0.3)
    final_accuracy_50 = accuracy(final_weights, x_test, y_test, 0.5)
    final_accuracy_10_transfer = accuracy(final_weights_transfer, x_test, y_test, 0.1)
    final_accuracy_30_transfer = accuracy(final_weights_transfer, x_test, y_test, 0.3)
    final_accuracy_50_transfer = accuracy(final_weights_transfer, x_test, y_test, 0.5)

    return np.array([final_accuracy_10, final_accuracy_30, final_accuracy_50, final_accuracy_10_transfer, final_accuracy_30_transfer, final_accuracy_50_transfer])

def print_final(accuracies, test_case, num_features, results):
    print("\n**********************************\nFirst Layer: " + str(
        test_case[0]) + "\nSecond Layer: " + str(test_case[1]) + "\nKeep Rate of: " + str(
        test_case[2]) + "\nLearning Rate of: " + str(test_case[3]) + "\nEpoch: " + str(
        test_case[4]) + "\nWith " + str(num_features) + " features: \n" + "The accuracies were:\n" + str(
        round(accuracies[2], 3)) + "/" +  str(
        round(accuracies[1], 3)) + "/" + str(
        round(accuracies[0], 3)) + "\n\n" +
        "We do transfer learning on this case!" + "\n"
        + "From " + str(
        round(results[2], 3)) + "/" +  str(
        round(results[1], 3)) + "/" + str(
        round(results[0], 3)) + " to " + str(
        round(results[5], 3)) + "/" +  str(
        round(results[4], 3)) + "/" + str(
        round(results[3], 3)))

# assumes x is already PCA-ed
def cross_validation_test(x, y, test_case, num_features):
    # first check on one iteration
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=TEST_SIZE, random_state=33)
    final_accuracy, test_case, num_features = process(test_case, x_train, x_test, y_train, y_test, num_features)
    # if good enough proceed to do 10-FOLD-CROSS-VALIDATION

    #  just for one run, to see the plot
    if final_accuracy > 0.5:
        print("Yeah! We got a good enough model!")

    elif final_accuracy > 0.9:
        kf = KFold(n_splits=K_FOLD_NUMBER, random_state = 22, shuffle = True)
        accuracies = np.array([0, 0, 0])
        for train_index, test_index in kf.split(x):
            x_train, x_test = x[train_index], x[test_index]
            y_train, y_test = y[train_index], y[test_index]
            accuracies = process_K_fold(test_case, x_train, x_test, y_train, y_test, num_features)/K_FOLD_NUMBER + accuracies
        results = transfer_test(x,y, test_case, num_features)

        print_final(accuracies, test_case, num_features, results)

def plot(epochs, train_errors, test_errors, acc):
    sns.set(style="whitegrid")
    sns.set(rc={"figure.figsize": (10, 6)})
    sns.set(font_scale=1.5)
    print(epochs)
    print(len(test_errors))

    # I just put test_errors for train because I thought that this might solve the ValueError: All arrays must be of the same length error
    sns.lineplot(x=range(epochs), y=train_errors, label='Train')
    sns.lineplot(x=range(epochs), y=test_errors, label='Test')
    plt.legend()
    plt.title(f"Mean Squared Error (acc: {round(acc, 3)})")
    plt.xlabel("Epochs")
    plt.ylabel("MSE")
    plt.show()

df = pd.read_csv('databases/wdbc_unsplit_norm.csv')
y = pd.get_dummies(df.label).values
y = y[:, 1]

y = y.reshape(-1, 1)

x = df.drop('label', 1)

# comment out when running main for!
# x = np.hstack((x, np.ones((x.shape[0], 1))))
# x_reduced = pca(x.T, 2)
# x_reduced = np.hstack((x_reduced, np.ones((x_reduced.shape[0], 1))))
# cross_validation_test(x_reduced,y, (8,18,0.95,0.3,2000), 2)
# hidden_size_1_arr = [8]
# hidden_size_2_arr = [8]
# keep_rate_arr = [0.95, 1]
# learning_rate_arr = [0.1]
# epochs_arr = [6000]
# hidden_size_1_arr = [8]
# hidden_size_2_arr = [8]
# keep_rate_arr = [0.95, 1]
# learning_rate_arr = [0.1]
# epochs_arr = [6000]

# hidden_size_1_arr = [2, 4, 8, 16]
# hidden_size_2_arr = [0, 2, 4, 8]
# keep_rate_arr = [0.95, 1]
# learning_rate_arr = [0.001, 0.01, 0.1]
# epochs_arr = [2000, 4000, 8000, 10000]

hidden_size_1_arr = [8]
hidden_size_2_arr = [0]
keep_rate_arr = [1]
learning_rate_arr = [0.1]
epochs_arr = [100000]

combined = [(h1, h2, kr, lr, e) for h1 in hidden_size_1_arr for h2 in hidden_size_2_arr for kr in keep_rate_arr for lr in learning_rate_arr for e in epochs_arr]

num_features = 10
x_reduced = pca(x.T, num_features)
x_reduced = np.hstack((x_reduced, np.ones((x_reduced.shape[0], 1))))

cross_validation_test(x_reduced, y, combined[0], num_features)

# for num_features in [10]:
#     x_reduced = pca(x.T, num_features)
#     x_reduced = np.hstack((x_reduced, np.ones((x_reduced.shape[0], 1))))
#     Parallel(n_jobs=12)(delayed(cross_validation_test)(x_reduced, y, test_case, num_features) for test_case in combined)


"""empty message

Revision ID: 3b52ec217d5c
Revises: 477cd18a5fa8
Create Date: 2019-05-29 03:05:03.622388

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '3b52ec217d5c'
down_revision = '477cd18a5fa8'
branch_labels = None
depends_on = None

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('projectslog', sa.Column('group_id', sa.String(), nullable=False))
    op.add_column('projectslog', sa.Column('supervisor_id', sa.String(), nullable=False))
    op.create_foreign_key(None, 'projectslog', 'groups', ['group_id'], ['id'])
    op.create_foreign_key(None, 'projectslog', 'supervisors', ['supervisor_id'], ['id'])
    # ### end Alembic commands ###

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(None, 'projectslog', type_='foreignkey')
    op.drop_constraint(None, 'projectslog', type_='foreignkey')
    op.drop_column('projectslog', 'supervisor_id')
    op.drop_column('projectslog', 'group_id')
    # ### end Alembic commands ###


from typing import List

class Solution:
    def arrayStringsAreEqual(self, word1: List[str], word2: List[str]) -> bool:
        w1 = ""
        w2 = ""
        for w in word1:
            w1 += w
        for w in word2:
            w2 += w
        return w1 == w2

word1 = ["ab", "c"]
word2 = ["a", "bc"]

word1 = ["a", "cb"]
word2 = ["ab", "c"]
print(Solution.arrayStringsAreEqual(Solution, word1, word2))

1662 Check If Two String Arrays are Equivalent.py


import cv2

from pyzbar import pyzbar

def detect_qrcode(frame):

    barcodes = pyzbar.decode(frame)
    qrCodeDetector = cv2.QRCodeDetector()
    decodedText, points, _ = qrCodeDetector.detectAndDecode(frame)
    for barcode in barcodes:
        (x, y, w, h) = barcode.rect
        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2)
        barcodeData = barcode.data.decode("utf-8")
        barcodeType = barcode.type

        text = "{} ({})".format(barcodeData, barcodeType)
        cv2.putText(frame, text, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX,
                    0.5, (0, 0, 255), 2)

    cv2.imshow("Image", frame)
    cv2.waitKey()

if __name__ == '__main__':

    img = cv2.imread("/media/mensa/Data/Task/TrolleyDetection/train_data/new_train_data/IMG_20200401_141536.jpg")
    detect_qrcode(frame=img)


#!/usr/bin/env python
# coding: utf-8

# **Task 07: Querying RDF(s)**

# In[1]:

github_storage = "https://raw.githubusercontent.com/FacultadInformatica-LinkedData/Curso2021-2022/master/Assignment4/course_materials"

# Leemos el fichero RDF de la forma que lo hemos venido haciendo

# In[4]:

from rdflib import Graph, Namespace, Literal
from rdflib.namespace import RDF, RDFS
from rdflib.plugins.sparql import prepareQuery

g = Graph()
g.namespace_manager.bind('ns', Namespace("http://somewhere#"), override=False)
g.namespace_manager.bind('vcard', Namespace("http://www.w3.org/2001/vcard-rdf/3.0#"), override=False)
g.parse(github_storage+"/rdf/example6.rdf", format="xml")

# **TASK 7.1: List all subclasses of "Person" with RDFLib and SPARQL**

# In[202]:

# TO DO
ns = Namespace("http://somewhere#")

q1 = prepareQuery('''
  SELECT ?subcls WHERE {
    ?subcls rdfs:subClassOf+ ns:Person
  }
  ''', initNs = {
      "ns":ns,
      "rdfs":RDFS
  }
)

# Visualize the results
for r in g.query(q1):
    print(r.subcls)

# In[203]:

for s,p,o in g.triples((None, RDFS.subClassOf, ns.Person)):
    print(s,p,o)
    for s1,p1,o1 in g.triples((None, p, s)):
        print(s1,p1,o1)

# **TASK 7.2: List all individuals of "Person" with RDFLib and SPARQL (remember the subClasses)**
#

# In[204]:

for s,p1,o in g.triples((None, RDF.type, ns.Person)):
        print(s,p,o)

for s, p, o in g.triples((None, RDFS.subClassOf, ns.Person)):
    for s1,p1,o1 in g.triples((None, RDF.type, s)):
        print(s1,p1,o1)

# Visualize the results

# In[205]:

q2 = prepareQuery('''
  SELECT  ?names WHERE {
    ?class rdfs:subClassOf* ns:Person .
    ?names a ?class .
  }
  ''', initNs = {
      "ns":ns,
      "rdfs":RDFS,
  }
)

# Visualize the results
for r in g.query(q2):
    print(r.names)

# **TASK 7.3: List all individuals of "Person" and all their properties including their class with RDFLib and SPARQL**
#

# In[206]:

q3 = prepareQuery('''
  SELECT ?names ?prop ?val WHERE {
    ?cls rdfs:subClassOf* ns:Person .
    ?names a ?cls .
    ?names ?prop ?val

  }
  ''', initNs = {
      "ns":ns,
      "rdfs":RDFS,
  }
)

# Visualize the results
for r in g.query(q3):
  print(r.names, r.prop, r.val)
# Visualize the results

# In[207]:

def individuals_and_prop(sub, prop, cls):
    for s1, p1, o1 in g.triples((sub, prop, cls)):
        for s2, p2, o2 in g.triples((s1, None, None)):
            print(s2,p2,o2)

individuals_and_prop(None, RDF.type, ns.Person)
for s, p, o in g.triples((None, RDFS.subClassOf, ns.Person)):
    individuals_and_prop(None, RDF.type, s)

# In[ ]:


import os
import sys
import time
import warnings
from pathlib import Path
from typing import Optional

import lightning as L
import torch
import random
import ahocorasick

from pill_each.spatial_lora import lora_pill

# support running without installing as a package
wd = Path(__file__).parent.parent.resolve()
sys.path.append(str(wd))

from model_spatial import LLaMA_spatial, LLaMASpatialConfig
from lit_llama import LLaMA
from lit_llama.utils import EmptyInitOnDevice, lazy_load, llama_model_lookup, quantization
from lit_llama.tokenizer import Tokenizer
from scripts.prepare_alpaca import generate_prompt
from prepare import gen_sentence, prepare, POI_Find_Dict

# from datasets import load_dataset

lora_r = 8
lora_alpha = 16
lora_dropout = 0.05
max_seq_length = 256
which_devices = [5]
spatial_mark = True

@torch.no_grad()
def not_spatial_generate(
    model: LLaMA,
    idx: torch.Tensor,
    max_new_tokens: int,
    *,
    max_seq_length: Optional[int] = None,
    temperature: float = 1.0,
    top_k: Optional[int] = None,
    eos_id: Optional[int] = None,
) -> torch.Tensor:
    """Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.

    The implementation of this function is modified from A. Karpathy's nanoGPT.

    Args:
        model: The model to use.
        idx: Tensor of shape (T) with indices of the prompt sequence.
        max_new_tokens: The number of new tokens to generate.
        max_seq_length: The maximum sequence length allowed.
        temperature: Scales the predicted logits by 1 / temperature
        top_k: If specified, only sample among the tokens with the k highest probabilities
        eos_id: If specified, stop generating any more token once the <eos> token is triggered
    """
    # create an empty tensor of the expected final shape and fill in the current tokens
    T = idx.size(0)
    T_new = T + max_new_tokens
    if max_seq_length is None:
        max_seq_length = min(T_new, model.config.block_size)

    device, dtype = idx.device, idx.dtype
    # create an empty tensor of the expected final shape and fill in the current tokens
    empty = torch.empty(T_new, dtype=dtype, device=device)
    empty[:T] = idx
    idx = empty
    input_pos = torch.arange(0, T, device=device)

    if idx.device.type == "xla":
        import torch_xla.core.xla_model as xm

        xm.mark_step()

    # generate max_new_tokens tokens
    for _ in range(max_new_tokens):
        x = idx.index_select(0, input_pos).view(1, -1)

        # forward
        logits = model(x, max_seq_length, input_pos)
        logits = logits[0, -1] / temperature

        # optionally crop the logits to only the top k options
        if top_k is not None:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits = torch.where(logits < v[[-1]], -float("Inf"), logits)

        probs = torch.nn.functional.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, num_samples=1).to(dtype=dtype)

        # advance
        input_pos = input_pos[-1:] + 1

        if idx.device.type == "xla":
            xm.mark_step()

        # concatenate the new generation
        idx = idx.index_copy(0, input_pos, idx_next)

        # if <eos> token is triggered, return the output (stop generation)
        if idx_next == eos_id:
            return idx[:input_pos]  # include the EOS token

    return idx

@torch.no_grad()
def generate(
        model: LLaMA_spatial,
        idx: torch.Tensor,
        poi_mask: torch.Tensor,
        spatial_pill: torch.Tensor,
        spatial_scope: torch.Tensor,
        max_new_tokens: int,
        *,
        max_seq_length: Optional[int] = None,
        temperature: float = 1.05,
        top_k: Optional[int] = None,
        eos_id: Optional[int] = None,
) -> (torch.Tensor, torch.Tensor):
    """Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.

    The implementation of this function is modified from A. Karpathy's nanoGPT.

    Args:
        model: The model to use.
        idx: Tensor of shape (T) with indices of the prompt sequence.
        max_new_tokens: The number of new tokens to generate.
        max_seq_length: The maximum sequence length allowed.
        temperature: Scales the predicted logits by 1 / temperature
        top_k: If specified, only sample among the tokens with the k highest probabilities
        eos_id: If specified, stop generating any more token once the <eos> token is triggered
    """
    # create an empty tensor of the expected final shape and fill in the current tokens
    T = idx.size(0)
    T_new = T + max_new_tokens
    if max_seq_length is None:
        max_seq_length = min(T_new, model.config.block_size)
    max_seq_length = max_seq_length + max_new_tokens

    device = idx.device
    # create an empty tensor of the expected final shape and fill in the current tokens
    empty_1 = torch.empty(T_new, dtype=idx.dtype, device=device)
    empty_2 = torch.zeros(T_new, dtype=poi_mask.dtype, device=device)
    empty_1[:T] = idx
    empty_2[:T] = poi_mask
    idx = empty_1
    poi_mask = empty_2
    input_pos = torch.arange(0, T, device=device)
    # spatial_pill_ = spatial_pill
    # spatial_pill = torch.zeros((T_new, model.config.s_embd), dtype=spatial_pill_.dtype, device=device)
    # spatial_pill[:T] = spatial_pill_
    if idx.device.type == "xla":
        import torch_xla.core.xla_model as xm

        xm.mark_step()

    # generate max_new_tokens tokens
    for _ in range(max_new_tokens):
        x = idx.index_select(0, input_pos).view(1, -1)
        pm = poi_mask.index_select(0, input_pos).view(1, -1)
        # forward
        logits, coord = model(x, pm, spatial_pill.unsqueeze(0),
                              spatial_scope.unsqueeze(0), max_seq_length)
        logits = logits[0, -1] / temperature

        # optionally crop the logits to only the top k options
        if top_k is not None:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits = torch.where(logits < v[[-1]], -float("Inf"), logits)

        probs = torch.nn.functional.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, num_samples=1).to(dtype=idx.dtype)

        # advance
        input_pos = torch.arange(0, input_pos.shape[0]+1, device=device)

        if idx.device.type == "xla":
            xm.mark_step()

        # concatenate the new generation
        idx = idx.index_copy(0, input_pos[-1], idx_next)
        # spatial_pill = spatial_pill.index_copy(0, input_pos, coord[0, -1:, :])

        # if <eos> token is triggered, return the output (stop generation)
        if idx_next == eos_id:
            return idx[:input_pos[-1]]  # include the EOS token

    return idx, spatial_pill

def generate_sentence(instruction, poi_token, poi_list, tokenizer, device):
    places = [poi_list[token][1] for token in poi_token]
    places = " , ".join(places)
    prompt = instruction + places
    encoded = tokenizer.encode(prompt, bos=True, eos=False, device=device)
    return encoded

def generate_main(
        tokenizer_path: str = "../checkpoints/lit-llama/tokenizer.model",
        out_dir: str = "../out/pill/poi/",
        trainset_dir: str = "../data/spatial_dataset/poi_train.pt",
        validset_dir: str = "../data/spatial_dataset/poi_test.pt",
        poi_dir: str = "../data/spatial_dataset/poi_list.pt",
        lora_path: Path = Path("../out/lora/alpaca/lit-llama-lora-finetuned.pth"),
        checkpoint_path: Path = Path("../checkpoints/lit-llama/7B/lit-llama.pth"),
        spatial=spatial_mark
):
    lora_path = Path("../out/pill/poi/lora-spatial-pill-finetuned.pth")
    fabric = L.Fabric(accelerator="cuda", devices=which_devices, precision="bf16-true")
    fabric.launch()
    # fabric.seed_everything(1337 + fabric.global_rank)

    if fabric.global_rank == 0:
        os.makedirs(out_dir, exist_ok=True)

    # train_data, val_data = load_datasets(data_dir=data_dir)
    data = torch.load(trainset_dir)
    data.extend(torch.load(validset_dir))
    poi_list = torch.load(poi_dir)
    poi_finder = POI_Find_Dict(poi_list)
    llama_config = LLaMASpatialConfig.from_name("7B")
    tokenizer = Tokenizer(Path(tokenizer_path))
    data, tmp = prepare(data, poi_list, tokenizer,
                        max_seq_length=max_seq_length,
                        stage="generate")
    #valid_data, _ = prepare(valid_data, poi_list, tokenizer, padded_vocab_size=llama_config.padded_vocab_size)
    llama_config.max_poi_len, llama_config.bbox = tmp

    print("Loading model ...", file=sys.stderr)
    t0 = time.time()
    if spatial:
        with fabric.init_module(), lora_pill(r=lora_r, alpha=lora_alpha, dropout=lora_dropout, enabled=True):
            model = LLaMA_spatial(llama_config)
            # strict=False because missing keys due to LoRA weights not contained in checkpoint state
            if checkpoint_path is not None:
                with lazy_load(checkpoint_path) as pretrained_checkpoint:
                    model.load_state_dict(pretrained_checkpoint, strict=False)
                    print("Load from checkpoint:{}.".format(checkpoint_path))
            if lora_path is not None:
                with lazy_load(lora_path) as lora_checkpoint:
                    model.load_state_dict(lora_checkpoint, strict=False)
                    print("Load from lora:{}.".format(lora_path))
    else:
        with lazy_load(checkpoint_path) as checkpoint:
            name = llama_model_lookup(checkpoint)
            with fabric.init_module(empty_init=True):
                model = LLaMA.from_name(name)
            model.load_state_dict(checkpoint)
    print(f"Time to load model: {time.time() - t0:.02f} seconds.", file=sys.stderr)

    device = fabric.device
    model.eval()

    command = ""
    while command != "quit":
        try:
            ix = int(command)
        except:
            ix = random.randint(0, len(data) - 1)
            print("random index: {}".format(ix))

        lang_input = data[ix]["language_inputs"].type(torch.int64).to(device)
        lang_label = data[ix]["language_labels"].type(torch.int64).to(device)
        poi_mask = data[ix]["poi_mask"].type(torch.bool).to(device)
        raw_spatial = data[ix]["raw_spatial"].type(torch.cuda.BFloat16Tensor).to(device)
        spatial_scope = data[ix]["spatial_scopes"].type(torch.int64).to(device)

        sentence = tokenizer.decode(lang_label.view(-1))
        print("correct answer: \n{}".format(sentence))
        if spatial:
            output, coords = generate(
                model,
                idx=lang_input,
                poi_mask=poi_mask,
                spatial_pill=raw_spatial,
                spatial_scope=spatial_scope,
                max_seq_length=256,
                max_new_tokens=12,
            )
        else:
            output = not_spatial_generate(model, lang_input, max_new_tokens=12)
            model.reset_cache()
        output = tokenizer.decode(output.view(-1))
        print("predicted answer: \n{}".format(output))
        command = input()

class WordExtract:
    def __init__(self, keywords):
        self.keywords = keywords  # type(list)  list id equals to keyword id
        self.automat = self.build_automat()

    def build_automat(self):
        auto = ahocorasick.Automaton()
        for index, keyword in enumerate(self.keywords):
            auto.add_word(keyword, (index, keyword))
        auto.make_automaton()
        return auto

    def extract_keywords(self, text):
        keyword_matches = set()
        for end_index, (keyword_index, original_keyword) in self.automat.iter(text):
            start_index = end_index - len(original_keyword) + 1
            keyword_matches.add((keyword_index, start_index, end_index))
        return keyword_matches

def extract_sample():
    sample_text = "这是一个示例文本，其中包含一些关键词，如Python、关键字、文本处理等如Python。"
    sample_keywords = ["关键词", "Python", "文本处理"]

    WE = WordExtract(sample_keywords)
    result = WE.extract_keywords(sample_text)  # extract word
    print("文本中的关键词:", result)

if __name__ == "__main__":
    from jsonargparse import CLI

    torch.set_float32_matmul_precision("high")
    warnings.filterwarnings(
        # Triggered internally at ../aten/src/ATen/EmptyTensor.cpp:31
        "ignore",
        message="ComplexHalf support is experimental and many operators don't support it yet"
    )
    warnings.filterwarnings(
        # Triggered in bitsandbytes/autograd/_functions.py:298
        "ignore",
        message="MatMul8bitLt: inputs will be cast from torch.bfloat16 to float16 during quantization",
    )
    CLI(generate_main)


from tkinter import*
from tkinter.colorchooser import*
from tkinter.simpledialog import*
from PIL import ImageGrab

mycolor = 'black'

def mouseClick(event):
    global x1,y1,x2,y2
    x1=event.x
    y1=event.y

def mouseDrop(event):
    global x1,y1,x2,y2,penWidth,penColor
    x2=event.x
    y2=event.y
    if shape=="line":
        canvas.create_line(x1,y1,x2,y2,width=penWidth,fill=penColor)
    elif shape=="oval":
        canvas.create_oval(x1,y1,x2,y2,width=penWidth,fill=penColor,outline=penColor)
    elif shape=="rectangle":
        canvas.create_rectangle(x1,y1,x2,y2,width=penWidth,fill=penColor,outline=penColor)

def getColor():
    global penColor
    color=askcolor()
    penColor=color[1]

def getWidth():
    global penWidth
    penWidth=askinteger("선 두께","선 두께(1~10)를 입력하세요",minvalue=1,maxvalue=10)

def paint(event):
    x1,y1=(event.x),(event.y)
    x2,y2=(event.x+2),(event.y-2)
    canvas.create_oval(x1,y1,x2,y2, fill=mycolor)

def drawLine():
    global shape
    shape="line"

def drawOval():
    global shape
    shape="oval"

def drawRectangle():
    global shape
    shape="rectangle"

def eraser():
    canvas.delete("all")

def save(event):#화면 캡쳐 해서 파일에 저장
    x = window.winfo_rootx()
    y = window.winfo_rooty()#창의 위치 얻어오기
    w = window.winfo_width() + x
    h = window.winfo_height() + y#창의 크기 얻어오기

    box = (x, y, w, h)
    img=ImageGrab.grab(box)#창의 크기만큼만 이미지저장
    saveas='손예은_.bmp'
    img.save(saveas)

window=None
canvas=None
x1,y1,x2,y2=None,None,None,None #선의 시작점과 끝점
penColor='black' #선의 기본 색-검정
penWidth=5 #선의 기본 굵기-5
fillColor=None
shape="paint"

#if shape=="line":

#if shape=="oval":
#if shape=="rectangle":
if __name__=="__main__":
    window=Tk()
    window.title("그림판")
    canvas=Canvas(window,height=850,width=850)

canvas.bind("<Button-1>",mouseClick)
canvas.bind("<ButtonRelease-1>",mouseDrop)
canvas.bind('<B1-Motion>',paint)
canvas.bind("<Button-2>",eraser)
window.bind("<space>", save) #스페이스를 누루면 캡쳐 후 저장
canvas.pack()

mainMenu=Menu(window)
window.config(menu=mainMenu)

shapeMenu=Menu(mainMenu)
mainMenu.add_cascade(label="모양",menu=shapeMenu)
shapeMenu.add_command(label="선",command=drawLine)
shapeMenu.add_command(label="원",command=drawOval)
shapeMenu.add_command(label="사각형",command=drawRectangle)

fileMenu=Menu(mainMenu)
mainMenu.add_cascade(label="설정",menu=fileMenu)
fileMenu.add_command(label="색상 선택",command=getColor)
fileMenu.add_command(label="선 두께 결정",command=getWidth)
fileMenu.add_command(label="지우기",command=eraser)

window.mainloop()