# -*- coding: utf-8 -*-
from __future__ import unicode_literals, print_function

import ast
import inspect
import io
import os
import re
import sys
import textwrap
import token
from typing import List
import unittest
from time import time

import astroid
import six
from asttokens import util, ASTTokens

from . import tools

try:
  from astroid.nodes.utils import Position as AstroidPosition
except Exception:
  AstroidPosition = ()

class TestMarkTokens(unittest.TestCase):
  maxDiff = None

  # We use the same test cases to test both nodes produced by the built-in `ast` module, and by
  # the `astroid` library. The latter derives TestAstroid class from TestMarkTokens. For checks
  # that differ between them, .is_astroid_test allows to distinguish.
  is_astroid_test = False
  module = ast

  def create_mark_checker(self, source, verify=True):
    atok = self.create_asttokens(source)
    checker = tools.MarkChecker(atok)

    # The last token should always be an ENDMARKER
    # None of the nodes should contain that token
    assert atok.tokens[-1].type == token.ENDMARKER
    if atok.text:  # except for empty files
      for node in checker.all_nodes:
        assert node.last_token.type != token.ENDMARKER

    if verify:
      checker.verify_all_nodes(self)
    return checker

  @staticmethod
  def create_asttokens(source):
    return ASTTokens(source, parse=True)

  def print_timing(self):
    # Print the timing of mark_tokens(). This doesn't normally run as a unittest, but if you'd like
    # to see timings, e.g. while optimizing the implementation, run this to see them:
    #
    #     nosetests -m print_timing -s tests.test_mark_tokens tests.test_astroid
    #
    # pylint: disable=no-self-use
    import timeit
    print("mark_tokens", sorted(timeit.repeat(
      setup=textwrap.dedent(
        '''
        import ast, asttokens
        source = "foo(bar(1 + 2), 'hello' + ', ' + 'world')"
        atok = asttokens.ASTTokens(source)
        tree = ast.parse(source)
        '''),
      stmt='atok.mark_tokens(tree)',
      repeat=3,
      number=1000)))

  def test_mark_tokens_simple(self):
    source = tools.read_fixture('astroid', 'module.py')
    m = self.create_mark_checker(source)

    # Line 14 is: [indent 4] MY_DICT[key] = val
    self.assertEqual(m.view_nodes_at(14, 4), {
      "Name:MY_DICT",
      "Subscript:MY_DICT[key]",
      "Assign:MY_DICT[key] = val"
    })

    # Line 35 is: [indent 12] raise XXXError()
    self.assertEqual(m.view_nodes_at(35, 12), {'Raise:raise XXXError()'})
    self.assertEqual(m.view_nodes_at(35, 18), {'Call:XXXError()', 'Name:XXXError'})

    # Line 53 is: [indent 12] autre = [a for (a, b) in MY_DICT if b]
    self.assertEqual(m.view_nodes_at(53, 20), {'ListComp:[a for (a, b) in MY_DICT if b]'})
    self.assertEqual(m.view_nodes_at(53, 21), {'Name:a'})
    if self.is_astroid_test:
      self.assertEqual(m.view_nodes_at(53, 23), {'Comprehension:for (a, b) in MY_DICT if b'})
    else:
      self.assertEqual(m.view_nodes_at(53, 23), {'comprehension:for (a, b) in MY_DICT if b'})

    # Line 59 is: [indent 12] global_access(local, val=autre)
    self.assertEqual(m.view_node_types_at(59, 12), {'Name', 'Call', 'Expr'})
    self.assertEqual(m.view_nodes_at(59, 26), {'Name:local'})
    if self.is_astroid_test:
      self.assertEqual(m.view_nodes_at(59, 33), {'Keyword:val=autre'})
    else:
      self.assertEqual(m.view_nodes_at(59, 33), {'keyword:val=autre'})
    self.assertEqual(m.view_nodes_at(59, 37), {'Name:autre'})

  def test_mark_tokens_multiline(self):
    source = (
"""(    # line1
a,      # line2
b +     # line3
  c +   # line4
  d     # line5
)""")
    m = self.create_mark_checker(source)

    self.assertIn('Name:a', m.view_nodes_at(2, 0))
    self.assertEqual(m.view_nodes_at(3, 0),  {
      'Name:b',
      'BinOp:b +     # line3\n  c',
      'BinOp:b +     # line3\n  c +   # line4\n  d',
    })

    all_text = {m.atok.get_text(node) for node in m.all_nodes}
    self.assertEqual(all_text, {
      source,
      'a', 'b', 'c', 'd',
      # All other expressions preserve newlines and comments but are parenthesized.
      'b +     # line3\n  c',
      'b +     # line3\n  c +   # line4\n  d',
    })
    self.assertIn('Tuple:' + source, m.view_nodes_at(1, 0))

  def verify_fixture_file(self, path):
    source = tools.read_fixture(path)
    m = self.create_mark_checker(source, verify=False)
    tested_nodes = m.verify_all_nodes(self)

    exp_index = (0 if six.PY2 else 1) + (3 if self.is_astroid_test else 0)
    # For ast on Python 3.9, slices are expressions, we handle them and test them.
    if not self.is_astroid_test and issubclass(ast.Slice, ast.expr):
      exp_index += 1
    exp_tested_nodes = self.expect_tested_nodes[path][exp_index]
    self.assertEqual(tested_nodes, exp_tested_nodes)

  # There is not too much need to verify these counts. The main reason is: if we find that some
  # change reduces the count by a lot, it's a red flag that the test is now covering fewer nodes.
  expect_tested_nodes = {
    #                                   AST                  | Astroid
    #                                   Py2   Py3  Py3+slice | Py2   Py3
    'astroid/__init__.py':            ( 4,    4,   4,          4,    4,   ),
    'astroid/absimport.py':           ( 4,    3,   3,          4,    3,   ),
    'astroid/all.py':                 ( 21,   23,  23,         21,   23,  ),
    'astroid/clientmodule_test.py':   ( 75,   67,  67,         69,   69,  ),
    'astroid/descriptor_crash.py':    ( 30,   28,  28,         30,   30,  ),
    'astroid/email.py':               ( 3,    3,   3,          1,    1,   ),
    'astroid/format.py':              ( 64,   61,  61,         62,   62,  ),
    'astroid/module.py':              ( 185,  174, 174,        171,  171, ),
    'astroid/module2.py':             ( 248,  253, 255,        240,  253, ),
    'astroid/noendingnewline.py':     ( 57,   59,  59,         57,   63,  ),
    'astroid/notall.py':              ( 15,   17,  17,         15,   17,  ),
    'astroid/recursion.py':           ( 6,    6,   6,          4,    4,   ),
    'astroid/suppliermodule_test.py': ( 20,   17,  17,         18,   18,  ),
  }

  # This set of methods runs verifications for the variety of syntax constructs used in the
  # fixture test files.
  # pylint: disable=multiple-statements
  def test_fixture1(self): self.verify_fixture_file('astroid/__init__.py')
  def test_fixture2(self): self.verify_fixture_file('astroid/absimport.py')
  def test_fixture3(self): self.verify_fixture_file('astroid/all.py')
  def test_fixture4(self): self.verify_fixture_file('astroid/clientmodule_test.py')
  def test_fixture5(self): self.verify_fixture_file('astroid/descriptor_crash.py')
  def test_fixture6(self): self.verify_fixture_file('astroid/email.py')
  def test_fixture7(self): self.verify_fixture_file('astroid/format.py')
  def test_fixture8(self): self.verify_fixture_file('astroid/module.py')
  def test_fixture9(self): self.verify_fixture_file('astroid/module2.py')
  def test_fixture10(self): self.verify_fixture_file('astroid/noendingnewline.py')
  def test_fixture11(self): self.verify_fixture_file('astroid/notall.py')
  def test_fixture12(self): self.verify_fixture_file('astroid/recursion.py')
  def test_fixture13(self): self.verify_fixture_file('astroid/suppliermodule_test.py')

  def test_deep_recursion(self):
    # This testcase has 1050 strings joined with '+', which causes naive recursions to fail with
    # 'maximum recursion depth exceeded' error. We actually handle it just fine, but we can't use
    # to_source() on it because it chokes on recursion depth. So we test individual nodes.
    source = tools.read_fixture('astroid/joined_strings.py')

    if self.is_astroid_test:
      if getattr(astroid, '__version__', '1') >= '2':
        # Astroid 2 no longer supports this; see
        # https://github.com/PyCQA/astroid/issues/557#issuecomment-396004274
        self.skipTest('astroid-2.0 does not support this')

      # Astroid < 2 does support this with optimize_ast set to True
      astroid.MANAGER.optimize_ast = True
      try:
        m = self.create_mark_checker(source, verify=False)
      finally:
        astroid.MANAGER.optimize_ast = False

      self.assertEqual(len(m.all_nodes), 4)     # This is the result of astroid's optimization
      self.assertEqual(m.view_node_types_at(1, 0), {'Module', 'Assign', 'AssignName'})
      const = next(n for n in m.all_nodes if isinstance(n, astroid.nodes.Const))
      # TODO: Astroid's optimization makes it impossible to get the right start-end information
      # for the combined node. So this test fails. To avoid it, don't set 'optimize_ast=True'. To
      # fix it, astroid would probably need to record the info from the nodes it's combining. Or
      # astroid could avoid the need for the optimization by using an explicit stack like we do.
      #self.assertEqual(m.atok.get_text_range(const), (5, len(source) - 1))
    else:
      m = self.create_mark_checker(source, verify=False)
      self.assertEqual(len(m.all_nodes), 2104)
      self.assertEqual(m.view_node(m.all_nodes[-1]),
                       "Constant:'F1akOFFiRIgPHTZksKBAgMCLGTdGNIAAQgKfDAcgZbj0odOnUA8GBAA7'")
      self.assertEqual(m.view_node(m.all_nodes[-2]),
                       "Constant:'Ii0uLDAxLzI0Mh44U0gxMDI5JkM0JjU3NDY6Kjc5Njo7OUE8Ozw+Oz89QTxA'")
      self.assertEqual(m.view_node(m.all_nodes[1053]),
                       "Constant:'R0lGODlhigJnAef/AAABAAEEAAkCAAMGAg0GBAYJBQoMCBMODQ4QDRITEBkS'")
      self.assertEqual(m.view_node(m.all_nodes[1052]),
                       "BinOp:'R0lGODlhigJnAef/AAABAAEEAAkCAAMGAg0GBAYJBQoMCBMODQ4QDRITEBkS'\r\n" +
                       "     +'CxsSEhkWDhYYFQ0aJhkaGBweGyccGh8hHiIkIiMmGTEiHhQoPSYoJSkqKDcp'")

      binop = next(n for n in m.all_nodes if n.__class__.__name__ == 'BinOp')
      self.assertTrue(m.atok.get_text(binop).startswith("'R0l"))
      self.assertTrue(m.atok.get_text(binop).endswith("AA7'"))

    assign = next(n for n in m.all_nodes if n.__class__.__name__ == 'Assign')
    self.assertTrue(m.atok.get_text(assign).startswith("x = ("))
    self.assertTrue(m.atok.get_text(assign).endswith(")"))

  def test_slices(self):
    # Make sure we don't fail on parsing slices of the form `foo[4:]`.
    source = "(foo.Area_Code, str(foo.Phone)[:3], str(foo.Phone)[3:], foo[:], bar[::2, :], bar2[:, ::2], [a[:]][::-1])"
    m = self.create_mark_checker(source)
    self.assertIn("Tuple:" + source, m.view_nodes_at(1, 0))
    self.assertEqual(m.view_nodes_at(1, 1),
                     { "Attribute:foo.Area_Code", "Name:foo" })
    self.assertEqual(m.view_nodes_at(1, 16),
                     { "Subscript:str(foo.Phone)[:3]", "Call:str(foo.Phone)", "Name:str"})
    self.assertEqual(m.view_nodes_at(1, 36),
                     { "Subscript:str(foo.Phone)[3:]", "Call:str(foo.Phone)", "Name:str"})
    # Slice and ExtSlice nodes are wrong, and in particular placed with parents. They are not very
    # important, so we skip them here.
    self.assertEqual({n for n in m.view_nodes_at(1, 56) if 'Slice:' not in n},
                     { "Subscript:foo[:]", "Name:foo" })
    self.assertEqual({n for n in m.view_nodes_at(1, 64) if 'Slice:' not in n and 'Tuple:' not in n},
                     { "Subscript:bar[::2, :]", "Name:bar" })

  def test_adjacent_strings(self):
    source = """
foo = 'x y z' \\
'''a b c''' "u v w"
bar = ('x y z'   # comment2
       'a b c'   # comment3
       'u v w'
      )
"""
    m = self.create_mark_checker(source)
    node_name = 'Const' if self.is_astroid_test else 'Constant'
    self.assertEqual(m.view_nodes_at(2, 6), {
      node_name + ":'x y z' \\\n'''a b c''' \"u v w\""
    })
    self.assertEqual(m.view_nodes_at(4, 7), {
      node_name + ":'x y z'   # comment2\n       'a b c'   # comment3\n       'u v w'"
    })

  def test_print_function(self):
    # This testcase imports print as function (using from __future__). Check that we can parse it.
    # verify_all_nodes doesn't work on Python 2 because the print() call parsed in isolation
    # is viewed as a Print node since it doesn't see the future import
    source = tools.read_fixture('astroid/nonregr.py')
    m = self.create_mark_checker(source, verify=six.PY3)

    # Line 16 is: [indent 8] print(v.get('yo'))
    self.assertEqual(m.view_nodes_at(16, 8),
                     { "Call:print(v.get('yo'))", "Expr:print(v.get('yo'))", "Name:print" })
    self.assertEqual(m.view_nodes_at(16, 14), {"Call:v.get('yo')", "Attribute:v.get", "Name:v"})

  # To make sure we can handle various hard cases, we include tests for issues reported for a
  # similar project here: https://bitbucket.org/plas/thonny

  if not six.PY2:
    def test_nonascii(self):
      # Test of https://bitbucket.org/plas/thonny/issues/162/weird-range-marker-crash-with-non-ascii
      # Only on PY3 because Py2 doesn't support unicode identifiers.
      for source in (
        "℘·2=1+a℘·b+a℘·2b",  # example from https://github.com/python/cpython/issues/68382
        "sünnikuupäev=str((18+int(isikukood[0:1])-1)//2)+isikukood[1:3]",
        "sünnikuupaev=str((18+int(isikukood[0:1])-1)//2)+isikukood[1:3]"):
        m = self.create_mark_checker(source)
        self.assertEqual(m.view_nodes_at(1, 0), {
          "Module:%s" % source,
          "Assign:%s" % source,
          "%s:%s" % ("AssignName" if self.is_astroid_test else "Name", source.split("=")[0]),
        })

  if sys.version_info[0:2] >= (3, 6):
    # f-strings are only supported in Python36. We don't handle them fully, for a couple of
    # reasons: parsed AST nodes are not annotated with correct line and col_offset (see
    # https://bugs.python.org/issue29051), and there are confusingly two levels of tokenizing.
    # Meanwhile, we only parse to the level of JoinedStr, and verify that.
    def test_fstrings(self):
      for source in (
        '(f"He said his name is {name!r}.",)',
        "f'{function(kwarg=24)}'",
        'a = f"""result: {value:{width}.{precision}}"""',
        """[f"abc {a['x']} def"]""",
        "def t():\n  return f'{function(kwarg=24)}'"):
        self.create_mark_checker(source)

    def test_adjacent_joined_strings(self):
        source = """
foo = f'x y z' \\
f'''a b c''' f"u v w"
bar = ('x y z'   # comment2
       'a b c'   # comment3
       f'u v w'
      )
"""
        m = self.create_mark_checker(source)
        self.assertEqual(m.view_nodes_at(2, 6), {
            "JoinedStr:f'x y z' \\\nf'''a b c''' f\"u v w\""
        })
        self.assertEqual(m.view_nodes_at(4, 7), {
            "JoinedStr:'x y z'   # comment2\n       'a b c'   # comment3\n       f'u v w'"
        })

  def test_splat(self):
    # See https://bitbucket.org/plas/thonny/issues/151/debugger-crashes-when-encountering-a-splat
    source = textwrap.dedent("""
      arr = [1,2,3,4,5]
      def print_all(a, b, c, d, e):
          print(a, b, c, d ,e)
      print_all(*arr)
    """)
    m = self.create_mark_checker(source)
    self.assertEqual(m.view_nodes_at(5, 0),
        { "Expr:print_all(*arr)", "Call:print_all(*arr)", "Name:print_all" })
    if not six.PY2 or self.is_astroid_test:
      self.assertEqual(m.view_nodes_at(5, 10), { "Starred:*arr" })
    self.assertEqual(m.view_nodes_at(5, 11), { "Name:arr" })

  def test_paren_attr(self):
    # See https://bitbucket.org/plas/thonny/issues/123/attribute-access-on-parenthesized
    source = "(x).foo()"
    m = self.create_mark_checker(source)
    self.assertEqual(m.view_nodes_at(1, 1), {"Name:x"})
    self.assertEqual(m.view_nodes_at(1, 0),
                     {"Module:(x).foo()", "Expr:(x).foo()", "Call:(x).foo()", "Attribute:(x).foo"})

  def test_conditional_expr(self):
    # See https://bitbucket.org/plas/thonny/issues/108/ast-marker-crashes-with-conditional
    source = "a = True if True else False\nprint(a)"
    m = self.create_mark_checker(source)
    name_a = 'AssignName:a' if self.is_astroid_test else 'Name:a'
    const_true = ('Const:True' if self.is_astroid_test else
                  'Name:True' if six.PY2 else
                  'Constant:True')
    self.assertEqual(m.view_nodes_at(1, 0),
                     {name_a, "Assign:a = True if True else False", "Module:" + source})
    self.assertEqual(m.view_nodes_at(1, 4),
                     {const_true, 'IfExp:True if True else False'})
    if six.PY2:
      self.assertEqual(m.view_nodes_at(2, 0), {"Print:print(a)"})
    else:
      self.assertEqual(m.view_nodes_at(2, 0), {"Name:print", "Call:print(a)", "Expr:print(a)"})

  def test_calling_lambdas(self):
    # See https://bitbucket.org/plas/thonny/issues/96/calling-lambdas-crash-the-debugger
    source = "y = (lambda x: x + 1)(2)"
    m = self.create_mark_checker(source)
    self.assertEqual(m.view_nodes_at(1, 4), {'Call:(lambda x: x + 1)(2)'})
    self.assertEqual(m.view_nodes_at(1, 15), {'BinOp:x + 1', 'Name:x'})
    if self.is_astroid_test:
      self.assertEqual(m.view_nodes_at(1, 0), {'AssignName:y', 'Assign:'+source, 'Module:'+source})
    else:
      self.assertEqual(m.view_nodes_at(1, 0), {'Name:y', 'Assign:' + source, 'Module:' + source})

  def test_comprehensions(self):
    # See https://bitbucket.org/plas/thonny/issues/8/range-marker-doesnt-work-correctly-with
    for source in (
      "[(key, val) for key, val in ast.iter_fields(node)]",
      "((key, val) for key, val in ast.iter_fields(node))",
      "{(key, val) for key, val in ast.iter_fields(node)}",
      "{key: val for key, val in ast.iter_fields(node)}",
      "[[c for c in key] for key, val in ast.iter_fields(node)]"):
      self.create_mark_checker(source)

  def test_trailing_commas(self):
    # Make sure we handle trailing commas on comma-separated structures (e.g. tuples, sets, etc.)
    for source in (
      "(a,b,)",
      "[c,d,]",
      "{e,f,}",
      "{h:1,i:2,}"):
      self.create_mark_checker(source)

  def test_tuples(self):
    def get_tuples(code):
      m = self.create_mark_checker(code)
      return [m.atok.get_text(n) for n in m.all_nodes if n.__class__.__name__ == "Tuple"]

    self.assertEqual(get_tuples("a,"), ["a,"])
    self.assertEqual(get_tuples("(a,)"), ["(a,)"])
    self.assertEqual(get_tuples("(a),"), ["(a),"])
    self.assertEqual(get_tuples("((a),)"), ["((a),)"])
    self.assertEqual(get_tuples("(a,),"), ["(a,),", "(a,)"])
    self.assertEqual(get_tuples("((a,),)"), ["((a,),)", "(a,)"])
    self.assertEqual(get_tuples("()"), ["()"])
    self.assertEqual(get_tuples("(),"), ["(),", "()"])
    self.assertEqual(get_tuples("((),)"), ["((),)", "()"])
    self.assertEqual(get_tuples("((),(a,))"), ["((),(a,))", "()", "(a,)"])
    self.assertEqual(get_tuples("((),(a,),)"), ["((),(a,),)", "()", "(a,)"])
    self.assertEqual(get_tuples("((),(a,),),"), ["((),(a,),),", "((),(a,),)", "()", "(a,)"])
    self.assertEqual(get_tuples('((foo, bar),)'), ['((foo, bar),)', '(foo, bar)'])
    self.assertEqual(get_tuples('(foo, bar),'), ['(foo, bar),', '(foo, bar)'])
    self.assertEqual(get_tuples('def foo(a=()): ((x, (y,)),) = ((), (a,),),'), [
      '()', '((x, (y,)),)', '(x, (y,))', '(y,)', '((), (a,),),', '((), (a,),)', '()', '(a,)'])
    self.assertEqual(get_tuples('def foo(a=()): ((x, (y,)),) = [(), [a,],],'), [
      '()', '((x, (y,)),)', '(x, (y,))', '(y,)', '[(), [a,],],', '()'])

  def test_dict_order(self):
    # Make sure we iterate over dict keys/values in source order.
    # See https://github.com/gristlabs/asttokens/issues/31
    source = 'f({1: (2), 3: 4}, object())'
    self.create_mark_checker(source)

  def test_del_dict(self):
    # See https://bitbucket.org/plas/thonny/issues/24/try-del-from-dictionary-in-debugging-mode
    source = "x = {4:5}\ndel x[4]"
    m = self.create_mark_checker(source)
    self.assertEqual(m.view_nodes_at(1, 4), {'Dict:{4:5}'})
    if self.is_astroid_test:
      self.assertEqual(m.view_nodes_at(1, 5), {'Const:4'})
    else:
      self.assertEqual(m.view_nodes_at(1, 5), {'Constant:4'})
    self.assertEqual(m.view_nodes_at(2, 0), {'Delete:del x[4]'})
    self.assertEqual(m.view_nodes_at(2, 4), {'Name:x', 'Subscript:x[4]'})

  if not six.PY2:
    def test_bad_tokenless_types(self):
      # Cases where _get_text_positions_tokenless is incorrect in 3.8.
      source = textwrap.dedent("""
        def foo(*, name: str):  # keyword-only argument with type annotation
          pass

        f(*(x))  # ast.Starred with parentheses
      """)
      self.create_mark_checker(source)

    def test_return_annotation(self):
      # See https://bitbucket.org/plas/thonny/issues/9/range-marker-crashes-on-function-return
      source = textwrap.dedent("""
        def liida_arvud(x: int, y: int) -> int:
          return x + y
      """)
      m = self.create_mark_checker(source)
      self.assertEqual(m.view_nodes_at(2, 0),
        {'FunctionDef:def liida_arvud(x: int, y: int) -> int:\n  return x + y'})
      if self.is_astroid_test:
        self.assertEqual(m.view_nodes_at(2, 16),   {'Arguments:x: int, y: int', 'AssignName:x'})
      else:
        self.assertEqual(m.view_nodes_at(2, 16),   {'arguments:x: int, y: int', 'arg:x: int'})
      self.assertEqual(m.view_nodes_at(2, 19),   {'Name:int'})
      self.assertEqual(m.view_nodes_at(2, 35),   {'Name:int'})
      self.assertEqual(m.view_nodes_at(3, 2),    {'Return:return x + y'})

  def test_keyword_arg_only(self):
    # See https://bitbucket.org/plas/thonny/issues/52/range-marker-fails-with-ridastrip-split
    source = "f(x=1)\ng(a=(x),b=[y])"
    m = self.create_mark_checker(source)
    self.assertEqual(m.view_nodes_at(1, 0),
                     {'Name:f', 'Call:f(x=1)', 'Expr:f(x=1)', 'Module:' + source})
    self.assertEqual(m.view_nodes_at(2, 0),
                     {'Name:g', 'Call:g(a=(x),b=[y])', 'Expr:g(a=(x),b=[y])'})
    self.assertEqual(m.view_nodes_at(2, 11), {'Name:y'})
    if self.is_astroid_test:
      self.assertEqual(m.view_nodes_at(1, 2), {'Keyword:x=1'})
      self.assertEqual(m.view_nodes_at(1, 4), {'Const:1'})
      self.assertEqual(m.view_nodes_at(2, 2), {'Keyword:a=(x)'})
      self.assertEqual(m.view_nodes_at(2, 8), {'Keyword:b=[y]'})
    else:
      self.assertEqual(m.view_nodes_at(1, 2), {'keyword:x=1'})
      self.assertEqual(m.view_nodes_at(1, 4), {'Constant:1'})
      self.assertEqual(m.view_nodes_at(2, 2), {'keyword:a=(x)'})
      self.assertEqual(m.view_nodes_at(2, 8), {'keyword:b=[y]'})

  def test_decorators(self):
    # See https://bitbucket.org/plas/thonny/issues/49/range-marker-fails-with-decorators
    source = textwrap.dedent("""
      @deco1
      def f():
        pass
      @deco2(a=1)
      def g(x):
        pass

      @deco3()
      def g(x):
        pass

      @deco4
      class C:
        pass
    """)
    m = self.create_mark_checker(source)
    # The `arguments` node has bogus positions here (and whenever there are no arguments). We
    # don't let that break our test because it's unclear if it matters to anything anyway.
    self.assertIn('FunctionDef:@deco1\ndef f():\n  pass', m.view_nodes_at(2, 0))
    self.assertEqual(m.view_nodes_at(2, 1), {'Name:deco1'})
    if self.is_astroid_test:
      self.assertEqual(m.view_nodes_at(5, 0), {
        'FunctionDef:@deco2(a=1)\ndef g(x):\n  pass',
        'Decorators:@deco2(a=1)'
      })
    else:
      self.assertEqual(m.view_nodes_at(5, 0), {'FunctionDef:@deco2(a=1)\ndef g(x):\n  pass'})
    self.assertEqual(m.view_nodes_at(5, 1), {'Name:deco2', 'Call:deco2(a=1)'})

    self.assertEqual(m.view_nodes_at(9, 1), {'Name:deco3', 'Call:deco3()'})

  def test_with(self):
    source = "with foo: pass"
    m = self.create_mark_checker(source)
    self.assertEqual(m.view_node_types_at(1, 0), {"Module", "With"})
    self.assertEqual(m.view_nodes_at(1, 0), {
      "Module:with foo: pass",
      "With:with foo: pass",
    })

    source = textwrap.dedent(
      '''
      def f(x):
        with A() as a:
          log(a)
          with B() as b, C() as c: log(b, c)
        log(x)
      ''')
    # verification fails on Python2 which turns `with X, Y` turns into `with X: with Y`.
    m = self.create_mark_checker(source, verify=six.PY3)
    self.assertEqual(m.view_nodes_at(5, 4), {
      'With:with B() as b, C() as c: log(b, c)'
    })
    self.assertEqual(m.view_nodes_at(3, 2), {
      'With:  with A() as a:\n    log(a)\n    with B() as b, C() as c: log(b, c)'
    })
    with_nodes = [n for n in m.all_nodes if n.__class__.__name__ == 'With']
    self.assertEqual({m.view_node(n) for n in with_nodes}, {
      'With:with B() as b, C() as c: log(b, c)',
      'With:  with A() as a:\n    log(a)\n    with B() as b, C() as c: log(b, c)',
    })

  def test_one_line_if_elif(self):
    source = """
if 1: a
elif 2: b
    """
    self.create_mark_checker(source)

  def test_statements_with_semicolons(self):
    source = """
a; b; c(
  17
); d # comment1; comment2
if 2: a; b; # comment3
if a:
  if b: c; d  # comment4
    """
    m = self.create_mark_checker(source)
    self.assertEqual(
      [m.atok.get_text(n) for n in m.all_nodes if util.is_stmt(n)],
      ['a', 'b', 'c(\n  17\n)', 'd', 'if 2: a; b', 'a', 'b',
       'if a:\n  if b: c; d', 'if b: c; d', 'c', 'd'])

  def test_complex_numbers(self):
    source = """
1
-1
j  # not a complex number, just a name
1j
-1j
1+2j
3-4j
1j-1j-1j-1j
    """
    self.create_mark_checker(source)

  def test_parens_around_func(self):
    source = textwrap.dedent(
      '''
      foo()
      (foo)()
      (lambda: 0)()
      (lambda: ())()
      (foo)((1))
      (lambda: ())((2))
      x = (obj.attribute.get_callback() or default_callback)()
      ''')
    m = self.create_mark_checker(source)
    self.assertEqual(m.view_nodes_at(2, 0), {"Name:foo", "Expr:foo()", "Call:foo()"})
    self.assertEqual(m.view_nodes_at(3, 1), {"Name:foo"})
    self.assertEqual(m.view_nodes_at(3, 0), {"Expr:(foo)()", "Call:(foo)()"})
    self.assertEqual(m.view_nodes_at(4, 0), {"Expr:(lambda: 0)()", "Call:(lambda: 0)()"})
    self.assertEqual(m.view_nodes_at(5, 0), {"Expr:(lambda: ())()", "Call:(lambda: ())()"})
    self.assertEqual(m.view_nodes_at(6, 0), {"Expr:(foo)((1))", "Call:(foo)((1))"})
    self.assertEqual(m.view_nodes_at(7, 0), {"Expr:(lambda: ())((2))", "Call:(lambda: ())((2))"})
    self.assertEqual(m.view_nodes_at(8, 4),
                     {"Call:(obj.attribute.get_callback() or default_callback)()"})
    self.assertIn('BoolOp:obj.attribute.get_callback() or default_callback', m.view_nodes_at(8, 5))

  def test_complex_slice_and_parens(self):
    source = 'f((x)[:, 0])'
    self.create_mark_checker(source)

  if six.PY3:
    def test_sys_modules(self):
      """
      Verify all nodes on source files obtained from sys.modules.
      This can take a long time as there are many modules,
      so it only tests all modules if the environment variable
      ASTTOKENS_SLOW_TESTS has been set.
      """
      from .test_astroid import AstroidTreeException

      modules = list(sys.modules.values())
      if not os.environ.get('ASTTOKENS_SLOW_TESTS'):
        modules = modules[:20]

      start = time()
      for module in modules:
        # Don't let this test (which runs twice) take longer than 13 minutes
        # to avoid the travis build time limit of 30 minutes
        if time() - start > 13 * 60:
          break

        try:
          filename = inspect.getsourcefile(module)
        except Exception:  # some modules raise weird errors
          continue

        if not filename:
          continue

        filename = os.path.abspath(filename)
        print(filename)
        try:
          with io.open(filename) as f:
            source = f.read()
        except OSError:
          continue

        if self.is_astroid_test and (
            # Astroid fails with a syntax error if a type comment is on its own line
            re.search(r'^\s*# type: ', source, re.MULTILINE)
        ):
          print('Skipping', filename)
          continue

        try:
          self.create_mark_checker(source)
        except AstroidTreeException:
          # Astroid sometimes fails with errors like:
          #     AttributeError: 'TreeRebuilder' object has no attribute 'visit_typealias'
          # See https://github.com/gristlabs/asttokens/actions/runs/6015907789/job/16318767911?pr=110
          # Should be fixed in the next astroid release:
          #     https://github.com/pylint-dev/pylint/issues/8782#issuecomment-1669967220
          # Note that this exception is raised before asttokens is even involved,
          # it's purely an astroid bug that we can safely ignore.
          continue

  if six.PY3:
    def test_dict_merge(self):
      self.create_mark_checker("{**{}}")

    def test_async_def(self):
      self.create_mark_checker("""
async def foo():
  pass

@decorator
async def foo():
  pass
""")

    def test_async_for_and_with(self):
      # Can't verify all nodes because in < 3.7
      # async for/with outside of a function is invalid syntax
      m = self.create_mark_checker("""
async def foo():
  async for x in y: pass
  async with x as y: pass
  """, verify=False)
      assert m.view_nodes_at(3, 2) == {"AsyncFor:async for x in y: pass"}
      assert m.view_nodes_at(4, 2) == {"AsyncWith:async with x as y: pass"}

    def test_await(self):
      # Can't verify all nodes because in astroid
      # await outside of an async function is invalid syntax
      m = self.create_mark_checker("""
async def foo():
  await bar
  """, verify=False)
      assert m.view_nodes_at(3, 2) == {"Await:await bar", "Expr:await bar"}

  if sys.version_info >= (3, 8):
    def test_assignment_expressions(self):
      # From https://www.python.org/dev/peps/pep-0572/
      self.create_mark_checker("""
# Handle a matched regex
if (match := pattern.search(data)) is not None:
    # Do something with match
    pass

# A loop that can't be trivially rewritten using 2-arg iter()
while chunk := file.read(8192):
   process(chunk)

# Reuse a value that's expensive to compute
[y := f(x), y**2, y**3]

# Share a subexpression between a comprehension filter clause and its output
filtered_data = [y for x in data if (y := f(x)) is not None]

y0 = (y1 := f(x))  # Valid, though discouraged

foo(x=(y := f(x)))  # Valid, though probably confusing

def foo(answer=(p := 42)):  # Valid, though not great style
    ...

def foo(answer: (p := 42) = 5):  # Valid, but probably never useful
    ...

lambda: (x := 1) # Valid, but unlikely to be useful

(x := lambda: 1) # Valid

lambda line: (m := re.match(pattern, line)) and m.group(1) # Valid

if any((comment := line).startswith('#') for line in lines):
    print("First comment:", comment)

if all((nonblank := line).strip() == '' for line in lines):
    print("All lines are blank")

partial_sums = [total := total + v for v in values]
""")

  if sys.version_info >= (3, 10):
    def test_match_case(self):
      m = self.create_mark_checker("""
if 0:
  match x:
    case ast.BinOp():
      if z:
        pass
    case cls(a,b) if y:
      pass
    case _:
      match y:
        case 1:
          pass
        case _:
          pass
""")
      self.assertEqual(m.view_nodes_at(10, 6), {
        'Match:'
        '      match y:\n'
        '        case 1:\n'
        '          pass\n'
        '        case _:\n'
        '          pass',
      })

  def parse_snippet(self, text, node):
    """
    Returns the parsed AST tree for the given text, handling issues with indentation and newlines
    when text is really an extracted part of larger code.
    """
    # If text is indented, it's a statement, and we need to put in a scope for indents to be valid
    # (using textwrap.dedent is insufficient because some lines may not indented, e.g. comments or
    # multiline strings). If text is an expression but has newlines, we parenthesize it to make it
    # parsable.
    # For expressions and statements, we add a dummy statement '_' before it because if it's just a
    # string contained in an astroid.Const or astroid.Expr it will end up in the doc attribute and be
    # a pain to extract for comparison
    # For starred expressions, e.g. `*args`, we wrap it in a function call to make it parsable.
    # For slices, e.g. `x:`, we wrap it in an indexing expression to make it parsable.
    indented = re.match(r'^[ \t]+\S', text)
    if indented:
      return self.module.parse('def dummy():\n' + text).body[0].body[0]
    if util.is_starred(node):
      return self.module.parse('f(' + text + ')').body[0].value.args[0]
    if util.is_slice(node):
      return self.module.parse('a[' + text + ']').body[0].value.slice
    if util.is_expr(node):
      return self.module.parse('_\n(' + text + ')').body[1].value
    if util.is_module(node):
      return self.module.parse(text)
    return self.module.parse('_\n' + text).body[1]

  def test_assert_nodes_equal(self):
    """
    Checks that assert_nodes_equal actually fails when given different nodes
    """

    def check(s1, s2):
      n1 = self.module.parse(s1)
      n2 = self.module.parse(s2)
      with self.assertRaises(AssertionError):
        self.assert_nodes_equal(n1, n2)

    check('a', 'b')
    check('a*b', 'a+b')
    check('a*b', 'b*a')
    check('(a and b) or c', 'a and (b or c)')
    check('a = 1', 'a = 2')
    check('a = 1', 'a += 1')
    check('a *= 1', 'a += 1')
    check('[a for a in []]', '[a for a in ()]')
    check("for x in y: pass", "for x in y: fail")
    check("1", "1.0")
    check("foo(a, b, *d, c=2, **e)",
          "foo(a, b, *d, c=2.0, **e)")
    check("foo(a, b, *d, c=2, **e)",
          "foo(a, b, *d, c=2)")
    check('def foo():\n    """xxx"""\n    None',
          'def foo():\n    """xx"""\n    None')

  nodes_classes = ast.AST
  context_classes = [ast.expr_context] # type: List[util.AstNode]
  iter_fields = staticmethod(ast.iter_fields)

  def assert_nodes_equal(self, t1, t2):
    # Ignore the context of each node which can change when parsing
    # substrings of source code. We just want equal structure and contents.
    for context_classes_group in self.context_classes:
      if isinstance(t1, context_classes_group):
        self.assertIsInstance(t2, context_classes_group)
        break
    else:
      self.assertEqual(type(t1), type(t2))

    if isinstance(t1, AstroidPosition):
      # Ignore the lineno/col_offset etc. from astroid
      return

    if isinstance(t1, (list, tuple)):
      self.assertEqual(len(t1), len(t2))
      for vc1, vc2 in zip(t1, t2):
        self.assert_nodes_equal(vc1, vc2)
    elif isinstance(t1, self.nodes_classes):
      self.assert_nodes_equal(
        list(self.iter_fields(t1)),
        list(self.iter_fields(t2)),
      )
    else:
      # Weird bug in astroid that collapses spaces in docstrings sometimes maybe
      if self.is_astroid_test and isinstance(t1, six.string_types):
        t1 = re.sub(r'^ +$', '', t1, flags=re.MULTILINE)
        t2 = re.sub(r'^ +$', '', t2, flags=re.MULTILINE)

      self.assertEqual(t1, t2)


