import random

from absl.testing import parameterized
import numpy as np

from tensorflow.python.client import session
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import tensor as tensor_lib
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gradients_impl
from tensorflow.python.ops.ragged import ragged_factory_ops
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensor
from tensorflow.python.platform import benchmark
from tensorflow.python.platform import googletest
from tensorflow.python.util import nest

def make_placeholder(t):
  return array_ops.placeholder_with_default(t, None)

def rebuild_ragged_tensor_with_value_rowids(rt, feed_dict=None, sess=None):
  """Returns a copy of `rt`, built using `from_value_rowids`.

  This ensures that RaggedTensor._cached_value_rowids is populated, which
  triggers a different code-path for converting ragged tensors to tensors.

  If `feed_dict` and `sess` are specified, then build the new `RaggedTensor`
  using placeholder tensors, and populate a feed dictionary that can be used
  to feed the placeholders.

  Args:
    rt: The RaggedTensor to copy.
    feed_dict: If specified, then build the new `RaggedTensor` using
      placeholders, and populate this dict with entries to feed those
      placeholders.
    sess: A session used to evaluate tensors; required if feed_dict is
      specified.

  Returns:
    A copy of `rt`, built using `from_value_rowids`.
  """
  if isinstance(rt, ragged_tensor.RaggedTensor):
    values = rebuild_ragged_tensor_with_value_rowids(rt.values, feed_dict, sess)
    rowids = rt.value_rowids()
    nrows = rt.nrows()
    if feed_dict is not None:
      rowids_ph = make_placeholder(rowids)
      nrows_ph = make_placeholder(nrows)
      feed_dict[rowids_ph] = sess.run(rowids)
      feed_dict[nrows_ph] = sess.run(nrows)
      rowids, nrows = rowids_ph, nrows_ph
    return ragged_tensor.RaggedTensor.from_value_rowids(values, rowids, nrows)
  else:
    if feed_dict is not None:
      rt_ph = make_placeholder(rt)
      feed_dict[rt_ph] = sess.run(rt)
      rt = rt_ph
    return rt

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

  def testDocStringExamples(self):
    """Example from ragged_to_tensor.__doc__."""
    rt = ragged_factory_ops.constant([[9, 8, 7], [], [6, 5], [4]])
    dt = rt.to_tensor()
    self.assertAllEqual(dt, [[9, 8, 7], [0, 0, 0], [6, 5, 0], [4, 0, 0]])

  @parameterized.named_parameters(
      # Simple 2D ragged tensors (with one ragged dimension)
      {
          'testcase_name': 'shape_2xN',
          'rt_input': [[0, 1, 2], [], [3]],
          'expected': [[0, 1, 2], [0, 0, 0], [3, 0, 0]]
      },
      {
          'testcase_name': 'shape_2xN_default_0D',
          'rt_input': [[0, 1, 2], [], [3]],
          'default': 5,
          'expected': [[0, 1, 2], [5, 5, 5], [3, 5, 5]]
      },
      {
          'testcase_name': 'empty_first_row',
          'rt_input': [[], [], [3, 4], []],
          'expected': [[0, 0], [0, 0], [3, 4], [0, 0]]
      },
      {
          'testcase_name': 'empty_last_row',
          'rt_input': [[0, 1, 2], [], [3], []],
          'expected': [[0, 1, 2], [0, 0, 0], [3, 0, 0], [0, 0, 0]]
      },
      {
          'testcase_name': 'shape_4xN',
          'rt_input': [[1, 2, 3], [], [4], [5, 6]],
          'expected': [[1, 2, 3], [0, 0, 0], [4, 0, 0], [5, 6, 0]]
      },
      {
          'testcase_name': 'shape_4xN_default_0D',
          'rt_input': [[1, 2, 3], [], [4], [5, 6]],
          'default': 9,
          'expected': [[1, 2, 3], [9, 9, 9], [4, 9, 9], [5, 6, 9]]
      },
      {
          'testcase_name': 'shape_2xN_already_dense',
          'rt_input': [[6, 7, 8], [9, 10, 11]],
          'expected': [[6, 7, 8], [9, 10, 11]],
      },
      {
          'testcase_name': 'shape_2xN_string_already_dense',
          'rt_input': [[b'a', b'b', b'c'],
                       [b'd', b'e', b'antidisestablishmentarianism']],
          'ragged_rank': 1,
          'expected': [[b'a', b'b', b'c'],
                       [b'd', b'e', b'antidisestablishmentarianism']],
      },
      # 3D ragged tensors with two ragged dimensions
      {
          'testcase_name': 'shape_4xNxM',
          'rt_input': [[[1, 2], [], [3, 4]], [], [[5]], [[6, 7], [8]]],
          'expected': [
              [[1, 2], [0, 0], [3, 4]],  #
              [[0, 0], [0, 0], [0, 0]],  #
              [[5, 0], [0, 0], [0, 0]],  #
              [[6, 7], [8, 0], [0, 0]],  #
          ]
      },
      {
          'testcase_name': 'shape_4xNxM_default_0D',
          'rt_input': [[[1, 2], [], [3, 4]], [], [[5]], [[6, 7], [8]]],
          'default': 9,
          'expected': [
              [[1, 2], [9, 9], [3, 4]],  #
              [[9, 9], [9, 9], [9, 9]],  #
              [[5, 9], [9, 9], [9, 9]],  #
              [[6, 7], [8, 9], [9, 9]],  #
          ]
      },
      {
          'testcase_name': 'shape_1xNx1_default_0D',
          'rt_input': [[[1], [2], [3]]],
          'ragged_rank': 1,
          'default': 0,
          'expected': [[[1], [2], [3]]],
      },
      {
          'testcase_name': 'shape_2xNx2_already_dense',
          'rt_input': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15], [16, 17]]],
          'ragged_rank': 1,
          'expected': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15], [16, 17]]],
      },
      {
          'testcase_name': 'shape_2xNx2_already_dense_default_1D',
          'rt_input': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15], [16, 17]]],
          'ragged_rank': 1,
          'default': [31, 32],
          'expected': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15], [16, 17]]],
      },
      {
          'testcase_name': 'shape_2xNx2_string_already_dense',
          'rt_input': [[[b'a', b'b'], [b'c', b'd'], [b'e', b'f']],
                       [[b'g', b'jalapeno'], [b'kangaroo', b'llama'],
                        [b'manzana', b'nectar']]],
          'ragged_rank': 1,
          'expected': [[[b'a', b'b'], [b'c', b'd'], [b'e', b'f']],
                       [[b'g', b'jalapeno'], [b'kangaroo', b'llama'],
                        [b'manzana', b'nectar']]],
      },
      # 3D ragged tensors with one ragged dimension
      {
          'testcase_name': 'shape_4xNx1_default_1D',
          'rt_input': [[[1], [2], [3]], [], [[4]], [[5], [6]]],
          'ragged_rank': 1,
          'default': [9],
          'expected': [[[1], [2], [3]],
                       [[9], [9], [9]],
                       [[4], [9], [9]],
                       [[5], [6], [9]]]
      },
      {
          'testcase_name': 'shape_2xNx2_default_0D',
          'rt_input': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15]]],
          'ragged_rank': 1,
          'default': 2,
          'expected': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15], [2, 2]]],
      },
      {
          'testcase_name': 'shape_2xNx2_default_1D',
          'rt_input': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15]]],
          'ragged_rank': 1,
          'default': [2, 3],
          'expected': [[[6, 7], [8, 9], [10, 11]],
                       [[12, 13], [14, 15], [2, 3]]],
      },
      # 4D ragged tensors with 3 ragged dimensions
      {
          'testcase_name': 'shape_1xNxMxK_default_0D',
          'rt_input': [[[[1], [2]], [], [[3]]]],
          'default': 9,
          'expected': [[[[1], [2]], [[9], [9]], [[3], [9]]]],
      },
      # Broadcast default
      {
          'testcase_name': 'shape_2xNx2x2_default_2x1',
          'rt_input': [[[[1, 2], [3, 4]]], []],
          'ragged_rank': 1,
          'default': [[5], [6]],
          'expected': [[[[1, 2], [3, 4]]],
                       [[[5, 5], [6, 6]]]],
      },
      {
          'testcase_name': 'shape_2xNx2x2_default_1x2',
          'rt_input': [[[[1, 2], [3, 4]]], []],
          'ragged_rank': 1,
          'default': [[5, 6]],
          'expected': [[[[1, 2], [3, 4]]],
                       [[[5, 6], [5, 6]]]],
      },
      # Explicit shape
      {
          'testcase_name': 'shape_4xN_with_crop',
          'rt_input': [[0, 1, 2, 3], [], [4], []],
          'shape': [2, 3],
          'expected': [[0, 1, 2], [0, 0, 0]],
      },
      {
          'testcase_name': 'shape_2xN_with_pad',
          'rt_input': [[1, 2], [3]],
          'shape': [3, 3],
          'expected': [[1, 2, 0], [3, 0, 0], [0, 0, 0]],
      },
      {
          'testcase_name': 'shape_4xN_with_crop_and_pad',
          'rt_input': [[0, 1, 2, 3], [], [4], []],
          'shape': [2, 8],
          'expected': [[0, 1, 2, 3, 0, 0, 0, 0],
                       [0, 0, 0, 0, 0, 0, 0, 0]],
      },
      {
          'testcase_name': 'shape_4xN_with_tuple_shape',
          'rt_input': [[0, 1, 2, 3], [], [4], []],
          'shape': (2, 3),
          'expected': [[0, 1, 2], [0, 0, 0]],
      },
      {
          'testcase_name': 'shape_4xN_with_tensorshape_shape',
          'rt_input': [[0, 1, 2, 3], [], [4], []],
          'shape': tensor_shape.TensorShape([2, 3]),
          'expected': [[0, 1, 2], [0, 0, 0]],
      },
      {
          'testcase_name': 'shape_4xN_with_partial_shape',
          'rt_input': [[0, 1, 2, 3], [], [4], []],
          'shape': tensor_shape.TensorShape([2, None]),
          'expected': [[0, 1, 2, 3], [0, 0, 0, 0]],
      },
      # Empty tensors
      {
          'testcase_name': 'shape_0xN',
          'rt_input': [],
          'ragged_rank': 1,
          'expected': [],
          'expected_shape': [0, 0],
      },
      {
          'testcase_name': 'shape_0xNxM',
          'rt_input': [],
          'ragged_rank': 2,
          'expected': [],
          'expected_shape': [0, 0, 0],
      },
      # {
      #     'testcase_name': 'shape_0xNx2',
      #     'rt_input': [],
      #     'ragged_rank': 1,
      #     'inner_shape': [2],
      #     'expected': [],
      #     'expected_shape': [0, 0, 2],
      # },
      {
          'testcase_name': 'shape_2xN_empty',
          'rt_input': [[], []],
          'expected': [[], []],
          'expected_shape': [2, 0],
      },
  )  # pyformat: disable
  def testRaggedTensorToTensor(self,
                               rt_input,
                               expected,
                               ragged_rank=None,
                               inner_shape=None,
                               default=None,
                               shape=None,
                               expected_shape=None):
    rt1 = ragged_factory_ops.constant(
        rt_input, ragged_rank=ragged_rank, inner_shape=inner_shape)
    rt2 = rebuild_ragged_tensor_with_value_rowids(rt1)
    for rt in [rt1, rt2]:
      for use_placeholder in [False, True]:
        if use_placeholder:
          if default is not None:
            default = make_placeholder(default)
          rt = nest.map_structure(make_placeholder, rt, expand_composites=True)
        dt = rt.to_tensor(default_value=default, shape=shape)
        self.assertIsInstance(dt, tensor_lib.Tensor)
        self.assertEqual(rt.dtype, dt.dtype)
        if shape is not None:
          self.assertTrue(dt.shape.is_compatible_with(shape))
        else:
          self.assertTrue(dt.shape.is_compatible_with(rt.shape))
        if expected_shape is not None:
          expected = np.ndarray(expected_shape, buffer=np.array(expected))
        self.assertAllEqual(dt, expected)

  @parameterized.parameters([
      {
          'rt_input': [[1, 2, 3]],
          'default': 'a',
          'error_type': TypeError,
          'error': r'Expected int32|Cannot convert',
      },
      {
          'rt_input': [[1, 2, 3]],
          'default': [0],
          'error': r'default_value\.shape=.* and '
                   r'rt_input\.flat_values\.shape=.* are incompatible: '
                   r'default_value\.rank = 1  must be less than '
                   r'rt_input\.flat_values\.rank = 1'
      },
      {
          'rt_input': [[[1, 2], [3, 4]], [[5, 6]]],
          'ragged_rank': 1,
          'default': [7, 8, 9],
          'error': r'default_value\.shape.* and '
                   r'rt_input\.flat_values\.shape.* are incompatible: '
                   r'default_value\.shape\[-1\] = 3 but '
                   r'rt_input\.flat_values\.shape\[-1\] = 2'
      },
      {
          'rt_input': [[1, 2, 3]],
          'shape': [3, 3, 3],
          'error': r'rt_input\.shape and shape=\[.,.,.\] are incompatible: '
                   r'rt_input\.rank = 2 but shape\.rank = 3'
      },
      {
          'rt_input': [[[1, 2, 3]]],
          'ragged_rank': 1,
          'shape': [1, 1, 4],
          'error': r'rt_input\.shape and shape=\[1,1,4\] are incompatible: '
                   r'rt_input\.shape\[2\] = 3 but shape\[2\] = 4'
      },
  ])
  def testError(self,
                rt_input,
                error,
                error_type=(ValueError, errors.InvalidArgumentError),
                default=None,
                ragged_rank=None,
                shape=None):

    rt = ragged_factory_ops.constant(rt_input, ragged_rank=ragged_rank)
    with self.assertRaisesRegex(error_type, error):
      self.evaluate(rt.to_tensor(default_value=default, shape=shape))
    rt_placeholder = nest.map_structure(
        make_placeholder, rt, expand_composites=True)
    with self.assertRaisesRegex(error_type, error):
      self.evaluate(
          rt_placeholder.to_tensor(default_value=default, shape=shape))

  def test_shape_limit_shape_is_tensor(self):
    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
    actual = input_data.to_tensor(
        shape=constant_op.constant([2, 3], dtype=dtypes.int64))
    self.assertAllEqual(actual, [[0, 1, 2], [0, 0, 0]])
    self.assertEqual(actual.shape.as_list(), [2, 3])

  def test_shape_limit_shape_is_tensor_unknown_rank(self):
    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
    actual = input_data.to_tensor(
        shape=constant_op.constant(-1, dtype=dtypes.int64))
    self.assertAllEqual(
        actual, [[0, 1, 2, 3], [0, 0, 0, 0], [4, 0, 0, 0], [0, 0, 0, 0]])
    self.assertTrue(actual.shape.is_compatible_with([4, 4]))

  def test_shape_limit_shape_is_tensor_unknown_dim(self):
    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
    actual = input_data.to_tensor(
        shape=constant_op.constant([2, -1], dtype=dtypes.int64))
    self.assertAllEqual(actual, [[0, 1, 2, 3], [0, 0, 0, 0]])
    self.assertTrue(actual.shape.is_compatible_with([2, None]))

  def test_shape_limit_shape_is_tensor_int32(self):
    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
    actual = input_data.to_tensor(
        shape=constant_op.constant([2, 3], dtype=dtypes.int32))
    self.assertAllEqual(actual, [[0, 1, 2], [0, 0, 0]])
    self.assertEqual(actual.shape.as_list(), [2, 3])

  def test_shape_expand_first_dim(self):
    input_data = ragged_factory_ops.constant([[0, 1, 2], [], [3]])
    actual = input_data.to_tensor(shape=[4, 4])
    self.assertAllEqual(
        actual, [[0, 1, 2, 0], [0, 0, 0, 0], [3, 0, 0, 0], [0, 0, 0, 0]])
    self.assertEqual(actual.shape.as_list(), [4, 4])

  def test_value_transposed(self):
    # Check that transposed data is not an issue.
    my_value = array_ops.transpose(
        constant_op.constant([[0, 1, 2, 3], [4, 5, 6, 7]]))
    input_data = RaggedTensor.from_value_rowids(
        values=my_value,
        value_rowids=constant_op.constant([0, 1, 2, 3], dtype=dtypes.int64),
        nrows=constant_op.constant(4, dtype=dtypes.int64),
        validate=True)
    self.assertAllEqual(input_data, [[[0, 4]], [[1, 5]], [[2, 6]], [[3, 7]]])

  def test_broadcast_default(self):
    # The dense dimension here is 2 x 2
    input_data = ragged_factory_ops.constant([[[[1, 2], [3, 4]]], []],
                                             ragged_rank=1)
    # This placeholder has a 2 x 1 dimension.
    default_value = make_placeholder([[5], [6]])
    actual = input_data.to_tensor(default_value=default_value)
    expected = [[[[1, 2], [3, 4]]], [[[5, 5], [6, 6]]]]
    self.assertAllEqual(actual, expected)

  def test_broadcast_default_no_placeholder(self):
    input_data = ragged_factory_ops.constant([[[[1, 2], [3, 4]]], []],
                                             ragged_rank=1)
    # default_value has a 2 x 1 dimension.
    default_value = constant_op.constant([[5], [6]], shape=None)
    actual = input_data.to_tensor(default_value=default_value)
    expected = [[[[1, 2], [3, 4]]], [[[5, 5], [6, 6]]]]
    self.assertAllEqual(actual, expected)

  def test_shape_expand_second_dim(self):
    input_data = ragged_factory_ops.constant([[0, 1, 2], [], [3], []])
    actual = input_data.to_tensor(shape=[3, 4])
    self.assertAllEqual(actual, [[0, 1, 2, 0], [0, 0, 0, 0], [3, 0, 0, 0]])

  @parameterized.parameters(
      ([2, 3, 4], None, [2, 3, 4]),
      ([2, 3, 4], [None, None, None], [2, 3, 4]),
      ([2, 3, 4], [None, 3, None], [2, 3, 4]),
      ([2, 3, 4], [None, 3, 4], [2, 3, 4]),
      ([2, 3, 4], [2, 3, 4], [2, 3, 4]),
      )
  def test_preserve_shape_roundtrip(
      self, input_shape, to_tensor_shape, expected_shape):
    tensor = array_ops.zeros(input_shape)
    ragged_from_tensor = RaggedTensor.from_tensor(tensor, ragged_rank=2)
    recovered_tensor = ragged_from_tensor.to_tensor(shape=to_tensor_shape)
    self.assertAllEqual(tensor.shape.as_list(), expected_shape)
    self.assertAllEqual(ragged_from_tensor.shape.as_list(), expected_shape)
    self.assertAllEqual(recovered_tensor.shape.as_list(), expected_shape)

  def test_empty_tensor_with_shape(self):
    input_data = RaggedTensor.from_value_rowids(
        values=constant_op.constant([], dtype=dtypes.int64),
        value_rowids=constant_op.constant([], dtype=dtypes.int64),
        nrows=constant_op.constant(2, dtype=dtypes.int64),
        validate=True)
    actual = input_data.to_tensor(default_value=3, shape=[2, 3])
    self.assertAllEqual(actual, [[3, 3, 3], [3, 3, 3]])

  # pylint: disable=bad-whitespace
  @parameterized.named_parameters([
      dict(
          testcase_name = '2d_default_shape',
          shape         = None,
          rt_value      = [[1, 2, 3], [4], [5, 6]],
          rt_grad       = [[9, 8, 7], [6], [3, 2]],
          default_value = 0,
          default_grad  = sum([5, 4, 1]),
          output_value  = [[1, 2, 3], [4, 0, 0], [5, 6, 0]],
          output_grad   = [[9, 8, 7], [6, 5, 4], [3, 2, 1]]),
      dict(
          testcase_name = '2d_pad',
          shape         = [4, 4],
          rt_value      = [[1, 2, 3], [4], [5, 6]],
          rt_grad       = [[9, 8, 7], [5], [1, 0]],
          default_value = 0,
          default_grad  = sum([6, 4, 3, 2, 1, 2, 3, 4, 5, 6]),
          output_value  = [
              [1, 2, 3, 0], [4, 0, 0, 0], [5, 6, 0, 0], [0, 0, 0, 0]],
          output_grad   = [
              [9, 8, 7, 6], [5, 4, 3, 2], [1, 0, 1, 2], [3, 4, 5, 6]]),
      dict(
          testcase_name = '2d_pad_and_crop',
          shape         = [5, 3],
          rt_value      = [[1, 2, 3], [4], [5, 6, 7, 8, 9], [8]],
          rt_grad       = [[9, 8, 7], [6], [3, 2, 1, 0, 0], [2]],
          default_value = 0,
          default_grad  = sum([5, 4, 3, 4, 5, 6, 7]),
          output_value  = [
              [1, 2, 3], [4, 0, 0], [5, 6, 7], [8, 0, 0], [0, 0, 0]],
          output_grad   = [
              [9, 8, 7], [6, 5, 4], [3, 2, 1], [2, 3, 4], [5, 6, 7]]),
      dict(
          testcase_name = '3d_rrank_2',
          shape         = [2, 2, 2],
          rt_value      = [[[9, 8, 7], [6]], [[5, 4]]],
          rt_grad       = [[[1, 2, 0], [3]], [[5, 6]]],
          default_value = 3,
          default_grad  = sum([4, 7, 8]),
          output_value  = [[[9, 8], [6, 3]], [[5, 4], [3, 3]]],
          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
      dict(
          testcase_name = '3d_rrank_1_with_0d_default',
          ragged_rank   = 1,
          shape         = [2, 2, 2],
          rt_value      = [[[9, 8], [7, 6]], [[5, 4]]],
          rt_grad       = [[[1, 2], [3, 4]], [[5, 6]]],
          default_value = 3,
          default_grad  = sum([7, 8]),
          output_value  = [[[9, 8], [7, 6]], [[5, 4], [3, 3]]],
          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
      dict(
          testcase_name = '3d_rrank_1_with_1d_default',
          ragged_rank   = 1,
          shape         = [2, 2, 2],
          rt_value      = [[[9, 8], [7, 6]], [[5, 4]]],
          rt_grad       = [[[1, 2], [3, 4]], [[5, 6]]],
          default_value = [3, 2],
          default_grad  = [7, 8],
          output_value  = [[[9, 8], [7, 6]], [[5, 4], [3, 2]]],
          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
      dict(
          testcase_name = '3d_rrank_1_with_1d_broadcast_default',
          ragged_rank   = 1,
          shape         = [2, 2, 2],
          rt_value      = [[[9, 8], [7, 6]], [[5, 4]]],
          rt_grad       = [[[1, 2], [3, 4]], [[5, 6]]],
          default_value = [3],
          default_grad  = [7 + 8],
          output_value  = [[[9, 8], [7, 6]], [[5, 4], [3, 3]]],
          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
      dict(
          testcase_name = '4d_rrank_1_with_2d_default',
          ragged_rank   = 1,
          shape         = [3, 3, 2, 1],
          rt_value      = [[[[9], [8]], [[7], [6]]], [[[5], [4]]]],
          rt_grad       = [[[[1], [2]], [[3], [4]]], [[[7], [8]]]],
          default_value = [[3], [2]],
          default_grad  = [[5 + 9 + 2 + 4 + 6 + 8], [6 + 1 + 3 + 5 + 7 + 9]],
          output_value  = [[[[9], [8]], [[7], [6]], [[3], [2]]],
                           [[[5], [4]], [[3], [2]], [[3], [2]]],
                           [[[3], [2]], [[3], [2]], [[3], [2]]]],
          output_grad   = [[[[1], [2]], [[3], [4]], [[5], [6]]],
                           [[[7], [8]], [[9], [1]], [[2], [3]]],
                           [[[4], [5]], [[6], [7]], [[8], [9]]]]),
      dict(
          testcase_name = '4d_rrank_1_with_with_0d_default',
          ragged_rank   = 1,
          shape         = [3, 3, 2, 1],
          rt_value      = [[[[9], [8]], [[7], [6]]], [[[5], [4]]]],
          rt_grad       = [[[[1], [2]], [[3], [4]]], [[[7], [8]]]],
          default_value = 3,
          default_grad  = 5 + 9 + 2 + 4 + 6 + 8 + 6 + 1 + 3 + 5 + 7 + 9,
          output_value  = [[[[9], [8]], [[7], [6]], [[3], [3]]],
                           [[[5], [4]], [[3], [3]], [[3], [3]]],
                           [[[3], [3]], [[3], [3]], [[3], [3]]]],
          output_grad   = [[[[1], [2]], [[3], [4]], [[5], [6]]],
                           [[[7], [8]], [[9], [1]], [[2], [3]]],
                           [[[4], [5]], [[6], [7]], [[8], [9]]]]),
      dict(
          testcase_name = 'zero_size',
          shape         = [0, 0],
          rt_value      = [[9, 8], [7, 6, 5], [4]],
          rt_grad       = [[0, 0], [0, 0, 0], [0]],
          default_value = 3,
          default_grad  = 0,
          output_value  = [],
          output_grad   = [])
  ])  # pyformat: disable
  def test_gradient(self,
                    shape,
                    rt_value,
                    rt_grad,
                    default_value,
                    default_grad,
                    output_value,
                    output_grad,
                    ragged_rank=None):
    """Tests that ragged_to_dense generates the right gradient.

    Args:
      shape: The `shape` arg for `ragged_to_dense`.
      rt_value: The `rt_input` arg for `ragged_to_dense`.
      rt_grad: The expected gradient for `rt_value`.  Corresponds 1:1 with
        `rt_value`.
      default_value: The `default_value` arg for `ragged_to_dense`.
      default_grad: The expected gradient for `default_value`.  Corresponds 1:1
        with `default_value`.
      output_value: The expected output of `ragged_to_dense`.
      output_grad: The gradient for the output (used to generate the gradients
        `rt_grad` and `default_grad`).  Corresponds 1:1 with `output_value`.
      ragged_rank: Ragged rank for `rt_value`.
    """
    rt_value = ragged_factory_ops.constant(
        rt_value, dtype=dtypes.float32, ragged_rank=ragged_rank)
    rt_grad = ragged_factory_ops.constant(
        rt_grad, dtype=dtypes.float32, ragged_rank=ragged_rank)
    default_value = constant_op.constant(default_value, dtype=dtypes.float32)
    default_grad = constant_op.constant(default_grad, dtype=dtypes.float32)
    output_value = constant_op.constant(
        output_value, dtype=dtypes.float32, shape=shape)
    output_grad = constant_op.constant(
        output_grad, dtype=dtypes.float32, shape=shape)
    shape = tensor_shape.as_shape(shape)

    # There are different code paths for ragged_to_dense, depending on whether
    # the RaggedTensor was created from row_splits or value_rowids.  Make sure
    # that we test both.
    for partition_type in ['row_splits', 'value_rowids']:
      rt_val = self.rt_with_partition_type(rt_value, partition_type)
      if context.executing_eagerly():
        self._test_gradient_helper(rt_val, default_value, shape, output_grad,
                                   output_value, rt_grad, default_grad)
      else:
        # There are different code paths when computing the gradient for
        # default_value, depending on whether shape info is statically
        # available; make sure that we test all code paths.
        for shape_info in ['known', 'unknown_dims', 'unknown_rank']:
          rt_val = self.wrap_in_placeholder(rt_val, shape_info)
          default_val = self.wrap_in_placeholder(default_value, shape_info)
          shape_val = self.wrap_in_placeholder(shape, shape_info)
          self._test_gradient_helper(rt_val, default_val, shape_val,
                                     output_grad, output_value, rt_grad,
                                     default_grad)

  def _test_gradient_helper(self, rt_val, default_val, shape_val, output_grad,
                            expected_output_val, expected_rt_grad,
                            expected_default_grad):
    if context.executing_eagerly():
      with backprop.GradientTape() as tape:
        tape.watch([rt_val, default_val])
        out = rt_val.to_tensor(default_val, shape=shape_val)
        actual_rt_grad, actual_default_grad = tape.gradient(
            out, (rt_val, default_val), output_gradients=output_grad)
    else:
      out = rt_val.to_tensor(default_val, shape=shape_val)
      actual_rt_grad, actual_default_grad = gradients_impl.gradients(
          ys=out, xs=(rt_val, default_val), grad_ys=output_grad)

    self.assertAllClose(out, expected_output_val)
    self.assertIsInstance(actual_rt_grad, RaggedTensor)
    self.assertAllClose(actual_rt_grad, expected_rt_grad)
    self.assertAllClose(actual_default_grad, expected_default_grad)

  def rt_with_partition_type(self, rt, partition_type):
    if isinstance(rt, tensor_lib.Tensor):
      return rt
    if partition_type == 'row_splits':
      return rt
    if partition_type == 'value_rowids':
      return ragged_tensor.RaggedTensor.from_value_rowids(
          self.rt_with_partition_type(rt.values, partition_type),
          rt.value_rowids(), rt.nrows())
    raise AssertionError('Unexpected partition_type %r' % partition_type)

  def wrap_in_placeholder(self, arg, shape_info):
    """Wraps `arg` in a placeholder to limit static shape info.

    Args:
      arg: The value to wrap.  A Tensor, RaggedTensor, or TensorShape.
      shape_info: One of ['known', 'unknown_dims', 'unknown_rank'].

    Returns:
      * If shape_info is 'known': returns `arg`.
      * If shape_info is 'unknown_dims': returns a placeholder wrapping `arg`
        where the dimension sizes are unknown.  If `arg` is a TensorShape,
        then convert it to a vector first.  If `arg` is a RaggedTensor, then
        wrap the flat_values.
      * If shape_info is 'unknown_rank': returns a placeholder wrapping `arg`
        where the rank is unknown.  If `arg` is a TensorShape, then convert it
        to a vector first.  If `arg` is a RaggedTensor, then wrap the
        flat_values.
    """
    if shape_info == 'known':
      return arg
    if isinstance(arg, ragged_tensor.RaggedTensor):
      return arg.with_flat_values(
          self.wrap_in_placeholder(arg.flat_values, shape_info))
    if isinstance(arg, tensor_shape.TensorShape):
      if arg.ndims is None:
        return arg
      arg = constant_op.constant(arg.as_list())
    if shape_info == 'unknown_rank':
      return array_ops.placeholder_with_default(arg, None)
    if shape_info == 'unknown_dims':
      return array_ops.placeholder_with_default(arg, [None] * arg.shape.rank)
    raise AssertionError('Unexpected shape_info %r' % shape_info)

  def test_shape_is_list_including_tensor_element(self):
    rt = ragged_factory_ops.constant([[1, 2, 3], [4], [5, 6]])
    result = rt.to_tensor(shape=[2, constant_op.constant(2)])
    self.assertAllEqual(result, [[1, 2], [4, 0]])

class RaggedToDenseBenchmark(googletest.Benchmark):

  # Configurations to test.  See `run_benchmark` for config param docs.
  CONFIGS = [
      {'shape': [10, 10]},
      {'shape': [10, 1000]},
      {'shape': [1000, 10]},
      {'shape': [1000, 10], 'fill': [1, 0.95]},  # Mostly full.
      {'shape': [1000, 10], 'fill': [1, 0.05]},  # Mostly empty.
      {'shape': [1000, 10], 'dtype': dtypes.string},
      {'shape': [1000, 10], 'dtype': dtypes.int64},
      {'shape': [100, 100]},
      {'shape': [50, 50, 32]},
      {'shape': [100, 100, 100], 'min_iters': 100},
      {'shape': [1000, 1000], 'min_iters': 100},
      {'shape': [10, 10, 10, 10, 10]},
      {'shape': [10, 10, 10, 10, 10], 'ragged_rank': 1},
      {'shape': [10, 10, 10, 10, 10], 'ragged_rank': 2},
      {'shape': [50, 50, 32], 'ragged_rank': 1, 'default_shape': [32]},
      {'shape': [200, 50, 32], 'ragged_rank': 1, 'default_shape': [32]}
  ]  # pyformat: disable

  def run_benchmark(self,
                    shape=(100, 100),
                    ragged_rank=None,
                    dtype=dtypes.float32,
                    fill=None,
                    default_shape=(),
                    output_shape=None,
                    min_iters=1000):
    """Run a benchmark with the specified configuration parameters.

    Args:
      shape: Bounding box for the input ragged tensor.
      ragged_rank: Ragged rank for the input ragged tensor.  Defaults to
        `len(shape)-1`.
      dtype: Data type for the input ragged tensor.
      fill: How full each dimension should be (0-1).  Corresponds 1:1 with
        `shape`.  Defaults to 0.8 for each dimension.
      default_shape: Shape for the default (padding) value.
      output_shape: Output shape -- ragged tensor will be padded or cropped to
        this shape.
      min_iters: Minimum iterations for benchmark.
    """
    if ragged_rank is None:
      ragged_rank = len(shape) - 1
    if fill is None:
      fill = [0.8 for _ in shape]

    # Build the inputs for the op.
    rt_input = self._generateRaggedTensor(shape, ragged_rank, dtype, fill)
    default_value = constant_op.constant(
        self._generateRaggedTensor(default_shape, 0, dtype), dtype=dtype)

    mbs = np.prod(shape) / (2**20)
    with session.Session(config=benchmark.benchmark_config()) as sess:
      extras = {
          'shape': shape,
          'ragged_rank': ragged_rank,
          'dtype': dtype,
          'fill': fill,
          'default_shape': default_shape
      }
      rt = ragged_factory_ops.constant(rt_input, dtype, ragged_rank=ragged_rank)

      # Inputs for with_splits:
      splits_rt_placeholder = ragged_factory_ops.placeholder(
          dtype, ragged_rank, shape[ragged_rank + 1:])
      splits_feed_dict = {splits_rt_placeholder: sess.run(rt)}

      # Inputs for with_rowids:
      rowids_feed_dict = {}
      rowids_rt_placeholder = rebuild_ragged_tensor_with_value_rowids(
          rt, rowids_feed_dict, sess)

      # Common arguments for benchmarks:
      run_op_benchmark_kwargs = dict(
          sess=sess,
          store_memory_usage=True,
          min_iters=min_iters,
          burn_iters=max(5, min_iters // 10),
          mbs=mbs,
          extras=extras)

      ragged_to_tensor_with_splits = splits_rt_placeholder.to_tensor(
          default_value=default_value)
      self.run_op_benchmark(
          op_or_tensor=ragged_to_tensor_with_splits.op,
          name='ragged_to_tensor_with_splits',
          feed_dict=splits_feed_dict,
          **run_op_benchmark_kwargs)

      ragged_to_tensor_with_rowids = rowids_rt_placeholder.to_tensor(
          default_value=default_value)
      self.run_op_benchmark(
          op_or_tensor=ragged_to_tensor_with_rowids.op,
          name='ragged_to_tensor_with_rowids',
          feed_dict=rowids_feed_dict,
          **run_op_benchmark_kwargs)

  def _generateRaggedTensor(self, shape, ragged_rank, dtype, fill=None, axis=0):
    if axis == len(shape):
      value = random.random()
      if dtype == dtypes.string:
        value = str(value)
      if dtype.is_integer:
        value = int(value * 1000)
      return value
    if axis == 0 or axis > ragged_rank:
      slice_size = shape[axis]
    else:
      slice_size = (np.random.geometric(fill[axis], shape[axis]) == 1).sum()
    return [
        self._generateRaggedTensor(shape, ragged_rank, dtype, fill, axis + 1)
        for _ in range(slice_size)
    ]

  def benchmark_ragged_to_dense(self):
    random.seed(5)
    for config in self.CONFIGS:
      self.run_benchmark(**config)

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


import sys
import random
from random import shuffle
from tqdm import tqdm
import os.path
import urllib
import urllib.error
import urllib.request
import imghdr
from ssl import CertificateError
from socket import timeout as TimeoutError
from socket import error as SocketError
import time
import pickle
import hashlib
import itertools

random.seed(42)

min_size = 7000
# api_url = "http://www.image-net.org/api/text/imagenet.synset.geturls?wnid=%s"
api_url = "http://www.image-net.org/api/text/imagenet.synset.geturls.getmapping?wnid=%s"

def hash(content):
    m = hashlib.md5()
    m.update(content)
    hex_value = m.hexdigest()
    return hex_value

class DownloadError(Exception):
    """Base class for exceptions in this module."""
    def __init__(self, message=""):
        self.message = message

def download(url, timeout, retry, sleep=0.8):
    """Downloads a file at given URL."""
    count = 0
    while True:
        try:
            f = urllib.request.urlopen(url, timeout=timeout)
            if f is None:
                raise DownloadError('Cannot open URL' + url)
            content = f.read()
            f.close()
            break
        except (urllib.error.HTTPError, CertificateError) as e:
            if isinstance(e, urllib.error.HTTPError):
                if e.code >= 400: raise DownloadError("No Access")
            count += 1
            if count > retry:
                raise DownloadError("Retry limit")
        except (TimeoutError, SocketError, IOError) as e:
            count += 1
            if count > retry:
                raise DownloadError("Retry limit")
            time.sleep(sleep)
    return content

def get_urls(id_to_synset_dict, limit=-1):
    urls = []

    print("Downloading urls from %s" % api_url)

    progress_bar = tqdm(total=len(id_to_synset_dict))
    for synset_id, i in id_to_synset_dict.items():
        url = api_url % synset_id
        content = download(url, timeout=15, retry=5, sleep=0.1)
        content = content.decode('ISO-8859-1')
        content = content.split('\n')[:-1] # last line is blank

        if limit is not None:
            shuffle(content)
            content = content[:limit]

        content = [(i, *c.split()) for c in content]
        urls += content
        progress_bar.set_postfix({"num_urls": len(urls)})
        progress_bar.update(1)

    return urls

def find_extension(content):
    extension = imghdr.what('', content) #check if valid image
    if extension == "jpeg":
        extension = "jpg"
    if extension == None:
        raise DownloadError()
    return extension

def write_file(content, file_name):
    with open(file_name, 'wb') as f:
        f.write(content)

if __name__ == "__main__":
    if len(sys.argv) == 1:
        print("""Specify number of images:
        > python download_imagenet_images.py [num_imgs]""")
        exit()

    # Id to Class information
    with open('./synset_id_to_class', 'r') as f:
        id_to_class = {}
        for l in f:
            c, i = l.strip().split()
            id_to_class[i] = c

    n_imgs = int(sys.argv[1])
    if os.path.exists('valid_urls.pkl'):
        with open('valid_urls.pkl', 'rb') as f:
            urls = pickle.load(f)
    else:
        urls = get_urls(id_to_class, limit=None)

        print(f"Storing urls in 'valid_urls.pkl'")
        with open('valid_urls.pkl', 'wb') as f:
            pickle.dump(urls, f)

    shuffle(urls)
    print(urls[:10])

    with open('flickr_empty.png', 'rb') as f:
        empty = f.read()

    # Hash of empty image.
    empty_hex = hash(empty)

    os.makedirs('images', exist_ok=True)

    print("Starting image download")
    progress_bar = tqdm(total=n_imgs)
    i = j = 0
    good = [True] * len(urls)
    while i < n_imgs and j < len(urls):
        try:
            cl      = urls[j][0]
            name    = urls[j][1]
            url     = urls[j][2]

            response = download(url, 2, 2)

            if hash(response) == empty_hex:
                raise DownloadError()

            ext = find_extension(response)
            file_name = os.path.join("images", cl + "_" + name + "." + ext)
            write_file(response, file_name)
            progress_bar.update(1)
            j += 1
            i += 1

        except DownloadError as d:
            # print("Download error", d)
            good[j] = False
            j += 1

        except Exception as e:
            # print("Exception", e, e.__class__.__name__)
            good[j] = False
            j += 1

    progress_bar.close()
    print(f"i: {i}, j: {j}, skipped: {j - i}")

    len_before = len(urls)
    urls = list(itertools.compress(urls, good))
    len_after = len(urls)

    print("Updating valid urls to only contain urls, that actually gave a good result")
    print(f"Went from {len_before} to {len_after} elements")
    with open('valid_urls.pkl', 'wb') as f:
        pickle.dump(urls, f)


import os

from SourceFile import SourceFile

class HandlersShortcutFile(SourceFile):
    '''handlers.yml -> handlers/main.yml'''
    def compile_to(self, target_root):
        role_path = self.parent.corresponding_target_folder
        handlers_dir_path = os.path.join(target_root, role_path, 'handlers')
        self._mkdir(handlers_dir_path)
        path = os.path.join(handlers_dir_path, 'main.yml')
        self._copy_file_to(path)


import os

import logging
import telegram
import openai
from flask import Flask, redirect, request, url_for
from dotenv import load_dotenv

load_dotenv()

# Initialize Flask app
app = Flask(__name__)

# Set environment variables
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
TELEGRAM_API_KEY = os.getenv("TELEGRAM_API_KEY")

# Set other constants
MODEL_NAME = "text-davinci-003"
TEMPERATURE = 0.6

# Set OpenAI API key
openai.api_key = OPENAI_API_KEY

try:
    # Initialize Telegram bot
    bot = telegram.Bot(token=TELEGRAM_API_KEY)
except Exception as e:
    # Raise error if initialization failed
    raise ValueError(f"Error initializing Telegram bot: {e}")

@app.route('/start', methods=['POST'])
def start():
    update = telegram.Update.de_json(request.get_json(force=True), bot)
    chat_id = update.message.chat_id
    username = update.message.from_user.username
    bot.send_message(chat_id=chat_id, text=f"Привет, {username}! Я бот. Твой chat_id: {chat_id}")
    return 'ok'

@app.route('/animal', methods=['POST'])
def index():
    # Get update from Telegram or form data
    update = telegram.Update.de_json(request.get_json(force=True), bot)
    animal = request.form.get('animal', update.message.text)

    try:
        # Call OpenAI API to generate response
        response = openai.Completion.create(
            model=MODEL_NAME,
            prompt=generate_prompt(animal),
            temperature=TEMPERATURE,
        )

        # Extract text from top suggestion
        result = response.choices[0].text
    except Exception as e:
        # Handle errors by returning empty result
        result = ''

    # Respond to user message in Telegram and redirect to index
    update.message.reply_text(result)
    return redirect(url_for('index', result=result))

def generate_prompt(animal):
    return f"{animal.capitalize()}"

app.logger.setLevel(logging.DEBUG)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 8000)))


class Pessoa:
    def __init__(self, idade=0, altura=0):
        self.idade = idade
        self.altura = altura

    def __repr__(self):
        return f"Idade: {self.idade}\nAltura: {self.altura}"

i = int(input("Quantidade de pessoas: "))
for i in range(i):
    idade = int(input("Idade: "))
    altura = float(input("Altura: "))
    pessoa = Pessoa(idade, altura)
    print(pessoa)


from time import sleep
from datetime import datetime
from sh import gphoto2 as gp
import signal, os, subprocess
import shutil
shot_date = datetime.now().strftime("%Y-%m-%d")
shot_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

triggerCommand = ["--capture-image-and-download"]
save_folder = "image-"+shot_date
def createSaveFolder():
    try:
        os.mkdir(save_folder)
    except:
      pass
    os.chdir(save_folder)

def captureImages():
    gp(triggerCommand)

def renameFiles():
            if os.path.isfile("capt0000.jpg"):
                print('lol')
            try:
                shutil.move("/home/pi/capt0000.jpg", f"/home/pi/{save_folder}/{shot_time}.jpg")
            except Exception as ex:
                print(ex)

def get_shot():
   captureImages()
   createSaveFolder()
   renameFiles()

#get_shot()


"""
code by Tae Hwan Jung(Jeff Jung) @graykode, Derek Miller @dmmiller612, modify by shwei
Reference: https://github.com/jadore801120/attention-is-all-you-need-pytorch
           https://github.com/JayParks/transformer
"""

"""
代码位置：
手把手教你用Pytorch代码实现Transformer模型（超详细的代码解读）
https://ugirc.blog.csdn.net/article/details/120394042

"""

# ====================================================================================================
# 数据构建
import math
import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as Data

device = 'cpu'
# device = 'cuda'

# transformer epochs
epochs = 100
# epochs = 1000

# 这里我没有用什么大型的数据集，而是手动输入了两对德语→英语的句子
# 还有每个字的索引也是我手动硬编码上去的，主要是为了降低代码阅读难度
# S: Symbol that shows starting of decoding input
# E: Symbol that shows starting of decoding output
# P: Symbol that will fill in blank sequence if current batch data size is short than time steps
sentences = [
    # 德语和英语的单词个数不要求相同
    # enc_input                dec_input           dec_output
    ['ich mochte ein bier P', 'S i want a beer .', 'i want a beer . E'],
    ['ich mochte ein cola P', 'S i want a coke .', 'i want a coke . E']
]

# 德语和英语的单词要分开建立词库
# Padding Should be Zero
src_vocab = {'P': 0, 'ich': 1, 'mochte': 2, 'ein': 3, 'bier': 4, 'cola': 5}
src_idx2word = {i: w for i, w in enumerate(src_vocab)}
src_vocab_size = len(src_vocab)

tgt_vocab = {'P': 0, 'i': 1, 'want': 2, 'a': 3, 'beer': 4, 'coke': 5, 'S': 6, 'E': 7, '.': 8}
idx2word = {i: w for i, w in enumerate(tgt_vocab)}
tgt_vocab_size = len(tgt_vocab)

src_len = 5  # （源句子的长度）enc_input max sequence length
tgt_len = 6  # dec_input(=dec_output) max sequence length

# Transformer Parameters
d_model = 512  # Embedding Size（token embedding和position编码的维度）
d_ff = 2048  # FeedForward dimension (两次线性层中的隐藏层 512->2048->512，线性层是用来做特征提取的），当然最后会再接一个projection层
d_k = d_v = 64  # dimension of K(=Q), V（Q和K的维度需要相同，这里为了方便让K=V）
n_layers = 6  # number of Encoder of Decoder Layer（Block的个数）
n_heads = 8  # number of heads in Multi-Head Attention（有几套头）

# ==============================================================================================
# 数据构建

def make_data(sentences):
    """把单词序列转换为数字序列"""
    enc_inputs, dec_inputs, dec_outputs = [], [], []
    for i in range(len(sentences)):
        enc_input = [[src_vocab[n] for n in sentences[i][0].split()]]  # [[1, 2, 3, 4, 0], [1, 2, 3, 5, 0]]
        dec_input = [[tgt_vocab[n] for n in sentences[i][1].split()]]  # [[6, 1, 2, 3, 4, 8], [6, 1, 2, 3, 5, 8]]
        dec_output = [[tgt_vocab[n] for n in sentences[i][2].split()]]  # [[1, 2, 3, 4, 8, 7], [1, 2, 3, 5, 8, 7]]

        enc_inputs.extend(enc_input)
        dec_inputs.extend(dec_input)
        dec_outputs.extend(dec_output)

    return torch.LongTensor(enc_inputs), torch.LongTensor(dec_inputs), torch.LongTensor(dec_outputs)

enc_inputs, dec_inputs, dec_outputs = make_data(sentences)

class MyDataSet(Data.Dataset):
    """自定义DataLoader"""

    def __init__(self, enc_inputs, dec_inputs, dec_outputs):
        super(MyDataSet, self).__init__()
        self.enc_inputs = enc_inputs
        self.dec_inputs = dec_inputs
        self.dec_outputs = dec_outputs

    def __len__(self):
        return self.enc_inputs.shape[0]

    def __getitem__(self, idx):
        return self.enc_inputs[idx], self.dec_inputs[idx], self.dec_outputs[idx]

loader = Data.DataLoader(MyDataSet(enc_inputs, dec_inputs, dec_outputs), 2, True)

# ====================================================================================================
# Transformer模型

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        """
        x: [seq_len, batch_size, d_model]
        """
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

def get_attn_pad_mask(seq_q, seq_k):
    # pad mask的作用：在对value向量加权平均的时候，可以让pad对应的alpha_ij=0，这样注意力就不会考虑到pad向量
    """这里的q,k表示的是两个序列（跟注意力机制的q,k没有关系），例如encoder_inputs (x1,x2,..xm)和encoder_inputs (x1,x2..xm)
    encoder和decoder都可能调用这个函数，所以seq_len视情况而定
    seq_q: [batch_size, seq_len]
    seq_k: [batch_size, seq_len]
    seq_len could be src_len or it could be tgt_len
    seq_len in seq_q and seq_len in seq_k maybe not equal
    """
    batch_size, len_q = seq_q.size()  # 这个seq_q只是用来expand维度的
    batch_size, len_k = seq_k.size()
    # eq(zero) is PAD token
    # 例如:seq_k = [[1,2,3,4,0], [1,2,3,5,0]]
    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)  # [batch_size, 1, len_k], True is masked
    return pad_attn_mask.expand(batch_size, len_q, len_k)  # [batch_size, len_q, len_k] 构成一个立方体(batch_size个这样的矩阵)

def get_attn_subsequence_mask(seq):
    """建议打印出来看看是什么的输出（一目了然）
    seq: [batch_size, tgt_len]
    """
    attn_shape = [seq.size(0), seq.size(1), seq.size(1)]
    # attn_shape: [batch_size, tgt_len, tgt_len]
    subsequence_mask = np.triu(np.ones(attn_shape), k=1)  # 生成一个上三角矩阵
    subsequence_mask = torch.from_numpy(subsequence_mask).byte()
    return subsequence_mask  # [batch_size, tgt_len, tgt_len]

# ==========================================================================================
class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()

    def forward(self, Q, K, V, attn_mask):
        """
        Q: [batch_size, n_heads, len_q, d_k]
        K: [batch_size, n_heads, len_k, d_k]
        V: [batch_size, n_heads, len_v(=len_k), d_v]
        attn_mask: [batch_size, n_heads, seq_len, seq_len]
        说明：在encoder-decoder的Attention层中len_q(q1,..qt)和len_k(k1,...km)可能不同
        """
        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)  # scores : [batch_size, n_heads, len_q, len_k]
        # mask矩阵填充scores（用-1e9填充scores中与attn_mask中值为1位置相对应的元素）
        scores.masked_fill_(attn_mask, -1e9)  # Fills elements of self tensor with value where mask is True.

        attn = nn.Softmax(dim=-1)(scores)  # 对最后一个维度(v)做softmax
        # scores : [batch_size, n_heads, len_q, len_k] * V: [batch_size, n_heads, len_v(=len_k), d_v]
        context = torch.matmul(attn, V)  # context: [batch_size, n_heads, len_q, d_v]
        # context：[[z1,z2,...],[...]]向量, attn注意力稀疏矩阵（用于可视化的）
        return context, attn

