import itertools

from absl.testing import parameterized
import numpy as np

from tensorflow.core.framework import full_type_pb2
from tensorflow.python import tf2
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import indexed_slices
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.framework import test_util
from tensorflow.python.framework.weak_tensor import WeakTensor
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_math_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import tensor_array_ops
from tensorflow.python.ops import variables
from tensorflow.python.ops import weak_tensor_ops  # pylint: disable=unused-import
from tensorflow.python.ops import weak_tensor_test_util
from tensorflow.python.ops.ragged import ragged_factory_ops
from tensorflow.python.platform import googletest

_convert_to_input_type = weak_tensor_test_util.convert_to_input_type
_get_weak_tensor = weak_tensor_test_util.get_weak_tensor

@test_util.run_all_in_graph_and_eager_modes
class ReduceTest(test_util.TensorFlowTestCase, parameterized.TestCase):
  # Test unary ops with optional dtype arg.

  @parameterized.parameters(
      ("WeakTensor", WeakTensor),
      ("Python", WeakTensor),
      ("NumPy", tensor.Tensor),
      ("Tensor", tensor.Tensor),
  )
  def testReduceAllDims(self, input_type, result_type):
    test_input = _convert_to_input_type(
        [[1, 2, 3], [4, 5, 6]], input_type, dtypes.int32
    )
    with test_util.device(use_gpu=True):
      res = math_ops.reduce_sum(test_input)
      self.assertIsInstance(res, result_type)
      self.assertEqual(self.evaluate(res), 21)

  def testReduceExtendType(self):
    test_in = np.random.randn(1000, 1000).astype(np.float32)
    in_f32 = _get_weak_tensor(test_in, dtypes.float32)
    in_bfl6 = math_ops.cast(test_in, dtypes.bfloat16)

    out_f32 = self.evaluate(math_ops.reduce_sum(in_f32))
    out_bf16 = self.evaluate(math_ops.reduce_sum(in_bfl6))
    expected = math_ops.cast(out_f32, dtypes.bfloat16)

    self.assertAllClose(out_bf16, expected, 1e-3)

  def testCountNonzero(self):
    # simple case
    x = _get_weak_tensor([[0, -2, 0], [4, 0, 0]], dtypes.int32)
    self.assertEqual(self.evaluate(math_ops.count_nonzero(x)), 2)

    # boolean input
    x = math_ops.not_equal(x, 0)
    self.assertEqual(self.evaluate(math_ops.count_nonzero(x)), 2)

    # would overflow if int8 would be used for internal calculations
    x = 2 * np.ones(512, dtype=np.int8)
    self.assertEqual(self.evaluate(math_ops.count_nonzero(x)), 512)

  @parameterized.parameters(
      ("WeakTensor", WeakTensor),
      ("Python", WeakTensor),
      ("NumPy", tensor.Tensor),
      ("Tensor", tensor.Tensor),
  )
  def testReduceExplicitAxes(self, input_type, result_type):
    x = _convert_to_input_type([[1, 2, 3], [4, 5, 6]], input_type, dtypes.int32)
    with test_util.device(use_gpu=True):
      for axis in (0, -2):
        res = math_ops.reduce_sum(x, axis=axis)
        self.assertIsInstance(res, result_type)
        self.assertAllEqual(res, [5, 7, 9])
      for axis in (1, -1):
        res = math_ops.reduce_sum(x, axis=axis)
        self.assertIsInstance(res, result_type)
        self.assertAllEqual(res, [6, 15])
      for axis in (None, (0, 1), (1, 0), (-1, 0), (0, -1), (-2, 1), (1, -2),
                   (-1, -2), (-2, -1)):
        res = math_ops.reduce_sum(x, axis=axis)
        self.assertIsInstance(res, result_type)
        self.assertEqual(self.evaluate(res), 21)

  def testReduceInvalidAxis(self):
    if context.executing_eagerly():
      # The shape check is in run a graph construction time. In eager mode,
      # it misses the check, magically return result given wrong shape.
      return
    x = _get_weak_tensor([[1, 2, 3], [4, 5, 6]], dtype=np.int32)
    axis = np.array([[0], [1]])
    with self.assertRaisesRegex(ValueError, "must be at most rank 1"):
      math_ops.reduce_sum(x, axis)

  def testReduceVar(self):
    x = _get_weak_tensor([[0, 0, 0], [0, 0, 0]], dtype=dtypes.float32)
    self.assertAllClose(self.evaluate(math_ops.reduce_variance(x)), 0)
    self.assertAllClose(
        self.evaluate(math_ops.reduce_variance(x, axis=0)), [0, 0, 0])

    x = _get_weak_tensor([[1, 2, 1, 1], [1, 1, 0, 1]])
    with self.assertRaisesRegex(TypeError, "must be either real or complex"):
      math_ops.reduce_variance(x)

    x = _get_weak_tensor([[1.0, 2.0, 1.0, 1.0], [1.0, 1.0, 0.0, 1.0]])
    self.assertEqual(self.evaluate(math_ops.reduce_variance(x)), 0.25)
    x_np = np.array([[1, 2, 1, 1], [1, 1, 0, 1]], "float32")
    self.assertEqual(np.var(x_np), 0.25)
    self.assertEqual(self.evaluate(math_ops.reduce_variance(x_np)), 0.25)

    x = ragged_factory_ops.constant([[5., 1., 4., 1.], [], [5., 9., 2.], [5.],
                                     []])
    self.assertAllClose(math_ops.reduce_variance(x, axis=0), [0., 16., 1., 0.])

  def testReduceVarComplex(self):
    # Ensure that complex values are handled to be consistent with numpy
    complex_ys = [([0 - 1j, 0 + 1j], dtypes.float64),
                  (np.array([0 - 1j, 0 + 1j], "complex64"), dtypes.float32),
                  (np.array([0 - 1j, 0 + 1j], "complex128"), dtypes.float64)]
    for y, dtype in complex_ys:
      y_result = math_ops.reduce_variance(y)
      self.assertEqual(np.var(y), 1.0)
      self.assertEqual(self.evaluate(y_result), 1.0)
      self.assertEqual(y_result.dtype, dtype)

  def testReduceStd(self):
    x = _get_weak_tensor([[0, 0, 0], [0, 0, 0]], dtypes.float32)
    self.assertAllClose(self.evaluate(math_ops.reduce_std(x)), 0)
    self.assertAllClose(
        self.evaluate(math_ops.reduce_std(x, axis=0)), [0, 0, 0])

    x = _get_weak_tensor([[1, 2, 1, 1], [1, 1, 0, 1]])
    with self.assertRaisesRegex(TypeError, "must be either real or complex"):
      math_ops.reduce_std(x)

    x = [[1., 2., 1., 1.], [1., 1., 0., 1.]]
    res = math_ops.reduce_std(x)
    self.assertEqual(self.evaluate(res), 0.5)
    self.assertIsInstance(res, WeakTensor)
    x_np = np.array(x)
    self.assertEqual(np.std(x_np), 0.5)
    self.assertEqual(self.evaluate(math_ops.reduce_std(x_np)), 0.5)
    self.assertIsInstance(math_ops.reduce_std(x_np), tensor.Tensor)

    x = ragged_factory_ops.constant([[5., 1., 4., 1.], [], [5., 9., 2.], [5.],
                                     []])
    self.assertAllClose(math_ops.reduce_std(x, axis=0), [0., 4., 1., 0.])

  def testReduceStdComplex(self):
    # Ensure that complex values are handled to be consistent with numpy
    complex_ys = [([0 - 1j, 0 + 1j], dtypes.float64),
                  (np.array([0 - 1j, 0 + 1j], "complex64"), dtypes.float32),
                  (np.array([0 - 1j, 0 + 1j], "complex128"), dtypes.float64)]
    for y, dtype in complex_ys:
      y_result = math_ops.reduce_std(y)
      self.assertEqual(np.std(y), 1.0)
      self.assertEqual(self.evaluate(y_result), 1.0)
      self.assertEqual(y_result.dtype, dtype)

@test_util.run_all_in_graph_and_eager_modes
class LogSumExpTest(test_util.TensorFlowTestCase):

  def testReduceLogSumExp(self):
    for dtype in [np.float16, np.float32, np.double]:
      x_np = np.random.rand(5, 5).astype(dtype)
      with test_util.use_gpu():
        y_tf_np = math_ops.reduce_logsumexp(x_np)
        y_np = np.log(np.sum(np.exp(x_np)))
        self.assertAllClose(y_tf_np, y_np)

  def testReductionIndices(self):
    for dtype in [np.float16, np.float32, np.double]:
      x_np = np.random.rand(5, 5).astype(dtype)
      with test_util.use_gpu():
        y_tf = math_ops.reduce_logsumexp(x_np, axis=[0])
        y_np = np.log(np.sum(np.exp(x_np), axis=0))
        self.assertShapeEqual(y_np, y_tf)
        y_tf_np = self.evaluate(y_tf)
        self.assertAllClose(y_tf_np, y_np)

  def testReductionIndices2(self):
    for dtype in [np.float16, np.float32, np.double]:
      x_np = np.random.rand(5, 5).astype(dtype)
      with test_util.use_gpu():
        y_tf = math_ops.reduce_logsumexp(x_np, axis=0)
        y_np = np.log(np.sum(np.exp(x_np), axis=0))
        self.assertShapeEqual(y_np, y_tf)
        y_tf_np = self.evaluate(y_tf)
        self.assertAllClose(y_tf_np, y_np)

  def testKeepDims(self):
    for dtype in [np.float16, np.float32, np.double]:
      x_np = np.random.rand(5, 5).astype(dtype)
      with test_util.use_gpu():
        y_tf_np = math_ops.reduce_logsumexp(x_np, keepdims=True)
        self.assertEqual(y_tf_np.shape.rank, x_np.ndim)
        y_np = np.log(np.sum(np.exp(x_np), keepdims=True))
        self.assertAllClose(y_tf_np, y_np)

  def testOverflow(self):
    x = [1000, 1001, 1002, 1003]
    for dtype in [np.float32, np.double]:
      x_np = np.array(x, dtype=dtype)
      max_np = np.max(x_np)
      with self.assertRaisesRegex(RuntimeWarning,
                                  "overflow encountered in exp"):
        out = np.log(np.sum(np.exp(x_np)))
        if out == np.inf:
          raise RuntimeWarning("overflow encountered in exp")

      with test_util.use_gpu():
        x_tf = _get_weak_tensor(x_np, shape=x_np.shape)
        y_tf_np = math_ops.reduce_logsumexp(x_tf)
        y_np = np.log(np.sum(np.exp(x_np - max_np))) + max_np
        self.assertAllClose(y_tf_np, y_np)

  def testUnderflow(self):
    x = [-1000, -1001, -1002, -1003]
    for dtype in [np.float32, np.double]:
      x_np = np.array(x, dtype=dtype)
      max_np = np.max(x_np)
      with self.assertRaisesRegex(RuntimeWarning,
                                  "divide by zero encountered in log"):
        out = np.log(np.sum(np.exp(x_np)))
        if out == -np.inf:
          raise RuntimeWarning("divide by zero encountered in log")

      with test_util.use_gpu():
        x_tf = _get_weak_tensor(x_np, shape=x_np.shape)
        y_tf_np = math_ops.reduce_logsumexp(x_tf)
        y_np = np.log(np.sum(np.exp(x_np - max_np))) + max_np
        self.assertAllClose(y_tf_np, y_np)

  def testInfinity(self):
    with test_util.use_gpu():
      res = math_ops.reduce_logsumexp(-np.inf)
      self.assertEqual(-np.inf, self.evaluate(res))

@test_util.run_all_in_graph_and_eager_modes
class RoundTest(test_util.TensorFlowTestCase):

  def testRounding(self):
    x = np.arange(-5.0, 5.0, .25)
    for dtype in [np.float32, np.double, np.int32]:
      x_np = np.array(x, dtype=dtype)
      with test_util.device(use_gpu=True):
        x_tf = _get_weak_tensor(x_np, shape=x_np.shape)
        y_tf = math_ops.round(x_tf)
        y_tf_np = self.evaluate(y_tf)
        y_np = np.round(x_np)
        self.assertAllClose(y_tf_np, y_np, atol=1e-2)

class SignTest(test_util.TensorFlowTestCase):

  def test_complex_sign_gradient(self):
    with context.eager_mode():
      x = math_ops.complex(1., 1.)
      with backprop.GradientTape() as t:
        t.watch(x)
        y = math_ops.sign(x)
      self.assertAllClose(
          t.gradient(y, x), math_ops.complex(0.353553, -0.353553))