#!/usr/bin/python
# Author:       F Massimo, implemented as function by K Cassou
# Date:         2018-02-10, 2020-02-12, 2020-06-21
# Purpose:      set of function for LPA simulation with SMILEI
# Source:       Python 3 (python2)
#####################################################################

### loading system module
from __future__ import (division, print_function, absolute_import, unicode_literals)
import os, sys
import warnings

# for execution on local SERA0X computers unit
host = os.popen('hostname -f').read()[:-1]
host_list = ['sera01.lal.in2p3.fr','sera03.lal.in2p3.fr','sera03.lal.in2p3.fr','sera04.lal.in2p3.fr']
if host in host_list:
    sys.path.append('/silver/PALLAS/simulations/smilei/postprocess/')

# set the current path for happi download path
module_filename = os.path.basename(__file__)
module_filepath = __file__
module_dirpath = module_filepath[:-len(module_filename)]
os.chdir(module_dirpath)

#print(module_dirpath)

# happi import KC : to be modified to enough to install happi, wrong path
try :
    import happi
except ImportError :
    try:
        os.system('svn checkout https://github.com/SmileiPIC/Smilei.git/trunk/happi/')
        print('last version of happi has been downloaded with in:', os.path.abspath(module_dirpath))
    except IOError:
        print('download the happi module from github: https://github.com/SmileiPIC/Smilei.git and copy it in `lpa2` module directory')
    pass