class MultiHeadAttention(nn.Module):
    """这个Attention类可以实现:
    Encoder的Self-Attention
    Decoder的Masked Self-Attention
    Encoder-Decoder的Attention
    """

    def __init__(self):
        super(MultiHeadAttention, self).__init__()
        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=False)  # q,k必须维度相同，不然无法做点积
        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=False)
        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=False)
        self.fc = nn.Linear(n_heads * d_v, d_model, bias=False)

    def forward(self, input_Q, input_K, input_V, attn_mask):
        """
        input_Q: [batch_size, len_q, d_model]
        input_K: [batch_size, len_k, d_model]
        input_V: [batch_size, len_v(=len_k), d_model]
        attn_mask: [batch_size, seq_len, seq_len]
        """
        residual, batch_size = input_Q, input_Q.size(0)
        # 下面的多头的参数矩阵是放在一起做线性变换的，然后再拆成多个头，这是工程实现的技巧
        # B: batch_size, S:seq_len, D: dim
        # (B, S, D) -proj-> (B, S, D_new) -split-> (B, S, Head, W) -trans-> (B, Head, S, W)
        #           线性变换               拆成多头

        # Q: [batch_size, n_heads, len_q, d_k]
        Q = self.W_Q(input_Q).view(batch_size, -1, n_heads, d_k).transpose(1, 2)
        # K: [batch_size, n_heads, len_k, d_k] # K和V的长度一定相同，维度可以不同
        K = self.W_K(input_K).view(batch_size, -1, n_heads, d_k).transpose(1, 2)
        # V: [batch_size, n_heads, len_v(=len_k), d_v]
        V = self.W_V(input_V).view(batch_size, -1, n_heads, d_v).transpose(1, 2)

        # 因为是多头，所以mask矩阵要扩充成4维的
        # attn_mask: [batch_size, seq_len, seq_len] -> [batch_size, n_heads, seq_len, seq_len]
        attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1)

        # context: [batch_size, n_heads, len_q, d_v], attn: [batch_size, n_heads, len_q, len_k]
        context, attn = ScaledDotProductAttention()(Q, K, V, attn_mask)
        # 下面将不同头的输出向量拼接在一起
        # context: [batch_size, n_heads, len_q, d_v] -> [batch_size, len_q, n_heads * d_v]
        context = context.transpose(1, 2).reshape(batch_size, -1, n_heads * d_v)
        # 再做一个projection
        output = self.fc(context)  # [batch_size, len_q, d_model]
        return nn.LayerNorm(d_model).to(device)(output + residual), attn