@test_util.run_all_in_graph_and_eager_modes
class ReciprocalNoNanTest(test_util.TensorFlowTestCase):

  allowed_dtypes = [dtypes.float32, dtypes.float64, dtypes.complex128]

  def testBasic(self):
    for dtype in self.allowed_dtypes:
      x = _get_weak_tensor([1.0, 2.0, 0.0, 4.0], dtype=dtype)

      y = math_ops.reciprocal_no_nan(x)

      target = _get_weak_tensor([1.0, 0.5, 0.0, 0.25], dtype=dtype)

      self.assertAllEqual(y, target)
      self.assertEqual(y.dtype.base_dtype, target.dtype.base_dtype)

  def testInverse(self):
    for dtype in self.allowed_dtypes:
      x = np.random.choice([0, 1, 2, 4, 5], size=(5, 5, 5))
      x = _get_weak_tensor(x, dtype=dtype)

      y = math_ops.reciprocal_no_nan(math_ops.reciprocal_no_nan(x))

      self.assertAllClose(y, x)
      self.assertEqual(y.dtype.base_dtype, x.dtype.base_dtype)

class EqualityTest(test_util.TensorFlowTestCase, parameterized.TestCase):

  @test_util.run_all_in_graph_and_eager_modes
  def testEqualityNone(self):
    x = _get_weak_tensor([1.0, 2.0, 0.0, 4.0], dtype=dtypes.float32)
    self.assertNotEqual(x, None)
    self.assertNotEqual(None, x)
    self.assertFalse(math_ops.tensor_equals(x, None))
    self.assertTrue(math_ops.tensor_not_equals(x, None))

  @parameterized.named_parameters(
      (f"-is_equals={is_equals}-float_literal_type={type(float_literal)}"  # pylint: disable=g-complex-comprehension
       f"-float_literal={float_literal}", is_equals, float_literal)
      for float_literal in [4.6, np.float32(4.6), 4.4, np.float32(4.4)]
      for is_equals in [True, False])
  def testEqualityNoDowncast(self, is_equals, float_literal):
    if (tf2.enabled() and isinstance(float_literal, np.float32) or
        not tf2.enabled() and isinstance(float_literal, float)):
      # TODO(b/199262800): Remove this skip
      self.skipTest("There is a bug in type promotion.")
    if is_equals:
      op = math_ops.tensor_equals
    else:
      op = math_ops.tensor_not_equals
    x = _get_weak_tensor(4)
    try:
      result = op(x, float_literal)
      if isinstance(result, tensor.Tensor):
        result = self.evaluate(result)
    except TypeError:
      # Throwing a TypeError is OK
      return
    self.assertEqual(result, not is_equals)

@test_util.run_all_in_graph_and_eager_modes
class ErfcinvTest(test_util.TensorFlowTestCase):

  def testErfcinv(self):
    values = _get_weak_tensor(
        np.random.uniform(0.1, 1.9, size=int(1e4)).astype(np.float32)
    )
    approx_id = math_ops.erfc(math_ops.erfcinv(values))
    self.assertAllClose(values, self.evaluate(approx_id))

@test_util.run_all_in_graph_and_eager_modes
class ArgMaxMinTest(test_util.TensorFlowTestCase):

  def _generateRandomWeakTensor(self, dtype, shape):
    if dtype.is_integer:
      array = np.random.default_rng().integers(
          low=dtype.min, high=dtype.max, size=shape, endpoint=True)
      return _get_weak_tensor(array, dtype=dtype)
    else:
      array = np.random.default_rng().uniform(low=-1.0, high=1.0, size=shape)
      return _get_weak_tensor(array, dtype=dtype)

  def _getValidDtypes(self):
    return (dtypes.float32, dtypes.float64, dtypes.int32, dtypes.int64)

  def testArgMax(self):
    shape = (24, 8)
    for dtype in self._getValidDtypes():
      tf_values = self._generateRandomWeakTensor(dtype, shape)
      np_values = self.evaluate(tf_values)
      for axis in range(0, len(shape)):
        np_max = np.argmax(np_values, axis=axis)
        tf_max = math_ops.argmax(tf_values, axis=axis)
        self.assertAllEqual(tf_max, np_max)

  def testArgMaxReturnsFirstOccurence(self):
    for dtype in self._getValidDtypes():
      values = _get_weak_tensor(
          [[10, 11, 15, 15, 10], [12, 12, 10, 10, 12]], dtype=dtype
      )
      self.assertAllEqual(
          math_ops.argmax(values, axis=1),
          np.argmax(self.evaluate(values), axis=1))

      # Long tensor to ensure works with multithreading/GPU
      values = array_ops.zeros(shape=(193681,), dtype=dtype)
      self.assertAllEqual(math_ops.argmax(values), 0)

  def testArgMaxUint16(self):
    shape = (24, 8)
    for dtype in self._getValidDtypes():
      tf_values = self._generateRandomWeakTensor(dtype, shape)
      np_values = self.evaluate(tf_values)
      for axis in range(0, len(shape)):
        np_max = np.argmax(np_values, axis=axis)
        tf_max = math_ops.argmax(
            tf_values, axis=axis, output_type=dtypes.uint16)
        self.assertAllEqual(tf_max, np_max)

  def testArgMin(self):
    shape = (24, 8)
    for dtype in self._getValidDtypes():
      tf_values = self._generateRandomWeakTensor(dtype, shape)
      np_values = self.evaluate(tf_values)
      for axis in range(0, len(shape)):
        np_min = np.argmin(np_values, axis=axis)
        tf_min = math_ops.argmin(tf_values, axis=axis)
        self.assertAllEqual(tf_min, np_min)

  def testArgMinReturnsFirstOccurence(self):
    for dtype in self._getValidDtypes():
      values = _get_weak_tensor(
          [[10, 11, 15, 15, 10], [12, 12, 10, 10, 12]], dtype=dtype
      )
      self.assertAllEqual(
          math_ops.argmin(values, axis=1),
          np.argmin(self.evaluate(values), axis=1))

      # Long tensor to ensure works with multithreading/GPU
      values = array_ops.zeros(shape=(193681,), dtype=dtype)
      self.assertAllEqual(math_ops.argmin(values), 0)

class CastTest(test_util.TensorFlowTestCase):

  def testCastWithFullType(self):

    @def_function.function
    def test_fn():
      ta = tensor_array_ops.TensorArray(dtypes.int32, size=1)
      h = math_ops.cast(ta.flow, dtypes.variant)

      t = full_type_pb2.FullTypeDef(
          type_id=full_type_pb2.TFT_PRODUCT,
          args=[full_type_pb2.FullTypeDef(type_id=full_type_pb2.TFT_ARRAY)])
      h.op.experimental_set_type(t)

      ta = tensor_array_ops.TensorArray(dtypes.int32, flow=h)
      ta = ta.write(0, _get_weak_tensor(1))
      return ta.stack()

    self.assertAllEqual(self.evaluate(test_fn()), [1])

@test_util.run_all_in_graph_and_eager_modes
class BinaryOpsTest(test_util.TensorFlowTestCase):

  def testRHSDispatchingAndErrorRaising(self):
    if context.executing_eagerly():
      error = ValueError
      error_message = r"Attempt to convert a value .* with an unsupported type"
    else:
      error = TypeError
      error_message = r"Failed to convert elements of .* to Tensor"

    class RHSReturnsTrue:

      def __radd__(self, other):
        return True

      def __rmod__(self, other):
        return False

    a = array_ops.ones([1], dtype=dtypes.int32) + RHSReturnsTrue()
    self.assertEqual(a, True)

    a = _get_weak_tensor(5, dtype=dtypes.int32) + RHSReturnsTrue()
    self.assertEqual(a, True)

    a = array_ops.ones([1], dtype=dtypes.float32) % RHSReturnsTrue()
    self.assertEqual(a, False)

    a = _get_weak_tensor(5, dtype=dtypes.float32) % RHSReturnsTrue()
    self.assertEqual(a, False)

    class RHSRaisesError:

      def __radd__(self, other):
        raise TypeError("RHS not implemented")

    with self.assertRaisesRegex(error, error_message):
      a = array_ops.ones([1], dtype=dtypes.int32) + RHSRaisesError()
      self.evaluate(a)

    with self.assertRaisesRegex(error, error_message):
      a = _get_weak_tensor([1], dtype=dtypes.int32) + RHSRaisesError()
      self.evaluate(a)

    class RHSReturnsNotImplemented:

      def __radd__(self, other):
        return NotImplemented

    with self.assertRaisesRegex(error, error_message):
      a = array_ops.ones([1], dtype=dtypes.int32) + RHSReturnsNotImplemented()
      self.evaluate(a)

      a = _get_weak_tensor([1], dtype=dtypes.int32) + RHSReturnsNotImplemented()
      self.evaluate(a)

    class RHSNotImplemented:
      pass

    with self.assertRaisesRegex(error, error_message):
      a = array_ops.ones([1], dtype=dtypes.int32) + RHSNotImplemented()
      self.evaluate(a)

      a = _get_weak_tensor([1], dtype=dtypes.int32) + RHSNotImplemented()
      self.evaluate(a)

@test_util.run_all_in_graph_and_eager_modes
class ScalarMulTest(test_util.TensorFlowTestCase):

  def testAcceptsRefs(self):
    if context.executing_eagerly():
      var = resource_variable_ops.ResourceVariable(10, name="var")
    else:
      var = variables.Variable(10)
    result = math_ops.scalar_mul(3, var)
    init = variables.global_variables_initializer()
    with test_util.device(use_gpu=True):
      self.evaluate(init)
      self.assertEqual(30, self.evaluate(result))

  def testAcceptsIndexedSlices(self):
    values = constant_op.constant([2, 3, 5, 7, 0, -1], shape=[3, 2])
    indices = constant_op.constant([0, 2, 5])
    # Test that patched scalar_mul works with IndexedSlices.
    x = math_ops.scalar_mul(-3, indexed_slices.IndexedSlices(values, indices))
    with test_util.device(use_gpu=True):
      self.assertAllEqual(
          self.evaluate(x.values), [[-6, -9], [-15, -21], [0, 3]]
      )
      self.assertAllEqual(self.evaluate(x.indices), [0, 2, 5])

class ComparisonOps(parameterized.TestCase, test_util.TensorFlowTestCase):

  def test_math_equal(self):
    self.assertAllEqual(math_ops.equal(1, constant_op.constant(1)), True)
    self.assertAllEqual(
        math_ops.equal(np.int_(1), constant_op.constant(1)), True
    )
    self.assertAllEqual(
        math_ops.equal(
            constant_op.constant(1, dtypes.float32),
            constant_op.constant(1, dtypes.int32),
        ),
        True,
    )

  def test_math_maximum(self):
    # Test math_ops.maximum.
    self.assertAllEqual(math_ops.maximum(1, constant_op.constant(2)), 2)
    self.assertAllEqual(
        math_ops.maximum(np.int_(1), constant_op.constant(1.5, dtypes.float32)),
        np.array(1.5, np.float32),
    )
    self.assertAllEqual(
        math_ops.maximum(
            constant_op.constant(5, dtypes.float32),
            constant_op.constant(1, dtypes.int32),
        ),
        5,
    )

    # Test gen_math_ops.maximum.
    self.assertAllEqual(gen_math_ops.maximum(1, constant_op.constant(2)), 2)
    self.assertAllEqual(
        gen_math_ops.maximum(
            np.int_(1), constant_op.constant(1.5, dtypes.float32)
        ),
        np.array(1.5, np.float32),
    )
    self.assertAllEqual(
        gen_math_ops.maximum(
            constant_op.constant(5, dtypes.float32),
            constant_op.constant(1, dtypes.int32),
        ),
        5,
    )

  def test_math_minimum(self):
    # Test math_ops.minimum.
    self.assertAllEqual(math_ops.minimum(1, constant_op.constant(2)), 1)
    self.assertAllEqual(
        math_ops.minimum(np.int_(1), constant_op.constant(1.1, dtypes.float32)),
        1,
    )
    self.assertAllEqual(
        math_ops.minimum(
            constant_op.constant(5, dtypes.float32),
            constant_op.constant(-1, dtypes.int32),
        ),
        -1,
    )

    # Test gen_math_ops.minimum.
    self.assertAllEqual(gen_math_ops.minimum(1, constant_op.constant(2)), 1)
    self.assertAllEqual(
        gen_math_ops.minimum(
            np.int_(1), constant_op.constant(1.1, dtypes.float32)
        ),
        1,
    )
    self.assertAllEqual(
        gen_math_ops.minimum(
            constant_op.constant(5, dtypes.float32),
            constant_op.constant(-1, dtypes.int32),
        ),
        -1,
    )

