from absl.testing import absltest
from absl.testing import parameterized
from flax import linen as nn
from flax.core import pop
from flax.linen import partitioning as flax_partitioning
import jax
import numpy as np

from flaxformer.components import embedding
from flaxformer.components import initializers

class EmbedTest(parameterized.TestCase):

  def test_embedder_raises_exception_for_incorrect_input_type(self):
    """Tests that inputs are integers and that an exception is raised if not."""
    embed = embedding.Embed(num_embeddings=10, features=5)
    inputs = np.expand_dims(np.arange(5, dtype=np.int64), 1)
    variables = embed.init(jax.random.PRNGKey(0), inputs)
    bad_inputs = inputs.astype(np.float32)
    with self.assertRaisesRegex(
        ValueError, 'Input type must be an integer or unsigned integer.'):
      _ = embed.apply(variables, bad_inputs)

  @parameterized.named_parameters(
      {
          'testcase_name': 'with_ones',
          'init_fn': jax.nn.initializers.ones,
          'num_embeddings': 10,
          'features': 5,
          'matrix_sum': 5 * 10,
      }, {
          'testcase_name': 'with_zeros',
          'init_fn': jax.nn.initializers.zeros,
          'num_embeddings': 10,
          'features': 5,
          'matrix_sum': 0,
      })
  def test_embedding_initializes_correctly(self, init_fn, num_embeddings,
                                           features, matrix_sum):
    """Tests if the Embed class initializes with the requested initializer."""
    embed = embedding.Embed(
        num_embeddings=num_embeddings,
        features=features,
        embedding_init=init_fn)
    inputs = np.expand_dims(np.arange(5, dtype=np.int64), 1)
    variables = embed.init(jax.random.PRNGKey(0), inputs)
    embedding_matrix = variables['params']['embedding']
    self.assertEqual(int(np.sum(embedding_matrix)), matrix_sum)

  def test_embedding_matrix_shape(self):
    """Tests that the embedding matrix has the right shape."""
    num_embeddings = 10
    features = 5
    embed = embedding.Embed(num_embeddings=num_embeddings, features=features)
    inputs = np.expand_dims(np.arange(features, dtype=np.int64), 1)
    variables = embed.init(jax.random.PRNGKey(0), inputs)
    embedding_matrix = variables['params']['embedding']
    self.assertEqual((num_embeddings, features), embedding_matrix.shape)

  def test_embedding_attend(self):
    """Tests that attending with ones returns sum of embedding vectors."""
    features = 5
    embed = embedding.Embed(num_embeddings=10, features=features)
    inputs = np.array([[1]], dtype=np.int64)
    variables = embed.init(jax.random.PRNGKey(0), inputs)
    query = np.ones(features, dtype=np.float32)
    result = embed.apply(variables, query, method=embed.attend)
    expected = np.sum(variables['params']['embedding'], -1)
    np.testing.assert_array_almost_equal(result, expected)

  def test_embedding_axis_names(self):
    rules = [
        ('my_batch_dim', 'data'),
        ('embed', None),
        ('vocab', None),
    ]
    with flax_partitioning.axis_rules(rules):
      embed = embedding.Embed(num_embeddings=10, features=5)
      inputs = np.array([1], dtype=np.int64)
      embed.init(
          jax.random.PRNGKey(0), inputs, input_axis_names=('my_batch_dim',))

  def test_embedding_axis_names_ctor(self):
    rules = [
        ('my_batch_dim', 'data'),
        ('embed', None),
        ('vocab', None),
    ]
    with flax_partitioning.axis_rules(rules):
      embed = embedding.Embed(
          num_embeddings=10, features=5, input_axis_names=('my_batch_dim',))
      inputs = np.array([1], dtype=np.int64)
      embed.init(jax.random.PRNGKey(0), inputs)

  def test_embedding_axis_names_call_overrides(self):
    rules = [
        ('my_batch_dim', 'data'),
        ('embed', None),
        ('vocab', None),
    ]
    with flax_partitioning.axis_rules(rules):
      embed = embedding.Embed(
          num_embeddings=10, features=5, input_axis_names=('other_batch_dim',))
      inputs = np.array([1], dtype=np.int64)
      embed.init(
          jax.random.PRNGKey(0), inputs, input_axis_names=('my_batch_dim',))