# Pytorch中的Linear只会对最后一维操作，所以正好是我们希望的每个位置用同一个全连接网络
class PoswiseFeedForwardNet(nn.Module):
    def __init__(self):
        super(PoswiseFeedForwardNet, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(d_model, d_ff, bias=False),
            nn.ReLU(),
            nn.Linear(d_ff, d_model, bias=False)
        )

    def forward(self, inputs):
        """
        inputs: [batch_size, seq_len, d_model]
        """
        residual = inputs
        output = self.fc(inputs)
        return nn.LayerNorm(d_model).to(device)(output + residual)  # [batch_size, seq_len, d_model]

class EncoderLayer(nn.Module):
    def __init__(self):
        super(EncoderLayer, self).__init__()
        self.enc_self_attn = MultiHeadAttention()
        self.pos_ffn = PoswiseFeedForwardNet()

    def forward(self, enc_inputs, enc_self_attn_mask):
        """E
        enc_inputs: [batch_size, src_len, d_model]
        enc_self_attn_mask: [batch_size, src_len, src_len]  mask矩阵(pad mask or sequence mask)
        """
        # enc_outputs: [batch_size, src_len, d_model], attn: [batch_size, n_heads, src_len, src_len]
        # 第一个enc_inputs * W_Q = Q
        # 第二个enc_inputs * W_K = K
        # 第三个enc_inputs * W_V = V
        enc_outputs, attn = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs,
                                               enc_self_attn_mask)  # enc_inputs to same Q,K,V（未线性变换前）
        enc_outputs = self.pos_ffn(enc_outputs)
        # enc_outputs: [batch_size, src_len, d_model]
        return enc_outputs, attn