allowed_var_op_input_combinations = [
    (dtypes.uint8, 10),
    (dtypes.uint8, "weak_i64"),
    (dtypes.uint8, dtypes.uint8),
    (dtypes.uint16, 10),
    (dtypes.uint16, "weak_i64"),
    (dtypes.uint16, dtypes.uint8),
    (dtypes.uint16, dtypes.uint16),
    (dtypes.uint32, 10),
    (dtypes.uint32, "weak_i64"),
    (dtypes.uint32, dtypes.uint8),
    (dtypes.uint32, dtypes.uint32),
    (dtypes.uint64, 10),
    (dtypes.uint64, "weak_i64"),
    (dtypes.uint64, dtypes.uint32),
    (dtypes.uint64, dtypes.uint64),
    (dtypes.int8, 10),
    (dtypes.int8, "weak_i64"),
    (dtypes.int8, dtypes.int8),
    (dtypes.int16, 10),
    (dtypes.int16, "weak_i64"),
    (dtypes.int16, dtypes.uint8),
    (dtypes.int16, dtypes.int8),
    (dtypes.int16, dtypes.int16),
    (dtypes.int32, 10),
    (dtypes.int32, "weak_i64"),
    (dtypes.int32, dtypes.uint16),
    (dtypes.int32, dtypes.int16),
    (dtypes.int32, dtypes.int32),
    (dtypes.int64, 10),
    (dtypes.int64, "weak_i64"),
    (dtypes.int64, dtypes.uint32),
    (dtypes.int64, dtypes.int32),
    (dtypes.int64, dtypes.int64),
    (dtypes.bfloat16, 10),
    (dtypes.bfloat16, "weak_i64"),
    (dtypes.bfloat16, 1.0),
    (dtypes.bfloat16, "weak_f64"),
    (dtypes.bfloat16, dtypes.int32),
    (dtypes.bfloat16, dtypes.bfloat16),
    (dtypes.float16, 10),
    (dtypes.float16, "weak_i64"),
    (dtypes.float16, 1.0),
    (dtypes.float16, "weak_f64"),
    (dtypes.float16, dtypes.int32),
    (dtypes.float16, dtypes.float16),
    (dtypes.float32, 10),
    (dtypes.float32, "weak_i64"),
    (dtypes.float32, 1.0),
    (dtypes.float32, "weak_f64"),
    (dtypes.float32, dtypes.int32),
    (dtypes.float32, dtypes.float32),
    (dtypes.float64, 10),
    (dtypes.float64, "weak_i64"),
    (dtypes.float64, 1.0),
    (dtypes.float64, "weak_f64"),
    (dtypes.float64, dtypes.int32),
    (dtypes.float64, dtypes.float64),
    (dtypes.complex64, 10),
    (dtypes.complex64, "weak_i64"),
    (dtypes.complex64, 1.0),
    (dtypes.complex64, "weak_f64"),
    (dtypes.complex64, 1.0 + 2.0j),
    (dtypes.complex64, "weak_c128"),
    (dtypes.complex64, dtypes.int32),
    (dtypes.complex64, dtypes.complex64),
    (dtypes.complex128, 10),
    (dtypes.complex128, "weak_i64"),
    (dtypes.complex128, 1.0),
    (dtypes.complex128, "weak_f64"),
    (dtypes.complex128, 1.0 + 2.0j),
    (dtypes.complex128, "weak_c128"),
    (dtypes.complex128, dtypes.int32),
    (dtypes.complex128, dtypes.float64),
    (dtypes.complex128, dtypes.complex128),
]

disallowed_var_op_input_combinations = [
    (dtypes.uint8, 1.0),
    (dtypes.uint8, "weak_f64"),
    (dtypes.uint8, dtypes.int8),
    (dtypes.uint8, dtypes.uint16),
    (dtypes.uint16, 1.0),
    (dtypes.uint16, "weak_f64"),
    (dtypes.uint16, dtypes.int8),
    (dtypes.uint16, dtypes.uint32),
    (dtypes.uint32, 1.0),
    (dtypes.uint32, "weak_f64"),
    (dtypes.uint32, dtypes.int32),
    (dtypes.uint32, dtypes.uint64),
    (dtypes.uint64, 1.0),
    (dtypes.uint64, "weak_f64"),
    (dtypes.uint64, dtypes.int8),
    (dtypes.uint64, dtypes.float16),
    (dtypes.int8, 1.0),
    (dtypes.int8, "weak_f64"),
    (dtypes.int8, dtypes.int16),
    (dtypes.int8, dtypes.float32),
    (dtypes.int8, dtypes.complex64),
    (dtypes.int16, 1.0),
    (dtypes.int16, "weak_f64"),
    (dtypes.int16, dtypes.int32),
    (dtypes.int16, dtypes.float32),
    (dtypes.int16, dtypes.complex64),
    (dtypes.int32, 1.0),
    (dtypes.int32, "weak_f64"),
    (dtypes.int32, dtypes.int64),
    (dtypes.int32, dtypes.float16),
    (dtypes.int32, dtypes.complex64),
    (dtypes.int64, 1.0),
    (dtypes.int64, "weak_f64"),
    (dtypes.int64, dtypes.uint64),
    (dtypes.int64, dtypes.float16),
    (dtypes.int64, dtypes.complex64),
    (dtypes.bfloat16, 1.0 + 2.0j),
    (dtypes.bfloat16, "weak_c128"),
    (dtypes.bfloat16, dtypes.float16),
    (dtypes.bfloat16, dtypes.float32),
    (dtypes.float16, 1.0 + 2.0j),
    (dtypes.float16, "weak_c128"),
    (dtypes.float16, dtypes.bfloat16),
    (dtypes.float16, dtypes.float32),
    (dtypes.float64, 1.0 + 2.0j),
    (dtypes.float64, "weak_c128"),
    (dtypes.float64, dtypes.complex64),
    (dtypes.complex64, dtypes.float64),
    (dtypes.complex64, dtypes.complex128),
]

def _weak_tensor_from_str(s):
  if s == "weak_i64":
    return _get_weak_tensor(1, dtype=dtypes.int64)
  elif s == "weak_f64":
    return _get_weak_tensor(1.0, dtype=dtypes.float64)
  elif s == "weak_c128":
    return _get_weak_tensor(1.0 + 2.0j, dtype=dtypes.complex128)
  else:
    raise ValueError(f"Unsupported str: {s}")

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

  @parameterized.parameters(
      itertools.product(
          allowed_var_op_input_combinations,
          ("assign", "assign_add", "assign_sub")))
  def testAllowedDtypes(self, v_dtype_and_delta, op):
    v_dtype, delta = v_dtype_and_delta
    if isinstance(delta, dtypes.DType):
      delta = constant_op.constant(1, delta)
    elif isinstance(delta, str):
      delta = _weak_tensor_from_str(delta)

    var = resource_variable_ops.ResourceVariable(10, dtype=v_dtype)
    result = getattr(var, op)(delta)
    with test_util.device(use_gpu=True):
      self.assertEqual(result.dtype, v_dtype)

  @parameterized.parameters(
      itertools.product(
          disallowed_var_op_input_combinations,
          ("assign", "assign_add", "assign_sub")))
  def testDisallowedDtypes(self, v_dtype_and_delta, op):
    v_dtype, delta = v_dtype_and_delta
    if isinstance(delta, dtypes.DType):
      delta = constant_op.constant(1, delta)
    elif isinstance(delta, str):
      delta = _weak_tensor_from_str(delta)

    var = resource_variable_ops.ResourceVariable(10, dtype=v_dtype)
    with self.assertRaises(TypeError):
      _ = getattr(var, op)(delta)

@test_util.run_all_in_graph_and_eager_modes
class DivAndModTest(test_util.TensorFlowTestCase):

  def numpySafeFloorDivInt(self, x, y):
    z = x // y
    # Numpy produces 0 for INT_MIN/-1, but we expect an overflow to INT_MIN
    # so that (INT_MIN/-1) + (INT_MIN % -1) = INT_MIN + 0 = INT_MIN.
    z[(x == np.iinfo(x.dtype).min) & (y == -1)] = np.iinfo(x.dtype).min
    return z

  def numpySafeFloorModInt(self, x, y):
    # Numpy crashes with a FPE for INT_MIN % -1.
    z = self.numpySafeFloorDivInt(x, y)
    return x - z * y

  def numpySafeTruncateDivInt(self, x, y):
    z = self.numpySafeFloorDivInt(x, y)
    # Round up if non-zero remainder and inputs have opposite signs.
    z[(x != z * y) & ((x < 0) != (y < 0))] += 1
    return z

  def numpySafeTruncateModInt(self, x, y):
    # Numpy crashes with a FPE for INT_MIN % -1.
    z = self.numpySafeTruncateDivInt(x, y)
    return x - z * y

  def intEdgeTestData(self, dtype):
    """Edge-case test data for integer types."""
    # INT_MIN/-1 expected to produce signed-integer overflow,
    # INT_MIN/INT_MAX expected to work.
    nums = np.array(
        [np.iinfo(dtype).min, -1, 1, np.iinfo(dtype).max], dtype=dtype
    ).reshape([4, 1])
    divs = nums.reshape([1, 4])
    return nums, divs

  @test_util.disable_asan("Expected signed integer overflow.")
  @test_util.disable_ubsan("Expected signed integer overflow.")
  def testFloorDivModIntEdges(self):
    for dtype in [np.int32, np.int64]:
      x, y = self.intEdgeTestData(dtype)
      x_weak, y_weak = _get_weak_tensor(x), _get_weak_tensor(y)
      tf_floor_div = math_ops.floor_div(x_weak, y_weak)
      np_floor_div = self.numpySafeFloorDivInt(x, y)
      self.assertIsInstance(tf_floor_div, WeakTensor)
      self.assertAllEqual(tf_floor_div, np_floor_div)

      tf_floor_mod = math_ops.floormod(x_weak, y_weak)
      np_floor_mod = self.numpySafeFloorModInt(x, y)
      self.assertIsInstance(tf_floor_div, WeakTensor)
      self.assertAllEqual(tf_floor_mod, np_floor_mod)
      z = math_ops.add(math_ops.multiply(tf_floor_div, y_weak), tf_floor_mod)
      # x = floor_div(x, y) * y + floor_mod(x, y)
      self.assertIsInstance(z, WeakTensor)
      self.assertAllEqual(z, np.broadcast_to(x, z.shape))

  @test_util.disable_asan("Expected signed integer overflow.")
  @test_util.disable_ubsan("Expected signed integer overflow.")
  def testTruncateDivModIntEdges(self):
    for dtype in [np.int32, np.int64]:
      x, y = self.intEdgeTestData(dtype)
      x_weak, y_weak = _get_weak_tensor(x), _get_weak_tensor(y)
      tf_truncate_div = math_ops.truncatediv(x_weak, y_weak)
      np_truncate_div = self.numpySafeTruncateDivInt(x, y)
      self.assertIsInstance(tf_truncate_div, WeakTensor)
      self.assertAllEqual(tf_truncate_div, np_truncate_div)

      tf_truncate_mod = math_ops.truncatemod(x_weak, y_weak)
      np_truncate_mod = self.numpySafeTruncateModInt(x, y)
      self.assertIsInstance(tf_truncate_mod, WeakTensor)
      self.assertAllEqual(tf_truncate_mod, np_truncate_mod)
      z = math_ops.add(
          math_ops.multiply(tf_truncate_div, y_weak), tf_truncate_mod
      )
      self.assertIsInstance(z, WeakTensor)
      # x = truncatediv(x, y) * y + truncatemod(x, y)
      self.assertAllEqual(z, np.broadcast_to(x, z.shape))