# modules import
import statsmodels.api              # necessary to be located before statsmodels.stats import.
import statsmodels.stats as stats
import numpy as np
import scipy.constants as sc
from scipy.optimize import curve_fit
from scipy.signal import find_peaks, peak_widths
import matplotlib.pyplot as plt

#################### Inputs ##############################################

species_name   = "electronfromion"
lambda0        = 0.8e-6 # Wavelength of the laser

default_directory = os.path.abspath(os.getcwd())
home_directory     = os.path.expanduser("~")

# used to apply a filter in energy (m_e c^2 units, or Lorentz factor)
E_min          = 0.
E_max          = 500

chunk_size     = 100000000  #Chunck of particles treated simultaneously

horiz_axis_conversion_factor = 0.512 # to convert from Smilei units to MeV
hist_conversion_factor       = 1.    # if equal to 1, the charge is in pC

################# Fundamental Physical constants ########################
eps0   = sc.epsilon_0;  # Electric permittivity of vacuum, F/m
mu0    = sc.mu_0;       # Magnetic permittivity of vacuum, kg.m.A-2s-2
e      = sc.e;          # Elementary charge, C
EeV    = sc.eV;         # 1 eV = 1.6e-19 Joules
c      = sc.c;          # Lightspeed, m/s
me     = sc.m_e;        # Electron mass, kg
mp     = sc.m_p;        # Proton mass, kg
h      = sc.h;          # Planck's constant, J.s
hbar   = sc.hbar