class DecoderLayer(nn.Module):
    def __init__(self):
        super(DecoderLayer, self).__init__()
        self.dec_self_attn = MultiHeadAttention()
        self.dec_enc_attn = MultiHeadAttention()
        self.pos_ffn = PoswiseFeedForwardNet()

    def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask):
        """
        dec_inputs: [batch_size, tgt_len, d_model]
        enc_outputs: [batch_size, src_len, d_model]
        dec_self_attn_mask: [batch_size, tgt_len, tgt_len]
        dec_enc_attn_mask: [batch_size, tgt_len, src_len]
        """
        # dec_outputs: [batch_size, tgt_len, d_model], dec_self_attn: [batch_size, n_heads, tgt_len, tgt_len]
        dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs, dec_inputs,
                                                        dec_self_attn_mask)  # 这里的Q,K,V全是Decoder自己的输入
        # dec_outputs: [batch_size, tgt_len, d_model], dec_enc_attn: [batch_size, h_heads, tgt_len, src_len]
        dec_outputs, dec_enc_attn = self.dec_enc_attn(dec_outputs, enc_outputs, enc_outputs,
                                                      dec_enc_attn_mask)  # Attention层的Q(来自decoder) 和 K,V(来自encoder)
        dec_outputs = self.pos_ffn(dec_outputs)  # [batch_size, tgt_len, d_model]
        return dec_outputs, dec_self_attn, dec_enc_attn  # dec_self_attn, dec_enc_attn这两个是为了可视化的