@test_util.run_all_in_graph_and_eager_modes
class DivNoNanTest(test_util.TensorFlowTestCase, parameterized.TestCase):
  _SUPPORTED_DTYPES = [
      dtypes.int32,
      dtypes.int64,
      dtypes.float32,
      dtypes.float64,
      dtypes.complex128,
  ]

  @parameterized.parameters(*_SUPPORTED_DTYPES)
  def testBasic(self, dtype):
    if dtype.is_unsigned:
      nums = np.arange(0, 120, 3).reshape(40, 1)
      divs = np.arange(0, 48, 4).reshape(1, 12)
    elif dtype.is_integer:
      nums = np.arange(-120, 120, 3).reshape(80, 1)
      divs = np.arange(-48, 48, 4).reshape(1, 24)
    else:
      nums = np.arange(-10, 10, 0.25).reshape(80, 1)
      divs = np.arange(-3, 3, 0.25).reshape(1, 24)
    assert 0 in divs, "Bad test set-up"

    tf_nums = _get_weak_tensor(nums, dtype=dtype)
    tf_divs = _get_weak_tensor(divs, dtype=dtype)

    # Use tf versions for expected value to ensure inputs are identical
    np_nums = self.evaluate(tf_nums)
    np_divs = self.evaluate(tf_divs)
    np_result = np.true_divide(np_nums, np_divs)
    np_result[:, np_divs[0] == 0] = 0

    with test_util.use_gpu():
      tf_result = math_ops.div_no_nan(tf_nums, tf_divs)
      self.assertIsInstance(tf_result, WeakTensor)
      self.assertAllCloseAccordingToType(tf_result, np_result)

  @parameterized.product(
      type_x=_SUPPORTED_DTYPES + [float, int],
      type_y=_SUPPORTED_DTYPES + [float, int],
  )
  def testSameSupportedTypesAsDivide(self, type_x, type_y):
    def one(type_):
      if type_ is int:
        return 1
      elif type_ is float:
        return 1.0
      else:
        return _get_weak_tensor(1, dtype=type_)

    x = one(type_x)
    y = one(type_y)

    divide_raises = False
    try:
      divide_result = math_ops.divide(x, y)
    except TypeError:
      divide_raises = True

    if divide_raises:
      with self.assertRaises(TypeError):
        _ = math_ops.div_no_nan(x, y)
    else:
      divide_no_nan_result = math_ops.div_no_nan(x, y)
      self.assertIsInstance(divide_no_nan_result, WeakTensor)
      self.assertIsInstance(divide_result, WeakTensor)
      self.assertEqual(divide_no_nan_result.dtype, divide_result.dtype)
      self.assertAllEqual(divide_no_nan_result, divide_result)

  @parameterized.parameters(
      (dtypes.float32),
      (dtypes.float64),
      (dtypes.complex128),
  )
  def testSmall(self, dtype):
    # Choose values whose squared magnitude underflows to zero/subnormal.
    zero = _get_weak_tensor([0, 0, 0, 0], dtype=dtype)
    divs = _get_weak_tensor([1e-25, -1e-20, 1e-165, -1e-160], dtype=dtype)
    tf_result = math_ops.div_no_nan(zero, divs)

    # Results should always be exactly zero.
    self.assertAllEqual(tf_result, zero)
    self.assertIsInstance(tf_result, WeakTensor)

  @parameterized.parameters(
      (dtypes.float32),
      (dtypes.float64),
      (dtypes.complex128),
  )
  def testNonFiniteInNumerator(self, dtype):
    nums = _get_weak_tensor([np.nan, np.inf, np.NINF], dtype=dtype)
    zeros = _get_weak_tensor([0, 0, 0], dtype=dtype)
    ones = _get_weak_tensor([1, 1, 1], dtype=dtype)
    with test_util.use_gpu():
      tf_result_zeros = math_ops.div_no_nan(nums, zeros)
      self.assertAllEqual([0, 0, 0], tf_result_zeros)
      self.assertIsInstance(tf_result_zeros, WeakTensor)
      tf_result_ones = math_ops.div_no_nan(nums, ones)
      self.assertAllEqual(nums / ones, tf_result_ones)
      self.assertIsInstance(tf_result_ones, WeakTensor)

if __name__ == "__main__":
  ops.set_dtype_conversion_mode("all")
  googletest.main()


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

import threading
import time

from enum import Enum, unique
from urllib.parse import urlparse
from lxml import etree
from tempfile import NamedTemporaryFile
from pathlib import Path
from pathvalidate import sanitize_filename

from PySide6.QtCore import QObject

from scihub_eva.globals.preferences import *
from scihub_eva.utils.preferences_utils import *
from scihub_eva.utils.logging_utils import *
from scihub_eva.utils.api_utils import *

@unique
class SciHubAPIRampageType(Enum):
    # Raw query
    RAW = 0

    # Query with captcha
    WITH_CAPTCHA = 1

@unique
class SciHubAPIError(Enum):
    # Unknown error
    UNKNOWN = 0

    # Cannot find a valid PDF when fetching PDF URL
    NO_VALID_PDF = 1

    # Cannot download automatically due to captcha
    BLOCKED_BY_CAPTCHA = 2

    # Wrong captcha
    WRONG_CAPTCHA = 3

class SciHubAPI(QObject, threading.Thread):
    def __init__(
            self,
            logger,
            callback,
            scihub_url,
            sess,
            raw_query=None,
            query=None,
            rampage_type=None
    ):
        QObject.__init__(self)
        threading.Thread.__init__(self)

        self._logger = logger
        self._callback = callback

        self._scihub_url = scihub_url
        self._sess = sess
        self._raw_query = raw_query
        self._query = query
        self._rampage_type = rampage_type
        self._captcha_answer = None

    def __del__(self):
        self._sess.close()

    @property
    def raw_query(self):
        return self._raw_query

    @raw_query.setter
    def raw_query(self, raw_query):
        self._raw_query = raw_query

    @property
    def query(self):
        return self._query

    @query.setter
    def query(self, query):
        self._query = query

    @property
    def rampage_type(self):
        return self._rampage_type

    @rampage_type.setter
    def rampage_type(self, rampage_type):
        self._rampage_type = rampage_type

    @property
    def captcha_answer(self):
        return self._captcha_answer

    @captcha_answer.setter
    def captcha_answer(self, captcha_answer):
        self._captcha_answer = captcha_answer

    def fetch_pdf_url(self, query):
        self._logger.info(
            self.tr('Using Sci-Hub URL: ') +
            f'{self._scihub_url}')

        query_type = guess_query_type(query)
        self._logger.info(self.tr('Query type: ') + query_type.upper())

        pdf_url = query
        err = None

        if query_type != 'pdf':
            try:
                self._logger.info(self.tr('Fetching PDF URL ...'))

                pdf_url_response = self._sess.post(
                    self._scihub_url, data={'request': query}, verify=False,
                    timeout=Preferences.get_or_default(
                        NETWORK_TIMEOUT_KEY,
                        NETWORK_RETRY_TIMES_DEFAULT,
                        type=int) / 1000.0)

                if pdf_url_response.status_code != 200:
                    self._logger.error(self.tr('Error {}').format(
                        pdf_url_response.status_code))
                    self._logger.info(
                        self.tr('You may need check it manually.'))
                    err = SciHubAPIError.UNKNOWN
                else:
                    html = etree.HTML(pdf_url_response.content)
                    article = \
                        html.xpath('//div[@id="article"]/embed[1]') or \
                        html.xpath('//div[@id="article"]/iframe[1]') if \
                        html is not None else None

                    if article and len(article) > 0:
                        pdf_url = urlparse(article[0].attrib['src'])
                        response_url = urlparse(pdf_url_response.url)

                        if pdf_url.scheme == '':
                            pdf_url = pdf_url._replace(
                                scheme=response_url.scheme)

                        if pdf_url.netloc == '':
                            pdf_url = pdf_url._replace(
                                netloc=response_url.netloc)

                        pdf_url = pdf_url.geturl()
                        pdf_url_html = f'{pdf_url}'

                        self._logger.info(
                            self.tr('Got PDF URL: ') + pdf_url_html)
                    else:
                        err = SciHubAPIError.NO_VALID_PDF

                        self._logger.error(self.tr('Failed to get PDF URL!'))
                        self._logger.info(
                            self.tr('You may need check it manually.'))
            except Exception as e:
                err = SciHubAPIError.UNKNOWN

                self._logger.error(self.tr('Failed to get PDF URL!'))
                self._logger.info(self.tr('You may need check it manually.'))
                self._logger.error(str(e))

        return pdf_url, err

    def get_captcha_info(self, pdf_captcha_response):
        captcha_id, captcha_img_url = None, None

        html = etree.HTML(pdf_captcha_response.content)
        imgs = html.xpath('//img[@id="captcha"]')
        ids = html.xpath('//input[@name="id"]')

        if len(imgs) > 0 and len(ids) > 0:
            captcha_id = ids[0].attrib['value']
            captcha_img_url = urlparse(imgs[0].attrib['src'])
            response_url = urlparse(pdf_captcha_response.url)

            if captcha_img_url.scheme == '':
                captcha_img_url = captcha_img_url._replace(
                    scheme=response_url.scheme)

            if captcha_img_url.netloc == '':
                captcha_img_url = captcha_img_url._replace(
                    netloc=response_url.netloc)

            captcha_img_url = captcha_img_url.geturl()

        return captcha_id, captcha_img_url

    def download_captcha_img(self, captcha_img_url):
        captcha_img_file = NamedTemporaryFile(delete=False)
        captcha_img_file_path = Path(captcha_img_file.name)

        captcha_img_res = self._sess.get(captcha_img_url, stream=True)

        if captcha_img_res.status_code == 200:
            for chuck in captcha_img_res:
                captcha_img_file.write(chuck)

        captcha_img_file.flush()
        captcha_img_file.close()

        return captcha_img_file_path

    def fetch_pdf_with_captcha(self, pdf_captcha_response):
        pdf, err = None, None

        captcha_id, _ = self.get_captcha_info(pdf_captcha_response)

        pdf_response = self._sess.post(
            pdf_captcha_response.url,
            data={'answer': self._captcha_answer, 'id': captcha_id},
            verify=False,
            timeout=Preferences.get_or_default(
                NETWORK_TIMEOUT_KEY,
                NETWORK_RETRY_TIMES_DEFAULT,
                type=int) / 1000.0)

        if pdf_response.status_code != 200:
            self._logger.error(self.tr('Error {}').format(
                pdf_response.status_code))
            self._logger.info(self.tr('You may need check it manually.'))
            err = SciHubAPIError.UNKNOWN
        elif pdf_response.headers['Content-Type'] == 'application/pdf':
            self._logger.info(self.tr('Angel [CAPTCHA] down!'))
            pdf = pdf_response.content
        else:
            pdf = pdf_response
            err = SciHubAPIError.WRONG_CAPTCHA

        return pdf, err

    def fetch_pdf(self, pdf_url):
        self._logger.info(self.tr('Fetching PDF ...'))

        pdf, err = None, None

        try:
            pdf_response = self._sess.get(
                pdf_url, verify=False,
                timeout=Preferences.get_or_default(
                    NETWORK_TIMEOUT_KEY,
                    NETWORK_RETRY_TIMES_DEFAULT,
                    type=int) / 1000.0)

            if pdf_response.status_code != 200:
                self._logger.error(self.tr('Error {}').format(
                    pdf_response.status_code))
                self._logger.info(self.tr('You may need check it manually.'))
                err = SciHubAPIError.UNKNOWN
            elif pdf_response.headers['Content-Type'] == 'application/pdf':
                pdf = pdf_response.content
            elif pdf_response.headers['Content-Type'].startswith('text/html'):
                self._logger.warn(self.tr('Angel [CAPTCHA] is coming!'))
                err = SciHubAPIError.BLOCKED_BY_CAPTCHA
                pdf = pdf_response
            else:
                self._logger.error(self.tr('Unknown PDF Content-Type!'))
                self._logger.info(self.tr('You may need check it manually.'))
        except Exception as e:
            err = SciHubAPIError.UNKNOWN

            self._logger.error(self.tr('Failed to get PDF!'))
            self._logger.info(self.tr('You may need check it manually.'))
            self._logger.error(str(e))

        return pdf, err

    def save_pdf(self, pdf, filename):
        pdf_name_formatter = Preferences.get_or_default(
            FILE_FILENAME_PREFIX_FORMAT_KEY,
            FILE_FILENAME_PREFIX_FORMAT_DEFAULT)

        if not Preferences.get_or_default(
                FILE_OVERWRITE_EXISTING_FILE_KEY,
                FILE_OVERWRITE_EXISTING_FILE_DEFAULT,
                type=bool):
            pdf_name_formatter += '_' + str(round(time.time() * 1000000))

        pdf_metadata = get_pdf_metadata(pdf)
        query_type = guess_query_type(self._raw_query)

        if query_type in ['doi', 'pmid']:
            pdf_metadata['id'] = self._raw_query
        else:
            for patten in ['_{id}', '{id}_', '{id}']:
                pdf_name_formatter = pdf_name_formatter.replace(patten, '')

        pdf_name_formatter += '_' + filename if pdf_name_formatter else filename

        try:
            pdf_name = pdf_name_formatter.format(**pdf_metadata)
        except Exception as e:
            self._logger.error(
                self.tr('Unsupported filename keywords: ') + pdf_name_formatter)
            return

        pdf_name = sanitize_filename(pdf_name, replacement_text='-')
        pdf_path = Path(Preferences.get(FILE_SAVE_TO_DIR_KEY)) / pdf_name

        with open(pdf_path, 'wb') as fp:
            fp.write(pdf)

        pdf_link = f'{pdf_path.as_posix()}'

        self._logger.info(self.tr('Saved PDF as: ') + pdf_link)

    def rampage(self, query, rampage_type):
        if rampage_type == SciHubAPIRampageType.RAW:
            self._logger.info(LOGGER_SEP)
            self._logger.info(self.tr('Dealing with query: ') + query)

            # Fetch PDF URL
            pdf_url, err = self.fetch_pdf_url(query)
            if err is not None:
                return None, err

            # Fetch PDF
            pdf, err = self.fetch_pdf(pdf_url)
            if err == SciHubAPIError.BLOCKED_BY_CAPTCHA:
                return pdf, err
            elif err is not None:
                return None, err

            # Save PDF
            filename = urlparse(pdf_url).path[1:].split('/')[-1]
            self.save_pdf(pdf, filename)
        elif rampage_type == SciHubAPIRampageType.WITH_CAPTCHA:
            # Fetch PDF with Captcha
            pdf, err = self.fetch_pdf_with_captcha(query)
            if err == SciHubAPIError.WRONG_CAPTCHA:
                self._logger.error(
                    self.tr('Wrong captcha, failed to kill Angel [CAPTCHA]!'))
                return None, err

            # Save PDF
            filename = urlparse(query.url).path[1:].split('/')[-1]
            self.save_pdf(pdf, filename)

        return None, None

    def run(self):
        res, err = self.rampage(self._query, self._rampage_type)
        self._callback(res, err)