class MultiEmbedTest(parameterized.TestCase):

  def test_multi_embed_returns_correct_shape(self):
    """Tests that we can build a generic combined embedder."""
    features = 5
    embedders = {
        'token_embed': nn.Embed(num_embeddings=10, features=features),
        'segment_embed': nn.Embed(num_embeddings=2, features=features),
        'position_embed': nn.Embed(num_embeddings=12, features=features)
    }
    model = embedding.MultiEmbed(embedders)

    token_ids = np.array([[0, 1, 2], [3, 4, 5]], dtype=np.int64)
    segment_ids = np.array([[0, 1, 1], [0, 0, 1]], dtype=np.int64)
    position_ids = np.arange(3, dtype=np.int64)[None]
    output, _ = model.init_with_output(
        jax.random.PRNGKey(0),
        token_embed=token_ids,
        segment_embed=segment_ids,
        position_embed=position_ids)
    self.assertEqual(output.shape, token_ids.shape + (features,))

  @parameterized.named_parameters(
      {
          'testcase_name': 'with_sum_method',
          'method': embedding.EmbedCombineMethod.SUM,
          'features': 7,
          'expected_shape': (2, 3, 7),
      }, {
          'testcase_name': 'with_concat_method',
          'method': embedding.EmbedCombineMethod.CONCAT,
          'features': 7,
          'expected_shape': (2, 3, 14),
      })
  def test_multi_embed_combines_embeddings_correctly(self, method, features,
                                                     expected_shape):
    """Tests that embeddings are correctly summed or concatenated."""
    embedders = {
        'token_embed': nn.Embed(num_embeddings=10, features=features),
        'segment_embed': nn.Embed(num_embeddings=2, features=features)
    }
    model = embedding.MultiEmbed(embedders)

    token_ids = np.array([[0, 1, 2], [3, 4, 5]], dtype=np.int64)
    segment_ids = np.array([[0, 1, 1], [0, 0, 1]], dtype=np.int64)
    variables = model.init(
        jax.random.PRNGKey(0), token_embed=token_ids, segment_embed=segment_ids)
    embeddings = model.apply(
        variables,
        token_embed=token_ids,
        segment_embed=segment_ids,
        method=model.get_individual_embeddings)
    output = model.apply(
        variables,
        embeddings,
        combine_method=method,
        method=model.combine_embeddings)
    self.assertEqual(output.shape, expected_shape)

  def test_multi_embed_combine_embeddings_raises_exception(self):
    """Tests that an exception is raised for an invalid combine method."""
    model = embedding.MultiEmbed({
        'token_embed': nn.Embed(num_embeddings=10, features=5),
    })
    token_ids = np.array([[0, 1, 2], [3, 4, 5]], dtype=np.int64)
    variables = model.init(jax.random.PRNGKey(0), token_embed=token_ids)
    embeddings = {'token_embed': np.ones([2, 3, 5], dtype=np.float32)}
    with self.assertRaises(ValueError):
      _ = model.apply(
          variables,
          embeddings,
          combine_method='invalid_combine_method',
          method=model.combine_embeddings)

  def test_multi_embed_can_return_individual_embeddings(self):
    """Tests that MultiEmbed returns a dictionary with each embedded input."""
    features = 5
    embedders = {
        'token_embed': nn.Embed(num_embeddings=10, features=features),
        'segment_embed': nn.Embed(num_embeddings=2, features=features),
        'position_embed': nn.Embed(num_embeddings=12, features=features)
    }
    model = embedding.MultiEmbed(embedders)

    token_ids = np.array([[0, 1, 2], [3, 4, 5]], dtype=np.int64)
    segment_ids = np.array([[0, 1, 1], [0, 0, 1]], dtype=np.int64)
    position_ids = np.arange(3, dtype=np.int64)[None]
    variables = model.init(
        jax.random.PRNGKey(0),
        token_embed=token_ids,
        segment_embed=segment_ids,
        position_embed=position_ids)
    embeddings = model.apply(
        variables,
        token_embed=token_ids,
        segment_embed=segment_ids,
        position_embed=position_ids,
        method=model.get_individual_embeddings)
    self.assertIn('token_embed', embeddings)
    self.assertIn('segment_embed', embeddings)
    self.assertIn('position_embed', embeddings)

  def test_multi_embed_returns_nonempty_gradient(self):
    """Tests that we can capture a non-empty gradient from MultiEmbed."""
    features = 5
    embedders = {
        'token_embed': nn.Embed(num_embeddings=10, features=features),
        'segment_embed': nn.Embed(num_embeddings=2, features=features),
        'position_embed': nn.Embed(num_embeddings=12, features=features)
    }
    model = embedding.MultiEmbed(
        embedders, sow_intermediates=True, capture_gradients=True)
    token_ids = np.array([[0, 1, 2], [3, 4, 5]], dtype=np.int64)
    segment_ids = np.array([[0, 1, 1], [0, 0, 1]], dtype=np.int64)
    position_ids = np.arange(3, dtype=np.int64)[None]
    variables = model.init(
        jax.random.PRNGKey(0),
        token_embed=token_ids,
        segment_embed=segment_ids,
        position_embed=position_ids)
    self.assertContainsSubset(('grads',), variables)

    def fake_loss(variables, token_ids, segment_ids, position_ids):
      """Returns a loss."""
      output, _ = model.apply(
          variables,
          token_embed=token_ids,
          segment_embed=segment_ids,
          position_embed=position_ids,
          mutable=['grads'])
      return output.sum()

    grad_fn = jax.grad(fake_loss)
    grads_variables = grad_fn(variables, token_ids, segment_ids, position_ids)
    grads = grads_variables['grads']
    self.assertContainsSubset(('output_grad',), grads)
    self.assertNotAlmostEqual(grads['output_grad'].sum(), 0.0)

def _arange_init(key, shape, dtype):
  del key
  return np.arange(np.prod(shape), dtype=dtype).reshape(shape)