########## Physical constants ##########################################
omega0 = 2*np.pi*c/lambda0          #
onel   = lambda0/ (2*np.pi)         #
ncrit  = eps0*me*omega0**2/e**2;    # critical density (m^-3, not cm^-3)

######### useful functions ####################################

def lin_interp(x, y, i, half):
    return x[i] + (x[i+1] - x[i]) * ((half - y[i]) / (y[i+1] - y[i]))

def half_max_x(x, y):
    half = max(y)/2.0
    signs = np.sign(np.add(y, -half))
    zero_crossings = (signs[0:-2] != signs[1:-1])
    zero_crossings_i = np.where(zero_crossings)[0]
    return [lin_interp(x, y, zero_crossings_i[0], half),lin_interp(x, y, zero_crossings_i[1], half)]

def fwhm(x,y):
    hmx = half_max_x(x,y)
    return hmx[1]-hmx[0]

def gaussian(x, amp, xcenter, width):
    return amp * np.exp(-(x-xcenter)**2 / width**2)

def weighted_median(data, weights):
    """
    Args:
      data (list or numpy.array): data
      weights (list or numpy.array): weights
    """
    data, weights = np.array(data).squeeze(), np.array(weights).squeeze()
    s_data, s_weights = map(np.array, zip(*sorted(zip(data, weights))))
    midpoint = 0.5 * sum(s_weights)
    if any(weights > midpoint):
        w_median = (data[weights == np.max(weights)])[0]
    else:
        cs_weights = np.cumsum(s_weights)
        idx = np.where(cs_weights <= midpoint)[0][-1]
        if cs_weights[idx] == midpoint:
            w_median = np.mean(s_data[idx:idx+2])
        else:
            w_median = s_data[idx+1]
    return w_median

def mad(data, axis=None):
    """
    Compute *Median Absolute Deviation* of an array along given axis.
    """
    # Median along given axis, but *keeping* the reduced axis so that
    # result can still broadcast against a.
    med = np.median(data, axis=axis, keepdims=True)
    mad = np.median(np.absolute(data - med), axis=axis)  # MAD along given axis

    return mad

def weighted_std(data, weights):
    """
    Compute *weighted_standard Deviation* of an array along given axis.
    """
    d = stats.weightstats.DescrStatsW(data,weights)

    return d.std

def weighted_mean(data, weights):
    """
    Compute *weighted_mean* of an array along given axis.
    """
    d = stats.weightstats.DescrStatsW(data,weights)

    return d.mean

def weighted_mad(data, weights):
    """
    Compute *weighted_Median Absolute Deviation* of an array along given axis.
    """
    wmed = weighted_median(data,weights)
    wmad = weighted_median(np.absolute(data - wmed),weights)  # MAD along given axis

    return wmad

########## load data with happi ##############################

def loadData(directory=default_directory):
    """loading data in the simulation directory and return an object pointing to the various
    files, see smilei website"""
    S = happi.Open(directory, show = False,verbose = False)
    return S

######### extract laser var ##############################

def getMaxinMovingWindow(S,var="Env_E_abs"):
    """ return the max of var on axis (r=0) for all timestep available
    S : is the simulation output object return by happi.Open()
    var : check namelist ["Env_E_abs]
    return a numpy array - var.max() and the timestep vector [0:iteration_max]
    """
    # read all timestep Available
    ts = S.Probe(0,var).getTimesteps()
    data = np.max(S.Prove(0,var).getData(),axis=1)
    varmax = np.stack((ts,data),axis=0)
    return varmax

def getLasera0(S,timeStep,var='Env_E_abs'):
    """ return the max of var on axis (r=0) for the timestep
    S : is the simulation output object return by happi.Open()
    timeStep : timestep smilei unit
    var : check namelist ["Env_E_abs]
    return a numpy array - var.max() and the timestep vector [0:iteration_max]
    """
    return np.max(S.Probe(0,var,timeStep).getData()[0])

def getLaserWaist(S,timeStep,var='Env_E_abs'):
    """ return the laser waist of Env or field `var` at the iteration
    S : is the simulation output object return by happi.Open()
    timestep : simulation timestep
    var : check namelist ["Env_E_abs" or laser field] to be updated for AM geometry
    return the waist evaluated with Gaussian fit in code units (lamda_0/2pi)
    """
    temp = S.Probe(1,var,timeStep).getData()[0]
    x_max,y_max = np.unravel_index(np.argmax(temp),temp.shape)
    init_vals = [np.max(temp),y_max, 1.0]
    a_val = temp[x_max,:]
    y_val = np.arange(0,temp.shape[1],1)
    # gaussian fit
    best_vals, covar = curve_fit(gaussian, y_val, a_val, p0=init_vals)
    return best_vals[2]

def getLaserPulselength(S,timeStep,var='Env_E_abs'):
    """ return the laser pulse length of Env or field `var` at the iteration
    S : is the simulation output object return by happi.Open()
    timestep : simulation timestep
    var : check namelist ["Env_E_abs" or laser field] to be updated for AM geometry
    return the pulse length FWHM evaluated with Gaussian fit in code units (lamda_0/2pi)
    """
    temp = S.Probe(1,var,timeStep).getData()[0]
    x_max,y_max = np.unravel_index(np.argmax(temp),temp.shape)
    init_vals = [np.max(temp),x_max, 1.0]
    a_val = temp[:,y_max]
    x_val = np.arange(0,temp.shape[0],1)
    # gaussian fit slightly underestimate FWHM value
    best_vals, _ = curve_fit(gaussian, x_val, a_val, p0=init_vals)
    return best_vals[2]*2*np.sqrt(2*np.log(2))

######### extract plasma profile ############################

def plasmaProfile(S):
    """ return the electon plasma density  profile
    S : is the simulation output object return by happi.Open()
    return the numpy array - plasProfile (x,ne) e-/m^3
    """
    nc = S.namelist.ncrit
    ne = np.array(S.namelist.xh_values)*nc
    plasProfile = np.array((S.namelist.xh_points,ne))
    return plasProfile

def dopantProfile(S):
    """ return the electon dopan density profile
    S : is the simulation output object return by happi.Open()
    return the numpy array - plasProfile (x,nN2) N2/m^3
    """
    nc = S.namelist.ncrit
    nd = np.array(S.namelist.xd_values)*nc
    dopProfile = np.array((S.namelist.xd_points,nd))
    return dopProfile

######### extract beam parameter for one iteration ###########