__all__ = [
    'SciHubAPIRampageType',
    'SciHubAPIError',
    'SciHubAPI'
]


import numpy as np
import nibabel as nib
import os
import glob
import pandas as pd
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt

class Evaluate:
    def __init__(self, conf):
        self.config = conf

    def get_whole_tumor_mask(self, data):
        return data > 0

    def get_tumor_core_mask(self, data):
        return np.logical_or(data == 1, data == 4)

    def get_enhancing_tumor_mask(self, data):
        return data == 4

    def dice_coefficient(self, truth, prediction):
        return 2 * np.sum(truth * prediction)/(np.sum(truth) + np.sum(prediction))

    def main(self):
    #    header = ("WholeTumor", "TumorCore", "EnhancingTumor")
    #    masking_functions = (get_whole_tumor_mask, get_tumor_core_mask, get_enhancing_tumor_mask)
        header = ["MSlesions"]
        masking_functions = (self.get_whole_tumor_mask)
        rows = list()
        subject_ids = list()
        path_to_prediction = os.path.abspath("results/prediction/rev_"+str(self.config.rev)+"/prediction_"+self.config.data_set)
        print(path_to_prediction)
        if not os.path.exists(path_to_prediction):
            os.makedirs(path_to_prediction)
        for case_folder in os.listdir(path_to_prediction):
            print("Case: ",case_folder)
            case_folder = os.path.join(path_to_prediction,case_folder)
            if not os.path.isdir(case_folder):
                continue
            subject_ids.append(os.path.basename(case_folder))
            truth_file = os.path.join(case_folder, "truth.nii.gz")
            truth_image = nib.load(truth_file)
            truth = truth_image.get_data()
            prediction_file = os.path.join(case_folder, "prediction.nii.gz")
            prediction_image = nib.load(prediction_file)
            prediction = prediction_image.get_data()

            rows.append([self.dice_coefficient(masking_functions(truth), masking_functions(prediction))])

        df = pd.DataFrame.from_records(rows, columns=header, index=subject_ids)
        print(df)
        path = os.path.abspath("results/prediction/rev_"+str(self.config.rev))
        if not os.path.exists(path):
            os.makedirs(path)
            df.to_csv(os.path.join(path, self.config.data_set+"_scores.csv"))
        else:
            df.to_csv(os.path.join(path, self.config.data_set+"_scores.csv"))

        scores = dict()
        for index, score in enumerate(df.columns):
            values = df.values.T[index]
            scores[score] = values[pd.isnull(values) == False]

        plt.boxplot(list(scores.values()), labels=list(scores.keys()))
        plt.ylabel("Dice Coefficient")
        plt.savefig("validation_scores_boxplot.png")
        plt.close()

        # if os.path.exists("./training.log"):
        #     training_df = pd.read_csv("./training.log").set_index('epoch')
        #
        #     plt.plot(training_df['loss'].values, label='training loss')
        #     plt.plot(training_df['val_loss'].values, label='validation loss')
        #     plt.ylabel('Loss')
        #     plt.xlabel('Epoch')
        #     plt.xlim((0, len(training_df.index)))
        #     plt.legend(loc='upper right')
        #     plt.savefig(os.path.join(path,'loss_graph.png'))


"""
This script summarized DE results.
"""
import numpy as np
import pandas as pd
import session_info
from pyhere import here
from functools import lru_cache
from scipy.stats import mannwhitneyu

@lru_cache()
def get_annotation(feature, tissue):
    """
    Get the annotation file based on feature and tissue.
    """
    feat_lt = {"genes": "gene", "transcripts": "tx",
               "exons": "exon", "junctions": "jxn"}
    new_feature = feat_lt[feature]
    fn = here("input/counts/text_files_counts/_m",
              f"{tissue.lower()}/{new_feature}_annotation.txt")
    return pd.read_csv(fn, sep="\t")

@lru_cache()
def get_resid(feature, tissue):
    """
    Using feature and tissue to select residualized expression.
    """
    return pd.read_csv(here("interaction_sex_sz/interaction_model",
                            f"{tissue.lower()}/_m/{feature}",
                            "residualized_expression.tsv"),
                       sep="\t", index_col=0)\
             .transpose()

@lru_cache()
def get_pheno():
    return pd.read_csv(here("input/phenotypes/_m/phenotypes.csv"))\
             .loc[:, ["RNum", "Sex", "Dx", "Race", "Age"]]

@lru_cache()
def merge_resid(feature, tissue):
    return pd.merge(get_pheno(), get_resid(feature, tissue),
                    left_on="RNum", right_index=True)

@lru_cache()
def subset_resid(tissue, feature, sex):
    sex_dict = {"Male": "F", "Female": "M"}
    df = merge_resid(feature, tissue)
    ctl = df[(df["Dx"] == "Control") &
             (df["Sex"] == sex_dict[sex])].copy()
    sz = df[(df["Dx"] == "SCZD") &
            (df["Sex"] == sex_dict[sex])].copy()
    return ctl, sz

@lru_cache()
def get_de(feature, tissue, sex, fdr):
    """
    Using feature and tissue, select the correct DE results
    file and filter based on FDR (input).
    """
    df = pd.read_csv(here("interaction_sex_sz/by_sex_sz",
                          f"{tissue.lower()}/{sex.lower()}_analysis/_m",
                          f"{feature}/diffExpr_szVctl_full.txt"),
                     sep="\t", index_col=0)\
           .rename(columns={"gene_id":"gencodeID",
                            "gencodeGeneID": "gencodeID",
                            "gene_name": "Symbol"})
    df["ensemblID"] = df.gencodeID.str.replace("\\..*", "", regex=True)
    return df[(df["adj.P.Val"] < fdr)].copy()

@lru_cache()
def annotate_de(feature, tissue, sex, fdr):
    """
    Annotate DE features
    """
    cols = ["Feature", "seqnames", "gencodeID", "ensemblID",
            "Symbol", "logFC", "P.Value", "adj.P.Val", "SE"]
    return get_de(feature, tissue, sex, fdr)\
        .merge(get_annotation(feature, tissue),
               left_index=True, right_on="name")\
        .rename(columns={"name": "Feature"})\
        .loc[:, cols].sort_values("adj.P.Val")

@lru_cache()
def annotate_pvals(feature, tissue, sex, fdr):
    if fdr == 1:
        return annotate_de(feature, tissue, sex, fdr)
    else:
        df = annotate_de(feature, tissue, sex, fdr)
        ctl, sz = subset_resid(tissue, feature, sex)
        pval_df = []
        for feature_id in df.Feature:
            stat, pval = mannwhitneyu(ctl[feature_id], sz[feature_id])
            pval_df.append(pval)
        df["other_pval"] = pval_df
        return df

@lru_cache()
def extract_de(feature, tissue, fdr):
    ff = annotate_pvals(feature, tissue, "Female", fdr)
    mm = annotate_pvals(feature, tissue, "Male", fdr)
    ff["Sex"] = "Female"; mm["Sex"] = "Male"
    return ff, mm

@lru_cache()
def get_unique(feature, tissue, fdr):
    ff, mm = extract_de(feature, tissue, fdr)
    ff_genes = list(set(ff.Feature) - set(mm.Feature))
    mm_genes = list(set(mm.Feature) - set(ff.Feature))
    ff = ff[(ff["Feature"].isin(ff_genes))].copy()
    mm = mm[(mm["Feature"].isin(mm_genes))].copy()
    return ff, mm

@lru_cache()
def filter_pvals(feature, tissue, fdr):
    if fdr == 0.05:
        ff, mm = get_unique(feature, tissue, fdr)
        ff = ff[(ff["other_pval"] < fdr)].copy()
        mm = mm[(mm["other_pval"] < fdr)].copy()
    else:
        ff, mm = extract_de(feature, tissue, fdr)
    return pd.concat([ff, mm], axis=0)

@lru_cache()
def extract_features(tissue, fdr):
    # Extract DE from mash model
    genes = filter_pvals("genes", tissue, fdr)
    trans = filter_pvals("transcripts", tissue, fdr)
    exons = filter_pvals("exons", tissue, fdr)
    juncs = filter_pvals("junctions", tissue, fdr)
    genes["Type"] = "Gene"; trans["Type"] = "Transcript"
    exons["Type"] = "Exon"; juncs["Type"] = "Junction"
    return genes, trans, exons, juncs

def annotate_chrom(df):
    df.loc[:, "Chrom_Type"] = "Other"
    df.loc[df["seqnames"].isin(["chrX", "chrY"]), "Chrom_Type"] = "Allosome"
    df.loc[df["seqnames"].str.contains("chr\d+"), "Chrom_Type"] = "Autosome"
    df.loc[df["seqnames"] == "chrM", "Chrom_Type"] = "Mitochondria"
    return df

def get_DEGs_result_by_tissue(tissue, fdr=0.05):
    # Schizophrenia = 1; Control = -1
    genes, trans, exons, juncs = extract_features(tissue, fdr)
    df = pd.concat([genes, trans, exons, juncs])
    df.loc[:, "Direction"] = np.sign(df.logFC)
    df["Type"] = df.Type.astype("category")\
                   .cat.reorder_categories(["Gene", "Transcript",
                                            "Exon", "Junction"])
    df["Tissue"] = tissue
    return annotate_chrom(df)

def print_summary(tissue, fdr=0.05):
    df = get_DEGs_result_by_tissue(tissue, fdr)
    if tissue == "Caudate":
        w_mode = "w"
    else:
        w_mode = "a"
    statement = f"Significant DE (FDR < 0.05) in {tissue}"
    with open("summarize_results.log", mode=w_mode) as f:
        print(statement, file=f)
        for sex in ["Female", "Male"]:
            print(sex, file=f)
            genes = df[(df["Sex"] == sex) &
                       (df["Type"] == "Gene")].copy()
            trans = df[(df["Sex"] == sex) &
                       (df["Type"] == "Transcript")].copy()
            exons = df[(df["Sex"] == sex) &
                       (df["Type"] == "Exon")].copy()
            juncs = df[(df["Sex"] == sex) &
                       (df["Type"] == "Junction")].copy()
            for variable in ["Feature", "gencodeID"]:
                print(variable, file=f)
                gg = len(set(genes[variable]))
                tt = len(set(trans[variable]))
                ee = len(set(exons[variable]))
                jj = len(set(juncs[variable]))
                print(f"\nGene:\t\t{gg}\nTranscript:\t{tt}"+\
                      f"\nExon:\t\t{ee}\nJunction:\t{jj}\n", file=f)

def merge_data():
    bigdata1 = []; bigdata2 = [];
    for tissue in ["Caudate", "DLPFC", "Hippocampus"]:
        print_summary(tissue)
        data2 = get_DEGs_result_by_tissue(tissue, 0.05)
        data1 = get_DEGs_result_by_tissue(tissue, 1)
        bigdata1.append(data1); bigdata2.append(data2)
    return pd.concat(bigdata1), pd.concat(bigdata2)

def main():
    df1, df2 = merge_data()
    # Summary
    print("\nSummary:")
    gene = df2[(df2["Type"] == "Gene")].copy()
    ff = gene[(gene["Sex"] == "Female")].copy()
    mm = gene[(gene["Sex"] == "Male")].copy()
    print("Unique DEGs (Female): %d" %
          ff.drop_duplicates(subset="gencodeID").shape[0])
    print("Unique DEGs (Male): %d" %
          mm.drop_duplicates(subset="gencodeID").shape[0])
    print(gene.groupby(["Tissue", "Sex", "Direction"]).size())
    # Output
    ## Full data
    df1.sort_values(["Tissue", "Sex", "Type", "P.Value"])\
       .to_csv("differential_expression_schizophrenia_by_sex_4features.txt.gz",
               sep='\t', index=False)
    ## FDR significant
    df2.sort_values(["Tissue", "Sex", "Type", "P.Value"])\
       .to_csv("differential_expression_schizophrenia_by_sex_4features.sig.txt.gz",
               sep='\t', index=False)
    # Session infomation
    session_info.show()

if __name__ == '__main__':
    main()


# Nikko Rush
# 7/11/17