class EmbeddingTest(parameterized.TestCase):

  def test_add_position_embs(self):
    """Tests that positional embeddings are correctly applied."""
    positions = np.array([[[0, 1, 2], [0, 1, 2]]], dtype=np.int32)
    position_embedder = embedding.PositionEmbed(
        num_embeddings=10,
        features=2,
        dtype=np.float32,
        embedding_init=_arange_init,
    )

    variables = position_embedder.init(jax.random.PRNGKey(0), inputs=positions)
    output_embeddings = position_embedder.apply(variables, inputs=positions)

    np.testing.assert_array_equal(
        output_embeddings,
        [[[[0, 1], [2, 3], [4, 5]], [[0, 1], [2, 3], [4, 5]]]],
    )

  def test_add_position_embs_decoder(self):
    """Tests that position embeddings are correctly applied in decoding mode."""
    positions = np.array([[[0, 1, 2], [0, 1, 2]]], dtype=np.int32)
    position_embedder = embedding.PositionEmbed(
        num_embeddings=10,
        features=3,
        dtype=np.float32,
        embedding_init=_arange_init,
    )
    variables = position_embedder.init(
        jax.random.PRNGKey(0), inputs=positions, decode=True
    )
    state, params = pop(variables, 'params')

    output_embeddings, state = position_embedder.apply(
        {'params': params, **state},
        inputs=positions,
        decode=True,
        mutable=['cache'],
    )

    np.testing.assert_array_equal(output_embeddings, [[[0, 1, 2]]])
    np.testing.assert_array_equal(
        state['cache']['position_embedder_index'], [1]
    )

    # Test that repeated access increments the cache_index.
    output_embeddings, state = position_embedder.apply(
        {'params': params, **state},
        inputs=positions,
        decode=True,
        mutable=['cache'],
    )

    np.testing.assert_array_equal(output_embeddings, [[[3, 4, 5]]])
    np.testing.assert_array_equal(
        state['cache']['position_embedder_index'], [2]
    )

  def test_add_position_embs_batch_decoder(self):
    """Tests that position embeddings are correctly applied in decoding mode."""
    positions = np.array([[[0, 1, 2]], [[0, 1, 2]]], dtype=np.int32)
    position_embedder = embedding.PositionEmbed(
        num_embeddings=10,
        features=2,
        dtype=np.float32,
        embedding_init=_arange_init,
    )
    variables = position_embedder.init(
        jax.random.PRNGKey(0), inputs=positions, decode=True
    )
    state, params = pop(variables, 'params')
    state['cache']['position_embedder_index'] = np.array([0, 2])

    output_embeddings, state = position_embedder.apply(
        {'params': params, **state},
        inputs=positions,
        decode=True,
        mutable=['cache'],
    )

    np.testing.assert_array_equal(output_embeddings, [[[2, 3]], [[6, 7]]])
    np.testing.assert_array_equal(
        state['cache']['position_embedder_index'], [1, 3]
    )

    # Test that repeated access increments the cache_index.
    output_embeddings, state = position_embedder.apply(
        {'params': params, **state},
        inputs=positions,
        decode=True,
        mutable=['cache'],
    )

    np.testing.assert_array_equal(output_embeddings, [[[4, 5]], [[8, 9]]])
    np.testing.assert_array_equal(
        state['cache']['position_embedder_index'], [2, 4]
    )

  def test_add_sinusoidal_position_embs(self):
    """Tests that sinusoidal positional embeddings are applied."""
    positions = np.array([[0, 1]])
    sinusoid_embedder = embedding.FixedEmbed(
        embedding_init=initializers.sinusoidal(),
        features=5,
        max_length=10,
        dtype=np.float32)

    variables = sinusoid_embedder.init(jax.random.PRNGKey(0), inputs=positions)
    output_embeddings = sinusoid_embedder.apply(variables, inputs=positions)

    expected_output_embeddings = np.array(
        [[[0.0, 0.0, 1.0, 1.0, 0.0], [0.84147, 0.0001, 0.540302, 1.0, 0.0]]],
        dtype=np.float32)
    np.testing.assert_array_almost_equal(output_embeddings,
                                         expected_output_embeddings)

  def test_regression_add_position_embeddings_sine(self):
    sequence_length = 7
    hidden_dim = 5
    positions = np.arange(sequence_length)[None, :]
    embed = embedding.FixedEmbed(
        embedding_init=initializers.sinusoidal(),
        features=hidden_dim,
        max_length=32,
        dtype=np.float32)
    outputs, params = embed.init_with_output(
        jax.random.PRNGKey(0), inputs=positions)
    self.assertEqual(params, {})
    expected = np.array([
        [
            [0.0, 0.0, 1.0, 1.0, 0.0],
            [
                0.8414709568023682, 9.999999747378752e-05, 0.5403022766113281,
                1.0, 0.0
            ],
            [
                0.9092974066734314, 0.00019999999494757503, -0.416146844625473,
                1.0, 0.0
            ],
            [
                0.14112000167369843, 0.00029999998514540493,
                -0.9899924993515015, 0.9999999403953552, 0.0
            ],
            [
                -0.756802499294281, 0.00039999998989515007, -0.6536436080932617,
                0.9999999403953552, 0.0
            ],
            [
                -0.9589242935180664, 0.0004999999655410647, 0.28366219997406006,
                0.9999998807907104, 0.0
            ],
            [
                -0.279415488243103, 0.0005999999702908099, 0.9601702690124512,
                0.9999998211860657, 0.0
            ],
        ],
    ])
    np.testing.assert_allclose(outputs, expected, rtol=1e-5)

  def test_regression_add_position_embeddings_learned(self):
    sequence_length = 7
    hidden_dim = 5
    positions = np.arange(sequence_length)[None, :]
    embeds = embedding.PositionEmbed(
        num_embeddings=32,
        features=hidden_dim,
        embedding_init=jax.nn.initializers.normal(
            stddev=1e-6),  # Use learned embeds.
        dtype=np.float32)
    outputs, params = embeds.init_with_output(
        jax.random.PRNGKey(0), inputs=positions)
    param_shapes = jax.tree_map(lambda x: x.shape, params)
    self.assertEqual(param_shapes['params'], {
        'pos_embedding': (32, 5),
    })
    np.testing.assert_allclose(
        outputs, [[[
            4.486782358981145e-07, -3.5457077274259063e-07,
            1.0727929122822388e-07, 1.3890753791656607e-07,
            -7.792502287884417e-07
        ],
                   [
                       -2.3853303332543874e-07, 8.467901011499634e-07,
                       -8.650059157844225e-08, 3.407756707929366e-07,
                       -2.226069852895307e-07
                   ],
                   [
                       -1.1156219414942825e-07, -6.10807489920262e-07,
                       1.5079198192324839e-06, 1.714607265057566e-07,
                       -1.9881858861481305e-06
                   ],
                   [
                       2.980690965159738e-07, 4.950079102172822e-08,
                       1.9940485174174682e-07, 7.134963198041078e-07,
                       1.5179145975707797e-06
                   ],
                   [
                       7.398467261054975e-08, -4.966278197571228e-07,
                       8.895806757891478e-08, -4.5914887891740364e-07,
                       1.1332289204801782e-06
                   ],
                   [
                       1.3519083950086497e-06, -1.0249741535517387e-06,
                       -1.1988712458332884e-06, -6.748282288526752e-08,
                       1.18359389489342e-06
                   ],
                   [
                       -1.6119436168082757e-06, 1.0756450308235799e-07,
                       -8.505542155035073e-07, 1.1778743100876454e-06,
                       -1.000783186100307e-06
                   ]]],
        rtol=1e-5)