def getBeamParam(S,iteration,species_name="electronfromion",sort = False, E_min=50,E_max=520,chunk_size=100000000,print_flag=True,save_flag=False):
    """return beams paramater for the species_name of the Smilei simulation data
    iteration : timestep
    S : is the simulation output object return by happi.Open()
    species_name :  [electronfromion], electron
    E_min :         [0] energy filter min
    E_max :         [400] energy filter max
    printflag :     [True] print output on screen.
    saveflag :      [False] True to save the data in an csv file
     """
    ########## Read data from Track Particles Diag ############
    track_part = S.TrackParticles(species = species_name, sort = sort, chunksize=chunk_size)
    #print("Available timesteps = ",track_part.getAvailableTimesteps())
    dt_adim    = S.namelist.dt
    for particle_chunk in track_part.iterParticles(iteration, chunksize=chunk_size):
        # Read data
        #if print_flag==True:
        #    print(particle_chunk.keys())
        px           = particle_chunk["px"]
        py           = particle_chunk["py"]
        pz           = particle_chunk["pz"]
        x            = particle_chunk["x"]
        y            = particle_chunk["y"]
        z            = particle_chunk["z"]
        w            = particle_chunk["w"]
        p            = np.sqrt((px**2+py**2+pz**2))                # momentum
        E            = np.sqrt((1.+p**2))
        Nparticles   = np.size(w)
        if print_flag == True:                                  # Number of particles read
            print("Read ",Nparticles," particles from the file")
        total_weight = w.sum()
        Q            = total_weight* e * ncrit * onel**3 * 10**(12) # Total charge in pC
        if print_flag == True:
            print("Total charge before filter in energy= ",Q," pC")
        # Apply a filter on energy
        filter       = np.intersect1d( np.where( E > E_min )[0] ,  np.where( E < E_max )[0] )
        x            = x[filter]
        y            = y[filter]
        z            = z[filter]
        px           = px[filter]
        py           = py[filter]
        pz           = pz[filter]
        E            = E[filter]
        w            = w[filter]
        p            = p[filter]
        total_weight = w.sum()
        Q            = total_weight* e * ncrit * onel**3 * 10**(12) # Total charge in pC
        if print_flag == True:
            print("Total charge after filter in Energy = ",Q," pC")
            print("Filter energy limits: ",E_min,", ",E_max," (m_e c^2)")
        if total_weight > 0:
            #Compute mean values
            x_moy    = (x    *w).sum() / total_weight
            y_moy    = (y    *w).sum() / total_weight
            z_moy    = (z    *w).sum() / total_weight
            #px_moy   = (px   *w).sum() / total_weight
            py_moy   = (py   *w).sum() / total_weight
            pz_moy   = (pz   *w).sum() / total_weight
            # p_moy    = (p    *w).sum() / total_weight
            #Place center of mass at the center of the coordinates
            x  -= x_moy
            y  -= y_moy
            z  -= z_moy
            #px -= px_moy
            py -= py_moy
            pz -= pz_moy
            # Compute properties of the bunch
            x2_moy   = (x**2 *w).sum() / total_weight
            y2_moy   = (y**2 *w).sum() / total_weight
            z2_moy   = (z**2 *w).sum() / total_weight
            #px2_moy  = (px**2*w).sum() / total_weight
            py2_moy  = (py**2*w).sum() / total_weight
            pz2_moy  = (pz**2*w).sum() / total_weight
            ypy_moy  = (y*py *w).sum() / total_weight
            zpz_moy  = (z*pz *w).sum() / total_weight
            py2ovpx2 = (py**2/px**2*w).sum()/total_weight #divergence y squared
            pz2ovpx2 = (pz**2/px**2*w).sum()/total_weight

            # normalized rms emittances based on Floettmann K. "Some basic features of the beam emittance" PRSTA, vol. 6, 3 (2003)
            # https://link.aps.org/doi/10.1103/PhysRevSTAB.6.034202

            emittancey = ( py2_moy*y2_moy - ypy_moy**2 )
            emittancez = ( pz2_moy*z2_moy - zpz_moy**2 )
            if emittancey > 0:
                emittancey = np.sqrt(emittancey) * onel * 1e6 # [um]
            else:
                emittancey = 0.
            if emittancez > 0:
                emittancez = np.sqrt(emittancez) * onel * 1e6 # [um]
            else:
                emittancez = 0.

            rmssize_longitudinal = 2*np.sqrt(x2_moy) * onel * 1e6 # [micron]
            rmssize_y =            2*np.sqrt(y2_moy) * onel * 1e6 # [micron]
            rmssize_z =            2*np.sqrt(z2_moy) * onel * 1e6 # [micron]
            divergence_rms = np.sqrt( py2ovpx2 + pz2ovpx2 )

            # Statistic on energy distribution of particules
            E_mean = weighted_mean(E,w)*0.512
            E_med = weighted_median(E,w)*0.512
            dE_wrms = weighted_std(E,w)/weighted_mean(E,w)*100
            dE_rms = np.std(E)/weighted_mean(E,w)*100
            dE_mad = weighted_mad(E,w)/weighted_median(E,w)*100

            # print beam parameter
            if print_flag == True:
                print("")
                print("--------------------------------------------")
                print("")
                print(" Read \t\t\t\t\t\t", np.size(E)," particles")
                print( "[0] Iteration = \t\t\t\t", iteration)
                print( "[1] Simulation time = \t\t\t\t\t", iteration*dt_adim*onel/c*1e15," fs")
                print( "[2] E_mean = \t\t\t\t\t", E_mean," MeV")
                print( "[3] E_med = \t\t\t\t\t", E_med, "MeV")
                print( "[4] DeltaE_rms / E_mean = \t\t\t", dE_rms , " %.")
                print( "[5] E_mad /E_med  = \t\t\t\t", dE_mad, " %.")
                print( "[6] Total charge = \t\t\t\t", Q, " pC.")
                print( "[7] Emittance_y = \t\t\t\t", emittancey," mm-mrad")
                print( "[8] Emittance_z = \t\t\t\t", emittancez," mm-mrad")
                print( "[9] size_x = \t\t\t\t", rmssize_longitudinal,"um (RMS)")
                print( "[10] divergence_rms = \t\t\t\t", divergence_rms*1e3,"mrad")
                print( "")
                print( "--------------------------------------------")
                print( "")

            # beam paramater list for iteration timestep

            if Q > 0.0 :
                # [0] timestep
                # [1] time [fs]
                # [2] weighted mean energy   [MeV]
                # [3] weighted median value   [MeV]
                # [4] weighted RMS energy spread   [%]
                # [5] RMS energy spread   [%]
                # [6] MAD energy spread [%]
                # [7] charge [pC]
                # [8] normalized RMS y-emittance [um]
                # [9] normalized RMS z-emittance [um]
                # [10] bunch RMS length [um]
                # [11] bunch RMS sigy [um]
                # [12] bunch RMS sigz [um]
                # [13] RMS divergence [mrad]
                beamparam_dict = {"iteration":iteration,
                "time": iteration*dt_adim*onel/c*1e15,
                "energy_wmean": E_mean,
                "energy_wmedian": E_med,
                "energy_wrms": dE_wrms,
                "energy_rms": dE_rms,
                "energy_wmad": dE_mad,
                "charge": Q,
                "emittance_y": emittancey,
                "emittance_z": emittancez,
                "size_x_rms": rmssize_longitudinal,
                "size_y_rms" : rmssize_y,
                "size_z_rms" : rmssize_z,
                "divergence_rms": divergence_rms*1e3}
            else:
                print('no data in the filtered energy range')
                beamparam_dict = {"iteration":iteration,
                "time": iteration*dt_adim*onel/c*1e15,
                "energy_wmean": np.nan,
                "energy_wmedian": np.nan,
                "energy_wrms": np.nan,
                "energy_rms": np.nan,
                "energy_wmad": np.nan,
                "charge": np.nan,
                "emittance_y": np.nan,
                "emittance_z": np.nan,
                "size_x_rms": np.nan,
                "size_y_rms" : np.nan,
                "size_z_rms" : np.nan,
                "divergence_rms": np.nan}

            # save beam parameter in a file
            if save_flag == True:
                print( "data saved in npy file")
                filename = 'smilei-beamparam'+str(iteration)+'.npy'
                filepath = home_directory+'/'+filename
                np.save(filepath,beamparam_dict)
            return beamparam_dict

def getPartAvailableSteps(S,species_name="electronfromion",sort = False, chunk_size=10000000):
    """return available timesteps for the trackParticles"""
    return S.TrackParticles(species = species_name, sort = False, chunksize=chunk_size).getAvailableTimesteps()

def getBeamCharge(S,iteration,species_name="electronfromion",sort = False, E_min=10,E_max=520,chunk_size=10000000,print_flag=True):
    """return beam charge for the species_name of the Smilei simulation data at the timestep iteration
    iteration : timestep
    S : is the simulation output object return by happi.Open()
    species_name :  [electronfromion], electron
    E_min :         [10] energy filter min
    E_max :         [520] energy filter max
    printflag :     [True] print output on screen.
    Q : charge []
     """
    ########## Read data from Track Particles Diag ############
    track_part = S.TrackParticles(species = species_name, sort = sort, chunksize=chunk_size)
    test_part =  S.TrackParticles(species = species_name, timesteps=iteration, sort = sort, chunksize=chunk_size).getData()
    #print("Available timesteps = ",track_part.getAvailableTimesteps())
    px = 0.
    py = 0.
    pz = 0.
    w  = 0.
    if iteration is None:
        Q = 0.
        if print_flag == True:
            print("Iteration or timeStep is None type, return Q=", Q)
    elif test_part[int(iteration)]['w'].sum() < 0.1:
        Q = 0.
        if print_flag == True:
            print("no enough particles return Q=", Q)
    else:
        for particle_chunk in track_part.iterParticles(iteration, chunksize=chunk_size):
            # Read data
            px           += particle_chunk["px"]
            py           += particle_chunk["py"]
            pz           += particle_chunk["pz"]
            w            += particle_chunk["w"]
        p            = np.sqrt(px**2+py**2+pz**2)
        E            = np.sqrt(1.+p**2)
        Nparticles   = np.size(w)
        if Nparticles < 1.:
            Q= 0.
        if print_flag == True:                                  # Number of particles read
            print("Read ",Nparticles," particles from the file")
        total_weight = w.sum()
        Q            = total_weight* e * ncrit * onel**3 * 10**(12) # Total charge in pC
        if print_flag == True:
            print("Total charge before filter in energy= ",Q," pC")
            print("Filter energy limits: ",E_min,", ",E_max," (m_e c^2)")
        # Apply a filter on energy
        filter       = np.intersect1d( np.where( E > E_min )[0] ,  np.where( E < E_max )[0] )
        w            = w[filter]
        total_weight = w.sum()
        Q            = total_weight* e * ncrit * onel**3 * 10**(12) # Total charge in pC
        if print_flag==True:
            print("Total charge after filter in energy= ",Q," pC")
    return Q

def getInjectionTime(S,ts,probeVar='Rho_electronfromion',threshold = 5e-3,print_flag = False):
    """ return the injection timestep and longitudinal coordinate of the injection.
    The injection is defined by a threshold on the `electron_from_ion` density
    S : is the simulation output object return by happi.Open()
    ts : timestep vector [numpy array]
    threshold : value of e- from ionisation max density on axis -n_ei/ncrit  [smilei units]
    t : index of ts at which injection occcurs
    ti : injection timestep
    xi : injection longitudinal position [m]
    """
    dls = S.namelist.lambda_0/(2*np.pi)
    for t in range(len(ts)):
        rhoei = S.Probe(0,probeVar,ts[t]).getData()[0]
        if np.abs(rhoei.min())> threshold:
            ti = ts[t]
            xi = ts[t]*dls
            if print_flag == True :
                print('index:', t)
                print('injection time:',ti,'timestep')
                print('injection x:',xi,'mm')
            break
        else :
            ti = None
            xi = None
    return t,ti,xi