import matplotlib.pyplot as plt

# Plot the Time Series for the given idea
def plot_idea_timeseries(idea_name, idea_numbers, ts_matrix):
    ts_data = ts_matrix[idea_numbers[idea_name]]

    start = 1980
    end = 2015

    plt.plot(ts_data)
    plt.show()

def show_time_series(idea_numbers, ts_matrix):
    # Get a list of idea names
    ideas = list(idea_numbers.keys())

    display_help()

    quit = False
    while not quit:
        selection = input(">")
        if selection in ideas:
            plot_idea_timeseries(selection, idea_numbers, ts_matrix)
        elif selection == "l":
            print(ideas)
        elif selection == "q":
            quit = True
        else:
            display_help()

def display_help():
    print("Enter an idea name to view the time series, l to list ideas, q to quit, or h to view this help message")


import requests
from PIL import Image
from io import BytesIO
from urllib.robotparser import RobotFileParser
from json import load
from os.path import exists
import logging

robots = RobotFileParser()
robots.set_url("https://www.unsplash.com/robots.txt")
robots.read()
logging.basicConfig(
            format="%(asctime)s %(message)s",
            datefmt="%Y-%m-%d %I:%M:%S %p",
            level=logging.DEBUG
        )

next_page = "https://unsplash.com/napi/feeds/home"
with open("/home/stux/headers.json") as fout:
    headers = load(fout)

while next_page:
    logging.debug("Fetching %s" %next_page)
    data = requests.get(next_page, headers=headers).json()
    next_page = data.get('next_page', None)

    for pic in data.get('photos', []):
        if exists(pic['id']+".jpeg"):
            logging.info("%s already exists" %(pic['id']+".jpeg"))
            continue
        logging.debug("Fetching %s" %pic['id'])
        pic_data =  requests.get(pic['urls']['full'])
        pic_data = Image.open(BytesIO(pic_data.content))
        logging.info("Saving picture %s" %(pic['id']+".jpeg"))
        pic_data.save(pic['id']+".jpeg", "JPEG")


import os
from datetime import datetime, timedelta

import numpy as np
import pdb
import pandas as pd
import plotly.express as px
import streamlit as st

from src.frontend import get_config, get_kds
from src.frontend import create_gsheets_database_connection, run_query

conn = create_gsheets_database_connection()

dirname = os.path.dirname(__file__)
cfg = get_config(os.path.join("src", "config.yml"))

st.markdown("# Player Stats")

# Retrieve Data
sheet_url_match = st.secrets["private_gsheets_url_match"]
sheet_url_match_legacy = st.secrets["private_gsheets_url_match_legacy"]
sheet_url_player = st.secrets["private_gsheets_url_player"]
sheet_url_player_legacy = st.secrets["private_gsheets_url_player_legacy"]
sheet_url_roster = st.secrets["private_gsheets_url_roster"]

with st.spinner("Retrieving Player and Match Data..."):
    df_team = run_query(f'SELECT * FROM "{sheet_url_match}"', conn)
    # team_rows_legacy = run_query(f'SELECT * FROM "{sheet_url_match_legacy}"', conn)
    df_player = run_query(f'SELECT * FROM "{sheet_url_player}"', conn)
    # player_rows_legacy = run_query(f'SELECT * FROM "{sheet_url_player_legacy}"', conn)
    df_roster = run_query(f'SELECT * FROM "{sheet_url_roster}"', conn)

df_player = df_player.merge(
    df_team.loc[:, ["match_id", "date", "map"]],
    on="match_id",
    how="left",
)
df_player = df_player.merge(
    df_roster,
    on=['Riot_IDs', 'name', 'tag'],
    how='left'
)

df_player = df_player[~df_player['Riot_IDs'].isin(cfg["retired_players"])]

df_team["date"] = pd.to_datetime(df_team["date"], utc=True)
min_date = df_team.date.min()
max_date = df_team.date.max() + timedelta(days=1)

select_col1, select_col2, select_col3 = st.columns(3)
with select_col1:
    start_date = st.date_input("Start Date", min_date)
with select_col2:
    end_date = st.date_input("End Date", max_date)

# convert to datetimes for comparison
start_date = datetime(start_date.year, start_date.month, start_date.day)
end_date = datetime(end_date.year, end_date.month, end_date.day)

df_player["date"] = pd.to_datetime(df_player["date"])
df_player = df_player.loc[
    (df_player.date >= start_date) & (df_player.date <= end_date), :
]

int_cols = {col: "int64" for col in cfg["integer_columns"]}
float_cols = {col: "float" for col in cfg["float_columns"]}

df_player.loc[:, cfg["integer_columns"]] = df_player.loc[
    :, cfg["integer_columns"]
].fillna(0)
df_player.loc[:, cfg["float_columns"]] = df_player.loc[:, cfg["float_columns"]].fillna(
    0
)
df_player = df_player.astype(int_cols)
df_player = df_player.astype(float_cols)

# Calculate z-scores
df_player['KDA'] = (df_player['kills'] + df_player['assists']) / df_player['deaths']

st.header("General Player Statistics")

df_c = (
    df_player.loc[:, ["Riot_IDs", "clutch_wins", "clutch_counts"]]
    .fillna(0)
    .groupby("Riot_IDs")
    .sum()
    .reset_index()
)
df_c["clutch_percent"] = df_c["clutch_wins"] / df_c["clutch_counts"] * 100
df_c["clutch_percent"] = df_c["clutch_percent"].round(2)
general_cols = [
    "kills",
    "deaths",
    "assists",
    "first_bloods",
    "first_duels",
    "first_deaths",
    "bodyshots",
    "headshots",
    "legshots",
]
pivot_general = (
    df_player.loc[
        :,
        ["Riot_IDs"] + general_cols,
    ]
    .groupby("Riot_IDs")
    .sum()
    .rename(columns={"first_bloods": "first bloods"})
)
pivot_acs = (
    df_player.loc[:, ["Riot_IDs", "acs"]]
    .groupby("Riot_IDs")
    .mean()
    .rename(columns={"acs": "avg. acs"})
)
pivot_general = pd.merge(pivot_general, pivot_acs, left_index=True, right_index=True)
pivot_general = pd.merge(
    pivot_general,
    df_c.set_index("Riot_IDs").loc[:, "clutch_percent"],
    left_index=True,
    right_index=True,
)
match_counts = (
    df_player.groupby(["Riot_IDs"])
    .count()
    .agent.reset_index()
    .rename(columns={"agent": "match counts"})
    .set_index("Riot_IDs")
)
pivot_general = pd.merge(pivot_general, match_counts, left_index=True, right_index=True)
# Calculate per match stats
pivot_general["kills / match"] = pivot_general.kills / pivot_general["match counts"]
pivot_general["deaths / match"] = pivot_general.deaths / pivot_general["match counts"]
pivot_general["assists / match"] = pivot_general.assists / pivot_general["match counts"]
pivot_general["headshot %"] = (
    pivot_general["headshots"]
    / (
        pivot_general["headshots"]
        + pivot_general["bodyshots"]
        + pivot_general["legshots"]
    )
    * 100
)
pivot_general = pivot_general.rename(
    columns={"first_duels": "first duels", "first_deaths": "first deaths"}
)
display_cols = [
    "match counts",
    "kills",
    "kills / match",
    "deaths",
    "deaths / match",
    "assists",
    "assists / match",
    "headshot %",
    "first bloods",
    "first duels",
    "first deaths",
    "avg. acs",
]
# pivot_general.columns=['kills', 'deaths', 'assists', 'first_bloods', 'eco_score', 'acs']

kill_pivot_agent = (
    df_player.pivot_table(values="kills", columns="agent", index="Riot_IDs", aggfunc="sum")
    .fillna(0)
    .astype("int64")
)
kill_pivot_agent["Total"] = kill_pivot_agent.apply(lambda row: sum(row), axis=1)

st.header("Average K/D by Agent")

wavg_col = "Weight. Avg. K/D"
df_player_kd, kd_pivot = get_kds(df_player, wavg_col)
# Change order of columns
new_cols = list(kd_pivot.columns)
new_cols.reverse()
kd_pivot = kd_pivot.loc[:, new_cols]
st.write(kd_pivot.sort_values(by=wavg_col, ascending=False).round(4))
st.write(
    "Note: Weighted Average K/D controls for the amount of times a player plays as a certain agent"
)

# Calculate Ranks
teams = df_player_kd.team.unique()
team_rank_list = []
for team in teams:
    df_filter = df_player_kd[df_player_kd.team == team]
    team_perc_ranks = df_filter.groupby('Riot_IDs').mean(numeric_only=True)['K/D'].rank(pct=True).round(2).reset_index().rename(columns={'K/D':'Team Rank'})
    team_rank_list.append(team_perc_ranks)
df_player_team_perc = pd.concat(team_rank_list)
# Calculate percentile scores for each player across the entire dataset
df_player_global_perc = df_player_kd.groupby('Riot_IDs').mean(numeric_only=True)['K/D'].rank(pct=True).round(2).reset_index().rename(columns={'K/D':'Global Rank'})
df_ranks = df_player_team_perc.merge(df_player_global_perc, on=['Riot_IDs'])
st.header("Average K/D by Map")
kd_pivot_map = df_player_kd.pivot_table(
    values="K/D", columns="map", index="Riot_IDs"
).fillna(0)
st.write(kd_pivot_map.round(4))
# st.plotly_chart(px.imshow(kd_pivot_map.round(4), title="Average K/D by Map Viz"))

st.header('Individual Statistics')
player_list = list(df_player['Riot_IDs'].unique())
player_list.extend(['All'])
player_select = st.selectbox(label='Select Player', options=player_list)

if player_select == 'All':
    st.write(pivot_general.loc[:, display_cols].sort_values(by="kills", ascending=False))
else:
    disp_pivot = pivot_general.loc[pivot_general.index==player_select, display_cols].\
        sort_values(by="kills", ascending=False).merge(df_ranks.loc[:, ['Riot_IDs', 'Global Rank', 'Team Rank']], on='Riot_IDs', how='left')
    st.write(disp_pivot)
st.header("Shooting Breakdown")
df_headshot_ts = df_player.loc[
    df_player['Riot_IDs'] == player_select, ["Riot_IDs", "date", "bodyshots", "legshots", "headshots"]
].dropna()
shot_types = {'headshots': int, 'bodyshots': int, 'legshots': int}
df_headshot_ts = df_headshot_ts.astype(shot_types)
df_headshot_ts["headshot %"] = (
    df_headshot_ts["headshots"]
    / (
        df_headshot_ts["headshots"]
        + df_headshot_ts["bodyshots"]
        + df_headshot_ts["legshots"]
    )
    * 100
)
fig = px.histogram(
    df_headshot_ts,
    x="Riot_IDs",
    y=["legshots", "bodyshots", "headshots"],
    barmode="stack",
    barnorm="percent",
).update_xaxes(categoryorder="total descending")
st.plotly_chart(fig)
# fig_shots = px.line(df_headshot_ts, x="date", y="count", color='Riot_IDs', line_shape="linear")
# st.plotly_chart(fig_shots)


from typing import List

from copy import deepcopy

class Side(object):
    def __init__(self, height: int, width: int):
        self.__array = [[0 for i in range(width)] for k in range(height)]

    @property
    def value(self) -> List[List[int]]:
        return self.__array

    @property
    def width(self) -> int:
        return len(self.__array[0])

    @property
    def height(self) -> int:
        return len(self.__array)

    def __repr__(self):
        _str_board = [' '.join([str(self.__array[y][x]) for x in range(self.width)]) for y in range(self.height)]
        return "\n".join(_str_board)

    def __str__(self):
        return self.__repr__()

    def fill(self, side: List) -> any:
        if len(side) != self.height:
            raise Exception("Wrong height")
        if len(side[0]) != self.width:
            raise Exception("Wrong width")

        self._check_data(side)
        self.__array = deepcopy(side)

    def calculate_square(self) -> int:
        _square = 0
        for row in self.__array:
            for value in row:
                _square += value

        return _square

    def _check_data(self, side: List) -> None:
        for x in side:
            for value in x:
                if value != 0 and value != 1:
                    raise Exception("Wrong data")

    def rotate(self, direction=1) -> None:
        # direction = 1 (left)
        # direction = -1 (Right)

        _new_array = [[0 for i in range(self.height)] for k in range(self.width)]

        if direction == 1:
            self._rotate_left(_new_array, self.__array)
        elif direction == -1:
            self._rotate_right(_new_array, self.__array)
        else:
            raise Exception("Wrong direction")
        self.__array = _new_array

    @staticmethod
    def _rotate_left(new_array: List, old_array: List) -> None:
        i = 0
        for row in old_array:
            j = 0
            for _ in new_array:
                new_array[j][i] = row[len(row) - j - 1]
                j += 1
            i += 1

    @staticmethod
    def _rotate_right(new_array: List, old_array: List) -> None:
        i = 0
        for row in old_array[::-1]:
            j = 0
            for _ in new_array:
                new_array[j][i] = row[j]
                j += 1
            i += 1