class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.src_emb = nn.Embedding(src_vocab_size, d_model)  # token Embedding
        self.pos_emb = PositionalEncoding(d_model)  # Transformer中位置编码时固定的，不需要学习
        self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])

    def forward(self, enc_inputs):
        """
        enc_inputs: [batch_size, src_len]
        """
        enc_outputs = self.src_emb(enc_inputs)  # [batch_size, src_len, d_model]
        enc_outputs = self.pos_emb(enc_outputs.transpose(0, 1)).transpose(0, 1)  # [batch_size, src_len, d_model]
        # Encoder输入序列的pad mask矩阵
        enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs)  # [batch_size, src_len, src_len]
        enc_self_attns = []  # 在计算中不需要用到，它主要用来保存你接下来返回的attention的值（这个主要是为了你画热力图等，用来看各个词之间的关系
        for layer in self.layers:  # for循环访问nn.ModuleList对象
            # 上一个block的输出enc_outputs作为当前block的输入
            # enc_outputs: [batch_size, src_len, d_model], enc_self_attn: [batch_size, n_heads, src_len, src_len]
            enc_outputs, enc_self_attn = layer(enc_outputs,
                                               enc_self_attn_mask)  # 传入的enc_outputs其实是input，传入mask矩阵是因为你要做self attention
            enc_self_attns.append(enc_self_attn)  # 这个只是为了可视化
        return enc_outputs, enc_self_attns

class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        self.tgt_emb = nn.Embedding(tgt_vocab_size, d_model)  # Decoder输入的embed词表
        self.pos_emb = PositionalEncoding(d_model)
        self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])  # Decoder的blocks

    def forward(self, dec_inputs, enc_inputs, enc_outputs):
        """
        dec_inputs: [batch_size, tgt_len]
        enc_inputs: [batch_size, src_len]
        enc_outputs: [batch_size, src_len, d_model]   # 用在Encoder-Decoder Attention层
        """
        dec_outputs = self.tgt_emb(dec_inputs)  # [batch_size, tgt_len, d_model]
        dec_outputs = self.pos_emb(dec_outputs.transpose(0, 1)).transpose(0, 1).to(
            device)  # [batch_size, tgt_len, d_model]
        # Decoder输入序列的pad mask矩阵（这个例子中decoder是没有加pad的，实际应用中都是有pad填充的）
        dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs).to(device)  # [batch_size, tgt_len, tgt_len]
        # Masked Self_Attention：当前时刻是看不到未来的信息的
        dec_self_attn_subsequence_mask = get_attn_subsequence_mask(dec_inputs).to(
            device)  # [batch_size, tgt_len, tgt_len]

        # Decoder中把两种mask矩阵相加（既屏蔽了pad的信息，也屏蔽了未来时刻的信息）
        dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask + dec_self_attn_subsequence_mask),
                                      0).to(device)  # [batch_size, tgt_len, tgt_len]; torch.gt比较两个矩阵的元素，大于则返回1，否则返回0

        # 这个mask主要用于encoder-decoder attention层
        # get_attn_pad_mask主要是enc_inputs的pad mask矩阵(因为enc是处理K,V的，求Attention时是用v1,v2,..vm去加权的，要把pad对应的v_i的相关系数设为0，这样注意力就不会关注pad向量)
        #                       dec_inputs只是提供expand的size的
        dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs)  # [batc_size, tgt_len, src_len]

        dec_self_attns, dec_enc_attns = [], []
        for layer in self.layers:
            # dec_outputs: [batch_size, tgt_len, d_model], dec_self_attn: [batch_size, n_heads, tgt_len, tgt_len], dec_enc_attn: [batch_size, h_heads, tgt_len, src_len]
            # Decoder的Block是上一个Block的输出dec_outputs（变化）和Encoder网络的输出enc_outputs（固定）
            dec_outputs, dec_self_attn, dec_enc_attn = layer(dec_outputs, enc_outputs, dec_self_attn_mask,
                                                             dec_enc_attn_mask)
            dec_self_attns.append(dec_self_attn)
            dec_enc_attns.append(dec_enc_attn)
        # dec_outputs: [batch_size, tgt_len, d_model]
        return dec_outputs, dec_self_attns, dec_enc_attns

class Transformer(nn.Module):
    def __init__(self):
        super(Transformer, self).__init__()
        self.encoder = Encoder().to(device)
        self.decoder = Decoder().to(device)
        self.projection = nn.Linear(d_model, tgt_vocab_size, bias=False).to(device)

    def forward(self, enc_inputs, dec_inputs):
        """Transformers的输入：两个序列
        enc_inputs: [batch_size, src_len]
        dec_inputs: [batch_size, tgt_len]
        """
        # tensor to store decoder outputs
        # outputs = torch.zeros(batch_size, tgt_len, tgt_vocab_size).to(self.device)

        # enc_outputs: [batch_size, src_len, d_model], enc_self_attns: [n_layers, batch_size, n_heads, src_len, src_len]
        # 经过Encoder网络后，得到的输出还是[batch_size, src_len, d_model]
        enc_outputs, enc_self_attns = self.encoder(enc_inputs)
        # dec_outputs: [batch_size, tgt_len, d_model], dec_self_attns: [n_layers, batch_size, n_heads, tgt_len, tgt_len], dec_enc_attn: [n_layers, batch_size, tgt_len, src_len]
        dec_outputs, dec_self_attns, dec_enc_attns = self.decoder(dec_inputs, enc_inputs, enc_outputs)
        # dec_outputs: [batch_size, tgt_len, d_model] -> dec_logits: [batch_size, tgt_len, tgt_vocab_size]
        dec_logits = self.projection(dec_outputs)
        return dec_logits.view(-1, dec_logits.size(-1)), enc_self_attns, dec_self_attns, dec_enc_attns

model = Transformer().to(device)
# 这里的损失函数里面设置了一个参数 ignore_index=0，因为 "pad" 这个单词的索引为 0，这样设置以后，就不会计算 "pad" 的损失（因为本来 "pad" 也没有意义，不需要计算）
criterion = nn.CrossEntropyLoss(ignore_index=0)
optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.99)  # 用adam的话效果不好

# ====================================================================================================
for epoch in range(epochs):
    for enc_inputs, dec_inputs, dec_outputs in loader:
        """
        enc_inputs: [batch_size, src_len]
        dec_inputs: [batch_size, tgt_len]
        dec_outputs: [batch_size, tgt_len]
        """
        enc_inputs, dec_inputs, dec_outputs = enc_inputs.to(device), dec_inputs.to(device), dec_outputs.to(device)
        # outputs: [batch_size * tgt_len, tgt_vocab_size]
        outputs, enc_self_attns, dec_self_attns, dec_enc_attns = model(enc_inputs, dec_inputs)
        loss = criterion(outputs, dec_outputs.view(-1))  # dec_outputs.view(-1):[batch_size * tgt_len * tgt_vocab_size]
        print('Epoch:', '%04d' % (epoch + 1), 'loss =', '{:.6f}'.format(loss))

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