class HashEmbedTest(parameterized.TestCase):

  def test_hash_embedder(self):
    """Checks the parameters and return value shapes."""
    num_tables = 2
    num_embeddings_per_table = 32
    features = 16
    embedder = embedding.HashEmbed(
        features=features,
        num_embeddings_per_table=num_embeddings_per_table,
        num_tables=num_tables)

    key = jax.random.PRNGKey(0)
    batch_size = 2
    seq_len = 8
    ids = np.ones([batch_size, seq_len], dtype=np.int32)
    outputs, variables = embedder.init_with_output(key, ids)
    self.assertSequenceEqual(outputs.shape, (batch_size, seq_len, features))

    param_shapes = jax.tree_map(lambda x: list(x.shape), variables['params'])
    self.assertSameStructure(
        param_shapes, {
            'hash_embedder_table_0': {
                'embedding': [num_embeddings_per_table, features // num_tables]
            },
            'hash_embedder_table_1': {
                'embedding': [num_embeddings_per_table, features // num_tables]
            },
        })

  def test_hash_embedder_4d(self):
    """Checks the parameters and return value shapes."""
    num_tables = 2
    num_embeddings_per_table = 32
    features = 16
    embedder = embedding.HashEmbed(
        features=features,
        num_embeddings_per_table=num_embeddings_per_table,
        num_tables=num_tables)

    key = jax.random.PRNGKey(0)
    batch_size = 2
    seq_len = 8
    another_dim = 4
    ids = np.ones([batch_size, seq_len, another_dim], dtype=np.int32)
    outputs, variables = embedder.init_with_output(key, ids)
    self.assertSequenceEqual(outputs.shape,
                             (batch_size, seq_len, another_dim, features))

    param_shapes = jax.tree_map(lambda x: list(x.shape), variables['params'])
    self.assertSameStructure(
        param_shapes, {
            'hash_embedder_table_0': {
                'embedding': [num_embeddings_per_table, features // num_tables]
            },
            'hash_embedder_table_1': {
                'embedding': [num_embeddings_per_table, features // num_tables]
            },
        })

class NgramHashEmbedTest(parameterized.TestCase):

  @parameterized.product(
      batch_sizes=[(2,), (2, 3)],
      use_segment_ids=[False, True],
  )
  def test_hash_embedder(self, batch_sizes, use_segment_ids):
    """Checks the parameters and return value shapes."""
    num_tables = 2
    num_embeddings_per_table = 32
    features = 16
    embedder = embedding.NgramHashEmbed(
        ngram_orders=[1, 3, 4],
        padding_id=0,
        features=features,
        num_embeddings_per_table=num_embeddings_per_table,
        num_tables=num_tables)

    key = jax.random.PRNGKey(0)
    seq_len = 8
    ids = np.ones([*batch_sizes, seq_len], dtype=np.int32)
    segment_ids = (
        np.tile([[0] * 5 + [1] * 3],
                (*batch_sizes, 1)) if use_segment_ids else None)
    outputs, variables = (
        embedder.init_with_output(key, ids, segment_ids=segment_ids))
    self.assertSequenceEqual(outputs.shape, (*batch_sizes, seq_len, features))

    param_shapes = jax.tree_map(lambda x: list(x.shape), variables['params'])
    expected_table_shape = [num_embeddings_per_table, features // num_tables]
    self.assertSameStructure(
        param_shapes, {
            '1gram_hash_embed_table_0': {
                'embedding': expected_table_shape
            },
            '1gram_hash_embed_table_1': {
                'embedding': expected_table_shape
            },
            '3gram_hash_embed_table_0': {
                'embedding': expected_table_shape
            },
            '3gram_hash_embed_table_1': {
                'embedding': expected_table_shape
            },
            '4gram_hash_embed_table_0': {
                'embedding': expected_table_shape
            },
            '4gram_hash_embed_table_1': {
                'embedding': expected_table_shape
            },
        })

  @parameterized.product(
      batch_sizes=[(2,), (2, 3)],)
  def test_packing_correctness(self, batch_sizes):
    """Checks the parameters and return value shapes."""
    num_tables = 2
    num_embeddings_per_table = 32
    features = 16
    embedder = embedding.NgramHashEmbed(
        ngram_orders=[1, 3, 4],
        padding_id=0,
        features=features,
        num_embeddings_per_table=num_embeddings_per_table,
        num_tables=num_tables)

    key = jax.random.PRNGKey(0)
    seq_len = 8
    segment_ids = np.tile([[0] * 5 + [1] * 3], (*batch_sizes, 1))

    ids1 = np.ones([*batch_sizes, seq_len], dtype=np.int32)
    outputs1, variables = (
        embedder.init_with_output(key, ids1, segment_ids=segment_ids))

    # Run the embedder again, with the same IDs passed in for segment=0 IDs, but
    # different IDs for segment=1.
    ids2 = ids1 + segment_ids
    outputs2 = embedder.apply(variables, ids2, segment_ids=segment_ids)

    # Verify that the change to segment=1 didn't alter the outputs of segment=0.
    np.testing.assert_allclose(outputs1[..., :5, :], outputs2[..., :5, :])

class RotaryTest(absltest.TestCase):

  def test_rotary_embedding(self):
    """Checks the shape of rotary encodings."""
    batch = 2
    qlen = 3
    qheads = 4
    d = 2 * 5
    klen = 6
    kheads = 7
    maxlen = 8

    q = np.ones((batch, qlen, qheads, d))
    k = np.ones((batch, klen, kheads, d))
    cos = np.ones((maxlen, d))
    sin = np.ones((maxlen, d))
    out_q, out_k = embedding.apply_rotary_embedding(q, k, cos, sin)
    self.assertEqual(out_q.shape, q.shape)
    self.assertEqual(out_k.shape, k.shape)

  def test_rotary_embedding_multiquery(self):
    """Checks the shape of rotary encodings."""
    batch = 2
    qlen = 3
    qheads = 4
    d = 2 * 5
    klen = 6
    maxlen = 8

    q = np.ones((batch, qlen, qheads, d))
    k = np.ones((batch, klen, d))
    cos = np.ones((maxlen, d))
    sin = np.ones((maxlen, d))
    out_q, out_k = embedding.apply_rotary_embedding(q, k, cos, sin)
    self.assertEqual(out_q.shape, q.shape)
    self.assertEqual(out_k.shape, k.shape)

  def test_rotary_embedding_decode(self):
    """Checks the shape of rotary encodings."""
    batch = 2
    qlen = 1
    qheads = 4
    d = 2 * 5
    klen = 6
    maxlen = 8

    q = np.ones((batch, qlen, qheads, d))
    k = np.ones((batch, klen, d))
    cos = np.ones((maxlen, d))
    sin = np.ones((maxlen, d))
    rotary_index = np.ones((batch,), dtype=np.int32)
    out_q, out_k = embedding.apply_rotary_embedding(
        q, k, cos, sin, decode=True, rotary_index=rotary_index)
    self.assertEqual(out_q.shape, q.shape)
    self.assertEqual(out_k.shape, k.shape)

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


from collections import defaultdict
from math import comb

def take_training(courses):
    topic_count, difficulty_count = defaultdict(int), defaultdict(int)
    for topic, difficulty in courses:
        topic_count[topic] += 1
        difficulty_count[difficulty] += 1

    bad = sum((difficulty_count[difficulty] - 1) * (topic_count[topic] - 1)
              for topic, difficulty in courses)
    return comb(len(courses), 3) - bad

if __name__ == "__main__":
    results = []
    for _ in range(int(input())):
        inputs = []
        for _ in range(int(input())):
            inputs.append(list(map(int, input().split())))
        results.append(take_training(inputs))
    print(*results, sep="\n")


# Author: Alexander Shen
# Date: 11/28/2020
# Description: Replicate Focus/Domination board game with the win condition of capturing 6 pieces.

class Player:
    """
    Player class represents blueprint to create a Player object.
    Stores the name and color according to player's tuple.
    Tracks the number of captured and reserved pieces.
    Can add a capture and reserve piece.
    Can return all parameters in init method.
    """

    def __init__(self, profile):
        """ Initializes the name and color of the player along with their captured and reserved pieces."""
        self._player_name = profile[0]
        self._player_color = profile[1].upper()
        self._capture = 0
        self._reserve = 0

    def add_capture(self):
        """ Adds a captured piece to the player's collection"""
        self._capture += 1

    def add_reserve(self):
        """ Adds a reserved piece to the player's collection."""
        self._reserve += 1

    def use_reserve(self):
        """ Subtracts a reserved piece from the player's collection."""
        self._reserve -= 1

    def get_name(self):
        """ Returns a player's name. """
        return self._player_name

    def get_color(self):
        """ Returns a player's color."""
        return self._player_color

    def get_capture(self):
        """ Returns the number of captured pieces in a player's collection."""
        return self._capture

    def get_reserve(self):
        """ Returns the number of reserved pieces in a player's collection."""
        return self._reserve

class FocusGame:
    """
    FocusGame class creates a blueprint to create a FocusGame object.
    Creates the board and player objects.
    establishes the turn which is always the first player.
    """

    def __init__(self, first, second):
        """ Initializes the player objects, the board, and the first turn."""
        self._first_player = Player(first)
        self._second_player = Player(second)
        self._turn = self._first_player
        self._board = [
            [["R"], ["R"], ["G"], ["G"], ["R"], ["R"]],
            [["G"], ["G"], ["R"], ["R"], ["G"], ["G"]],
            [["R"], ["R"], ["G"], ["G"], ["R"], ["R"]],
            [["G"], ["G"], ["R"], ["R"], ["G"], ["G"]],
            [["R"], ["R"], ["G"], ["G"], ["R"], ["R"]],
            [["G"], ["G"], ["R"], ["R"], ["G"], ["G"]]
        ]

    def which_player(self, player_name):
        """ Given the player's name, return the object."""
        if player_name.upper() == self._first_player.get_name().upper():
            return self._first_player
        else:
            return self._second_player

    def change_turn(self, player):
        """ Changes the player's turn"""
        if player == self._first_player:
            self._turn = self._second_player
        else:
            self._turn = self._first_player

    def show_pieces(self, location):
        """ Returns a list of pieces within a given a tuple."""
        return self._board[location[0]][location[1]]

    def show_captured(self, player):
        """ Returns player's number of captured pieces."""
        return self.which_player(player).get_capture()

    def show_reserve(self, player):
        """ Returns player's number of reserved pieces"""
        return self.which_player(player).get_reserve()

    def print_board(self):
        """ Prints the board"""
        for lines in range(6):
            print(self._board[lines])

    def move_piece(self, player, select, move, number_pieces):
        """
        Move the desired amount of pieces to a location.
        It will call the check function to see if any input parameters are invalid.
        If the inputs are correct, then it will proceed to move the piece.
        If the movement causes a stack over 5, then it will call overflow and reduce it down to 5.
        It will also check winning conditions before displaying successfully moved.
        """
        if self.check(player, select, move, number_pieces):
            player_profile = self.which_player(player)
            piece_select = self._board[select[0]][select[1]]
            bottom_place = len(piece_select) - number_pieces
            for num in range(number_pieces):
                bottom_piece = piece_select[bottom_place]
                self._board[select[0]][select[1]].pop(bottom_place)
                self._board[move[0]][move[1]].append(bottom_piece)
            if len(self._board[move[0]][move[1]]) > 5:
                self.overflow(player, move)
                if player_profile.get_capture() == 6:
                    return player_profile.get_name() + " Wins"
            return "successfully moved"

        else:
            return False

    def reserved_move(self, player, move):
        """
        Move a piece from the player's reserve onto an existing stack.
        It will append a piece to the stack if the player's reserve is > 0.
        It will then check to see if the player has won, if the stack is over 5.
        """
        try:
            if self._board[move[0]][move[1]]:
                pass
        except IndexError:
            return False
        if move[0] < 0 or move[1] < 0:
            return False

        player_profile = self.which_player(player)
        if player_profile.get_reserve() > 0 and player_profile == self._turn:
            self._board[move[0]][move[1]].append(player_profile.get_color())
            player_profile.use_reserve()
            self.change_turn(player_profile)
            if len(self._board[move[0]][move[1]]) > 5:
                self.overflow(player, move)
                if player_profile.get_capture() == 6:
                    return player_profile.get_name() + " Wins"
            return "successfully moved"
        else:
            return False

    def overflow(self, player, move):
        """
        If the stack is over 5, then it will loop until it is not greater than 5.
        The bottom piece will be removed and added to the correct counter based on the player of the current turn.
        """
        player_profile = self.which_player(player)
        piece_move = self._board[move[0]][move[1]]
        while len(self._board[move[0]][move[1]]) > 5:
            bottom_piece = piece_move[0]
            self._board[move[0]][move[1]].pop(0)
            if bottom_piece == player_profile.get_color():
                player_profile.add_reserve()
            if bottom_piece != player_profile.get_color():
                player_profile.add_capture()

    def check(self, player, select, move, number_pieces):
        """
        Returns False if given the wrong move.
        This function is used to cover all the ways the piece can cause an invalid location.
        The player cannot move diagonal.
        The player cannot move an empty location.
        The player cannot source from or move to any location outside of the board.
        The player cannot move a piece that is not their color.
        The player cannot move to a location that does not match the number of pieces to be moved.
        The player cannot move if it not their turn.
        The player cannot move an incorrect number of pieces.
        """
        try:
            if self._board[move[0]][move[1]] and self._board[select[0]][select[1]]:
                pass
        except IndexError:
            return False
        piece_select = self._board[select[0]][select[1]]
        player_profile = self.which_player(player)
        if move[0] < 0 or move[1] < 0 or select[0] < 0 or select[1] < 0:
            return False
        if move[0] != select[0] and move[1] != select[1]:
            return False
        if not piece_select:
            return False
        if piece_select[len(piece_select) - 1] != player_profile.get_color():
            return False
        if number_pieces != abs(move[0] - select[0]) and move[1] == select[1]:
            return False
        if number_pieces != abs(move[1] - select[1]) and move[0] == select[0]:
            return False
        if player_profile != self._turn:
            return False
        if number_pieces > len(self._board[select[0]][select[1]]) or number_pieces <= 0:
            return False

        self.change_turn(player_profile)
        return True


"""
576. 出界的路径数
给定一个 m × n 的网格和一个球。球的起始坐标为 (i,j) ，你可以将球移到相邻的单元格内，或者往上、下、左、右四个方向上移动使球穿过网格边界。但是，你最多可以移动 N 次。找出可以将球移出边界的路径数量。答案可能非常大，返回 结果 mod 109 + 7 的值。

示例 1：

输入: m = 2, n = 2, N = 2, i = 0, j = 0
输出: 6
解释:

示例 2：

输入: m = 1, n = 3, N = 3, i = 0, j = 1
输出: 12
解释:

说明:

球一旦出界，就不能再被移动回网格内。
网格的长度和高度在 [1,50] 的范围内。
N 在 [0,50] 的范围内。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/out-of-boundary-paths
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。
"""
from functools import lru_cache

class Solution:
    def findPaths(self, m: int, n: int, maxMove: int, startRow: int, startColumn: int) -> int:
        """
        思路：记忆化递归
        1. 记录每个位置和剩余次数对应的出圈次数
        @param m:
        @param n:
        @param maxMove:
        @param startRow:
        @param startColumn:
        @return:
        """
        pos = [(0, 1), (0, -1), (1, 0), (-1, 0)]
        mod = 10 ** 9 + 7

        @lru_cache(None)
        def dfs(i, j, s):
            if s < 0: return 0
            if i < 0 or i == m or j < 0 or j == n: return 1
            ans = 0
            for dx, dy in pos:
                ans += dfs(i + dx, j + dy, s - 1)
                ans %= mod
            return ans

        res = dfs(startRow, startColumn, maxMove)
        dfs.cache_clear()
        return res % mod

if __name__ == '__main__':
    m = 2
    n = 2
    N = 2
    i = 0
    j = 0
    print(Solution().findPaths(m, n, N, i, j))


import sys
sys.path.insert(0, 'code/')
from process_data import *
from utils import *
from val import *
from bilstm_crf_model import *
import argparse
from keras.callbacks import ModelCheckpoint
from numpy.random import seed
from tensorflow import set_random_seed
seed(1)
set_random_seed(2)

if __name__ == '__main__':
    #Preprocessing
    data_makeup()
    print("############## Ready! ###############")
    get_train_data(data_dir='data/train/')

    parser = argparse.ArgumentParser(description='TRAIN')
    parser.add_argument('--num', type=int, default=0)
    parser.add_argument('--embed', type=int, default=300)
    parser.add_argument('--units', type=int, default=300)
    parser.add_argument('--epoch', type=int, default=4)
    parser.add_argument('--gpu', type=int, default=0)
    #parser.add_argument('--save', type=str)
    parser.add_argument('--batch', type=int, default=64)

    args = parser.parse_args()
    gpu_config(args.gpu)

    # Build model
    model, (train_x, train_y), (test_x, test_y) = bilstm_crf_model.create_model(args.embed, args.units)

    # only get the best single model
    filepath = 'model.h5'
    checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=True, mode='max')

    # Train
    model.fit(train_x, train_y, batch_size=args.batch, epochs=args.epoch, validation_data=[test_x, test_y], callbacks=[checkpoint])

    # Predict
    model_dir = 'model.h5'
    model, (word2idx, chunk_tags) = bilstm_crf_model.create_model(args.embed, args.units, train=False)
    model.load_weights(model_dir)

    test_dir = 'data/ruijin_round1_test_b_20181112/'
    submit_dir = 'submit/'
    test(test_dir, submit_dir, model, word2idx, chunk_tags)


import pendulum
from airflow.decorators import dag, task
from api_interface.get_stations_waqi import get_stations_waqi
from db.db_engine import get_db
from shared_models.stations_waqi import Waqi_Stations, Waqi_Stations_Temp
from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert

@dag(
    schedule="@daily",
    start_date=pendulum.datetime(2023, 1, 1, tz="UTC"),
    catchup=False,
    tags=["stations"],
)
def etl_stations_waqi():
    """This dag retrieves all stations from the World Air Quality Index
    project: https://aqicn.org/api/"""

    @task()
    def create_stations_temp():
        with get_db() as db:
            engine = db.get_bind()
            if engine.has_table("stations_waqi_temp"):
                Waqi_Stations_Temp.__table__.drop(db.get_bind())
            Waqi_Stations_Temp.__table__.create(db.get_bind())

    @task()
    def get_stations():
        stations = get_stations_waqi()
        return stations

    @task()
    def load_stations_temp(waqi_stations):
        with get_db() as db:
            db.execute(insert(Waqi_Stations_Temp), waqi_stations)
            db.commit()

    @task()
    def load_stations():
        with get_db() as db:
            insert_stmt = insert(Waqi_Stations).from_select(
                (
                    Waqi_Stations.station_id,
                    Waqi_Stations.station_name,
                    Waqi_Stations.latitude,
                    Waqi_Stations.longitude,
                    Waqi_Stations.request_datetime,
                    Waqi_Stations.data_datetime,
                ),
                select(Waqi_Stations_Temp),
            )

            upsert = insert_stmt.on_conflict_do_update(
                index_elements=[Waqi_Stations.station_id],
                set_={
                    Waqi_Stations.station_name: insert_stmt.excluded.station_name,
                    Waqi_Stations.latitude: insert_stmt.excluded.latitude,
                    Waqi_Stations.longitude: insert_stmt.excluded.longitude,
                    # Waqi_Stations.request_datetime:
                    # Waqi_Stations... insert_stmt.excluded.request_datetime,
                    Waqi_Stations.data_datetime: insert_stmt.excluded.data_datetime,
                },
            )
            db.execute(upsert)
            db.commit()

    task_1 = create_stations_temp()
    task_2 = get_stations()
    task_3 = load_stations_temp(task_2)
    task_4 = load_stations()
    task_1 >> task_2 >> task_3 >> task_4

etl_stations_waqi()


import os
from datetime import datetime

import Settings
import utils
from aio import Pipe, TagStreamParser

class PlayerFileSaverPipe(Pipe):
    def __init__(self, folder=""):
        super().__init__()
        self.file_dict = {}
        self.folder = folder
        if not os.path.exists(folder):
            os.makedirs(folder)

    def get_file(self, tag):
        player_id = utils.tag_to_player_id(tag)
        if player_id not in self.file_dict:
            file_name = self.folder + "/id" + player_id + ".txt"
            self.file_dict[player_id] = open(file_name, 'w')
        return self.file_dict[player_id]

    def next(self, data) -> object:
        copy = utils.csv(data) + "\n"
        self.get_file(data[TagStreamParser.TAG]).write(copy)
        return data

    def on_close(self):
        for file in self.file_dict.values():
            file.flush()
            file.close()


from flask import Flask, render_template, request
from flask_pymongo import PyMongo
from pymongo import MongoClient
import json

add = Flask(__name__)

add.config['MONGO_URI'] = "mongodb://localhost:27017/my_db"
add.config['MONGO_DBNAME'] = 'my_db'
mongo = PyMongo(add)
my_db = mongo.db

@add.route('/', methods=['GET','POST'])
def main():
	return render_template('add_index.html')

@add.route('/<b_id>', methods=['GET','POST'])
def response(b_id):
	# return "getting URL parameter " + b_id
	book_id = b_id
	book = my_db["book_details"]
	ret_dic = book.find_one({'p_key':int(book_id)})
	b_name = ret_dic.get("title")

	return render_template('add_index.html', b_name=b_name)

	if len(note_s) == 0:
		return render_template('add_index.html', err_msg="Please enter something. A note cannot be blank!")

@add.route('/response/<b_name>', methods=['GET','POST'])
def added(b_name):
	note_s = request.form.get("note_s")
	notes = my_db["notes"]
	book = my_db["book_details"]
	b_name = b_name
	# return "this is testing " + b_name
	temp_dic = {}
	temp_dic['note'] = note_s
	temp_dic['b_name'] = b_name
	x = notes.insert_one(temp_dic)
	return render_template('add_index.html', succ_msg="Note has been added successfully.", b_name=b_name)

# @app.route('/response', methods=['POST'])
# def response():
# 	keyword = request.form.get("keyword")
# 	note_title = request.form.get("note_title")
# 	note_s  = request.form.get("note_s")
# 	note_r  = request.form.get("note_r")
# 	book = my_db["book_details"]

# 	if keyword is not None:
# 		myquery = { "title": { "$regex": keyword } }
# 		out_key = book.find(myquery)
# 		list_key = []
# 		print(out_key)
# 		for i in out_key:
# 			list_key.append(i)

# 		le = len(list_key)

# 		return render_template('index.html', list_key=list_key, out_key=out_key, le=le)

# 	elif note_title is not None and note_s is not None:
# 		f = open("notes.txt","a")
# 		f.write("\n" + note_title + "|")
# 		f.write(note_s)
# 		f.close()
# 		return render_template('index.html', note_s="Note has been submitted.")

# 	elif note_r is not None:
# 		f = open("notes.txt", "r")
# 		flag = 0
# 		for i in f:
# 			list_notes = i.split("|")
# 			for j in list_notes:
# 				j.strip()
# 			print(list_notes)
# 			if note_r.strip() in list_notes:
# 				f.close()
# 				flag = 1
# 				return render_template('index.html', list_notes=list_notes)#, t_label = t_label, n_label = n_label)

# 		if flag == 0:
# 			return render_template('index.html', err_msg = "Sorry, there is no note with that title. Please try again.")

if __name__ == "__main__":
	add.debug = True
	add.run(host='0.0.0.0', port=5001)


def sudoku_grid_correct(sudoku: list):
    count = 0
    if row_correct(sudoku) == True:
        count += 1
    if column_correct(sudoku) == True:
        count += 1
    if block_correct(sudoku) == True:
        count += 1

    if count == 3:
        return True
    return False

def row_correct(sudoku):
    row = 0
    index = 0
    count = 0
    for i in range(len(sudoku)):
        for x in sudoku:
            for y in x:
                if y > 0:
                    if x.count(y) != 1:
                        count += 1
    if count > 0:
        return False
    return True

def column_correct(sudoku):
    count = 0
    newlist = []
    col = 0
    for i in range(len(sudoku)):
        for x in sudoku:
            newlist.append(x[col])
        for i in newlist:
            if i > 0:
                if newlist.count(i) != 1:
                    count += 1
        col += 1
        newlist = []
    if count > 0:
        return False
    return True

def block_correct(sudoku):
    newlist = []
    count = 0
    blocklist = [(0, 0), (0, 3), (0, 6), (3, 0), (3, 3), (3, 6), (6, 0), (6, 3), (6, 6)]
    for rowcol in blocklist:
        row = rowcol[0]
        col = rowcol[1]
        for i in range(3):
            for i in range(3):
                newlist.append(sudoku[row][col])
                col += 1
            row += 1
            col = rowcol[1]
        for i in newlist:
            if i > 0:
                if newlist.count(i) != 1:
                    count += 1
        newlist = []

    if count > 0:
        return False
    return True

if __name__ == "__main__":
    sudoku = [
        [2, 6, 7, 8, 3, 9, 5, 0, 4],
        [9, 0, 3, 5, 1, 0, 6, 0, 0],
        [0, 5, 6, 0, 0, 0, 8, 3, 9],
        [5, 1, 9, 0, 4, 6, 3, 2, 8],
        [8, 0, 2, 1, 0, 5, 7, 0, 6],
        [6, 7, 4, 3, 2, 0, 0, 0, 5],
        [0, 0, 0, 4, 5, 7, 2, 6, 3],
        [3, 2, 0, 0, 8, 0, 0, 5, 7],
        [7, 4, 5, 0, 0, 3, 9, 0, 1],
    ]
    print(sudoku_grid_correct(sudoku))


from habit.habit_model import Habit, HabitHistory, HabitNotCreated
import click

def complete(habitid, session):
    """
     When this method is called by the user through the command line interface
     it first checks if the habit is overdue. If it is, broken_habit_update
     method is called which calls the schedule() thus ensuring h.next_run is
     changed before h.can_complete is checked.
     This ensurer a user does not mark overdue habits as completed instead of
     broken.
     It should be noted that this method only works if the user tries to comp-
     lete the habit. that is to say broken fields will on be updated on the
     next attempt to complete a habit.
     """
    try:
        q = session.query(Habit)
        h = q.get(habitid)
        if h.overdue:
            h.broken_habit_update()  # works only if overdue is True
            hist = HabitHistory(habitid, h.activity, h.periodicity,
                                h.creation_date, h.completed, h.completed_at,
                                h.streak, h.date_broken, h.broken_count,
                                h.next_run)
            hist_id = session.add(hist)
            click.echo(f"too late to complete Habit {h.habitid}, habit broken"
                       f" fields have been updated instead.")
            session.commit()
            session.close()
            return hist_id

        elif h.can_complete:
            h.complete()  # works only if can_complete is True

            hist = HabitHistory(habitid, h.activity, h.periodicity,
                                h.creation_date, h.completed, h.completed_at,
                                h.streak, h.date_broken, h.broken_count,
                                h.next_run)
            hist_id = session.add(hist)
            click.echo(f"Habit {h.habitid} successfully completed!")
            session.commit()
            session.close()
            return hist_id
        else:
            click.echo(f"Habit {h.habitid} is not due for completion yet")

    except AttributeError:
        raise HabitNotCreated