class Detail(object):

    def __init__(self, height: int, width: int, depth: int, name: str = '', symbol: str = '1'):
        self.__object = [[[0 for k in range(depth)] for j in range(width)] for i in range(height)]
        self.__side: int = 2
        self.__rotation_steps: List[int] = []
        self.__name = name
        self.__symbol = symbol

    @property
    def name(self):
        return self.__name

    @property
    def value(self) -> List[List[List[int]]]:
        return self.__object

    @property
    def width(self) -> int:
        return len(self.__object[0])

    @property
    def height(self) -> int:
        return len(self.__object)

    @property
    def depth(self) -> int:
        return len(self.__object[0][0])

    @property
    def side_index(self) -> int:
        return self.__side

    @property
    def symbol(self):
        return self.__symbol

    def __repr__(self) -> str:
        return "\nName: %s\nside: %d\nSteps: %s\nDetail\n%s\n" % \
               (self.__name, self.__side, self.__rotation_steps, str(self.get_current_side()))

    def __str__(self) -> str:
        return self.__repr__()

    def rotate(self, direction=1) -> None:
        # direction = 1 - left,
        # direction = -1 - Right
        if direction != 1 and direction != -1:
            raise Exception("Wrong direction")

        last_step = [self.__rotation_steps[-1]] \
            if len(self.__rotation_steps) else []
        if len(last_step) and last_step[0] != direction:
            self.__rotation_steps = self.__rotation_steps[:-1]
        else:
            self.__rotation_steps.append(direction)
            total_length = len(self.__rotation_steps)
            if total_length >= 4:
                if self.__rotation_steps[-1] == \
                        self.__rotation_steps[-2] == \
                        self.__rotation_steps[-3] == \
                        self.__rotation_steps[-4]:
                    self.__rotation_steps = self.__rotation_steps[0: -4]

    def chose_side(self, side_index: int) -> None:
        if side_index < 0 or side_index > 5:
            raise Exception("Only 0..5 sides")

        self.__side = side_index

    def fill(self, array: List) -> None:
        if len(array) != self.height:
            raise Exception("Wrong height")
        if len(array[0]) != self.width:
            raise Exception("Wrong width")
        if len(array[0][0]) != self.depth:
            raise Exception("Wrong depth")

        self._check_data(array)
        self.__object = deepcopy(array)

    def _check_data(self, array: List[List[List[int]]]):
        for h in range(self.height):
            for w in range(self.width):
                for d in range(self.width):
                    if array[h][w][d] != 1 \
                            and array[h][w][d] != 0:
                        raise Exception("Wrong data")

    def get_side(self, side: int) -> Side:
        if side < 0 or side > 5:
            raise Exception("Side has to be 0..5")
        _return: Side = None
        if side == 0:
            _return = Side(self.width, self.depth)
            _return.fill(self.__object[0])
        elif side == 1:
            _return = Side(self.depth, self.width)
            _return.fill(self.__object[-1][::-1])
        elif side == 2:
            _return = Side(self.height, self.width)
            _return.fill([el[-1] for el in self.__object])
        elif side == 3:
            _return = Side(self.height, self.depth)
            _return.fill([[side[-1] for side in el[::-1]] for el in self.__object])
        elif side == 4:
            _return = Side(self.height, self.width)
            _return.fill([el[0][::-1] for el in self.__object])
        elif side == 5:
            _return = Side(self.height, self.depth)
            _return.fill([[side[0] for side in el] for el in self.__object])
        if side != 0 and side != 1:
            self._apply_rotation(_return)
        return _return

    def _apply_rotation(self, side: Side) -> None:
        for direction in self.__rotation_steps:
            side.rotate(direction)

    def get_current_side(self) -> Side:
        return self.get_side(self.__side)

    def get_opposite_side(self) -> Side:
        _opposite_sides = {
            0: 1,
            1: 0,
            2: 4,
            3: 5,
            4: 2,
            5: 3
        }

        return self.get_side(_opposite_sides[self.__side])

    @property
    def sides(self):
        return [2, 3, 4, 5]

def create_detail_from_data(name: str, symbol: str, data: List[List[List[int]]]) -> Detail:
    _return = Detail(len(data), len(data[0]), len(data[0][0]), name, symbol)
    _return.fill(data)
    return _return


from django.urls import include, path

app_name = "api"

# urls from apps
urlpatterns = [
    path(r'menu/', include('backend.menu.urls')),
    path(r'feedback/', include('backend.feedback.urls')),
    path(r'cart/', include('backend.cart.urls')),
    path(r'booking/', include('backend.booking.urls')),

    # path(r'auth/', include('rest_framework.urls', namespace='rest_framework')),
    # path(r'auth/', include('djoser.urls.authtoken')),
]


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

'''
Clone an undirected graph. Each node in the graph contains a label and a list of its neighbors.

OJ's undirected graph serialization:
Nodes are labeled uniquely.

We use # as a separator for each node, and , as a separator for node label and each neighbor of the node.
As an example, consider the serialized graph {0,1,2#1,2#2,2}.

The graph has a total of three nodes, and therefore contains three parts as separated by #.

First node is labeled as 0. Connect node 0 to both nodes 1 and 2.
Second node is labeled as 1. Connect node 1 to node 2.
Third node is labeled as 2. Connect node 2 to node 2 (itself), thus forming a self-cycle.
Visually, the graph looks like the following:

       1
      / \
     /   \
    0 --- 2
         / \
         \_/
'''
import collections

class UndirectedGraphNode:
    def __init__(self, x):
        self.label = x
        self.neighbors = []

class Solution:
    # @param node, a undirected graph node
    # @return a undirected graph node
    def cloneGraph1(self, node):
        # BFS
        if not node:
            return

        _queue = collections.deque([node])
        record = {node.label: UndirectedGraphNode(node.label)}

        while _queue:
            item = _queue.popleft()
            for nei in item.neighbors:
                if nei.label not in record:
                    record[nei.label] = UndirectedGraphNode(nei.label)
                    _queue.append(nei)
                record[item.label].neighbors.append(record[nei.label])
        return record[node.label]

    def cloneGraph(self, node):
        # DFS
        def dfs(node):
            if node.label in self.record:
                return self.record[node.label]

            self.record[node.label] = UndirectedGraphNode(node.label)
            for nei in node.neighbors:
                self.record[node.label].neighbors.append(dfs(nei))
            return self.record[node.label]

        if not node:
            return

        self.record = {}
        return dfs(node)


import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf

img = cv.imread(r"C:\Users\danie\Desktop\Python Practice\pythonProject\Lib\20.jpeg")

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data() #mnist= the stock dataset of handwritten digits

x_train = tf.keras.utils.normalize(x_train, axis = 1)
x_test = tf.keras.utils.normalize(x_test, axis = 1)

model = tf.keras.models.Sequential()

model.add(tf.keras.layers.Flatten(input_shape = (28,28)))
#28x28 input shape is number of pixels in each image

model.add(tf.keras.layers.Dense(units= 128, activation = tf.nn.relu))
model.add(tf.keras.layers.Dense(units= 128, activation = tf.nn.relu))
#connecting neuron layers from one layer to the next, more neurons = more computing power more sophistication more accuracy
#rectified linear unit : The Rectified Linear Unit is the most commonly used activation function in deep learning models.
# The function returns 0 if it receives any negative input, but for any positive value x it returns that value back.

model.add(tf.keras.layers.Dense(units = 10 , activation = tf.nn.softmax))
#creating output neurons. softmax scales percentages of each final neuron to sum to 1(100%) giving the probability score.

model.compile(optimizer = "adam", loss = "sparse_categorical_crossentropy", metrics = ["accuracy"])
model.fit(x_train, y_train, epochs = 3)
accuracy, loss = model.evaluate(x_test, y_test)
print(accuracy)
print(loss)

model.save('digits.model')

########################################################################################
def rescaleFrame(frame, scale = 1): #this function works  for images, video and live videeo
    width = 28
    height = 28

    dimensions = (width, height)

    return cv.resize(frame, dimensions, interpolation = cv.INTER_AREA)
########################################################################################

rescaled_image = rescaleFrame(img)
gray = cv.cvtColor(rescaled_image, cv.COLOR_BGR2GRAY)
blur = cv.GaussianBlur(gray, (5,5), cv.BORDER_DEFAULT)

canny = cv.Canny(blur, 125, 175)
cv.imshow('Blur', canny)

contours, hierarchies = cv.findContours(canny, cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)

img1 = canny
img1 = np.array([img1])
prediction = model.predict(img1)
print(f' the result is probably {np.argmax(prediction)}')  # argmax gives result with highest probability
plt.imshow(img[0], cmap=plt.cm.binary)
plt.show()

cv.waitKey(0)


import json
from types import MethodType

from browser import document, html, ajax, window  # type: ignore

import data
from data import START_POS, SCALE, PIXELS, Ability, abilities

def factory_activate_abl(box):
    ability_id = box.id.lstrip('box-')
    div = box.parent
    spec_id = div.id.lstrip('options-')
    spec = document[f'spec-{spec_id}']
    bar = [el for el in spec.children if el.id == f'abl-{ability_id}'][0]
    def _activate(ev):
        if bar.style.height == '0px':
            bar.style.height = ''
            bar.style.visibility = ''
        else:
            bar.style.height = '0px'
            bar.style.visibility = 'hidden'
    return _activate

def factory_activate_spec(box):
    div = box.parent
    abl_boxes = div.select('.ability-checkbox')
    spec_id = box.id.lstrip('box-')
    spec = document[f'spec-{spec_id}']

    def _activate(ev):
        if spec.style.height == '0px':
            spec.style.height = ''
            spec.style.visibility = ''
            for abl_box in abl_boxes:
                abl_box.disabled = False
        else:
            spec.style.height = '0px'
            spec.style.visibility = 'hidden'
            for abl_box in abl_boxes:
                abl_box.disabled = True
    return _activate

def factory_change_specs(inp):
    spec_name = inp.id.lstrip('count-')
    spec_obj = next(spec for spec in data.specs if spec.name == spec_name)
    spec_div = document[f'spec-{spec_name}']
    opt_div = document[f'options-{spec_name}']

    def _add(ev):
        max_id = min(int(inp.value) - 1, 3)

        specs = spec_div.children
        for spec in specs:
            if int(spec.id[-1]) <= max_id:
                spec.style.height = ''
                spec.style.visibility = ''
            else:
                spec.style.height = '0px'
                spec.style.visibility = 'hidden'
        for i in range(len(specs), max_id + 1):
            spec_div <= Spec(spec_obj, i)

        opts = list(filter(lambda x: isinstance(x, html.DIV), opt_div.children))
        for spec_opts in opts:
            if int(spec_opts.id[-1]) <= max_id:
                spec_opts.style.height = ''
                spec_opts.style.visibility = ''
            else:
                spec_opts.style.height = '0px'
                spec_opts.style.visibility = 'hidden'
        for i in range(len(opts), max_id + 1):
            opt_div <= SpecOpts(spec_obj, i)
        ability_boxes()
        spec_boxes()
        nicknames_inputs()
    return _add

def factory_change_nick(inp):
    spec_id = inp.id.lstrip('nick-')
    lbl = document[f'lbl-{spec_id}']

    def change_nick(ev):
        lbl.textContent = inp.value
    return change_nick

def ability_boxes():
    checkboxes = document.select('.ability-checkbox')
    for box in checkboxes:
        box.unbind('click')
        box.bind('click', factory_activate_abl(box))

def spec_boxes():
    checkboxes = document.select('.spec-checkbox')
    for box in checkboxes:
        box.unbind('click')
        box.bind('click', factory_activate_spec(box))

def spec_inps():
    spec_inputs = document.select('.spec-inp')
    for inp in spec_inputs:
        inp.unbind('input')
        inp.bind('input', factory_change_specs(inp))

def nicknames_inputs():
    inps = document.select('.spec-nick')
    for inp in inps:
        inp.unbind('input')
        inp.bind('input', factory_change_nick(inp))

def px_to_float(px: str) -> float:
    return float(px[:-2])

def px_to_time(px: str) -> str:
    time_stamp = int((px_to_float(px) - START_POS) / SCALE)
    return f'{time_stamp // 60}:{time_stamp % 60:02d}'

def float_to_time(n: float) -> str:
    time_stamp = int((n - START_POS) / SCALE)
    return f'{time_stamp // 60}:{time_stamp % 60:02d}'