def greedy_decoder(model, enc_input, start_symbol):
    """贪心编码
    For simplicity, a Greedy Decoder is Beam search when K=1. This is necessary for inference as we don't know the
    target sequence input. Therefore we try to generate the target input word by word, then feed it into the transformer.
    Starting Reference: http://nlp.seas.harvard.edu/2018/04/03/attention.html#greedy-decoding
    :param model: Transformer Model
    :param enc_input: The encoder input
    :param start_symbol: The start symbol. In this example it is 'S' which corresponds to index 4
    :return: The target input
    """
    enc_outputs, enc_self_attns = model.encoder(enc_input)
    dec_input = torch.zeros(1, 0).type_as(enc_input.data)
    terminal = False
    next_symbol = start_symbol
    while not terminal:
        # 预测阶段：dec_input序列会一点点变长（每次添加一个新预测出来的单词）
        dec_input = torch.cat([dec_input.to(device), torch.tensor([[next_symbol]], dtype=enc_input.dtype).to(device)],
                              -1)
        dec_outputs, _, _ = model.decoder(dec_input, enc_input, enc_outputs)
        projected = model.projection(dec_outputs)
        prob = projected.squeeze(0).max(dim=-1, keepdim=False)[1]
        # 增量更新（我们希望重复单词预测结果是一样的）
        # 我们在预测是会选择性忽略重复的预测的词，只摘取最新预测的单词拼接到输入序列中
        next_word = prob.data[-1]  # 拿出当前预测的单词(数字)。我们用x'_t对应的输出z_t去预测下一个单词的概率，不用z_1,z_2..z_{t-1}
        next_symbol = next_word
        if next_symbol == tgt_vocab["E"]:
            terminal = True
        # print(next_word)

    # greedy_dec_predict = torch.cat(
    #     [dec_input.to(device), torch.tensor([[next_symbol]], dtype=enc_input.dtype).to(device)],
    #     -1)
    greedy_dec_predict = dec_input[:, 1:]
    return greedy_dec_predict

# ==========================================================================================
# 预测阶段
enc_inputs, _, _ = next(iter(loader))
for i in range(len(enc_inputs)):
    greedy_dec_predict = greedy_decoder(model, enc_inputs[i].view(1, -1).to(device), start_symbol=tgt_vocab["S"])
    print(enc_inputs[i], '->', greedy_dec_predict.squeeze())
    print([src_idx2word[t.item()] for t in enc_inputs[i]], '->',
          [idx2word[n.item()] for n in greedy_dec_predict.squeeze()])


from __future__ import print_function
from builtins import range

import numpy as np
import pandas as pd
import sys
import argparse
import os

from sklearn.datasets import load_boston

from sklearn.linear_model import LinearRegression, Ridge, Lasso, LogisticRegression
from sklearn.svm import SVR
from sklearn.ensemble import RandomForestRegressor

from clr_regressors import KPlaneRegressor, CLRpRegressor, CLRcRegressor, RegressorEnsemble
from evaluate import evaluate_all

def preprocess_data(X, y):
  X -= np.min(X, axis=0, keepdims=True)
  X /= np.max(X, axis=0, keepdims=True) / 2.0
  X -= 1.0
  shuffle_idx = np.random.choice(X.shape[0], X.shape[0], replace=False)
  X = X[shuffle_idx]
  y = y[shuffle_idx]
  return X, y

if __name__ == '__main__':
  parser = argparse.ArgumentParser(description='Search parameters')
  parser.add_argument('--dataset', required=True)
  parser.add_argument('--seed', default=0, type=int)
  parser.add_argument('--n_jobs', default=1, type=int)
  parser.add_argument('--global_parallel', dest='global_parallel', action='store_true')
  args = parser.parse_args()
  np.random.seed(args.seed)

  if args.dataset == 'boston':
    boston = load_boston()
    X = boston.data
    y = boston.target
    constr_id = 8
  elif args.dataset == 'abalone':
    abalone_data = pd.read_csv('data/abalone.data', header=None)
    with open('data/abalone.names', 'r') as fin:
      abalone_descr = fin.read()
    X = pd.get_dummies(abalone_data.iloc[:,:-1], columns=[0]).as_matrix().astype(np.float)
    X = np.hstack((X, (np.digitize(X[:, 2], np.linspace(0.1, 0.2, 10)))[:,np.newaxis]))
    y = abalone_data.iloc[:, 8].as_matrix().astype(np.float)
    constr_id = 10
  elif args.dataset == 'auto-mpg':
    data = pd.read_csv('data/auto-mpg.data', header=None, sep='\s+', na_values='?')
    data = data.dropna()
    X = pd.get_dummies(data.iloc[:,1:-1], columns=[7]).as_matrix().astype(np.float)
    y = data[0].as_matrix().astype(np.float)
    constr_id = 5
  else:
    print("Dataset is not supported")
    sys.exit(0)
  X, y = preprocess_data(X, y)

  params = {
    'lr': [Ridge(alpha=1e-5), X, y],
    'ridge 10.0': [Ridge(alpha=10.0), X, y],
  }

  for C in [0.1, 1.0, 16.0, 32.0, 100.0, 128.0]:
    for g in ['auto', 0.25, 0.5, 1.0]:
      for eps in [2 ** (-8), 0.01, 0.25, 0.5]:
        params['svr C={}, g={}, eps={}'.format(C, g, eps)] = [
          SVR(C=C, gamma=g, epsilon=eps), X, y
        ]
  n_jobs = 1 if args.global_parallel else args.n_jobs
  for max_depth in [None, 10, 50]:
    for max_features in ['auto', 5]:
      for min_samples_split in [2, 10, 30]:
        for min_samples_leaf in [1, 10, 30]:
          params[
            'rf md={}, mf={}, mss={}, msl={}'.format(
              max_depth,max_features, min_samples_split, min_samples_leaf)
          ] = [RandomForestRegressor(
            n_estimators=100, max_depth=max_depth,
            max_features=max_features,
            min_samples_split=min_samples_split,
            min_samples_leaf=min_samples_leaf,n_jobs=n_jobs), X, y]

  # TODO: don't refit the model for weighted/not weighted evaluation
  for f in [True, False]:
    for w in [True, False]:
      for k in [2, 4, 6, 8]:
        for l in [0, 1, 10, 100]:
          params['kplane k={} l={} w={} f={}'.format(k, l, w, f)] = [
            KPlaneRegressor(k, l, weighted=w, fuzzy=f), X, y]
          params['CLR_p k={} l={} w={} f={}'.format(k, l, w, f)] = [
            CLRpRegressor(k, l, weighted=w, fuzzy=f), X, y]
          params['CLR_p LR=1.0 k={} l={} w={} f={}'.format(k, l, w, f)] = [
            CLRpRegressor(k, l, weighted=w, fuzzy=f, clf=LogisticRegression()), X, y]
          params['kplane k={} l={} w={} f={} ens=10'.format(k, l, w, f)] = [
            RegressorEnsemble(KPlaneRegressor(k, l, weighted=w, fuzzy=f)), X, y]
          params['CLR_p k={} l={} w={} f={} ens=10'.format(k, l, w, f)] = [
            RegressorEnsemble(CLRpRegressor(k, l, weighted=w, fuzzy=f)), X, y]
          params['CLR_p LR=1.0 k={} l={} w={} f={} ens=10'.format(k, l, w, f)] = [
            RegressorEnsemble(CLRpRegressor(k, l, weighted=w, fuzzy=f, clf=LogisticRegression())), X, y]

  for k in [2, 4, 6, 8]:
    for l in [0, 1, 10, 100]:
      params['CLR_c k={} l={}'.format(k, l)] = [
        CLRcRegressor(k, l, constr_id=constr_id), X, y]
      params['CLR_c k={} l={} ens=10'.format(k, l)] = [
        RegressorEnsemble(CLRcRegressor(k, l, constr_id=constr_id)), X, y]
      params['kplane k={} l={} w=size'.format(k, l)] = [
        KPlaneRegressor(k, l, weighted='size'), X, y]
      params['kplane k={} l={} w=size ens=10'.format(k, l)] = [
        RegressorEnsemble(KPlaneRegressor(k, l, weighted='size')), X, y]

  results = evaluate_all(
    params,
    file_name="results/{}-tmp1.csv".format(args.dataset),
    n_jobs=args.n_jobs,
    gl_parallel=args.global_parallel,
  )
  results = results.sort_values('test_mse_mean')

  add_params = {}
  # assuming ensembles are always best
  algos = [CLRpRegressor, KPlaneRegressor]
  algo_names = ['CLR_p', 'kplane']
  for algo, algo_name in zip(algos, algo_names):
    for idx in results.index:
      if algo_name in idx:
        k = int(idx.split()[1].split('=')[1])
        l = int(idx.split()[2].split('=')[1])
        w = idx.split()[3].split('=')[1] == 'True'
        f = idx.split()[4].split('=')[1] == 'True'
        for alpha in [0.01, 0.1, 1.0, 10.0, 100.0]:
          add_params[idx + ' Lasso {}'.format(alpha)] = [
            RegressorEnsemble(algo(k, l, weighted=w, fuzzy=f, clr_lr=Lasso(alpha))), X, y
          ]
          add_params[idx + ' Ridge {}'.format(alpha)] = [
            RegressorEnsemble(algo(k, l, weighted=w, fuzzy=f, clr_lr=Ridge(alpha))), X, y
          ]
        break
  for idx in results.index:
    if 'CLR_c' in idx:
      k = int(idx.split()[1].split('=')[1])
      l = int(idx.split()[2].split('=')[1])
      for alpha in [0.01, 0.1, 1.0, 10.0, 100.0]:
        add_params[idx + ' Lasso {}'.format(alpha)] = [
          RegressorEnsemble(CLRcRegressor(k, l, constr_id=constr_id, clr_lr=Lasso(alpha))), X, y
        ]
        add_params[idx + ' Ridge {}'.format(alpha)] = [
          RegressorEnsemble(CLRcRegressor(k, l, constr_id=constr_id, clr_lr=Ridge(alpha))), X, y
        ]
      break
  add_results = evaluate_all(
    add_params,
    file_name="results/{}-tmp2.csv".format(args.dataset),
    n_jobs=args.n_jobs,
    gl_parallel=args.global_parallel,
  )

  res_complete = results.append(add_results)
  res_complete.to_csv("results/{}.csv".format(args.dataset))
  os.remove("results/{}-tmp1.csv".format(args.dataset))
  os.remove("results/{}-tmp2.csv".format(args.dataset))


#!/usr/bin/python3

# assign the value of 12 to base and 16 to height.
base = 12
height = 16

# calculate
area = 0.5 * base * height

# print the value of area
print(area)


# https://leetcode.com/problems/insertion-sort-list/
# tags: #linked_list, #sorting
#
# Solution: Sort by swapping nodes
# Using insertion sort we need to remove current from its original position and insert it at its correct position:
# * Update current.next to behind which is the position before which current needs to be inserted.
# * Update next pointer of previous behind node (behind_prev.next) because current is now inserted before behind.
# * Update next pointer of previous current node (curr_prev.next) because current was removed previously from here.
# Time complexity: O(n^2), Space complexity O(1)
import sys
from typing import Optional

class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next

    def __str__(self):
        return str(self.val)

class Solution:
    def insertionSortList(self, head: Optional[ListNode]) -> Optional[ListNode]:
        dummy = ListNode(-(sys.maxsize - 1), head)
        curr_prev, current = head, head.next

        while current:
            curr_next = current.next

            if current.val > curr_prev.val:
                curr_prev = current
            else:
                behind_prev, behind = dummy, dummy.next

                while behind.val < current.val:
                    behind_prev, behind = behind, behind.next

                current.next = behind
                behind_prev.next = current
                curr_prev.next = curr_next

            current = curr_next

        return dummy.next

if __name__ == "__main__":
    sol = Solution()
    h1 = ListNode(4)
    h1.next = ListNode(2)
    h1.next.next = ListNode(1)
    h1.next.next.next = ListNode(3)
    h1 = sol.insertionSortList(h1)

    c = h1
    while c:
        print(c, end=" -> ")
        c = c.next
    print("", end="\n")