def getSpectrum(S,iteration_to_plot,species_name= "electronfromion",horiz_axis_name= "E", E_min=50, E_max = 640,plot_flag = False, print_flag = False, nbins_horiz = 200, normalized = False):
    """ return spectrum plot or data for a given timesteps
    S : smilei output data
    iteration_to_plot : timestep
    species_name : [electronfromion], electron
    horiz_axis_name : [E] can be px, p or E
    E_min : [50] min value considered in histogram for the horiz axis, in code units
    E_max : [640] max value considered in histogram for the horiz axis, in code units
    option : plot_flag, print_flag,
    nbins_horiz : binning energy histogram [200]
    normalized : normalization of the histogram [False]
    return spectrum data as numpy arrays  (horizontal axis (E, or p)), dQd(E,or p), Epeak, dQdE_max, Ewidth
    """

    #  horizontal axis limits (m_e c^2 units, or Lorentz factor)
    horiz_axis_min = E_min  # Max value considered in histogram for the horiz axis, in code units
    horiz_axis_max = E_max   # Min value considered in histogram for the horiz axis, in code units
    horiz_axis_conversion_factor = 0.512 # to convert from Smilei units to MeV
    hist_conversion_factor       = 1.    # if equal to 1, the charge is in pC
    energy_axis = np.zeros((1,nbins_horiz)) # initialization to avoid unbondedlocalerror.
    specData = np.zeros((1,nbins_horiz))
    Ewidth = 0.0
    Epeak = 0.0
    dQdE_max = 0.0

    ########## Read data from Track Particles Diag #####################
    sort = False
    chunk_size=100000000
    track_part = S.TrackParticles(species = species_name, sort = sort, chunksize=chunk_size)
    #print("Available timesteps = ",track_part.getAvailableTimesteps())

    for particle_chunk in track_part.iterParticles(iteration_to_plot, chunksize=chunk_size):
        # Read data
        #if print_flag==True:
        #    print(particle_chunk.keys())
        px           = particle_chunk["px"]
        py           = particle_chunk["py"]
        pz           = particle_chunk["pz"]
        x            = particle_chunk["x"]
        y            = particle_chunk["y"]
        z            = particle_chunk["z"]
        w            = particle_chunk["w"]
        p            = np.sqrt((px**2+py**2+pz**2))                # momentum
        E            = np.sqrt((1.+p**2))
        Nparticles   = np.size(w)                                    # Number of particles read
        if print_flag == True:
            print("Read ",Nparticles," particles from the file")
        total_weight = w.sum()
        Q            = total_weight* e * ncrit * onel**3 * 10**(12) # Total charge in pC
        if print_flag == True:
            print("Total charge before filter in energy= ",Q," pC")
        # Apply a filter on energy
        filter       = np.intersect1d( np.where( E > E_min )[0] ,  np.where( E < E_max )[0] )
        x            = x[filter]
        y            = y[filter]
        z            = z[filter]
        px           = px[filter]
        py           = py[filter]
        pz           = pz[filter]
        E            = E[filter]
        w            = w[filter]
        p            = p[filter]
        total_weight = w.sum()
        Q            = total_weight* e * ncrit * onel**3 * 10**(12) # Total charge in pC
        if print_flag == True:
            print("Total charge after filter in Energy = ",Q," pC")
            print("Filter energy limits: ",E_min,", ",E_max," (m_e c^2)")

        if Q > 0.0 :

            # Compute 1D histogram
            possible_axes_names =["x","y","z","px","py","pz","E"]
            axes                =[x,y,z,px,py,pz,E]

            if horiz_axis_name in possible_axes_names:
                horiz_axis = axes[possible_axes_names.index(horiz_axis_name)]
            else:
                print("Error, invalid axis")
                exit(0)

            hist1D, horiz_edges = np.histogram(horiz_axis, \
                                        bins=nbins_horiz, \
                                        range=[horiz_axis_min,horiz_axis_max], weights=w)
            #print(np.shape(horiz_edges))
            dhoriz_axis                    = abs(horiz_edges[1]-horiz_edges[0]) # bin size

            # histogram: integrated in dhoriz_axis and gives the total charge
            histogram_spectrum = hist1D*hist_conversion_factor/dhoriz_axis/horiz_axis_conversion_factor*e * ncrit * onel**3  * 10**(12)
            if normalized == True:
                histogram_spectrum = histogram_spectrum / histogram_spectrum[:].max()

            # horizontal axis
            horiz_edges = horiz_edges[0:-1]
            binx = dhoriz_axis*horiz_axis_conversion_factor
            horiz_edges = horiz_edges + 0.5*binx
            energy_axis = horiz_edges*horiz_axis_conversion_factor

            # Preparation for Plot
            if normalized == True:
                plot_title   = "Normalized histogram"
            else:
                plot_title   = 'dQ/d'+horiz_axis_name+" (pC/MeV)"

            #print np.shape(histogram_spectrum)
            #
            if print_flag == True:
                print('Bins size: dx = ',binx)

            histogram_spectrum[histogram_spectrum==0.]=float(np.nan)

            #if print_flag==True:
            #    print(len(energy_axis))

            specData = np.array((histogram_spectrum))

            # Plot
            if plot_flag == True:
                fig = plt.figure()
                fig.set_facecolor('w')

                plt.xlabel(horiz_axis_name+" (MeV)")
                plt.title(plot_title)

                #extnt = np.array([horiz_axis.min()*horiz_axis_conversion_factor, \
                #            horiz_axis.max()*horiz_axis_conversion_factor ])
                #print("Values extension for ",horiz_axis_name," (all particles):")
                #print(extnt)

                #extnt = np.array([horiz_axis_min*horiz_axis_conversion_factor, \
                #            horiz_axis_max*horiz_axis_conversion_factor])
                #print( "Values extension for ",horiz_axis_name," (particles included in the chosen horiz axis limits):")
                #print(extnt)

                plt.plot(energy_axis,histogram_spectrum)
                plt.xlim([horiz_axis_min*horiz_axis_conversion_factor,horiz_axis_max*horiz_axis_conversion_factor])

                #plt.savefig(home_directory+"/E_Spectrum.png",format='png')
                plt.show()

            # compute the full width half maximum using scipy.signal.findpeaks
            try :
                with warnings.catch_warnings():
                    warnings.filterwarnings('ignore', r'All-NaN slice encountered')
                    prom = (np.nanmax(specData)-np.nanmin(specData))*0.66 #factor might be adjusted
                    p , _  = find_peaks(specData,prominence=prom)
                if len(p)==0 :
                    Epeak = 0
                    Ewidth = 0
                    dQdE_max = 0
                else :
                    Epeak = energy_axis[p[0]]
                    dQdE_max = specData[p[0]]
                    Ewidth = binx*peak_widths(specData, p, rel_height=0.5)[0][0]
            except ValueError :
                Epeak = np.nan
                Ewidth = np.nan
                dQdE_max = np.nan

            if print_flag == True:
                print( "")
                print( "--------------------------------------------")
                print( "")
                print("beam Peak energy: \t",Epeak,"MeV")
                print("beam FWHM energy: \t",Ewidth,"MeV")
                print( "")
                print( "--------------------------------------------")
                print( "")
        # no charge in the energy range  Q = 0.
        else :
            energy_axis = np.nan*np.zeros((nbins_horiz))
            specData = np.nan*np.zeros((nbins_horiz))
            Epeak  = np.nan
            dQdE_max = np.nan
            Ewidth = np.nan

    return energy_axis, specData, Epeak, dQdE_max, Ewidth

def getPartParam(S,iteration,species_name="electronfromion",sort= False,chunk_size=100000000,E_min=25, E_max = 520,print_flag = True):
    """return x,y,z,px,py,pz,E,w,p for all particle at timesteps iteration within the filter"""
    track_part = S.TrackParticles(species = species_name,sort = sort,  chunksize=chunk_size)
    #print("Available timesteps = ",track_part.getAvailableTimesteps())

    for particle_chunk in track_part.iterParticles(iteration, chunksize=chunk_size):
        # Read data
        #if print_flag==True:
        #	print(particle_chunk.keys())
        px           = particle_chunk["px"]
        py           = particle_chunk["py"]
        pz           = particle_chunk["pz"]
        x            = particle_chunk["x"]
        y            = particle_chunk["y"]
        z            = particle_chunk["z"]
        w            = particle_chunk["w"]
        p            = np.sqrt((px**2+py**2+pz**2))                # momentum
        E            = np.sqrt((1.+p**2))
        Nparticles   = np.size(w)
        if print_flag==True:                                  # Number of particles read
            print("Read ",Nparticles," particles from the file")
        total_weight = w.sum()
        Q            = total_weight* e * ncrit * onel**3 * 10**(12) # Total charge in pC
        if print_flag==True:
            print("Total charge before filter in energy= ",Q," pC")
            print("Filter energy limits: ",E_min,", ",E_max," (m_e c^2)")
        # Apply a filter on energy
        filter       = np.intersect1d( np.where( E > E_min )[0] ,  np.where( E < E_max )[0] )
        x            = x[filter]
        y            = y[filter]
        z            = z[filter]
        px           = px[filter]
        py           = py[filter]
        pz           = pz[filter]
        E            = E[filter]
        w            = w[filter]
        p            = p[filter]
        total_weight = w.sum()
        Q            = total_weight* e * ncrit * onel**3 * 10**(12) # Total charge in pC
        if print_flag==True:
            print("Total charge after filter in energy= ",Q," pC")
    return np.array([x,y,z,px,py,pz,E,w,p])

def getPSxrms(S,iteration,species_name="electronfromion",sort= False,chunk_size=100000000,E_min=25, E_max = 520,print_flag = True):
    """return x,px for all particle at timesteps iteration within the filter"""
    track_part = S.TrackParticles(species = species_name,sort = sort,  chunksize=chunk_size)
    #print("Available timesteps = ",track_part.getAvailableTimesteps())

    for particle_chunk in track_part.iterParticles(iteration, chunksize=chunk_size):
        # Read data
        #if print_flag==True:
        #	print(particle_chunk.keys())
        px           = particle_chunk["px"]
        py           = particle_chunk["py"]
        pz           = particle_chunk["pz"]
        x            = particle_chunk["x"]
        w            = particle_chunk["w"]
        p            = np.sqrt((px**2+py**2+pz**2))                # momentum
        E            = np.sqrt((1.+p**2))
        Nparticles   = np.size(w)
        if print_flag==True:                                  # Number of particles read
            print("Read ",Nparticles," particles from the file")

        # Apply a filter on energy
        filter       = np.intersect1d( np.where( E > E_min )[0] ,  np.where( E < E_max )[0] )
        x            = x[filter]
        px           = px[filter]
        w            = w[filter]
        Nparticles   = np.size(w)
        if print_flag==True:                                  # Number of particles read
            print("After filtering",Nparticles," particles")

    return np.array([x,px,w])


import inspect

from src.utils import load_dataset, path_root
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.pyplot import figure

def load(methods,filtering='yg'):
    # methods = ['contrast', 'otsu', ... ]
    dfs = []
    for method in methods:
        dfs.append(load_dataset(method, filtering=filtering, aggregate=False))
    return dfs

def aggregat(df):
    aggregate_by = {'fitness': np.max}  # 'dice': np.max, 'jaccard': np.max, 'acc': np.max
    return df.groupby(['method', 'file'], as_index=False).agg(aggregate_by)

def describe(name):
    plt.xlabel("numer zdjęcia")
    plt.ylabel('dopasowanie')
    plt.savefig(path_root('plots', 'barplots', name))
    plt.show()

def bar_1_param(figsize):
    methods = ['contrast', 'otsu']
    dfs = load(methods)

    df1 = aggregat(dfs[0])

    N = df1.shape[0]
    ind = np.arange(N)
    width = 0.25
    plt.figure(figsize=figsize, dpi=80)
    bar1 = plt.bar(ind, df1.fitness, width, color='r')

    df1 = aggregat(dfs[1])
    bar2 = plt.bar(ind + width, df1.fitness, width, color='g')
    plt.xticks(ind + width / 2, [f'{i}' for i in range(1, 1 + N)])
    plt.legend((bar1, bar2), methods)
    plt.title('metody z jednym parametrem')
    describe(inspect.currentframe().f_code.co_name)

def bar_2_param(figsize):
    methods = ['median', 'midgrey','bernsen','mean']
    dfs = load(methods)

    df1 = aggregat(dfs[0])

    N = df1.shape[0]
    ind = np.arange(N)
    width = 0.15
    plt.figure(figsize=figsize, dpi=80)
    bar1 = plt.bar(ind, df1.fitness, width, color='r')

    df1 = aggregat(dfs[1])
    bar2 = plt.bar(ind + width, df1.fitness, width, color='g')

    df1 = aggregat(dfs[2])
    bar3 = plt.bar(ind + 2*width, df1.fitness, width, color='b')

    df1 = aggregat(dfs[3])
    bar4 = plt.bar(ind + 3*width, df1.fitness, width, color='y')

    plt.xticks(ind + 1.5*width , [f'{i}' for i in range(1, 1 + N)])
    plt.legend((bar1,bar2 ,bar3, bar4), methods)
    plt.title('metody z dwoma parametrami')

    describe(inspect.currentframe().f_code.co_name)

def bar_3_param(figsize):
    methods = ['sauvola', 'phansalkar','niblack']
    dfs = load(methods)

    df1 = aggregat(dfs[0])

    N = df1.shape[0]
    ind = np.arange(N)
    width = 0.15
    plt.figure(figsize=figsize, dpi=80)
    bar1 = plt.bar(ind, df1.fitness, width, color='r')

    df1 = aggregat(dfs[1])
    bar2 = plt.bar(ind + width, df1.fitness, width, color='g')

    df1 = aggregat(dfs[2])
    bar3 = plt.bar(ind + 2*width, df1.fitness, width, color='b')

    plt.xticks(ind + width , [f'{i}' for i in range(1, 1 + N)])
    plt.legend((bar1,bar2 ,bar3), methods)
    plt.title('metody z trzema parametrami')
    describe(inspect.currentframe().f_code.co_name)

figsize = (19,6,)
bar_1_param(figsize)
bar_2_param(figsize)
bar_3_param(figsize)


# python 另一种方式与 Mysql 进行交互： ORM框架
# 持久化层：在业务逻辑层和数据库层起到了衔接作用，他可以将内存中的数据模型转化为存储模型，或者将存储模型转化为内存中的数据模型。
# 业务逻辑层：业务对象（对象、属性、继承）--》 持久化层（ORM--》ODBC/JDBC）--》 数据库层（RDBMS：表、字段、索引））
# ORM(Object Relation Mapping) 对象关系映射:她是RDBMS和业务实体对象之间的一个映射，可以吧底层的RDBMS封装成业务实体对象，提供给业务逻辑使用。
from sqlalchemy import create_engine
from sqlalchemy import Column, String, Integer, Float
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import or_
from sqlalchemy import func