class Draggable(html.DIV):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def drag(self, ev):
        start_pos = ev.clientX

        def _stop_drag(ev):
            self.unbind('mousemove')
            self.unbind('mouseup')

        def _drag(ev):
            nonlocal start_pos
            delta = start_pos - ev.clientX
            start_pos = ev.clientX
            self.move(delta)

        self.bind('mouseup', _stop_drag)
        self.bind('mouseleave', _stop_drag)
        self.bind('mousemove', _drag)

    def move(self, delta: int, target=None):
        new_pos = self.offsetLeft - delta
        new_pos_right = new_pos + self.ability.cool_down
        obstacles_right = \
            False if not self._right else px_to_float(self._right.style.left) < new_pos_right
        obstacles_left = \
            False if not self._left else px_to_float(self._left.style.right) > new_pos

        if obstacles_left and delta > 0:
            self._left.move(SCALE, self)
        elif obstacles_right and delta < 0:
            self._right.move(-SCALE, self)
        elif new_pos <= PIXELS and new_pos >= START_POS:
            self._move(new_pos)
            if target:
                target.move(delta)

    def _move(self, new_pos: int):
        self.style.left = f'{new_pos}px'
        self.style.right = f'{new_pos + self.ability.cool_down}px'
        self.tooltip.text = px_to_time(self.style.left)

class AbilityBar(Draggable):
    """Draggable bar with ability name"""
    def __init__(self, ability: Ability, position: float, id: str):
        super().__init__(Class='ability-bar higher20')
        self.ability = ability
        self.pos = position
        self.id = f"abl-bar-{ability.name}{id}"

        self <= self._content()
        self.tooltip = html.SPAN(px_to_time(self.style.left), Class='tooltiptext')
        self <= self.tooltip

    @classmethod
    def new(cls, bar: html.DIV, ability: Ability):
        bar.ability = ability
        bar.drag = MethodType(cls.drag, bar)
        bar.move = MethodType(cls.move, bar)
        bar._move = MethodType(cls._move, bar)
        bar.tooltip = bar.select('.tooltiptext')[0]
        return bar

    def _content(self):
        self.style.left = f'{self.pos}px'
        self.style.right = f'{self.pos + self.ability.cool_down}px'
        self.style.width = self.ability.cool_down

        duration_style = {
            'max-width': self.ability.duration,
            'width': self.ability.duration,
            'background-color': self.ability.color_main,
        }
        duration = html.TD(
            self.ability.name,
            style=duration_style,
            Class='duration'
        )

        cd_length = self.ability.cool_down - self.ability.duration
        cd_style = {
            'width': cd_length,
            'background-color': self.ability.color_cd,
        }
        cd = html.TD(style=cd_style)

        tr = html.TR()
        tr <= duration
        tr <= cd
        tbody = html.TBODY()
        tbody <= tr
        ability_tbl = html.TABLE(Class='ability-tbl')
        ability_tbl <= tbody
        return ability_tbl

class AbilityRow(html.DIV):
    """Several draggable bars in one row"""
    def __init__(self, ability: Ability, id_: int):
        self.ability = ability
        abl_id = f'abl-{ability.name}{id_}'
        super().__init__(id=abl_id, Class='tooltip abl-bar')
        self._bars = [AbilityBar(ability, START_POS, '0')]
        self.bind_bars()

        self.btn_add = html.BUTTON('+', Class='btn-add', id=f'btn-{ability.name}{id_}')
        self.btn_add.unbind('click')
        self.btn_add.bind('click', self.add_bar())
        self <= self.btn_add
        self <= self._bars
        if not ability.check:
            self.style.visibility = 'hidden'
            self.style.height = '0px'

    @classmethod
    def new(cls, bar: html.DIV):
        bar.add_bar = MethodType(cls.add_bar, bar)
        bar.bind_bars = MethodType(cls.bind_bars, bar)

        ability_id = bar.id.lstrip('abl-')
        ability_name = ability_id[:-1]
        ability = next(abl for abl in abilities if abl.name == ability_name)
        bar.ability = ability
        abl_bars = bar.select('.ability-bar')
        bar._bars = [AbilityBar.new(abl_bar, ability) for abl_bar in abl_bars]
        bar.bind_bars()

        bar.btn_add = document[f'btn-{ability_id}']
        bar.btn_add.unbind('click')
        bar.btn_add.bind('click', bar.add_bar())
        return bar

    def add_bar(self):
        def adder(ev):
            bar_count = len(self._bars)
            last_bar = self._bars[-1]
            new_bar = AbilityBar(self.ability, px_to_float(last_bar.style.right), id=bar_count)
            self._bars.append(new_bar)
            self.bind_bars()
            self <= self._bars
        return adder

    def bind_bars(self):
        bar_count = len(self._bars)
        for i, bar in enumerate(self._bars):
            bar._left = self._bars[i - 1] if i > 0 else None
            bar._right = None if i == (bar_count - 1) else self._bars[i + 1]
            bar.unbind('mousedown')
            bar.bind('mousedown', bar.drag)

class Spec(html.DIV):
    """Division with several ability rows belonging to one specialization or player"""
    def __init__(self, spec: data.Spec, id: int):
        div_id = f'spec-{spec.name}{id}'
        super().__init__(id=div_id, Class=div_id.rstrip(str(id)))
        self.spec = spec
        self._id = id
        self._fill()
        self.inp.bind('input', self.change_nick)

    def _fill(self):
        name = self.spec.name
        self.inp = html.INPUT(
            id=f'nick-{name}{self._id}',
            Class='spec-nick',
            type='text',
            name=name,
            value=name,
            size=12
        )
        self <= self.inp
        for ability in self.spec.abilities:
            self <= AbilityRow(ability, self._id)
            self <= html.BR()

    def change_nick(self, ev):
        pass

    @classmethod
    def new(cls, div: html.DIV):
        return div

class SpecOpts(html.DIV):
    """Set of options for a spec - spec on/off, abilities on/off."""
    def __init__(self, spec: data.Spec, id: int):
        div_id = f'options-{spec.name}{id}'
        super().__init__(id=div_id)
        self.spec = spec
        self._id = id
        self._fill()

    def _fill(self):
        name = self.spec.name
        inp = html.INPUT(
            id=f'box-{name}{self._id}',
            Class='spec-checkbox',
            type='checkbox',
            name=f'box-{name}',
            checked=True
        )
        self <= inp
        self <= html.LABEL(html.B(name, id=f'lbl-{name}{self._id}'), For=f'box-{name}{self._id}')
        self <= html.BR()
        for ability in self.spec.abilities:
            name = ability.name
            inp = html.INPUT(
                id=f'box-{name}{self._id}',
                Class='ability-checkbox',
                type='checkbox',
                name=f'box-{name}{self._id}')
            if ability.check:
                inp.checked = True
            lbl = html.LABEL(name, For=f'box-{name}{self._id}')
            self <= inp
            self <= lbl
            self <= html.BR()

def phase_serialize():
    # Serialize phases to list of dicts.
    boss_phases = []
    phases = document['phases']
    for phase_div in phases.children:
        phase_name = phase_div.textContent.strip()
        try:
            phase = list(filter(lambda x: x['name'] == phase_name, boss_phases))[0]
        except IndexError:
            phase = dict(
                name=phase_name,
                color=phase_div.style['background-color'],
                intervals=[]
            )
        interval = [float_to_time(px_to_float(phase_div.style.left) - SCALE),
                    float_to_time(px_to_float(phase_div.style.right) - SCALE)]
        phase['intervals'].append(interval)
        if len(phase['intervals']) == 1:
            boss_phases.append(phase)
    return boss_phases

def casts_serialize():
    # Serialize boss casts to list of dicts.
    boss_casts = []
    boss_actions = document['boss-casts']
    for action_div in boss_actions.children:
        action = dict(
            name=action_div.text.strip(),
            color=action_div.children[0].style['background-color'],
            casts=[]
        )
        for cast in action_div.children:
            tooltip = cast.select('.tooltiptext')[0]
            text = tooltip.textContent.strip()
            time = text.split(' at')[1]
            action['casts'].append(time)
        boss_casts.append(action)
    return boss_casts

def layout_serialize():
    layout = []
    specs = document['specs']
    for spec_div in specs.children:
        spec = {
            'spec': {
                'name': spec_div.id.lstrip('spec-')
            },
            'players': []
        }
        for player in spec_div.children:
            inst = {
                'nick': player.select('input')[0].value,
                'is_visible': player.style.visibility != 'hidden',
                'abilities': []
            }
            for abl in player.select('.abl-bar'):
                ability = {
                    'casts': [],
                    'name': abl.id.lstrip('abl-')[:-1]
                }
                for a in abl.select('.ability-bar'):
                    tooltip = a.select('.tooltiptext')[0]
                    time = tooltip.textContent.strip()
                    ability['casts'].append(time)
                ability['is_visible'] = abl.style.visibility != 'hidden'
                inst['abilities'].append(ability)
            spec['players'].append(inst)
        layout.append(spec)
    # Without this line dragging brakes ^^.
    enable()
    return layout

def serialize(ev):
    seconds = (len(document.select('.timeline-ruler-tick')) /
               len(document.select('.timeline-ruler')))
    res = {
        'encounter': {
            'seconds': seconds,
            'name': document['boss-select'].selectedOptions[0].text,
            # 'phases': phase_serialize(),
            'boss_actions': casts_serialize()
        },
        'layout': layout_serialize()
    }
    send(res)

def send(d: dict):
    req = ajax.Ajax()
    req.open('POST', '/save_layout')
    req.set_header('content-type', 'application/json')
    _data = json.dumps(d)
    print(_data)
    req.send(_data)
    req.bind('complete', show_msg)

def show_msg(req):
    document['modal'].style.display = 'block'
    p = document['modal-content']

    if req.status == 200:
        msg = f'{window.location.hostname}/{req.text}'

    else:
        msg = f'Error {req.status}'

    p.textContent = msg

def modal_and_save():
    btn = html.BUTTON('Save', id='save-btn')
    btn.bind('click', serialize)
    document <= btn

    # modal = document['modal']
    span = document['close-modal']
    span.bind('click', close_modal)
    window.bind('keydown', esc_close_modal)
    # window.bind('click', close_modal)

def esc_close_modal(ev):
    if ev.keyCode == 27:
        document['modal'].style.display = 'none'

def close_modal(ev):
    document['modal'].style.display = 'none'

def lock_unlock():
    lock_btn = document['lock']

    def _lock_bind(ev):
        if lock_btn.textContent == 'Lock':
            lock()
            lock_btn.textContent = 'Unlock'
        else:
            unlock()
            lock_btn.textContent = 'Lock'

    lock_btn.bind('click', _lock_bind)

def lock():
    for el in document.select('.ability-bar'):
        el.unbind('mousedown')
    for el in document.select('input'):
        el.disabled = True

def unlock():
    enable()
    for el in document.select('input'):
        el.disabled = False

def enable():
    [AbilityRow.new(bar) for bar in document.select('.abl-bar')]

if __name__ == '__main__':
    ability_boxes()
    spec_boxes()
    spec_inps()
    nicknames_inputs()
    enable()
    modal_and_save()
    lock_unlock()


import os
import shutil
import uuid
import zipfile

from executeAutoTranslater import autoTranslate
from magicBeing import *

def makeZip(sourceDir, outFilename):
    zipf = zipfile.ZipFile(outFilename, "w", 8)
    pre_len = len(os.path.dirname(sourceDir))
    for parent, dirnames, filenames in os.walk(sourceDir):
        for filename in filenames:
            pathfile = os.path.join(parent, filename)
            arcname = pathfile[pre_len:].strip(os.path.sep)  # 相对路径
            zipf.write(pathfile, arcname)
    zipf.close()

def extractZip(targetFile, outPath):
    zipf = zipfile.ZipFile(targetFile, "r", 8)
    zipf.extractall(outPath)

def dododoSingle(filePath: str):
    with open(filePath,'r',encoding="utf-8") as f:
        context = f.read()

    a = context
    replaceLst = []

    while 'execute' in a:
        tag = uuid.uuid4()
        cmd = a[a.find("execute"):][:a[a.find("execute"):].find('\n')]
        a = a.replace(cmd,tag)
        replaceLst.append((cmd,tag))

    for cmd,tag in replaceLst:
        a = a.replace(tag,autoTranslate(cmd.decode("utf-8")).encode("utf-8"))

def forfile(pathIndex: str):
    pass

def __main__():
    while True:

        filepath, _ = formatipt("请输入MCFUNCTION文件或MCPACK文件地址(可以键入文件夹名)\n：",os.path.exists,"文件不存在或无法打开，请重新输入。",)

        if os.path.isdir:
            pass
        elif os.path.isfile:
            if filepath.endswith(".mcpack"):
                if os.path.exists(r"./temp/"):
                    shutil.rmtree(r"./temp/")
                os.makedirs(r"./temp/")
                extractZip(filepath,r"./temp/")

                shutil.rmtree(r"./temp/")

        prt('完成')

if __name__ == "__main__":
    __main__()