def prime(n):
    if n==1:
        return 0
    for i in range(2,n//2+1):
        if i%n==0:
            return 0
    return 1

x=int(input())
c=0
for n in range(1,x):
    if x%n==0:
        if prime(n):
            print(n,end=' ')
            c+=1
if c==0:
    print('-1')


r'''
{xrst_begin mixed_hes_fixed_obj_xam.py}
{xrst_spell
   laplace
}

ran_likelihood: Example and Test
################################

p(y|theta, u)
*************
In this example math:`y` given :math:`( \theta , u )`
is distributed normally with mean :math:`\bar{y} + u`
and variance :math:`\theta`; i.e.

.. math::

   - \log [ \B{p} ( y | \theta , u ) ]
   =
   \log \left[ \sqrt{ 2 \pi \theta } \right]
   +
   \frac{1}{2} ( y - \bar{y} - u )^2 / \theta

p(u|theta)
**********
In this example, the prior for :math:`u` given :math:`\theta`
is a normal with mean zero and
standard deviation :math:`\sigma`; i.e.

.. math::

   - \log [ \B{p} ( u | \theta ) ]
   =
   \log \left[ \sqrt{ 2 \pi \sigma^2 } \right]
   +
   \frac{1}{2} u^2 / \sigma^2

p(y|theta)
**********
For this example, Laplace approximation is equal to :math:`\B{p}(y|\theta)`
i.e, it is exact. Furthermore,

.. math::

   - \log[ \B{p}(y|\theta) ]
   =
   \log \left[ \sqrt{ 2 \pi ( \theta + \sigma^2) } \right]
   +
   \frac{1}{2} ( y - \bar{y} )^2 / ( \theta + \sigma^2 )

Derivative
**********
For this example there is no fixed effects likelihood or constraints.
Hence the derivative of the fixed effects objective,
w.r.t the fixed effects :math:`\theta`, is

.. math::

   \frac{1}{2} ( \theta + \sigma^2 )^{-1}
   -
   \frac{1}{2} ( y - \bar{y} )^2 ( \theta + \sigma^2 )^{-2}

Hessian
*******
Taking the derivative of the expression above
w.r.t the fixed effects :math:`\theta` we obtain the Hessian:

.. math::

   ( y - \bar{y} )^2 ( \theta + \sigma^2 )^{-3}
   -
   \frac{1}{2} ( \theta + \sigma^2 )^{-2}

{xrst_literal
   # BEGIN SOURCE
   # END SOURCE
}

{xrst_end mixed_hes_fixed_obj_xam.py}
'''
# BEGIN SOURCE
def hes_fixed_obj_xam() :
   import cppad_py
   import numpy
   ok         = True
   #
   y_bar  = 1.0 # mean of the data y
   y      = 2.0 # data that depends on random effects
   sigma  = 0.5 # standard deviation for the random effects
   #
   # value of theta and u at which we will record ran_likelihood( theta , u)
   theta_u  = numpy.array([ sigma*sigma, 0.0 ], dtype=float)
   #
   # independent variables during the recording
   atheta_u = cppad_py.independent(theta_u)
   #
   # split out theta and u
   atheta   = atheta_u[0]
   au       = atheta_u[1]
   #
   # - log[ p(y|theta,u) ] (dropping terms that are constant w.r.t. theta,u)
   atmp          = ( y  - y_bar - au )
   ap_y_theta_u  = 0.5 * atmp * atmp / atheta
   ap_y_theta_u += 0.5 * numpy.log( atheta )
   #
   # - log[ p(u|theta) ] (dropping terms that are constant w.r.t. theta,u)
   atmp          = au / sigma
   ap_u_theta    = 0.5 * atmp * atmp
   #
   # - log[ p(y|theta, u) p(u|theta) ]
   av_0 = ap_y_theta_u + ap_u_theta
   #
   # function mapping (theta,u) -> v
   av   = numpy.array( [ av_0 ] )
   r    = cppad_py.d_fun(atheta_u, av)
   #
   # mixed_obj
   theta = theta_u[0:1]
   u     = theta_u[1:2]
   mixed_obj = cppad_py.mixed(
      fixed_init     = theta        ,
      random_init    = u            ,
      ran_likelihood = r,
   )
   #
   # optimal random effects
   options    = 'String  sb    yes\n'     # suppress optimizer banner
   options   += 'Integer print_level 0\n' # suppress optimizer trace
   random_opt = mixed_obj.optimize_random(
      random_ipopt_options  = options    ,
      fixed_vec             = theta      ,
   )
   #
   # hes_fixed_obj_rcv
   hes_fixed_obj_rcv = cppad_py.sparse_rcv()
   mixed_obj.hes_fixed_obj(
      hes_fixed_obj_rcv        ,
      fixed_vec  = theta       ,
      random_opt = random_opt  ,
   )
   #
   theta = theta[0]
   term1 = (theta + sigma * sigma)
   term2 = (y - y_bar) * (y - y_bar)
   check = term2 / (term1 * term1 * term1)  - 0.5 / (term1 * term1 )
   #
   # check solution
   ok = ok and hes_fixed_obj_rcv.nr()     == 1
   ok = ok and hes_fixed_obj_rcv.nc()     == 1
   ok = ok and hes_fixed_obj_rcv.nnz()    == 1
   ok = ok and hes_fixed_obj_rcv.row()[0] == 0
   ok = ok and hes_fixed_obj_rcv.col()[0] == 0
   ok = ok and abs( hes_fixed_obj_rcv.val()[0] / check - 1.0 ) < 1e-8
   return ok
# END SOURCE
def test_hes_fixed_obj_xam() :
   assert hes_fixed_obj_xam()


# -*- coding: utf-8 -*-

import os

from pkg_resources import parse_version
import pkginfo

def _egg_info(path_to_egg='../../src/'):
    path_to_egg = os.path.join(
        os.path.dirname(__file__), path_to_egg)
    egg_info = pkginfo.Develop(path_to_egg)
    release = egg_info.version
    version = '%s.%s' % tuple(map(int, parse_version(release)[0:2]))
    return egg_info.name, egg_info.author, version, release

project, author, version, release = _egg_info()
copyright = '2009, %s' % author

# Extension
extensions = [
    'sphinx.ext.autodoc',
    'sphinx.ext.intersphinx',
    'sphinx.ext.todo',
    ]
intersphinx_mapping = {'http://docs.python.org/': None}

# Source
master_doc = 'index'
templates_path = ['_templates']
source_suffix = '.rst'
exclude_trees = []
pygments_style = 'sphinx'

# html build settings
html_theme = 'default'
html_static_path = ['_static']

# htmlhelp settings
htmlhelp_basename = '%sdoc' % project

# latex build settings
latex_documents = [
  ('index', '%s.tex' % project, u'%s Documentation' % project,
   author, 'manual'),
]


'3'
import signal
from neuron import h, gui
import numpy as np
import matplotlib.pyplot as plt
from open_pickle import read_from_pickle
import os
from tqdm import tqdm
from add_figure import add_figure
from glob import glob
import pickle
from calculate_F_factor import calculate_F_factor

path_single_traces=glob('data/traces_img/2017_05_08_A_0006/*pA.p')
path=path_single_traces[0]
I=int(path[path.rfind('/')+1:path.rfind('pA')])
# path1='/ems/elsc-labs/segev-i/moria.fridman/project/data_analysis_git/data_analysis/data/short_pulse_mean.p'
path='/ems/elsc-labs/segev-i/moria.fridman/project/data_analysis_git/data_analysis/data/short_pulse/mean_short_pulse_with_parameters.p'

I=-50

SPINE_START = 60
shrinkage_factor=1#1.0/0.7
resize_diam_by=1

CM=1#2/2
RM=15000#5684*2#*2
RA=100

start_fit= 2000
end_fit=4900#3960

print('the injection current is',I,flush=True)
folder_='/ems/elsc-labs/segev-i/moria.fridman/project/data_analysis_git/data_analysis/'
h.load_file("import3d.hoc")
h.load_file("nrngui.hoc")
h.load_file('stdlib.hoc')
h.load_file("stdgui.hoc")
# h.loadfile("stdrun.hoc")
do_calculate_F_factor=True
spine_type="mouse_spine"
if spine_type=="groger_spine":
    V_head=0.14
    spine_neck_diam=0.164
    spine_neck_L=0.782
elif spine_type=="mouse_spine":
    #mouse
    from math import pi
    head_area=0.37
    r_head=np.sqrt(head_area/(4*pi))
    spine_neck_L=0.73
    # HEAD_DIAM=0.667
    spine_neck_diam=0.25 #0.164/07=0.23
    spine_density=1.08
    V_head=4/3*pi*r_head**3

def SIGSEGV_signal_arises(signalNum, stack):
    print(f"{signalNum} : SIGSEGV arises")
    # Your code
signal.signal(signal.SIGSEGV, SIGSEGV_signal_arises)
def get_inj(T,I,V):
    #found the begining,end and median of the injection
    I_abs = np.abs(I)
    inj_start = np.where(I_abs > I_abs.max() / 4.0)[0][0] - 1
    inj_end = np.where(I_abs > I_abs.max() / 4.0)[0][-1]
    inj = np.median(I[inj_start:inj_end])
    return inj, T[inj_start], T[inj_end]

# the function that looking for the best parameter fit
#F_factor is the correction from the spine A(spine area)/area(without spine)
# F_factor = {} - take the F_factor for the specific segment
def change_model_pas(CM=1, RA = 250, RM = 20000.0, E_PAS = -70.0, F_factor = {}):
    #input the neuron property    h.dt = 0.1

    h.distance(0,0.5, sec=soma) # it isn't good beacause it change the synapse distance to the soma
    #h.distance(0, sec=soma)
    for sec in h.allsec(): ##check if we need to insert Ra,cm,g_pas,e_pas to the dendrit or just to the soma
        sec.Ra = RA
        sec.cm = CM#*(1.0/0.7)
        sec.g_pas = (1.0 / RM)#*(1.0/0.7)
        sec.e_pas = E_PAS
    for sec in cell.dend:
        for seg in sec: #count the number of segment and calclate g_factor and total dend distance,
            # how many segment have diffrent space larger then SPINE_START that decided
            if h.distance(seg) > SPINE_START:
                if do_calculate_F_factor:
                    F_factor=calculate_F_factor(cell,V_head,spine_neck_diam,spine_neck_L)
                else:
                    F_factor = 1.9#2.0 # F_factor[sec]
                seg.cm *= F_factor
                seg.g_pas *= F_factor

## e_pas is the equilibrium potential of the passive current
def plot_res(RM, RA, CM, save_name= "fit",print_full_graph=False):
    folder_="data/fit/"
    try: os.mkdir(folder_)
    except FileExistsError: pass
    if resize_diam_by!=1 and shrinkage_factor!=1:
        try: os.mkdir(folder_+'dend*'+str(round(resize_diam_by,2))+' &shrinkage by '+str(round(shrinkage_factor,2))+'_'+str(I)+'pA/')
        except FileExistsError: pass
        save_folder=folder_+'dend*'+str(round(resize_diam_by,2))+' &shrinkage by '+str(round(shrinkage_factor,2))+'_'+str(I)+'pA/'
    elif resize_diam_by!=1:
        try: os.mkdir(folder_+'dend*'+str(resize_diam_by)+'_'+str(I)+'pA/')
        except FileExistsError: pass
        save_folder=folder_+'dend*'+str(resize_diam_by)+'_'+str(I)+'pA/'
    elif shrinkage_factor!=1:
        try: os.mkdir(folder_+'F_shrinkage='+str(round(shrinkage_factor,2))+'_'+str(I)+'pA/')
        except FileExistsError: pass
        save_folder=folder_+'F_shrinkage='+str(round(shrinkage_factor,2))+'_'+str(I)+'pA/'
    else:
        try: os.mkdir(folder_+str(I)+'pA/')
        except FileExistsError: pass
        save_folder=folder_+str(I)+'pA/'

    # creat a clamp and record it for the chosen parameter
    ## save_name need to incloud the folder path
    change_model_pas(CM=CM, RA=RA, RM=RM, E_PAS = E_PAS)
    Vvec = h.Vector() #cerat vector to record on it
    Tvec = h.Vector() #cerat vector to record on it
    Vvec.record(soma(0.5)._ref_v) #where to recprd
    Tvec.record(h._ref_t) #when it record
    h.cvode.store_events(Vvec)

    h.run()
    npTvec = np.array(Tvec)
    npVec = np.array(Vvec)
    add_figure("fit "+str(I)+"pA\nRM="+str(round(RM,1))+",RA="+str(round(RA,1))+",CM="+str(round(CM,2)),'mS','mV')
    plt.plot(npTvec[start_fit:end_fit], npVec[start_fit:end_fit], color = 'r', linestyle ="--") #plot the recorded short_pulse
    plt.plot(T[start_fit:end_fit], V[start_fit:end_fit], color = 'green')
    plt.legend(['NEURON_sim','decay_to_fitting'])
    plt.savefig(save_folder+'/'+save_name+"_decay.png")
    plt.close()
    exp_V = V#[int(180.0 / h.dt):int(800.0 / h.dt)]
    npVec = npVec#[int(180.0 / h.dt):int(800.0 / h.dt)]
    npVec = npVec[:len(exp_V)]
    error_1 = np.sqrt(np.sum(np.power(np.mean(exp_V[:2000]) - np.mean(npVec[:2000]), 2)))  # error from mean rest
    error_2 = np.sqrt(np.sum(np.power(exp_V[start_fit:end_fit] - npVec[start_fit:end_fit], 2))/(end_fit-start_fit))  #  error for the decay
    error_3 = np.sqrt(np.sum(np.power(np.mean(exp_V[4100:4900]) - np.mean(npVec[4100:4900]), 2)))  # error for maximal voltage
    error_tot = np.sqrt(np.sum(np.power(exp_V - npVec, 2))/len(exp_V)) # mean square error

    print('error_total=',round(error_tot,3))
    print('error_decay=', round(error_2,3))
    print('error_mean_max_voltage=', round(error_3,3))
    print('error_from_rest=', round(error_1,3))
    if print_full_graph:
        add_figure('fit'+str(I)+'pA part ['+str(start_fit)+':'+str(end_fit)+']',short_pulse[0].units,short_pulse[0].units)
        plt.plot(T, V, color = 'k') #plot short_pulse data
        plt.plot(npTvec[:len(npVec)], npVec, color = 'r', linestyle ="--") #plot the recorded short_pulse
        plt.plot(T[start_fit:end_fit], V[start_fit:end_fit], color = 'green')
        plt.suptitle('error from full graph='+str(round(error_tot,3))+' and eddor from decay='+str(round(error_2,3)))
        plt.legend(['data','NEURON_sim','decay_to_fitting'])
        plt.savefig(save_folder+'/'+save_name+"_full_graph.png")
        plt.close()
    return save_folder

def efun(vals):
    #check the fitting
    # if the parameter incloud the fitting (not aqual to 1) check that the result is makes sense, if not return 1e6
    # if the result is make sense calculate the error between the record simulation and the initial data record
    ## *_IX is the parameter we play with them
    ## *_const is the basic parameters we return if the  result doesn't make sense
    if RM_IX != -1 :
        if vals.x[RM_IX] > 100000:
            return (1e6)
        RM = vals.x[RM_IX]
    else: RM = RM_const

    if CM_IX != -1:
        if vals.x[CM_IX] >2 :
            return (1e6)
        CM = vals.x[CM_IX]
    else:CM = CM_const

    if RA_IX != -1:
        if vals.x[RA_IX] > 300:
            return (1e6)
        RA = vals.x[RA_IX]
    else:RA = RA_const
    if (CM < 0.3 or RM < 2000 or RA <50):
        return 1e6
    # print('RA:',RA, '   CM:',CM, '   RM:',RM)

    change_model_pas(CM=CM, RA=RA, RM = RM, E_PAS = E_PAS)
    Vvec = h.Vector()
    Vvec.record(soma(0.5)._ref_v)

    h.run()
    npVec = np.array(Vvec)

    exp_V = V#[int(180.0/h.dt):int(800.0/h.dt)]
    npVec = npVec#[int(180.0/h.dt):int(800.0/h.dt)]
    npVec = npVec[:len(exp_V)]
    error_tot = np.sqrt(np.sum(np.power(exp_V - npVec, 2)))#/len(exp_V)) # mean square error

    error_1 = np.sqrt(np.sum(np.power(np.mean(exp_V[:2000]) - np.mean(npVec[:2000]), 2)))  # error from mean rest
    error_2 = np.sqrt(np.sum(np.power(exp_V[start_fit:end_fit] - npVec[start_fit:end_fit], 2))) #/(end_fit-start_fit)  #  error for the decay
    error_3 = np.sqrt(np.sum(np.power(np.mean(exp_V[4100:4900]) - np.mean(npVec[4100:4900]), 2)))  # error for maximal voltage

    return error_2 + (end_fit-start_fit)*error_3  #@# ask yoni if the calculation is right

class Cell: pass
def mkcell(fname):
    #def to read ACS file
    h('objref cell, tobj')
    loader = h.Import3d_GUI(None)
    loader.box.unmap()
    loader.readfile(fname)
    c = Cell()
    loader.instantiate(c)
    return c
def instantiate_swc(filename):
    h('objref cell, tobj')
    h.load_file('allen_model.hoc')
    h.execute('cell = new allen_model()')
    h.load_file(filename)
    nl = h.Import3d_SWC_read()
    nl.quiet = 1
    nl.input(filename)
    i3d = h.Import3d_GUI(nl, 0)
    i3d.instantiate(h.cell)
    return h.cell
#########################################
# build the model
######################################################

fname = "05_08_A_01062017_Splice_shrink_FINISHED_LABEL_Bluecell_spinec91.ASC"
# cell=instantiate_swc('/ems/elsc-labs/segev-i/moria.fridman/project/data_analysis_git/data_analysis/try1.swc')
cell =mkcell(fname)
print (cell)
sp = h.PlotShape()
sp.show(0)  # show diameters
## delete all the axons
for sec in cell.axon:
   h.delete_section(sec=sec)

soma= cell.soma[0]

#insert pas to all other section
for sec in tqdm(h.allsec()):
    sec.insert('pas') # insert passive property
    sec.nseg = int(sec.L/10)+1  #decide that the number of segment will be 21 with the same distances
for sec in h.allsec():
    sec.diam = sec.diam*resize_diam_by
    sec.L*=shrinkage_factor

# clamp = h.IClamp(cell.dend[82](0.996)) # insert clamp(constant potentientiol) at the soma's center
clamp = h.IClamp(soma(0.5)) # insert clamp(constant potentientiol) at the soma's center
clamp.amp = I/1000#-0.05 ## supopsed to be 0.05nA
if path in path_single_traces:
    start_inj=10500
    end_inj= 20400
    hz=0.1
    clamp.dur = (end_inj-start_inj)*hz
    clamp.delay = start_inj*hz
else:
    clamp.dur = 200
    clamp.delay = 296
######################################################
# load the data and see what we have
######################################################
short_pulse_dict = read_from_pickle(path)
short_pulse=short_pulse_dict['mean']
V = np.array(short_pulse[0])
short_pulse[1]=short_pulse[1].rescale('ms')
T = np.array(short_pulse[1])
T = T-T[0]
T=T
if path in path_single_traces:
    E_PAS = np.mean(V[:9000])
    end_fit=len(T)
else:
    E_PAS = np.mean(V[2945-500:2945])
    # E_PAS = np.mean(V[:2000])

h.tstop = (T[-1]-T[0])
h.v_init=E_PAS
h.dt = 0.1
h.steps_per_ms = h.dt

CM_IX = 2
RM_IX=0
RA_IX = 1

RM_const = 60000.0
RA_const = 100
CM_const = 1.0

print("free params:")

h.attr_praxis(1e-9,1000,0)
opt_vals = h.Vector(3)#3
opt_vals.x[RM_IX] =RM
opt_vals.x[RA_IX] = RA
opt_vals.x[CM_IX] = CM
change_model_pas(CM=CM, RA=RA, RM=RM, E_PAS=E_PAS)

imp = h.Impedance(sec=soma)
imp.loc(soma(0.5))
imp.compute(0)
imp.input(0)

plot_res(CM=CM, RM=RM, RA=RA, save_name="before")
# plot_res(CM=CM, RM=RM, RA=RA, save_name="before")

#plot_res(CM=0.97, RM=11853.6, RA=99.6, save_name="test")

print('the initial impadence is', imp.input(0))

# allway run the fitting 3 time to avoid stack in local minima
for i in range(3):
    RMSD = h.fit_praxis(efun,opt_vals)   #@# take too much time if the fitting isn't found
    print('RMSD:',RMSD)
    RM = opt_vals.x[RM_IX]
    RA = opt_vals.x[RA_IX]
    CM = opt_vals.x[CM_IX]

    print("RMSD", RMSD,", RM",  RM, ", RA",  RA, ", CM",  CM)
    if i==2:
        plot_res(CM=CM, RM=RM, RA=RA, save_name="_fit_after_" + str(i + 1), print_full_graph=True)
    else:
        save_folder=plot_res(CM=CM, RM=RM, RA=RA, save_name="_fit_after_" + str(i + 1))

    # imp.compute(0)
    # print('the impadence is',imp.input(0))
pickle.dump({
        "RM": RM,
        "RA": RA,
        "CM": CM
    }, open(save_folder+'/' + "final_result_dend*"+str(resize_diam_by)+".p", "wb"))
#


import json
import logging

logging.basicConfig(filename='app.log', filemode='w', format='%(name)s - %(levelname)s - %(message)s')

def unpack_json(file):
    try:
        with open(file) as json_file:
            data = json.load(json_file)
            return data
    except FileNotFoundError:
        logging.error(f"File {file} doesn't exists")

def parse_user(output_file, *args):
    data = []
    for items in args:
        json = unpack_json(items)
        if json != None:
            for item in json:
                if "name" in item.keys():
                    repeated = 0
                    if len(data) > 0:
                        for items in data:
                            if item["name"] == items["name"]:
                                repeated += 1
                    if repeated == 0:
                        data.append(item)
    return create_json(output_file, data)

def create_json(file, data):
    with open(file, "w") as json_file:
        json.dump(data, json_file, indent=4)

    return file

print(parse_user('user3.json' , 'user1.json', 'user2.json'))

sprint 06 ( With (in, out). Serialization )/Task 2/Task 2.py


import time
import threading

def subNum():
    global num  #每个线程中都获取这个全局变量
    lock.acquire()#锁定
    temp = num
    time.sleep(0.0000001)
    num=temp-1
    lock.release()#解锁

num=100

lock=threading.Lock()#创建同步锁

thread_list=[]
for i in range(100):
    t=threading.Thread(target=subNum)
    t.start()
    thread_list.append(t)

# for i in thread_list:#等待所有线程执行完毕
#     t.join()

print('Result:',num)

网络编程和并发/进程与线程/同步锁.py


import cv2
import omni2

values = [25, 25, 0, 65, 65, 65, 0, 170]
thrower = 100
speed = 40

cv2.namedWindow("Controller")

while True:
    k = cv2.waitKey(1) & 0xFF

    if k == ord("w"):
        print("Forward")
        omni2.forward(values, speed)
    if k == ord("s"):
        print("Backward")
        omni2.forward(values, speed * -1)
    if k == ord("d"):
        print("Right")
        omni2.rotate(values, speed * -1)
    if k == ord("a"):
        print("Left")
        omni2.rotate(values, speed)
    if k == ord("t"):
        print("Throw")
        omni2.startThrow(values, thrower)
    if k == ord("r"):
        print("Stop throw")
        omni2.endThrow(values)
    if k == ord("e"):
        print("Stop")
        omni2.stopAll(values)
    if k == ord("q"):
        print("Break")
        omni2.stopAll(values)
        break

cv2.destroyAllWindows()


"""
非 binary indexed tree 方法
"""
class NumMatrix_1(object):

    def __init__(self, matrix):
        lenM = len(matrix)
        if lenM == 0:
            return
        lenN = len(matrix[0])
        if lenN == 0:
            return

        self.m = matrix
        self.f = [[0 for x in xrange(lenN+1)] for y in xrange(lenM+1)]
        for i in xrange(1, lenM+1):
            for j in xrange(1, lenN+1):
                self.f[i][j] = self.f[i-1][j] + self.f[i][j-1] - self.f[i-1][j-1] + matrix[i-1][j-1]

    def update(self, row, col, val):
        lenM = len(self.f)
        lenN = len(self.f[0])
        for i in xrange(row+1, lenM):
            for j in xrange(col+1, lenN):
                self.f[i][j] += val-self.m[row][col]
        self.m[row][col] = val

    def sumRegion(self, row1, col1, row2, col2):

        return self.f[row2+1][col2+1] - self.f[row1][col2+1] - self.f[row2+1][col1] + self.f[row1][col1]

# Your NumMatrix object will be instantiated and called as such:
# obj = NumMatrix(matrix)
# obj.update(row,col,val)
# param_2 = obj.sumRegion(row1,col1,row2,col2)

"""
binary index tree 方法
"""

class NumMatrix_2(object):
    def __init__(self, matrix):
        self.lenM = len(matrix)
        if self.lenM == 0:
            return
        self.lenN = len(matrix[0])
        if self.lenN == 0:
            return

        self.arr = [[0 for x in xrange(self.lenN)] for y in xrange(self.lenM)]
        self.btree = [[0 for x in xrange(self.lenN + 1)]
                      for y in xrange(self.lenM + 1)]

        for i in xrange(self.lenM):
            for j in xrange(self.lenN):
                self.update(i, j, matrix[i][j])

    def update(self, row, col, val):
        diff = val - self.arr[row][col]
        self.arr[row][col] = val

        x = row + 1
        y = col + 1

        while x <= self.lenM:
            while y <= self.lenN:
                self.btree[x][y] += diff
                y += self.lowbit(y)
            x += self.lowbit(x)
            y = col + 1

    def sumRegion(self, row1, col1, row2, col2):
        return (self.prefixSum(row2, col2) - self.prefixSum(row2, col1 - 1) -
                self.prefixSum(row1 - 1, col2) +
                self.prefixSum(row1 - 1, col1 - 1))

    def prefixSum(self, row, col):
        sum = 0

        x = row + 1
        y = col + 1

        while x > 0:
            while y > 0:
                sum += self.btree[x][y]
                y -= self.lowbit(y)
            x -= self.lowbit(x)
            y = col + 1

        return sum

    def lowbit(self, x):
        return x & -x

# Your NumMatrix object will be instantiated and called as such:
# obj = NumMatrix(matrix)
# obj.update(row,col,val)
# param_2 = obj.sumRegion(row1,col1,row2,col2)