# 创建对象的基类：
Base = declarative_base()
# 定义Player对象：
class Player(Base):
  # 表名
  __tablename__ = 'player'

  # 表的结构：
  player_id = Column(Integer, primary_key=True, autoincrement=True)
  team_id = Column(Integer)
  player_name = Column(String(255))
  height = Column(Float(3,2))

# 初始化数据库连接
engine = create_engine('mysql+mysqlconnector://root:09090909@localhost:3306/wz')
# 创建 DBSession 类型：
DBSession = sessionmaker(bind=engine)

# 创建 session 对象：
session = DBSession()
# 创建Player对象：
new_player = Player(team_id=1003, player_name="约翰-科林斯", height=2.08)
# 增
# 添加 到session
session.add(new_player)
# 提交即保存到数据库
session.commit()

# 查
# to_dict() 方法到Base类中
def to_dict(self):
  return {
    c.name: getattr(self, c.name, None) for c in self.__table__.columns
  }
Base.to_dict = to_dict
# 查询身高 >= 2.08m 的球员
rows = session.query(Player).filter(Player.height >= 2.08).all()
# rows = session.query(Player).filter(Player.height >= 2.08, Player.height <= 2.10).all()
# rows = session.query(Player).filter(or_(Player.height >= 2.08, Player.height <= 2.10)).all()
# rows = session.query(Player.team_id, func.count(Player.player_id)).group_by(Player.team_id).having(func.count(Player.player_id) > 5).order_by(func.count(Player.player_id).asc()).all()
print([row.to_dict() for row in rows])

# 改
row = session.query(Player).filter(Player.player_name=='索恩-马克').first()
row.height = 2.17
session.commit()

# 删
row = session.query(Player).filter(Player.player_name=='约翰-科林斯').first()
session.delete(row)
session.commit()

# 关闭session
session.close()

# SQLAlchemy中，采用 Column 对字段进行定义，常用的数据类型：
#   Integer 整数型
#   Float   浮点类型
#   Decimal 定点类型
#   Boolean 布尔类型
#   Date    datetime.date 日期类型
#   Time    datetime.time 时间类型
#   String  字符串类型，使用时需要指定长度，区别于 Text 类型
#   Text    文本类型

#   除了指定 Column 的数据类型外，也可以指定 Column 的参数，对对象创建列约束
#   default       默认值
#   primary_key   是否为主键
#   unique        是否唯一
#   autoincrement 是否自动增长


from selenium import webdriver

# driver = webdriver.Chrome()

# 火狐浏览器 先记住密码
# 火狐浏览器 帮助---故障排除信息--配置文件夹
# profile_director = "配置文件夹的路径"
# # 实例化火狐浏览器配置项
# profile = webdriver.FirefoxProfile(profile_director)
# driver = webdriver.Firefox(profile)

# 谷歌浏览器

user_data_dir = r"--user-data-dir=C:\Users\windows-pc\AppData\Local\Google\Chrome\User Data"
# # # 配置谷歌浏览器加载项
options = webdriver.ChromeOptions()
options.add_argument(user_data_dir)

driver = webdriver.Chrome(options=options)

driver.get("http://oms-uat.jiangxi-isuzu.cn/#/login")


from typing import Dict, Union

from first_distance import FirstDistance
from first_time import FirstTime

class FirstPace(object):

    def __init__(self, minutes: int = 0, seconds: int = 0, length_unit: str = 'mile'):

        """
        Constructor

        :param minutes:
        :type minutes: int
        :param seconds:
        :type seconds: int
        :param length_unit:
        :type length_unit: str
        :return: instance of FirstPace
        :rtype: FirstPace
        """
        if not FirstDistance.is_valid_unit(unit=length_unit):
            raise ValueError('"{}" is not a valid length unit'.format(length_unit))
        self.time = FirstTime(minutes=minutes, seconds=seconds)
        self.length_unit = length_unit

    def __str__(self):

        return str(self.time) + ' min per ' + self.length_unit

    def to_json(self, output_unit: Union[str, None] = None) -> Dict:

        if output_unit and output_unit != self.length_unit:
            dist = FirstDistance(1.0, output_unit)
            seconds_time = self.to_time(dist, 'second')
            output_pace = FirstPace(seconds=round(seconds_time), length_unit=output_unit)
            return {'pace': str(output_pace), 'length_unit': output_unit, 'time': output_pace.time.to_json()}
        else:
            return {'pace': str(self), 'length_unit': self.length_unit, 'time': self.time.to_json()}

    def to_html(self, output_unit: Union[str, None] = None) -> str:

        if output_unit and output_unit != self.length_unit:
            dist = FirstDistance(1.0, output_unit)
            seconds_time = self.to_time(dist, 'second')
            output_pace = FirstPace(seconds=round(seconds_time), length_unit=output_unit)
            return '{} min per {}'.format(str(output_pace.time), output_unit)
        else:
            return '{} min per {}'.format(str(self.time), self.length_unit)

    @classmethod
    def from_string(cls, str_input: str):

        """
        Constructor: Instantiate FirstPace from a string input

        :param str_input: format - '0:MM:SS per unit'
        :type str_input: str
        :return: instance of FirstPace
        :rtype: FirstPace
        """
        tokens = str_input.split()

        p_time = FirstTime.from_string(string=tokens[0])  # pass the exception on
        length_unit = tokens[-1]

        if not FirstDistance.is_valid_unit(unit=tokens[-1]):
            raise ValueError('"{}" is not a valid length unit'.format(length_unit))

        return cls(minutes=p_time.seconds//60, seconds=p_time.seconds % 60, length_unit=length_unit)

    @classmethod
    def copy(cls, from_pace: 'FirstPace'):

        return cls.from_string(str(from_pace))

    def to_time(self, distance: FirstDistance, unit: str) -> float:

        """
        How much time will take to run a given distance with this pace

        :param distance: the distance
        :type distance: FirstDistance
        :param unit: the desired unit of the result
        :type unit: str
        :return: the time value for this unit
        :rtype: float
        """
        factor = distance.convert_to(unit=self.length_unit)
        seconds = self.time.total_seconds() * factor
        result_time = FirstTime(seconds=round(seconds))
        return result_time.convert_to(unit=unit)

    def to_distance(self, time: FirstTime, unit: str) -> float:

        """
        How far you run given duration with this pace

        :param time: the duration
        :type time: FirstTime
        :param unit: the desired unit of the result
        :type unit: str
        :return: the distance value for this unit
        :rtype: float
        """
        factor = time.total_seconds()/self.time.total_seconds()
        result_distance = FirstDistance(distance=factor, unit=self.length_unit)
        return result_distance.convert_to(unit=unit)

    @classmethod
    def from_time_distance(cls, time: FirstTime, distance: FirstDistance, unit: str = None):

        """
        Constructor: Initiate FirstPace from time/distance

        :param time:
        :type time: FirstTime
        :param distance:
        :type distance: FirstDistance
        :param unit: length unit
        :type unit: str
        :return: instance to FirstPace
        :rtype: FirstPace
        """
        if unit is None:
            unit = distance.unit
        factor = distance.convert_to(unit=unit)  # 400m with unit = mile will become ~0.25
        seconds = time.total_seconds() / factor  # 2 minutes for 400m will give ~(2*60)/0.25

        return cls(minutes=int(seconds//60), seconds=round(seconds % 60), length_unit=unit)

    def increment(self, seconds: int) -> None:

        """
        Increment the pace by number of seconds - for instructions like 'RP+15'

        :param seconds:
        :type seconds: int
        """
        self.time = FirstTime(seconds=self.time.seconds + seconds)

    def meters_per_second_delta(self, delta_in_seconds: int) -> float:

        """
        Convert to speed in m/s for tcx with delta for tolerance

        :param delta_in_seconds:
        :type delta_in_seconds: int
        :return: calculated speed in m/s
        :rtype: float
        """
        seconds = self.time.total_seconds() + delta_in_seconds
        meters = FirstDistance(distance=1.0, unit=self.length_unit).convert_to(unit='m')

        return meters / seconds


import numpy as np
import matplotlib.pyplot as plt

def f(n, r, primary):
 return primary * (1 + r/100)**n

x = np.linspace(1,99)
primary = 100
r = 5
y = f(x,r,primary)
ax1 = plt.plot(x,y)

print(type(ax1))
print(ax1)
plt.show()

#============================================

y1 = np.array([1.0,2.0,3.0,4.0,5.0])
y2 = np.array([1.2,1.6,3.1,4.2,4.8])
y3 = np.array([3.2,1.1,2.0,3.9,2.5])

fig,ax2 = plt.subplots()
print(type(ax2)) #
print(type(fig)) # fig representa a figura background
lines = ax2.plot(y1,'o-',y2,'x--',y3,'*-.')
ax2.set_title("Plot of the data y1,y2, and y3")
ax2.set_xlabel("x axis label")
ax2.set_ylabel("y axis label")
ax2.legend(("data y1","data y2","data y3"))
plt.show()

#==============================================================#

#plot1

x1 = np.array([0,1,2,3])
y2 = np.array([3,8,1,10])

fig,ax3 = plt.subplots(2,2)
ax3.set_title("Plot of the data y2")
ax3.set_xlabel("x axis label")
ax3.set_ylabel("y axis label")
ax3.legend(("data y2"))

plt.subplot(2,2,3)
plt.plot(x1,y2)
plt.title("SALES")

#plot2
x = np.array([0,1,2,3])
y = np.array([10,20,30,40])

plt.subplot(2,2,2)
plt.plot(x1,y2)
plt.title("INCOME")

plt.suptitle("MY SHOP")
plt.show()

plt.savefig("figure.1.png")


__author__ = 'diegopinheiro'

from common.attribute_converter import AttributeConverter
from common.attribute import Attribute
from common.float_converter_2 import FloatConverter2

class Rule():

    def __init__(self,
                 input_attributes=list(),
                 output_attributes=list(),
                 bit_string=list()):
        self.input_attributes = input_attributes
        self.output_attributes = output_attributes
        self.bit_string = bit_string
        self.antecedents = dict()
        self.consequents = dict()
        self.accuracy = 0
        self.coverage = 0
        self.initialize_rule()

    def initialize_rule(self):
        attribute_index = 0
        antecedents = dict()
        for attribute in self.input_attributes:
            antecedent = list()
            if attribute.type == Attribute.TYPE_DISCRETE:
                attribute_number_bits = AttributeConverter.get_number_representation(attribute)
                attribute_bits = self.bit_string[attribute_index:attribute_index+attribute_number_bits]
                for bit_position in range(0, len(attribute_bits)):
                    if attribute_bits[bit_position] == 1:
                        category_bits = [0] * len(attribute_bits)
                        category_bits[bit_position] = 1
                        antecedent.append(AttributeConverter.get_attribute_category(attribute,category_bits))
                antecedents[attribute] = antecedent
                attribute_index += attribute_number_bits

            elif attribute.type == Attribute.TYPE_CONTINUOUS:
                attribute_number_bits = FloatConverter2.get_number_representation()
                attribute_bits_low = FloatConverter2.get_value(self.bit_string[attribute_index:attribute_index+attribute_number_bits])
                attribute_bits_high = FloatConverter2.get_value(self.bit_string[attribute_index+attribute_number_bits:attribute_index+2*attribute_number_bits])

                antecedents[attribute] = [attribute_bits_low, attribute_bits_high]
                attribute_index += 2 * attribute_number_bits

        consequents = dict()
        for attribute in self.output_attributes:
            # attribute_number_bits = AttributeConverter.get_number_representation(attribute)
            attribute_bits = int(self.bit_string[attribute_index][0])
            # for bit_position in range(0, len(attribute_bits)):
            #     if attribute_bits[bit_position] == 1:
            #         consequents[attribute] = AttributeConverter.get_attribute_category(attribute, attribute_bits)
                    # break
            consequents[attribute] = attribute.categories[attribute_bits]
            attribute_index += 1

        self.antecedents = antecedents
        self.consequents = consequents

    def is_valid(self):
        if len(self.consequents.keys()) != len(self.output_attributes) or len(self.antecedents.keys()) != len(self.input_attributes):
            return False

        is_valid = True
        for a in self.antecedents.keys():
            if a.type == Attribute.TYPE_DISCRETE:
                if len(self.antecedents[a]) == 0:
                    is_valid = False
                    break
            elif a.type == Attribute.TYPE_CONTINUOUS:
                low_boundary = self.antecedents[a][0]
                high_boundary = self.antecedents[a][1]
                if low_boundary == float("Nan") or high_boundary == float("Nan"):
                    is_valid = False
                    break

                if low_boundary > high_boundary:
                    is_valid = False
                    break

        return is_valid

    def is_antecedents_match(self, register):
        is_applied = True
        for attribute in self.antecedents.keys():
            value = self.antecedents[attribute]
            if attribute.type == Attribute.TYPE_DISCRETE:
                if not value.__contains__(register[attribute.index]):
                    is_applied = False
                    break
            elif attribute.type == Attribute.TYPE_CONTINUOUS:
                low_boundary = value[0]
                high_boundary = value[1]
                desired = float(register[attribute.index])
                if not low_boundary <= desired <= high_boundary:
                    is_applied = False
                    break
        return is_applied

    def print_rule(self):
        rule_string = ""
        for attribute in self.antecedents:
            rule_string = rule_string.replace("?", "^")
            if attribute.type == Attribute.TYPE_DISCRETE:
                rule_string += attribute.name + " in " + str(self.antecedents[attribute]) + " ? "
            elif attribute.type == Attribute.TYPE_CONTINUOUS:
                rule_string += str(self.antecedents[attribute][0]) + " <= " + attribute.name + " <= " + str(self.antecedents[attribute][1]) + " ? "
        rule_string = rule_string.replace("?", "=>")

        for attribute in self.consequents:
            rule_string = rule_string.replace("?", "^")
            if attribute.type == Attribute.TYPE_DISCRETE:
                rule_string += attribute.name + " = " + str(self.consequents[attribute]) + " "
            elif attribute.type == Attribute.TYPE_CONTINUOUS:
                rule_string += str(self.consequents[attribute][0]) + " <= " + attribute.name + " <= " + str(self.consequents[attribute][1]) + " ? "
        # rule_string = rule_string.replace("?", "=>")

        # rule_string += " " + str(self.target_distribution)
        return rule_string


import argparse
import os
import pathlib
import shutil
import stat
import subprocess
import sys
import typing

from loguru import logger

aliases = {
    "gs": "git status",
    "gl": "git log --oneline",
}  # These don't have any actual useful use

GIT_TEMPLATE_URL = "https://github.com/Nipa-Code/python-project-template.git"

"""
NOTE: Missing to Configure pre-commit.
"""

def read_user_args(args: str) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Used to config the name of the project and possibly other stuff."
    )
    parser.add_argument(
        "-n",
        "--name",
        action="store",
        default="new_project",
        required=False,
        help="The name of new project.",
    )

    parser.add_argument(
        "-e",
        "--allow-only-empty",
        action="store_true",
        default=True,
        required=False,
        help="If True, will not create folders if directory is not empty.",
    )

    parser.add_argument(
        "-u",
        "--url",
        action="store",
        default=GIT_TEMPLATE_URL,
        required=False,
        help="The url of the template to clone. By default using some template of mine.",
    )

    return parser.parse_args(args)

def add_alias(name: str = "setup-project", command: str = None):
    """
    Add new alias to windows powershell to execute commands with simple shortcuts.
    :param name: The name of new alias.
    :param command: The command to execute.
    """
    if name is not None and command is not None:
        logger.info(f"adding alias {name} to {command}")
        subprocess.run(f"New-Alias {name} '{command}'", shell=True, capture_output=True)
        aliases[name] = command

def del_even_readonly(action, name, exc):
    """
    Edit file permissions and remove that file if it's read-only.
    :param name: The name of the file to be removed.
    """
    os.chmod(name, stat.S_IWRITE)
    os.remove(name)

def create_template(
    name: typing.Optional[str] = "new_project",
    allow_only_empty: typing.Optional[bool] = True,
    url: typing.Optional[str] = GIT_TEMPLATE_URL,
):
    """
    Function to create project templates automatically. Second param is kind of useless for now.
    -> create folders
    -> clone project template from Github
    -> initialize git
    -> initialize poetry
    :param name: The name of new project.
    :param allow_only_empty: If True, will not create folders if directory is not empty.
    :param url: The url of the project template, default is set on codebase.
    """
    if not os.listdir():  # check if path is empty. NOTE: add "not" if it is not there
        logger.info(
            f"Creating template for project: '{name}', from: '{url}', Directory is empty"
        )
        subprocess.run(f"mkdir {name}", shell=True, capture_output=True)
        logger.debug("Created project folder")

        subprocess.run(
            "git init",  # run init with the path of git
            shell=True,
            capture_output=True,
        )
        logger.info("Initialized git")

        subprocess.run(
            "git -b main",
            shell=True,
            capture_output=True,
        )
        logger.info("Switched branch to 'main'")

        subprocess.run(f"git clone {url}", shell=True, capture_output=True)
        logger.info("Cloned project template files")

        files = os.listdir("./python-project-template")
        files.__delitem__(files.index(".git"))  # delete .git from the files to move
        for file in files:
            shutil.move(f"./python-project-template/{file}", ".")

        # Delete the folder "python-project-template"
        if os.path.exists("./python-project-template"):
            shutil.rmtree("./python-project-template", onerror=del_even_readonly)

        logger.debug("Successfully pulled project files from Github and moved them.")

        logger.debug("Running 'poetry install', this may take a while...")
        subprocess.run("poetry install", shell=True, capture_output=True, timeout=180.0)
        logger.info("Installed poetry environment packages")

        subprocess.run(
            "git add .",
            shell=True,
            capture_output=True,
        )
        logger.debug("Added all files to git")

        subprocess.run(
            "git commit -m 'initial commit'",
            shell=True,
            capture_output=True,
        )
        logger.debug("Succesfully ran git commit, files are now saved to git")
        logger.info("Finished setup, ready to go!")
    else:
        logger.warning("Not an empty directory, aborting!")

if __name__ == "__main__":
    # read args when this file is invoked
    args = read_user_args(sys.argv[1:])
    create_template(args.name, allow_only_empty=True, url=args.url)


import jwt
import cloudinary.uploader
from rest_framework.response import Response
from rest_framework import status
from rest_framework.decorators import api_view
from drf_spectacular.utils import extend_schema
from django.shortcuts import get_object_or_404, render
from accounts.models import Accounts, Address
from artist.models import Artist, Event, Post
from artist.serializers import eventSerializer, postSerializer
from userapp.models import Booking, Order
from .serializers import addressSerializer, bookingSerializer, orderSerializer, paymentSerializer, userSerializer
# Create your views here.

@extend_schema(responses=userSerializer)
@api_view(['POST'])
def verify_token(request):
    print(request.headers,'kkkkkkkkkkkkkk')
    token = request.data['token']
    decoded = jwt.decode(token,'secret',alogrithms='HS256')
    admi_n = Accounts.objects.get(email = decoded.get('email'))
    serializer = userSerializer(admi_n,many=False)
    if admi_n:
        return Response(serializer.data)
    else:
        return Response({'status': 'Token Invalid'})

@extend_schema(responses=userSerializer(many=False))
@api_view(['GET'])
def user_profile(request,id):
    try:
        user = Accounts.objects.get(id=id)
        serializer = userSerializer(user,many=False)
        user_details = serializer.data
        return Response({'data':user_details}, status=status.HTTP_200_OK)
    except Accounts.DoesNotExist:
        return Response({'error': 'User not found'}, status=status.HTTP_404_NOT_FOUND)

@api_view(['POST'])
def userProfilePic(request, user_id):
    user = Accounts.objects.get(id=user_id)   # Assuming the artist is authenticated
    profile_picture = request.FILES["profile_img"]
    if profile_picture:
        upload_result = cloudinary.uploader.upload(
            profile_picture,
            folder='profiles'
        )
        print(upload_result, 'lllllllllllllll')
        profile_picture_url = upload_result['secure_url']
       # Update the profile_img field with the image URL
        user.profile_img = profile_picture_url
        user.save()
        return Response({"profile_picture_url": profile_picture_url})
    else:
        return Response({"message": "Unsuccessful"})

@api_view(['POST'])
def userCoverPic(request, user_id):
    user = Accounts.objects.get(id=user_id)   # Assuming the artist is authenticated
    cover_picture = request.FILES["cover_img"]
    if cover_picture:
        upload_result = cloudinary.uploader.upload(
            cover_picture,
            folder='profiles'
        )
        print(upload_result, 'lllllllllllllll')
        cover_picture_url = upload_result['secure_url']
       # Update the cover_img field with the image URL
        user.cover_img = cover_picture_url
        user.save()
        return Response({"cover_picture_url": cover_picture_url})
    else:
        return Response({"message": "Unsuccessful"})

@extend_schema(responses=userSerializer)
@api_view(['PUT'])
def user_profile_update(request,id):
    try:
        user = Accounts.objects.get(id=id)
        serializer = userSerializer(user,data=request.data)
        if serializer.is_valid():
            serializer.save()
            return Response("Profile updated Successfully")
        else:
            return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
    except Accounts.DoesNotExist:
        return Response("Artist bot found", status=status.HTTP_404_NOT_FOUND )

@extend_schema(responses=eventSerializer)
@api_view(['GET'])
def event_details(request,id):
    events = Event.objects.filter(conducting_artist_id=id)
    serializer = eventSerializer(events,many=True)
    return Response(serializer.data)

@extend_schema(responses=bookingSerializer)
@api_view(['PATCH'])
def booking_event(request):
    booking_serializer = bookingSerializer(data=request.data)
    if booking_serializer.is_valid() :
        booking_data = booking_serializer.validated_data
        event_id = booking_data['eventname'].id
        slot_no = booking_data['slot_no']
        # Get the event
        try:
            event = Event.objects.get(id=event_id)
        except Event.DoesNotExist:
            return Response({'error': 'Event not found'}, status=status.HTTP_404_NOT_FOUND)
        # Check if the total_slots is already zero
        if event.total_slots == 0:
            return Response({'error': 'No available slots'}, status=status.HTTP_400_BAD_REQUEST)
        # Calculate the new total_slots value
        new_total_slots = event.total_slots - slot_no
        if new_total_slots < 0:
            new_total_slots = 0
        # Update the event's total_slots
        event.total_slots = new_total_slots
        event.save()
        # Save the booking
        booking_serializer.save()
        return Response({'status': 'Event booked successfully'}, status=status.HTTP_201_CREATED)
    return Response(booking_serializer.errors, status=status.HTTP_400_BAD_REQUEST)

@extend_schema(responses=bookingSerializer)
@api_view(['GET'])
def bookedevent_list(request,id):
    booked_list = Booking.objects.filter(username=id)
    serializer = bookingSerializer(booked_list,many=True)
    return Response(serializer.data,status=status.HTTP_200_OK)

@extend_schema(responses=postSerializer)
@api_view(['GET'])
def artists_postlist(request,artist_id):
    posts = Post.objects.filter(artist_id = artist_id)
    serializer = postSerializer(posts,many=True)
    return Response(serializer.data,status=status.HTTP_200_OK)

# @extend_schema(responses=postSerializer)
# @api_view(['GET'])
# def buy_post(request,post_id):
#     try:
#         address = Address.objects.get(user=request.user)
#         post = Post.objects.get(id=post_id)
#         base_price = post.base_price
#         shipping_price = post.shipping_price
#         total_price = base_price + shipping_price
#         data = {
#             'post':post.id,'base_price':base_price,
#             'shipping_price':shipping_price,'total_price':total_price,
#             'address':address
#             }
#         return Response(data, status=status.HTTP_201_CREATED)
#     except Post.DoesNotExist:
#         return Response({'status':'Post not found'},status=status.HTTP_400_BAD_REQUEST)

@extend_schema(responses=addressSerializer)
@api_view(['POST'])
def add_address(request):
    serializer = addressSerializer(data=request.data)
    if serializer.is_valid():
        serializer.save(user=request.user)
        return Response("New address added succesfully", status=status.HTTP_201_CREATED)
    return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)

@extend_schema(responses=addressSerializer)
@api_view(['GET'])
def address_details(request,userid):
    adresses = Address.objects.filter(user_id=userid)
    serializer = addressSerializer(adresses,many=True)
    return Response({'data':serializer.data},status=status.HTTP_200_OK)

@extend_schema(responses=addressSerializer)
@api_view(['DELETE'])
def delete_address(request,id):
    address = get_object_or_404(Address,id=id)
    address.delete()
    return Response("Address deleted successfully")

@extend_schema(responses=orderSerializer)
@api_view(['GET'])
def order_list(request,user_id):
    order_list = Order.objects.filter(user_buyer=user_id)
    serializer = orderSerializer(order_list,many=True)
    return Response({'data':serializer.data},status=status.HTTP_200_OK)