import pprint
from typing import Text, List, Dict, Optional, Set, Tuple, Callable, Union

from absl import logging

from xls.common.xls_error import XlsError
from xls.dslx import ast
from xls.dslx import bit_helpers
from xls.dslx import deduce
from xls.dslx import dslx_builtins
from xls.dslx import extract_conversion_order
from xls.dslx.concrete_type import ArrayType
from xls.dslx.concrete_type import BitsType
from xls.dslx.concrete_type import ConcreteType
from xls.dslx.concrete_type import EnumType
from xls.dslx.concrete_type import FunctionType
from xls.dslx.concrete_type import TupleType
from xls.dslx.parametric_expression import ParametricExpression
from xls.dslx.parametric_instantiator import SymbolicBindings
from xls.dslx.span import PositionalError
from xls.dslx.span import Span
from xls.ir.python import bits as bits_mod
from xls.ir.python import fileno as fileno_mod
from xls.ir.python import function as ir_function
from xls.ir.python import function_builder
from xls.ir.python import lsb_or_msb
from xls.ir.python import number_parser
from xls.ir.python import package as ir_package
from xls.ir.python import source_location
from xls.ir.python import type as type_mod
from xls.ir.python import value as ir_value
from xls.ir.python import verifier as verifier_mod
from xls.ir.python.function_builder import BValue

class ParametricConversionError(XlsError):
  """Raised when we attempt to IR convert a parametric function."""

class ConversionError(PositionalError):
  """Raised on an issue converting to IR at a particular file position."""

def _int_to_bits(value: int, bit_count: int) -> bits_mod.Bits:
  """Converts a Python arbitrary precision int to a Bits type."""
  if bit_count <= 64:
    return bits_mod.UBits(value, bit_count) if value >= 0 else bits_mod.SBits(
        value, bit_count)
  return number_parser.bits_from_string(
      bit_helpers.to_hex_string(value, bit_count), bit_count=bit_count)

class _IrConverterFb(ast.AstVisitor):
  """An AST visitor that converts AST nodes into IR.

  Note that ASTs can only be converted to IR once they have been fully
  concretized; there is no parametric-function support in the IR text!

  Attributes:
    module: Module that we're converting IR for.
    node_to_ir: Mapping from AstNode to the IR that was used to emit it (as a
      FunctionBuilder BValue).
    symbolic_bindings: Mapping from parametric binding name (e.g. "N" in
      `fn [N: u32] id(x: bits[N]) -> bits[N]`) to its value in this conversion.
    node_to_type: Type information for the AstNodes determined during the type
      checking phase (that must precede IR conversion).
    _constant_deps: Externally-noted constant dependencies that this function
      has (as free variables); noted via add_constant_dep().
    emit_positions: Whether or not we should emit position data based on the AST
      node source positions.
    dslx_name: The name of the DSLX function that's currently being translated.
  """

  def __init__(self,
               package: ir_package.Package,
               module: ast.Module,
               node_to_type: deduce.NodeToType,
               emit_positions: bool,
               dslx_name: Optional[Text] = None):
    self.module = module
    self.node_to_type = node_to_type
    self.emit_positions = emit_positions
    self.package = package
    self.node_to_ir = {
    }  # type: Dict[ast.AstNode, Union[BValue, Tuple[int, BValue]]]
    self.symbolic_bindings = {}  # type: Dict[Text, int]
    self._constant_deps = []
    self.fb = None  # Optional[function_builder.FunctionBuilder]
    # TODO(leary): 2019-07-19 Create a way to get the file path from the module.
    self.fileno = self.package.get_or_create_fileno('fake_file.x')
    # Number of "counted for" nodes we've observed in this function.
    self.counted_for_count = 0
    self.last_expression = None  # Optional[ast.Expr]
    self.dslx_name = dslx_name

  def _extract_module_level_constants(self, m: ast.Module):
    """Populates `self.symbolic_bindings` with module-level constant values."""
    for constant in m.get_constants():
      if isinstance(constant.value, ast.Number):
        value = constant.value.get_value_as_int()
        logging.vlog(
            3, 'Found module-level constant %s for symbolic bindings: %s',
            constant.name.identifier, value)
        self.symbolic_bindings[constant.name.identifier] = value
    logging.vlog(2, 'Symbolic bindings now: %s', self.symbolic_bindings)

  def add_constant_dep(self, constant: ast.Constant) -> None:
    logging.vlog(2, 'Adding constant dep: %s', constant)
    self._constant_deps.append(constant)

  def _type_to_ir(self, concrete_type: ConcreteType) -> type_mod.Type:
    """Converts a concrete type to its corresponding IR representation."""
    assert isinstance(concrete_type, ConcreteType), concrete_type
    logging.vlog(4, 'Converting concrete type to IR: %s', concrete_type)
    if isinstance(concrete_type, ArrayType):
      element_type = self._type_to_ir(concrete_type.get_element_type())
      element_count = concrete_type.size
      if not isinstance(element_count, int):
        raise ValueError(
            'Expect array element count to be integer; got {!r}'.format(
                element_count))
      result = self.package.get_array_type(element_count, element_type)
      logging.vlog(
          4, 'Converted type to IR; concrete type: %s ir: %s element_count: %d',
          concrete_type, result, concrete_type.size)
      return result
    elif isinstance(concrete_type, BitsType) or isinstance(
        concrete_type, EnumType):
      return self.package.get_bits_type(concrete_type.get_total_bit_count())
    else:
      if not isinstance(concrete_type, TupleType):
        raise ValueError(
            'Expect type to be bits/enum, array, or tuple; got: '
            f'{concrete_type} ({concrete_type!r})')
      members = tuple(
          self._type_to_ir(m) for m in concrete_type.get_unnamed_members())
      return self.package.get_tuple_type(members)

  def _resolve_type_to_ir(self, node: ast.AstNode) -> type_mod.Type:
    concrete_type = self._resolve_type(node)
    assert isinstance(concrete_type, ConcreteType), concrete_type
    try:
      return self._type_to_ir(concrete_type)
    except ValueError as e:
      if 'Expect type to be' in str(e):
        raise ConversionError('Could not resolve type: {}'.format(e), node.span)
      if 'Expect array element count to be' in str(e):
        raise ConversionError('Could not resolve type: {}'.format(e), node.span)
      raise

  def _def(self, node: ast.AstNode, ir_func: Callable[..., BValue], *args,
           **kwargs) -> BValue:
    # TODO(leary): 2019-07-19 When everything is Python 3 we can switch to use a
    # single kwarg "span" after vararg with a normal pytype annotation.
    span = kwargs.pop('span', None)
    assert not kwargs
    assert isinstance(span,
                      (type(None), Span)), 'Expect span kwarg as Span or None.'

    if span is None and hasattr(node, 'span') and isinstance(node.span, Span):
      span = node.span
    loc = None
    if span is not None and self.emit_positions:
      start_pos = span.start
      lineno = fileno_mod.Lineno(start_pos.lineno)
      colno = fileno_mod.Colno(start_pos.colno)
      loc = source_location.SourceLocation(self.fileno, lineno, colno)

    try:
      ir = ir_func(*args, loc=loc)
    except TypeError:
      logging.error('Failed to call IR-creating function %s with args: %s',
                    ir_func, args)
      raise
    assert isinstance(
        ir, BValue), 'Expect ir_func to return a BValue; got {!r}'.format(ir)
    logging.vlog(4, 'Define node "%s" to be %s', node, ir)
    self.node_to_ir[node] = ir
    return ir

  def _def_alias(self, from_: ast.AstNode, to: ast.AstNode) -> BValue:
    self.node_to_ir[to] = self.node_to_ir[from_]
    return self._use(to)

  def _def_const(self, node: ast.AstNode, value: int, bit_count: int) -> BValue:
    bits = _int_to_bits(value, bit_count)
    ir = self._def(node, self.fb.add_literal_bits, bits)
    self.node_to_ir[node] = (value, ir)
    return ir

  def _is_const(self, node: ast.AstNode) -> bool:
    """Returns whether "node" corresponds to a known-constant (int) value."""
    record = self.node_to_ir[node]
    return isinstance(record, tuple)

  def _get_const(self, node: ast.AstNode) -> int:
    """Retrieves the known-constant (int) value associated with "node"."""
    record = self.node_to_ir[node]
    if not isinstance(record, tuple):
      raise ConversionError(
          'Expected a constant, but value does not appear constant.',
          node.get_span_or_fake())
    return record[0]

  def _use(self, node: ast.AstNode) -> BValue:
    record = self.node_to_ir[node]
    if isinstance(record, tuple):
      return record[1]
    else:
      return record

  def _resolve_dim(self, dim):
    while isinstance(dim, ParametricExpression):
      try:
        orig = dim
        dim = dim.evaluate(self.symbolic_bindings)
        logging.vlog(4, 'Evaluated dim %s to %s via %s', orig, dim,
                     self.symbolic_bindings)
      except KeyError:
        logging.vlog(4, 'Could not resolve %s dim %s via symbolic bindings %r',
                     dim, self.symbolic_bindings)
        return dim
    return dim

  def _resolve_type(self, node: ast.AstNode) -> ConcreteType:
    try:
      concrete_type = self.node_to_type[node]
    except deduce.TypeMissingError:
      raise ConversionError(
          f'Failed to convert to IR because type was missing for AST node: '
          f'{node} :: {node!r}', node.get_span_or_fake())

    assert isinstance(concrete_type, ConcreteType), concrete_type
    result = concrete_type.map_size(self._resolve_dim)
    logging.vlog(4, 'Resolved concrete type from %s to %s via %s',
                 concrete_type, result, self.symbolic_bindings)
    return result

  @ast.AstVisitor.no_auto_traverse
  def visit_TypeAnnotation(self, node: ast.TypeAnnotation):
    pass

  def visit_Ternary(self, node: ast.Ternary):
    self._def(node, self.fb.add_sel, self._use(node.test),
              self._use(node.consequent), self._use(node.alternate))

  def visit_Binop(self, node: ast.Binop):
    lhs_type = self.node_to_type[node.lhs]
    signed_input = isinstance(lhs_type, BitsType) and lhs_type.signed
    operator_to_f = {
        # Arithmetic.
        ast.Binop.ADD: self.fb.add_add,
        ast.Binop.SUB: self.fb.add_sub,
        ast.Binop.MUL: self.fb.add_smul if signed_input else self.fb.add_umul,
        ast.Binop.DIV: self.fb.add_udiv,
        # Comparisons.
        ast.Binop.EQ: self.fb.add_eq,
        ast.Binop.NE: self.fb.add_ne,
        ast.Binop.GE: self.fb.add_sge if signed_input else self.fb.add_uge,
        ast.Binop.GT: self.fb.add_sgt if signed_input else self.fb.add_ugt,
        ast.Binop.LE: self.fb.add_sle if signed_input else self.fb.add_ule,
        ast.Binop.LT: self.fb.add_slt if signed_input else self.fb.add_ult,
        # Shifts.
        ast.Binop.SHRL: self.fb.add_shrl,
        ast.Binop.SHLL: self.fb.add_shll,
        ast.Binop.SHRA: self.fb.add_shra,
        # Bitwise.
        ast.Binop.XOR: self.fb.add_xor,
        ast.Binop.AND: self.fb.add_and,
        ast.Binop.OR: self.fb.add_or,
        # Logical.
        ast.Binop.LOGICAL_AND: self.fb.add_and,
        ast.Binop.LOGICAL_OR: self.fb.add_or,
    }
    f = operator_to_f.get(node.operator.get_kind_or_keyword())
    if f:
      pass
    elif node.operator.kind == ast.Binop.CONCAT:
      # pylint: disable=g-long-lambda
      f = lambda x, y, loc: self.fb.add_concat([x, y], loc=loc)
    else:
      raise NotImplementedError(node.operator)

    self._def(node, f, self._use(node.lhs), self._use(node.rhs))

  def _next_counted_for_ordinal(self) -> int:
    result = self.counted_for_count
    self.counted_for_count += 1
    return result

  def _visit_matcher(self, matcher: ast.NameDefTree, index: Tuple[int, ...],
                     matched_value: BValue,
                     matched_type: ConcreteType) -> BValue:
    if matcher.is_leaf():
      leaf = matcher.get_leaf()
      if isinstance(leaf, ast.WildcardPattern):
        return self._def(matcher, self.fb.add_literal_bits,
                         bits_mod.UBits(1, 1))
      elif isinstance(leaf, (ast.Number, ast.EnumRef)):
        leaf.accept(self)
        return self._def(matcher, self.fb.add_eq, self._use(leaf),
                         matched_value)
      elif isinstance(leaf, ast.NameRef):
        result = self._def(matcher, self.fb.add_eq, self._use(leaf.name_def),
                           matched_value)
        self._def_alias(leaf.name_def, to=leaf)
        return result
      else:
        assert isinstance(
            leaf, ast.NameDef
        ), 'Expected leaf to be wildcard, number, or name; got: {!r}'.format(
            leaf)
        ok = self._def(leaf, self.fb.add_literal_bits, bits_mod.UBits(1, 1))
        self.node_to_ir[matcher] = self.node_to_ir[leaf] = matched_value
        return ok
    else:
      ok = self.fb.add_literal_bits(bits_mod.UBits(value=1, bit_count=1))
      for i, (element, element_type) in enumerate(
          zip(matcher.tree, matched_type.get_unnamed_members())):
        # Extract the element.
        member = self.fb.add_tuple_index(matched_value, i)
        cond = self._visit_matcher(element, index + (i,), member, element_type)
        ok = self.fb.add_and(ok, cond)
      return ok

  @ast.AstVisitor.no_auto_traverse
  def visit_Match(self, node: ast.Match):
    if (not node.arms or not node.arms[-1].patterns[0].is_irrefutable()):
      raise ConversionError(
          'Only matches with trailing irrefutable patterns are currently handled.',
          node.span)

    node.matched.accept(self)
    matched = self._use(node.matched)
    matched_type = self._resolve_type(node.matched)

    default_arm = node.arms[-1]
    assert len(default_arm.patterns) == 1, (
        'Multiple patterns in default arm is not yet implemented for IR '
        'conversion.')
    self._visit_matcher(default_arm.patterns[0], (len(node.arms) - 1,), matched,
                        matched_type)
    default_arm.expr.accept(self)

    arm_selectors = []
    arm_values = []
    for i, arm in enumerate(node.arms[:-1]):
      this_arm_selectors = []
      for pattern in arm.patterns:
        selector = self._visit_matcher(pattern, (i,), matched, matched_type)
        this_arm_selectors.append(selector)
      if len(this_arm_selectors) > 1:
        arm_selectors.append(self.fb.add_nary_or(this_arm_selectors))
      else:
        arm_selectors.append(this_arm_selectors[0])
      arm.expr.accept(self)
      arm_values.append(self._use(arm.expr))

    # So now we have the following representation of the match arms:
    #   match x {
    #     42  => blah
    #     64  => snarf
    #     128 => yep
    #     _   => burp
    #   }
    #
    #   selectors:     [x==42, x==64, x==128]
    #   values:        [blah,  snarf,    yep]
    #   default_value: burp
    self.node_to_ir[node] = self.fb.add_match_true(arm_selectors, arm_values,
                                                   self._use(default_arm.expr))
    self.last_expression = node

  def visit_Unop(self, node: ast.Unop):
    kind = node.operator.get_kind_or_keyword()
    if kind == ast.Unop.NEG:
      self._def(node, self.fb.add_neg, self._use(node.operand))
    elif kind == ast.Unop.INV:
      self._def(node, self.fb.add_not, self._use(node.operand))
    else:
      raise NotImplementedError(kind)

  def _visit_width_slice(self, node: ast.Index, width_slice: ast.WidthSlice,
                         lhs_type: ConcreteType) -> None:
    width_slice.start.accept(self)
    x = self.fb.add_shrl(self._use(node.lhs), self._use(width_slice.start))
    self._def(node, self.fb.add_bit_slice, x, 0,
              self._resolve_type(node).get_total_bit_count())

  def visit_Attr(self, node: ast.Attr) -> None:
    lhs_type = self.node_to_type[node.lhs]
    index = lhs_type.tuple_names.index(node.attr.identifier)
    self._def(node, self.fb.add_tuple_index, self._use(node.lhs), index)

  def visit_Index(self, node: ast.Index) -> None:
    lhs_type = self.node_to_type[node.lhs]
    if isinstance(lhs_type, TupleType):
      self._def(node, self.fb.add_tuple_index, self._use(node.lhs),
                self._get_const(node.index))
    elif isinstance(lhs_type, BitsType) and not isinstance(lhs_type, ArrayType):
      index_slice = node.index
      if isinstance(index_slice, ast.WidthSlice):
        return self._visit_width_slice(node, index_slice, lhs_type)
      assert isinstance(index_slice, ast.Slice), index_slice

      def extract_maybe_number(n: Optional[ast.Number]) -> Optional[int]:
        if n is None:
          return None
        return n.get_value_as_int()

      start = extract_maybe_number(index_slice.start)
      limit = extract_maybe_number(index_slice.limit)
      start, width = bit_helpers.resolve_bit_slice_indices(
          lhs_type.get_total_bit_count(), start, limit)
      self._def(node, self.fb.add_bit_slice, self._use(node.lhs), start, width)
    else:
      self._def(node, self.fb.add_array_index, self._use(node.lhs),
                self._use(node.index))

  def visit_Number(self, node: ast.Number):
    type_ = self._resolve_type(node)
    self._def_const(node, node.get_value_as_int(), type_.get_total_bit_count())

  @ast.AstVisitor.no_auto_traverse
  def visit_Constant(self, node: ast.Constant) -> None:
    node.value.accept(self)
    self._def_alias(node.value, to=node.name)

  @ast.AstVisitor.no_auto_traverse
  def visit_Array(self, node: ast.Array) -> None:
    members = []
    for member in node.members:
      member.accept(self)
      members.append(self._use(member))
    self._def(node, self.fb.add_array, members, members[0].get_type())

  # Note: need to traverse to define constants for members.
  def visit_ConstantArray(self, node: ast.ConstantArray) -> None:
    array_type = self._resolve_type(node)
    e_type = array_type.get_element_type()
    values = []
    for n in node.members:
      e = self._get_const(n)
      values.append(
          ir_value.Value(_int_to_bits(e, e_type.get_total_bit_count())))
    if node.has_ellipsis:
      while len(values) < array_type.size:
        values.append(values[-1])
    self._def(node, self.fb.add_literal_value,
              ir_value.Value.make_array(values))

  def _cast_to_array(self, node: ast.Cast, output_type: ConcreteType) -> None:
    bits = self._use(node.expr)
    slices = []
    element_bit_count = output_type.get_element_type().get_total_bit_count()
    # MSb becomes lowest-indexed array element.
    for i in range(0, output_type.get_total_bit_count(), element_bit_count):
      slices.append(self.fb.add_bit_slice(bits, i, element_bit_count))
    slices.reverse()
    element_type = self.package.get_bits_type(element_bit_count)
    self._def(node, self.fb.add_array, slices, element_type)

  def _cast_from_array(self, node: ast.Cast, output_type: ConcreteType) -> None:
    array = self._use(node.expr)
    array_type = self._resolve_type_to_ir(node.expr)
    pieces = []
    for i in range(array_type.get_size()):
      pieces.append(
          self.fb.add_array_index(
              array, self.fb.add_literal_bits(bits_mod.UBits(i, 32))))
    self._def(node, self.fb.add_concat, pieces)

  @ast.AstVisitor.no_auto_traverse
  def visit_Cast(self, node: ast.Cast) -> None:
    node.expr.accept(self)
    output_type = self._resolve_type(node)
    if isinstance(output_type, ArrayType):
      return self._cast_to_array(node, output_type)
    if not (isinstance(output_type, BitsType) or
            isinstance(output_type, EnumType)):
      raise NotImplementedError(
          'Cast can only handle bits output types; got: '
          f'{output_type} @ {node.span} ({output_type!r})')
    input_type = self._resolve_type(node.expr)
    if isinstance(input_type, ArrayType):
      return self._cast_from_array(node, output_type)
    new_bit_count = output_type.get_total_bit_count()
    input_type = self._resolve_type(node.expr)
    if new_bit_count < input_type.get_total_bit_count():
      self._def(node, self.fb.add_bit_slice, self._use(node.expr), 0,
                new_bit_count)
    else:
      signed_input = input_type.get_signedness()
      f = self.fb.add_signext if signed_input else self.fb.add_zeroext
      self._def(node, f, self._use(node.expr), new_bit_count)

  @ast.AstVisitor.no_auto_traverse
  def visit_XlsTuple(self, node: ast.XlsTuple) -> None:
    for o in node.members:
      o.accept(self)
    operands = tuple(self._use(o) for o in node.members)
    self._def(node, self.fb.add_tuple, operands)

  def _deref_struct_or_enum(
      self, node: Union[ast.Struct, ast.TypeDef, ast.Enum, ast.ModRef]
  ) -> Union[ast.Struct, ast.Enum]:
    while isinstance(node, ast.TypeDef):
      annotation = node.type_
      if not annotation.is_typeref():
        raise NotImplementedError(
            'Unhandled typedef for resolving to struct-or-enum: %s' %
            annotation)
      node = annotation.typeref.type_def

    if isinstance(node, (ast.Struct, ast.Enum)):
      return node

    assert isinstance(node, ast.ModRef), node
    imported_mod, _ = self.node_to_type.get_imports()[node.mod]
    td = imported_mod.get_typedef_by_name()[node.value_tok.value]
    # Recurse to resolve the typedef within the imported module.
    td = self._deref_struct_or_enum(td)
    assert isinstance(td, (ast.Struct, ast.Enum)), td
    return td

  def _deref_struct(self, node: Union[ast.Struct, ast.ModRef]) -> ast.Struct:
    result = self._deref_struct_or_enum(node)
    assert isinstance(result, ast.Struct), result
    return result

  def _deref_enum(self, node: Union[ast.Enum, ast.ModRef]) -> ast.Enum:
    result = self._deref_struct_or_enum(node)
    assert isinstance(result, ast.Enum), result
    return result

  @ast.AstVisitor.no_auto_traverse
  def visit_SplatStructInstance(self, node: ast.SplatStructInstance) -> None:
    node.splatted.accept(self)
    orig = self._use(node.splatted)
    updates = {}
    for k, e in node.members:
      e.accept(self)
      updates[k] = self._use(e)
    struct = self._deref_struct(node.struct)

    members = []
    for i, k in enumerate(struct.member_names):
      if k in updates:
        members.append(updates[k])
      else:
        members.append(self.fb.add_tuple_index(orig, i))
    self._def(node, self.fb.add_tuple, members)

  @ast.AstVisitor.no_auto_traverse
  def visit_StructInstance(self, node: ast.StructInstance) -> None:
    operands = []
    struct = self._deref_struct(node.struct)
    for _, m in node.get_ordered_members(struct):
      m.accept(self)
      operands.append(self._use(m))
    operands = tuple(operands)
    self._def(node, self.fb.add_tuple, operands)

  def _is_constant_zero(self, node: ast.AstNode) -> bool:
    return isinstance(node, ast.Number) and node.get_value_as_int() == 0

  @ast.AstVisitor.no_auto_traverse
  def visit_For(self, node: ast.For) -> None:
    node.init.accept(self)

    def query_const_range_call() -> int:
      """Returns trip count if this is a `for ... in range(CONST)` construct."""
      range_callee = (
          isinstance(node.iterable, ast.Invocation) and
          isinstance(node.iterable.callee, ast.NameRef) and
          node.iterable.callee.identifier == 'range')
      if not range_callee:
        raise ConversionError(
            'For-loop is of an unsupported form for IR conversion; only a '
            "'range(0, const)' call is supported, found non-range callee.",
            node.span)
      if len(node.iterable.args) != 2:
        raise ConversionError(
            'For-loop is of an unsupported form for IR conversion; only a '
            "'range(0, const)' call is supported, found inappropriate number "
            'of arguments.', node.span)
      if not self._is_constant_zero(node.iterable.args[0]):
        raise ConversionError(
            'For-loop is of an unsupported form for IR conversion; only a '
            "'range(0, const)' call is supported, found inappropriate number "
            'of arguments.', node.span)
      arg = node.iterable.args[1]
      arg.accept(self)
      if not self._is_const(arg):
        raise ConversionError(
            'For-loop is of an unsupported form for IR conversion; only a '
            "'range(const)' call is supported, did not find a const value "
            f'for {arg} ({arg!r}).', node.span)
      return self._get_const(arg)

    # TODO(leary): We currently only support counted loops of the form:
    #
    #   for (i, ...): (u32, ...) in range(N) {
    #      ...
    #   }
    trip_count = query_const_range_call()

    logging.vlog(3, 'Converting for-loop @ %s', node.span)
    body_converter = _IrConverterFb(
        self.package,
        self.module,
        self.node_to_type,
        emit_positions=self.emit_positions,
        dslx_name=self.dslx_name)
    body_converter.symbolic_bindings = dict(self.symbolic_bindings)
    body_fn_name = ('__' + self.fb.name + '_counted_for_{}_body').format(
        self._next_counted_for_ordinal()).replace('.', '_')
    body_converter.fb = function_builder.FunctionBuilder(
        body_fn_name, self.package)
    flat = node.names.flatten1()
    assert len(
        flat
    ) == 2, 'Expect an induction binding and loop carry binding; got {!r}'.format(
        flat)

    # Add the induction value.
    assert isinstance(
        flat[0], ast.NameDef
    ), 'Induction variable was not a NameDef: {0} ({0!r})'.format(flat[0])
    body_converter.node_to_ir[flat[0]] = body_converter.fb.add_param(
        flat[0].identifier.encode('utf-8'), self._resolve_type_to_ir(flat[0]))

    # Add the loop carry value.
    if isinstance(flat[1], ast.NameDef):
      body_converter.node_to_ir[flat[1]] = body_converter.fb.add_param(
          flat[1].identifier.encode('utf-8'), self._resolve_type_to_ir(flat[1]))
    else:
      # For tuple loop carries we have to destructure names on entry.
      carry_type = self._resolve_type_to_ir(flat[1])
      carry = body_converter.node_to_ir[flat[1]] = body_converter.fb.add_param(
          '__loop_carry', carry_type)
      body_converter._visit_matcher(  # pylint: disable=protected-access
          flat[1], (), carry, self._resolve_type(flat[1]))

    # Free variables are suffixes on the function parameters.
    freevars = node.body.get_free_variables(node.span.start)
    freevars = freevars.drop_defs(lambda x: isinstance(x, ast.BuiltinNameDef))
    for name_def in freevars.get_name_defs():
      type_ = self.node_to_type[name_def]
      if isinstance(type_, FunctionType):
        continue
      logging.vlog(3, 'Converting freevar name: %s', name_def)
      body_converter.node_to_ir[name_def] = body_converter.fb.add_param(
          name_def.identifier.encode('utf-8'),
          self._resolve_type_to_ir(name_def))

    node.body.accept(body_converter)
    body_function = body_converter.fb.build()
    logging.vlog(3, 'Converted body function: %s', body_function.name)

    stride = 1
    invariant_args = tuple(
        self._use(name_def)
        for name_def in freevars.get_name_defs()
        if not isinstance(self.node_to_type[name_def], FunctionType))
    self._def(node, self.fb.add_counted_for, self._use(node.init), trip_count,
              stride, body_function, invariant_args)

  def _get_mangled_name(self, function_name: Text, free_keys: Set[Text],
                        m: ast.Module,
                        symbolic_bindings: Optional[SymbolicBindings]) -> Text:
    """Returns mangled name of function 'name' w/ given parametric bindings."""
    symbolic_binding_keys = set(k for k, _ in symbolic_bindings or ())
    if free_keys > symbolic_binding_keys:
      raise ValueError(
          'Not enough symbolic bindings to convert function {!r}; need {!r} got {!r}'
          .format(function_name, free_keys, symbolic_binding_keys))
    mod_name = m.name.replace('.', '_')
    if not symbolic_bindings:
      return '__{}__{}'.format(mod_name, function_name)
    suffix = '_'.join(str(self._resolve_dim(v)) for _, v in symbolic_bindings)
    return '__{}__{}__{}'.format(mod_name, function_name, suffix)

  def _get_invocation_bindings(self,
                               invocation: ast.Invocation) -> SymbolicBindings:
    """Returns the symbolic bindings of the invocation.

    We must provide the current evaluation context (module name, function name,
    symbolic bindings) in order to retrieve the correct symbolic bindings to use
    in the invocation.

    Args:
      invocation: Invocation that the bindings are being retrieved for.

    Returns:
      The symbolic bindings for the given invocation.
    """
    # We only consider function symbolic bindings for invocations.
    # The typechecker doesn't care about module-level constants.
    module_level_constants = {
        c.name.identifier
        for c in self.module.get_constants()
        if isinstance(c.value, ast.Number)
    }
    relevant_symbolic_bindings = tuple(
        (k, v)
        for k, v in self.symbolic_bindings.items()
        if k not in module_level_constants)
    key = (self.module.name, self.dslx_name, relevant_symbolic_bindings)
    logging.vlog(2, 'Invocation %s symbolic bindings: %r key: %r', invocation,
                 invocation.symbolic_bindings, key)
    return invocation.symbolic_bindings.get(key, ())

  def _get_callee_identifier(self, node: ast.Invocation) -> Text:
    logging.vlog(3, 'Getting callee identifier for invocation: %s', node)
    if isinstance(node.callee, ast.NameRef):
      callee_name = node.callee.identifier
      m = self.module
    elif isinstance(node.callee, ast.ModRef):
      m = self.node_to_type.get_imports()[node.callee.mod][0]
      callee_name = node.callee.value_tok.value
    else:
      raise NotImplementedError('Callee not currently supported @ {}'.format(
          node.span))
    try:
      function = m.get_function(callee_name)
    except KeyError:
      # For e.g. builtins that are not in the module we just provide the name
      # directly.
      return callee_name
    if not function.is_parametric():
      return self._get_mangled_name(function.name.identifier,
                                    function.get_free_parametric_keys(), m,
                                    None)
    resolved_symbolic_bindings = self._get_invocation_bindings(node)
    logging.vlog(2, 'Node %s @ %s symbolic bindings %r', node, node.span,
                 resolved_symbolic_bindings)
    assert resolved_symbolic_bindings, node
    return self._get_mangled_name(function.name.identifier,
                                  function.get_free_parametric_keys(), m,
                                  resolved_symbolic_bindings)

  def _def_map_with_builtin(self, parent_node: ast.Invocation,
                            node: ast.NameRef, arg: ast.AstNode,
                            symbolic_bindings: SymbolicBindings) -> BValue:
    """Makes the specified builtin available to the package."""
    mangled_name = self._get_mangled_name(node.name_def.identifier, set(),
                                          self.module, symbolic_bindings)

    arg = self._use(arg)
    if mangled_name not in self.package.get_function_names():
      fb = function_builder.FunctionBuilder(mangled_name, self.package)
      param = fb.add_param('arg', arg.get_type().get_element_type())
      builtin_name = node.name_def.identifier
      assert builtin_name in dslx_builtins.UNARY_BUILTIN_NAMES, dslx_builtins.UNARY_BUILTIN_NAMES
      fbuilds = {'clz': fb.add_clz, 'ctz': fb.add_ctz}
      assert set(fbuilds.keys()) == dslx_builtins.UNARY_BUILTIN_NAMES, set(
          fbuilds.keys())
      fbuilds[builtin_name](param)
      fb.build()
    return self._def(parent_node, self.fb.add_map, arg,
                     self.package.get_function(mangled_name))

  def _visit_map(self, node: ast.Invocation) -> BValue:
    for arg in node.args[:-1]:
      arg.accept(self)
    arg = self._use(node.args[0])
    fn_node = node.args[1]

    if isinstance(fn_node, ast.NameRef):
      map_fn_name = fn_node.name_def.identifier
      if map_fn_name in dslx_builtins.PARAMETRIC_BUILTIN_NAMES:
        return self._def_map_with_builtin(node, fn_node, node.args[0],
                                          self._get_invocation_bindings(node))
      else:
        lookup_module = self.module
        fn = lookup_module.get_function(map_fn_name)
    elif isinstance(fn_node, ast.ModRef):
      map_fn_name = fn_node.value_tok.value
      imports = self.node_to_type.get_imports()
      lookup_module, _ = imports[fn_node.mod]
      fn = lookup_module.get_function(map_fn_name)
    else:
      raise NotImplementedError(
          'Unhandled function mapping: {!r}'.format(fn_node))

    node_sym_bindings = self._get_invocation_bindings(node)
    mangled_name = self._get_mangled_name(fn.name,
                                          fn.get_free_parametric_keys(),
                                          lookup_module, node_sym_bindings)

    return self._def(node, self.fb.add_map, arg,
                     self.package.get_function(mangled_name))

  def _visit_update(self, node: ast.Invocation, args: Tuple[BValue,
                                                            ...]) -> BValue:
    array, target_index, update_value = args
    return self._def(node, self.fb.add_array_update, array, target_index,
                     update_value)

  def _visit_bitwise_reduction(self, node: ast.Invocation, args: Tuple[BValue,
                                                                       ...],
                               called_name: Text) -> BValue:
    add_func = None
    if called_name == 'and_reduce':
      add_func = self.fb.add_and_reduce
    elif called_name == 'or_reduce':
      add_func = self.fb.add_or_reduce
    elif called_name == 'xor_reduce':
      add_func = self.fb.add_xor_reduce
    else:
      raise NotImplementedError(
          'Unknown bitwise reduction: {}'.format(called_name))
    return self._def(node, add_func, args[0])

  def _visit_bit_slice(self, node: ast.Invocation, args: Tuple[BValue,
                                                               ...]) -> BValue:
    lhs, _, width = args
    width_type = width.get_type()
    return self._def(node, self.fb.add_bit_slice, lhs,
                     self._get_const(node.args[1]), width_type.get_bit_count())

  def _visit_rev(self, node: ast.Invocation, args: Tuple[BValue,
                                                         ...]) -> BValue:
    arg, = args
    return self._def(node, self.fb.add_reverse, arg)

  def _visit_signex(self, node: ast.Invocation, args: Tuple[BValue,
                                                            ...]) -> BValue:
    lhs, rhs = args
    rhs_type = rhs.get_type()
    return self._def(node, self.fb.add_signext, lhs, rhs_type.get_bit_count())

  def _visit_clz(self, node: ast.Invocation, args: Tuple[BValue,
                                                         ...]) -> BValue:
    lhs, = args
    return self._def(node, self.fb.add_clz, lhs)

  def _visit_ctz(self, node: ast.Invocation, args: Tuple[BValue,
                                                         ...]) -> BValue:
    lhs, = args
    return self._def(node, self.fb.add_ctz, lhs)

  def _visit_one_hot(self, node: ast.Invocation, args: Tuple[BValue,
                                                             ...]) -> BValue:
    lhs, _ = args
    lsb_prio = self._get_const(node.args[1])
    return self._def(
        node, self.fb.add_one_hot, lhs,
        lsb_or_msb.LsbOrMsb.LSB if lsb_prio else lsb_or_msb.LsbOrMsb.MSB)

  def _visit_one_hot_sel(self, node: ast.Invocation,
                         args: Tuple[BValue, ...]) -> BValue:
    lhs, array = args
    array_type = array.get_type()
    rhs_elements = []
    for i in range(array_type.get_size()):
      rhs_elements.append(
          self.fb.add_array_index(
              array, self.fb.add_literal_bits(bits_mod.UBits(i, 32))))
    return self._def(node, self.fb.add_one_hot_sel, lhs, rhs_elements)

  def _visit_scmp(self, node: ast.Invocation, args: Tuple[BValue, ...],
                  which: Text) -> BValue:
    lhs, rhs = args
    return self._def(node, getattr(self.fb, 'add_{}'.format(which)), lhs, rhs)

  def _visit_trace(self, node: ast.Invocation, args: Tuple[BValue,
                                                           ...]) -> BValue:
    return self._def(node, self.fb.add_identity, args[0])

  @ast.AstVisitor.no_auto_traverse
  def visit_Invocation(self, node: ast.Invocation):
    called_name = self._get_callee_identifier(node)

    def accept_args() -> Tuple[BValue, ...]:
      for arg in node.args:
        arg.accept(self)
      return tuple(self._use(arg) for arg in node.args)

    if called_name == 'fail!':
      args = accept_args()
      assert len(node.args) == 1, ('Expect fail! builtin to only accept a '
                                   'single argument; got: {!r}'.format(args))
      self._def(node, self.fb.add_identity, args[0])
    elif called_name == 'update':
      self._visit_update(node, accept_args())
    elif called_name == 'signex':
      self._visit_signex(node, accept_args())
    elif called_name == 'clz':
      self._visit_clz(node, accept_args())
    elif called_name == 'ctz':
      self._visit_ctz(node, accept_args())
    elif called_name == 'map':
      # Map needs special arg handling, so we handle that inside.
      self._visit_map(node)
    elif called_name == 'one_hot':
      self._visit_one_hot(node, accept_args())
    elif called_name == 'one_hot_sel':
      self._visit_one_hot_sel(node, accept_args())
    elif called_name == 'bit_slice':
      self._visit_bit_slice(node, accept_args())
    elif called_name == 'rev':
      self._visit_rev(node, accept_args())
    elif called_name in ('sgt', 'sge', 'slt', 'sle'):
      self._visit_scmp(node, accept_args(), called_name)
    elif called_name in ('and_reduce', 'or_reduce', 'xor_reduce'):
      self._visit_bitwise_reduction(node, accept_args(), called_name)
    elif called_name == 'trace':
      self._visit_trace(node, accept_args())
    else:
      try:
        f = self.package.get_function(called_name)
      except Exception as e:
        # TODO(leary): Switch to new pybind11 more-specific exception.
        raise ConversionError(
            'Failed to get function from invocation: {}'.format(e), node.span)
      self._def(node, self.fb.add_invoke, accept_args(), f)

  def visit_ConstRef(self, node: ast.ConstRef) -> None:
    self._def_alias(node.name_def, to=node)

  def visit_NameRef(self, node: ast.NameRef) -> None:
    self._def_alias(node.name_def, node)

  def visit_EnumRef(self, node: ast.EnumRef) -> None:
    enum = self._deref_enum(node.enum)
    value = enum.get_value(node.value_tok.value)
    value.accept(self)
    self._def_alias(from_=value, to=node)

  @ast.AstVisitor.no_auto_traverse
  def visit_Let(self, node: ast.Let):
    node.rhs.accept(self)
    if node.name_def_tree.is_leaf():
      self._def_alias(node.rhs, to=node.name_def_tree.get_leaf())
      node.body.accept(self)
      self._def_alias(node.body, node)
    else:
      # Walk the tree performing tuple_index operations to get to the binding
      # levels desired.

      names = [self._use(node.rhs)]  # List[BValue]

      def walk(x: ast.NameDefTree, level: int, index: int) -> None:
        """Invoked at each level of the name def tree.

        Binds the name in the name def tree to the corresponding value being
        pattern matched.

        Args:
          x: The current level of the NameDefTree.
          level: Level in the NameDefTree (root is 0).
          index: Index of node in the current tree level (e.g. leftmost is 0).
        """
        del names[level:]
        names.append(
            self._def(
                x,
                self.fb.add_tuple_index,
                names[-1],
                index,
                span=(x.get_leaf().span if x.is_leaf() else x.span)))
        if x.is_leaf():
          self._def_alias(x, x.get_leaf())

      node.name_def_tree.do_preorder(walk)
      node.body.accept(self)
      self._def_alias(node.body, to=node)

    if self.last_expression is None:
      self.last_expression = node.body

  @ast.AstVisitor.no_auto_traverse
  def visit_Param(self, node: ast.Param):
    self._def(node.name, self.fb.add_param, node.name.identifier,
              self._resolve_type_to_ir(node.type_))

  def _visit_Function(
      self, node: ast.Function,
      symbolic_bindings: Optional[SymbolicBindings]) -> ir_function.Function:
    self.symbolic_bindings = {} if symbolic_bindings is None else dict(
        symbolic_bindings)
    self._extract_module_level_constants(self.module)
    # We use a function builder for the duration of converting this
    # ast.Function. When it's done being built, we drop the reference to it (by
    # setting self.fb to None).
    self.fb = function_builder.FunctionBuilder(
        self._get_mangled_name(node.name.identifier,
                               node.get_free_parametric_keys(), self.module,
                               symbolic_bindings), self.package)
    try:
      for param in node.params:
        param.accept(self)

      for parametric_binding in node.parametric_bindings:
        logging.vlog(4, 'Resolving parametric binding %s', parametric_binding)

        sb_value = self.symbolic_bindings[parametric_binding.name.identifier]
        value = self._resolve_dim(sb_value)
        assert isinstance(value, int), \
            'Expect integral parametric binding; got {!r}'.format(value)
        self._def_const(
            parametric_binding, value,
            self._resolve_type(parametric_binding.type_).get_total_bit_count())
        self._def_alias(parametric_binding, to=parametric_binding.name)

      for dep in self._constant_deps:
        dep.accept(self)
      del self._constant_deps[:]

      node.body.accept(self)

      last_expression = self.last_expression or node.body
      if isinstance(last_expression, ast.NameRef):
        self._def(last_expression, self.fb.add_identity,
                  self._use(last_expression))
      f = self.fb.build()
      logging.vlog(3, 'Built function: %s', f.name)
      verifier_mod.verify_function(f)
      return f
    finally:
      self.fb = None

  @ast.AstVisitor.no_auto_traverse
  def visit_Function(
      self,
      node: ast.Function,
      symbolic_bindings: Optional[SymbolicBindings] = None,
      module: Optional[ast.Module] = None) -> ir_function.Function:
    return self._visit_Function(node, symbolic_bindings)

  def get_text(self) -> Text:
    return self.package.dump_ir()

def _convert_one_function(package: ir_package.Package,
                          module: ast.Module,
                          function: ast.Function,
                          node_to_type: deduce.NodeToType,
                          symbolic_bindings: Optional[SymbolicBindings] = None,
                          emit_positions: bool = True) -> Text:
  """Converts a single function into its emitted text form.

  Args:
    package: IR package we're converting the function into.
    module: Module we're converting a function within.
    function: Function we're converting.
    node_to_type: Type information about module from the typechecking phase.
    symbolic_bindings: Parametric bindings to use during conversion, if this
      function is parametric.
    emit_positions: Whether to emit position information into the IR based on
      the AST's source positions.

  Returns:
    The converted IR function text.
  """
  function_by_name = module.get_function_by_name()
  constant_by_name = module.get_constant_by_name()
  converter = _IrConverterFb(
      package,
      module,
      node_to_type,
      emit_positions=emit_positions,
      dslx_name=function.name.identifier)

  freevars = function.body.get_free_variables(
      function.span.start).get_name_def_tups()
  logging.vlog(3, 'Free variables for function: %s', freevars)
  for identifier, name_def in freevars:
    if identifier in function_by_name or isinstance(name_def,
                                                    ast.BuiltinNameDef):
      pass
    elif identifier in constant_by_name:
      converter.add_constant_dep(constant_by_name[identifier])
    else:
      raise NotImplementedError(identifier)

  symbolic_binding_keys = set(k for k, _ in symbolic_bindings or ())
  f_parametric_keys = function.get_free_parametric_keys()
  if f_parametric_keys > symbolic_binding_keys:
    raise ValueError(
        'Not enough symbolic bindings to convert function {!r}; need {!r} got {!r}'
        .format(function.name.identifier, f_parametric_keys,
                symbolic_binding_keys))

  logging.vlog(3, 'Converting function: %s; symbolic bindings: %s', function,
               symbolic_bindings)
  f = converter.visit_Function(function, symbolic_bindings)
  return f.dump_ir(recursive=False)

def convert_module_to_package(
    module: ast.Module,
    node_to_type: deduce.NodeToType,
    emit_positions: bool = True) -> ir_package.Package:
  """Converts the contents of a module to IR form.

  Args:
    module: Module to convert.
    node_to_type: Concrete type information used in conversion.
    emit_positions: Whether to emit positional metadata into the output IR.

  Returns:
    The IR package that corresponds to this module.
  """
  emitted = []  # type: List[Text]
  package = ir_package.Package(module.name)
  order = extract_conversion_order.get_order(module, node_to_type.get_imports())
  logging.vlog(3, 'Convert order: %s', pprint.pformat(order))
  for record in order:
    emitted.append(
        _convert_one_function(
            package,
            record.m,
            record.f,
            node_to_type,
            symbolic_bindings=record.bindings,
            emit_positions=emit_positions))

  verifier_mod.verify_package(package)
  return package

def convert_module(module: ast.Module,
                   node_to_type: deduce.NodeToType,
                   emit_positions: bool = True) -> Text:
  """Same as convert_module_to_package, but converts to IR text."""
  return convert_module_to_package(module, node_to_type,
                                   emit_positions).dump_ir()

def convert_one_function(module: ast.Module,
                         entry_function_name: Text,
                         node_to_type: deduce.NodeToType,
                         emit_positions: bool = True) -> Text:
  """Returns function named entry_function_name in module as IR text."""
  package = ir_package.Package(module.name)
  _convert_one_function(
      package,
      module,
      module.get_function(entry_function_name),
      node_to_type,
      emit_positions=emit_positions)
  return package.dump_ir()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author：Janet Chou
with open('data.txt','a') as f1:
    with open('file_list_20170430_new.txt', 'r') as f:
        flag_list=list()
        count_d=0
        for i in f:
            flag = i.split(',', 4)[1]
            MD5 = i.split(',', 4)[2]
            if flag == 'n' and count_d<582:
                flag_list.append(1)
                count_d+=1
                f1.write(i)


# -*- coding: utf-8 -*-
"""
Created on Wed Apr 20 2022
Modified Wed Jan 4 2023: cleaning up

@author: abrielmann

"""
import os
import pandas as pd
pd.options.mode.chained_assignment = None  # disable SettingWithCopyWarning

# %% ---------------------------------------------------------
# Specify directories; settings
# ------------------------------------------------------------
os.chdir('..')
homeDir = os.getcwd()
# IDs (not including excluded subjs):
participantList = ['q22fa', 'eunhf', 'fax28', '1z0ca', '0xmq6',
                   'z19vf', 'pu7uh', 's4tsr', 'n0htt', 'zouzq', 'nqfzn',
                   'oveao', 'mh0hn', 'mbx6w', '46vpz', 'sdx2b', 'jmzs9',
                   'jeu26', 'hvjjw', 'xj6bc', '3l6s4', 'chb6q', 'kwqcc',
                   'czgyu', 'nyp52', 'pshnw', 'faow1', 'kuoa6', '7df33',
                   'mrx1q', 'quzfn', 're1sg', 'j49v1', '19bq6', 'r4ron',
                   '8kyd7', 'nq3sp', 'rsw55', '37o95', 'wtr4q', 'ar5vx',
                   '6yaxg', 'jv7jg', 'ms8r4', '520tj', '9gjbe',
                   'njzz4', 'rq280', 'tg4o2', 'fpt0a', 'k02z1', 'u0pc9',
                   'hxal4', 'h647a', 'r1g8p', '89w87', 'zv0ou', 'qrbxe',
                   '6r8qk']
participantList.sort()
dataDir = homeDir + '/'
plot = False
plotIndividuals = True  # plots for each participant's rating distribution

# %% ---------------------------------------------------------
# Load data
# ------------------------------------------------------------
rawDf = pd.read_csv(dataDir + 'merged_participantInfo.csv')
df = rawDf[rawDf['subj'].isin(participantList)]

# %% ---------------------------------------------------------
# demographics; use print() to display
# ------------------------------------------------------------
df.age.describe()
df.education.value_counts()
df.nationality.value_counts()
df.dogPresent.value_counts()
df.dogPast.value_counts()
df.imgBrowseDur.value_counts()
df.imgBrowseLike.value_counts()
df.dogLiking.value_counts()
df.boredom.value_counts()

# %%---------------------------------------------------------
# related to rating behavior; use print() to display
# ------------------------------------------------------------
df.corrRatingSliderPos.describe()
df.ratingCorrWithAvg.describe()


# -*- coding: utf-8 -*-
# __author__ = 'zouxuan'
from biocluster.agent import Agent
from biocluster.tool import Tool
import os
import subprocess
from biocluster.core.exceptions import OptionError
import re  # add re by guhaidong 20171025
from mbio.packages.taxon.mask_taxon import mask_taxon  # add by guhaidong 20171025

class PermanovaAgent(Agent):
    """
    脚本permanova.pl
    version v1.0
    author: zouxuan
    last_modified:2017.10.19
    """

    def __init__(self, parent):
        super(PermanovaAgent, self).__init__(parent)
        options = [
            {"name": "abu_table", "type": "infile",
             "format": "meta.otu.otu_table, meta.otu.tax_summary_dir, toolapps.table"},
            # modify by zhouxuan 20170623 小工具的模块是指定toolapps.table这个文件类型的
            # {"name": "level", "type": "string", "default": "otu"},
            {"name": "envtable", "type": "infile", "format": "meta.otu.group_table"},
            {"name": "permutations", "type": "int", "default": 999},
            {"name": "dis_method", "type": "string", "default": "bray"},
            {"name": "binary", "type": "string", "default": "false"}
        ]
        self.add_option(options)

    def check_options(self):
        """
        重写参数检查
        """
        if not self.option('abu_table').is_set:
            raise OptionError('必须提供丰度表', code="32703101")
        if not self.option('envtable').is_set:
            raise OptionError('必须提供环境因子表', code="32703102")

    def set_resource(self):
        """
        设置所需资源
        """
        self._cpu = 2
        self._memory = '3G'

    def end(self):
        result_dir = self.add_upload_dir(self.output_dir)
        result_dir.add_relpath_rules([
        ])
        # print self.get_upload_files()
        super(PermanovaAgent, self).end()

class PermanovaTool(Tool):
    def __init__(self, config):
        super(PermanovaTool, self).__init__(config)
        self._version = '1.0'
        self.perl_path = "program/perl-5.24.0/bin/perl"
        self.script1_path = self.config.PACKAGE_DIR + "/beta_diversity/permanova.pl"
        self.R_path = 'program/R-3.3.1/bin/Rscript'

    def run_permanova(self):
        otu_table = self.option("abu_table").prop['path']
        self.name_to_name = mask_taxon(otu_table, self.work_dir + "/tmp_mask_otu.xls")  # add by guhaidong 20171025
        otu_table = self.work_dir + '/tmp_mask_otu.xls'  # add by guhaidong 20171025
        cmd = self.perl_path
        cmd += ' %s -i %s -env %s -o %s  -pe %s -d_method %s -binary %s' % (
            self.script1_path, otu_table, self.option("envtable").prop['path'],
            self.work_dir + '/permanova.xls', self.option("permutations"), self.option("dis_method"),
            self.option("binary"))
        self.logger.info(cmd)
        command = self.add_command('cmd', cmd)
        command.run()
        self.wait(command)
        if command.return_code == 0:
            self.logger.info("permanova的r文件生成成功")
        else:
            self.set_error("permanova的r文件生成失败", code="32703101")
            raise Exception("permanova的r文件生成失败")
        cmd_1 = '%s %s' % (self.R_path, self.work_dir + "/cmd.r")
        self.logger.info(cmd_1)
        command1 = self.add_command('cmd_1', cmd_1)
        command1.run()
        self.wait(command1)
        if command1.return_code == 0:
            self.logger.info("R程序计算permanova成功")
        else:
            self.set_error("R程序计算permanova失败", code="32703102")
            raise Exception("R程序计算permanova失败")

    def dashrepl(self, matchobj):
        """
        add func by guhaidong 20171025
        """
        return self.name_to_name[matchobj.groups()[0]]

    def add_taxon(self, old_result, taxon_result):
        """
        add func by guhaidong 20171025
        description: 将旧注释的名称，根据词典替换成新注释名称
        """
        with open(old_result, "r") as f, open(taxon_result, "w") as w:
            # w.write(old_result)
            for i in f.readlines():
                # line = i.strip()
                new_line = re.sub(r"(name\d+)", self.dashrepl, i)
                w.write(new_line)

    def run(self):
        super(PermanovaTool, self).run()
        self.run_permanova()
        self.set_output()

    def set_output(self):
        linksites = os.path.join(self.output_dir, 'permanova.xls')
        if os.path.exists(linksites):
            os.remove(linksites)
        os.link(self.work_dir + '/permanova.xls', linksites)
        self.end()


import csv
from time import time
from lzss import inflate, deflate

pairs = [
    #['corpus/customer.xml', 'results/lzss-customer.csv'],
    ['corpus/bible.txt',    'results/lzss-bible.csv'],
    ['corpus/image.jpg',    'results/lzss-image.csv'],
    #['corpus/zeros.txt',    'results/lzss-zeros.csv'],
]

W = [2**i-1 for i in [6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]]
L = [2**i-1 for i in [8]]

for input_file, results_file in pairs:
    b = open(input_file, mode="rb").read()
    out = csv.writer(open(results_file, mode="w"))
    out.writerow(["W", "L", "Compression Ratio", "Encode Time", "Decode Time"])

    for w in W:
        for l in L:
            print("{input_file} {w} {l}".format(input_file=input_file, w=w, l=l))
            deflate_t0 = time()
            x = deflate(b, W=w, L=l)
            deflate_t1 = time()
            inflate_t0 = time()
            inflate(x, W=w, L=l)
            inflate_t1 = time()
            x.fill()
            out.writerow([
                w,
                l,
                (len(x) // 8) / len(b),
                deflate_t1 - deflate_t0,
                inflate_t1 - inflate_t0,
            ])


import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import open3d as o3d

pcd = o3d.io.read_point_cloud("体素栅格.pcd")
print(pcd)  # 输出点云点的个数
# 计算法线，搜索半径1cm，只考虑邻域内的30个点
pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.01, max_nn=30))
print(np.asarray(pcd.normals))
np.savetxt('体素栅格法向量.txt', np.asarray(pcd.normals))
DD = np.asarray(pcd.normals)

# DD = np.loadtxt('体素栅格法向量.txt')
df2 = pd.DataFrame(DD, columns=['Nx', 'Ny', 'Nz'])
print(df2)
df2.plot(kind='kde')
plt.show()


# -*- coding: utf-8 -*-
"""
@Time : 2023/4/26 16:04
@Author : Chao Zhu
"""
import torch
import torch.nn as nn
from torch.nn.utils import weight_norm

class Chomp1d(nn.Module):
    def __init__(self, chomp_size):
        super(Chomp1d, self).__init__()
        self.chomp_size = chomp_size

    def forward(self, x):
        return x[:, :, :-self.chomp_size].contiguous()

class TemporalBlock(nn.Module):
    def __init__(self, n_inputs, n_outputs, kernel_size, stride, dilation, padding, dropout=0.2):
        super(TemporalBlock, self).__init__()
        """
        构成TCN的核心Block, 原作者在图中称为Residual block, 是因为它存在残差连接.
        但注意, 这个模块包含了2个Conv1d.
        :param n_inputs: int, 输入通道数或者特征数
        :param n_outputs: int, 输出通道数或者特征数
        :param kernel_size: int, 卷积核尺寸
        :param stride: int, 步长, 在TCN固定为1
        :param dilation: int, 膨胀系数. 与这个Residual block(或者说, 隐藏层)所在的层数有关系.
                                例如, 如果这个Residual block在第1层, dilation = 2**0 = 1;
                                      如果这个Residual block在第2层, dilation = 2**1 = 2;
                                      如果这个Residual block在第3层, dilation = 2**2 = 4;
                                      如果这个Residual block在第4层, dilation = 2**3 = 8 ......
        :param padding: int, 填充系数. 与kernel_size和dilation有关.
        :param dropout: float, dropout比率
        """
        self.conv1 = weight_norm(nn.Conv1d(n_inputs, n_outputs, kernel_size,
                                           stride=stride, padding=padding, dilation=dilation))

        # 因为 padding 的时候, 在序列的左边和右边都有填充, 所以要裁剪
        self.chomp1 = Chomp1d(padding)
        self.relu1 = nn.ReLU()
        self.dropout1 = nn.Dropout(dropout)

        self.conv2 = weight_norm(nn.Conv1d(n_outputs, n_outputs, kernel_size,
                                           stride=stride, padding=padding, dilation=dilation))
        self.chomp2 = Chomp1d(padding)
        self.relu2 = nn.ReLU()
        self.dropout2 = nn.Dropout(dropout)

        self.net = nn.Sequential(self.conv1, self.chomp1, self.relu1, self.dropout1,
                                 self.conv2, self.chomp2, self.relu2, self.dropout2)

        # 1×1的卷积. 只有在进入Residual block的通道数与出Residual block的通道数不一样时使用.
        # 一般都会不一样, 除非num_channels这个里面的数, 与num_inputs相等. 例如[5,5,5], 并且num_inputs也是5
        self.downsample = nn.Conv1d(n_inputs, n_outputs, 1) if n_inputs != n_outputs else None

        # 在整个Residual block中有非线性的激活. 这个容易忽略!
        self.relu = nn.ReLU()
        self.init_weights()

    def init_weights(self):
        self.conv1.weight.data.normal_(0, 0.01)
        self.conv2.weight.data.normal_(0, 0.01)
        if self.downsample is not None:
            self.downsample.weight.data.normal_(0, 0.01)

    def forward(self, x):
        out = self.net(x)
        res = x if self.downsample is None else self.downsample(x)
        return self.relu(out + res)

class TemporalConvNet(nn.Module):
    def __init__(self, num_inputs, num_channels, kernel_size=2, dropout=0.2):
        super(TemporalConvNet, self).__init__()
        """
            :param num_inputs: int,  输入通道数或者特征数
            :param num_channels: list, 每层的hidden_channel数. 例如[5,12,3], 代表有3个block,
                                    block1的输出channel数量为5;
                                    block2的输出channel数量为12;
                                    block3的输出channel数量为3.
            :param kernel_size: int, 卷积核尺寸
            :param dropout: float, drop_out比率
        """
        layers = []
        num_levels = len(num_channels)
        # 可见，如果num_channels=[5,12,3]，那么
        # block1的dilation_size=1
        # block2的dilation_size=2
        # block3的dilation_size=4
        for i in range(num_levels):
            dilation_size = 2 ** i
            in_channels = num_inputs if i == 0 else num_channels[i-1]
            out_channels = num_channels[i]
            layers += [TemporalBlock(in_channels, out_channels, kernel_size, stride=1, dilation=dilation_size,
                                     padding=(kernel_size-1) * dilation_size, dropout=dropout)]

        self.network = nn.Sequential(*layers)
        self.fc = nn.Linear(12, 1)

    def forward(self, batch, device):
        src = batch['park_x'].to(device)  # [batch_size, length, channels]
        src = src.permute(0, 2, 1)  # [batch_size, channels, length]
        out_last = self.network(src)  # [batch_size, 1, length]
        out = self.fc(out_last)  # [batch_size, 1, 1]
        return out


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Apr 13 11:34:20 2022

@author: chris
"""
import numpy as np
from scipy.interpolate import interp1d

from model.eddington import ERDF
from model.helper import make_array, invert_function

def physics_model(physics_name, log_m_c, initial_guess, bounds,
                  fixed_m_c=True, eddington_erdf_params=None):
    '''
    Return physics model that relates SMF and HMF, including model function,
    model name, initial guess and physical parameter bounds for fitting,
    that related SMF and HMF.
    The model function parameters are left free and can be obtained via fitting.
    Three models implemented:
        none                : no feedback adjustment
        sn                  : supernova feedback
        both                : supernova and black hole feedback
        quasar              : black hole growth model
        eddington_free_ERDF : bolometric luminosity model using free ERDF
        eddington           : bolometric luminosity model using fixed ERDF
                              (need erdf parameter as additional input)
    '''
    if physics_name not in ['none', 'stellar', 'stellar_blackhole', 'quasar',
                            'eddington_free_m_c_free_ERDF',
                            'eddington_free_ERDF', 'eddington']:
        raise NameError('physics_name not known.')

    if physics_name == 'none':
        return(NoFeedback(log_m_c, initial_guess[:1],
                          bounds[:, :1]))

    if physics_name == 'eddington':
        if eddington_erdf_params is None:
            raise ValueError('Eddington model with fixed ERDF needs ERDF '
                             'parameter passed via eddington_erdf_params '
                             'argument.')
        return(QuasarLuminosity(log_m_c, initial_guess[:2], bounds[:, :2],
               eddington_erdf_params))

    if fixed_m_c:
        if physics_name == 'stellar':
            physics = StellarFeedback(log_m_c, initial_guess[:-1],
                                      bounds[:, :-1])
        elif physics_name == 'stellar_blackhole':
            physics = StellarBlackholeFeedback(log_m_c, initial_guess,
                                               bounds)
        elif physics_name == 'quasar':
            physics = QuasarGrowth(log_m_c, initial_guess,
                                   bounds)
        elif physics_name == 'eddington_free_ERDF':
            physics = QuasarLuminosity_free_ERDF(log_m_c, initial_guess,
                                                 bounds)
        elif physics_name == 'eddington':
            if eddington_erdf_params is None:
                raise ValueError('Eddington model with fixed ERDF needs ERDF '
                                 'parameter passed via eddington_erdf_params '
                                 'argument.')
            physics = QuasarLuminosity(log_m_c, initial_guess[:-2],
                                    bounds[:, :-2], eddington_erdf_params)

    else:
        # add log_m_c to initial guess and bounds
        initial_guess = np.insert(initial_guess, 0, log_m_c)
        log_m_c_range = 3  # +- range in which log_m_c is expected
        bounds = np.insert(bounds, 0, [10, log_m_c+log_m_c_range], axis=1)

        if physics_name == 'stellar':
            physics = StellarFeedback_free_m_c(log_m_c,
                                               initial_guess[:-1],
                                               bounds[:, :-1])
        elif physics_name == 'stellar_blackhole':
            physics = StellarBlackholeFeedback_free_m_c(log_m_c,
                                                        initial_guess,
                                                        bounds)
        elif physics_name == 'quasar':
            physics = QuasarGrowth_free_m_c(log_m_c, initial_guess,
                                            bounds)
        elif physics_name == 'eddington_free_ERDF':
            physics = QuasarLuminosity_free_m_c_free_ERDF(log_m_c,
                                                          initial_guess,
                                                          bounds)
        elif physics_name in 'eddington':
            raise NotImplementedError('physics model with free m_c but fixed '
                                      'ERDF not implemented.')
    return(physics)

################ MODEL WITH FREE CRITICAL MASS ################################

class StellarBlackholeFeedback_free_m_c(object):
    def __init__(self, log_m_c, initial_guess, bounds):
        '''
        Model for stellar + black hole feedback with free m_c.

        '''
        self.name = 'stellar_blackhole_free_m_c'
        self.log_m_c = log_m_c               # critical mass for feedback
        self.initial_guess = initial_guess   # initial guess for least_squares
        # fit
        # parameter (A, alpha, beta) bounds
        self.bounds = bounds

        # parameter used compared to the unconstrained model
        self.parameter_used = [0,1,2,3]
        self.fixed_m_c_flag = False

        # max halo mass (just used to avoid overflows)
        self._upper_m_h = 50
        # latest parameter used
        self._current_parameter = None
        self.log_m_h_function = None

    def calculate_log_quantity(self, log_m_h, log_m_c, log_A, alpha, beta):
        '''
        Calculate observable quantity from input halo mass and model parameter.
        '''
        log_m_h = make_array(log_m_h)
        log_m_h = self._check_overflow(log_m_h)

        sn, bh = self._variable_substitution(log_m_h, log_m_c, alpha, beta)
        log_quantity = np.empty_like(bh)

        # deal with bh becoming infinite sometimes
        inf_mask = np.isfinite(sn+bh)
        sn = sn[inf_mask]
        bh = bh[inf_mask]

        log_quantity[inf_mask] = log_A + log_m_h[inf_mask] - np.log10(sn + bh)
        log_quantity[np.logical_not(inf_mask)] = np.inf
        return(log_quantity)

    def calculate_log_halo_mass(self, log_quantity, log_m_c, log_A, alpha,
                                beta, num=500):
        '''
        Calculate halo mass from input quantity quantity and model parameter.
        Do this by calculating table of quantity for halo masses near critical
        mass and linearly interpolating inbetween (and extrapolating beyond
        range). Once interpolated function is created, it's saved and called
        for new calculations until new parameter are passed.
        num controls number of points that halo mass is initially calculated for,
        higher values make result more accurate but also more expensive.

        '''
        if alpha == beta == 0:
            return(log_quantity - log_A + np.log10(2))

        if self._current_parameter == [log_m_c, log_A, alpha, beta]:
            pass
        else:
            self._current_parameter = [log_m_c, log_A, alpha, beta]

            # create lookup tables of halo mass and quantities
            log_m_h_lookup = self._make_m_h_space(log_m_c, log_A, alpha,
                                                  beta, num=num)
            log_quantity_lookup = StellarBlackholeFeedback_free_m_c.\
                calculate_log_quantity(
                    self, log_m_h_lookup, log_m_c, log_A,
                    alpha, beta)

            # use lookup table to create callable spline
            self.log_m_h_function = interp1d(log_quantity_lookup,
                                             log_m_h_lookup,
                                             bounds_error=False,
                                             fill_value=([-np.inf], [np.inf]))

        log_m_h = self.log_m_h_function(log_quantity)
        log_m_h = make_array(log_m_h)
        log_m_h = self._check_overflow(log_m_h)  # deal with very large m_h
        return(log_m_h)

    def calculate_log_halo_mass_alternative(self, log_quantity, log_m_c, log_A,
                                            alpha, beta, xtol=0.1):
        '''
        Calculate halo mass from input observable quantity and model paramter.
        This is done by numerical inverting the relation for the quantity as
        function of halo mass (making this a very expensive operation).
        Absolute tolerance for inversion can be adjusted using xtol.
        This method is more accurate, but more time consuming (also has
        problems if beta is close to 1).
        '''
        log_m_h = invert_function(func=self.calculate_log_quantity,
                                  fprime=self.calculate_dlogquantity_dlogmh,
                                  fprime2=self.calculate_d2logquantity_dlogmh2,
                                  x0_func=self._initial_guess,
                                  y=log_quantity,
                                  args=(log_m_c, log_A, alpha, beta),
                                  xtol=xtol)
        return(log_m_h)

    def calculate_dlogquantity_dlogmh(self, log_m_h, log_m_c, log_A, alpha,
                                      beta):
        '''
        Calculate d/d(log m_h) log_quantity. High mass end for beta near
        one treated as special case, where value apporaches zero.
        '''
        log_m_h = make_array(log_m_h)
        log_m_h = self._check_overflow(log_m_h)

        sn, bh = self._variable_substitution(log_m_h, log_m_c, alpha, beta)
        first_derivative = np.empty_like(bh)

        # deal with bh becoming infinite sometimes
        inf_mask = np.isfinite(bh+sn)
        sn = sn[inf_mask]
        bh = bh[inf_mask]

        first_derivative[inf_mask] = 1 - (-alpha * sn + beta * bh) / (sn + bh)
        first_derivative[np.logical_not(inf_mask)] = 0
        return(first_derivative)

    def calculate_d2logquantity_dlogmh2(self, log_m_h, log_m_c, log_A, alpha,
                                        beta):
        '''
        Calculate d^2/d(log m_h)^2 log_quantity. High mass end for beta near
        one treated as special case, where value apporaches zero.
        '''
        log_m_h = make_array(log_m_h)
        log_m_h = self._check_overflow(log_m_h)

        x = 10**((alpha+beta) * (log_m_h - log_m_c))
        second_derivative = np.empty_like(x)

        # deal with bh becoming infinite sometimes
        inf_mask = np.isfinite(x)
        numerator = -np.log(10) * (alpha+beta)**2 * x[inf_mask]
        denominator = (x[inf_mask]+1)**2

        second_derivative[inf_mask] = numerator/denominator
        second_derivative[np.logical_not(inf_mask)] = 0
        return(second_derivative)

    def _variable_substitution(self, log_m_h, log_m_c, alpha, beta):
        '''
        Transform input quanitities to more easily handable quantities.
        '''
        log_m_h = self._check_overflow(log_m_h)

        ratio = log_m_h - log_m_c         # m_h/m_c in log space
        log_stellar = - alpha * ratio     # log of sn feedback contribution
        if beta == 0:
            log_black_hole = np.full_like(log_stellar, 0)
        else:
            log_black_hole = beta * ratio     # log of bh feedback contribution
        return(10**log_stellar, 10**log_black_hole)

    def _initial_guess(self, log_quantity, log_m_c, log_A, alpha, beta):
        '''
        Calculate initial guess for halo mass (for function inversion). Do this
        by calculating the high and low mass end approximation of the relation.
        '''
        # turnover quantity, where dominating feedback changes
        log_q_turn = log_A - np.log10(2) + log_m_c

        if log_quantity < log_q_turn:
            x0 = (log_quantity - log_A + alpha * log_m_c) / (1 + alpha)
        else:
            x0 = (log_quantity - log_A - beta * log_m_c) / (1 - beta)
        return(x0)

    def _check_overflow(self, log_m_h):
        '''
        Control for overflows by checking if halo mass exceeds upper limit.
        Also deal with Nans.
        '''
        if np.isnan(log_m_h).any():
            return(np.full_like(log_m_h, np.nan))
        if np.any(log_m_h > self._upper_m_h):
            if np.isscalar(log_m_h):
                log_m_h = np.inf
            else:
                log_m_h[log_m_h > self._upper_m_h] = np.inf
        return(log_m_h)

    def _calculate_feedback_regimes(self, log_m_c, log_A, alpha, beta,
                                   log_epsilon=-2, output='halo_mass'):
        '''
        Calculate halo mass where one of the feedbacks is strongly dominating
        in the sense that (M_h/M_c)^-alpha > epsilon * (M_h/M_c)^beta and
        the other way around. Returns 3 values: [log_m_c, log_m_sn, log_m_ bh].
        If output='halo_mass' (default) return transition
        halo masses, if output='quantity' return quantity values (for this
        halo mass and parameter).
        '''
        # check relative contribution of feedbacks using sn = epsilon*bh to
        # get limiting mass where either contribution drops below epsilon,
        # this is log_m_lim = -1/(alpha+beta) * log(epsilon) + log_m_c

        # mass where sn feedback is dominating
        log_m_h_sn = 1/(alpha+beta)  * log_epsilon + log_m_c
        # mass where bh feedback is dominating
        log_m_h_bh = -1/(alpha+beta) * log_epsilon + log_m_c

        if output == 'halo_mass':
            return(np.array([log_m_c, log_m_h_bh, log_m_h_sn]))
        elif output == 'quantity':
            # calculate corresponding quantity values
            log_q_c = StellarBlackholeFeedback_free_m_c.calculate_log_quantity(
                                                   self, log_m_c, log_m_c,
                                                   log_A, alpha, beta)[0]
            log_q_sn = StellarBlackholeFeedback_free_m_c.calculate_log_quantity(
                                                   self, log_m_h_sn, log_m_c,
                                                   log_A, alpha, beta)[0]
            log_q_bh = StellarBlackholeFeedback_free_m_c.calculate_log_quantity(
                                                   self, log_m_h_bh, log_m_c,
                                                   log_A, alpha, beta)[0]
            return(np.array([log_q_c, log_q_sn, log_q_bh]))
        else:
            raise ValueError('output must be either \'halo_mass\' or '
                             '\'quantity\'.')

    def _make_m_h_space(self, log_m_c, log_A, alpha, beta, num, log_epsilon=-3):
        '''
        Create array of m_h points arranged so that the points are dense where
        q(m_h) changes quickly and sparse where they aren't.
        num controls number of points that are calculated.
        epsilon controls how far out points are created (epsilon is the
        ration between the strength of the two feedback types, i.e.
        sn/bh = epsilon).

        '''

        # calculate halo masses where one of the feedbacks strongly dominates
        _, log_m_bh, log_m_sn = StellarBlackholeFeedback_free_m_c.\
                                 _calculate_feedback_regimes(
                                    self, log_m_c, log_A, alpha, beta,
                                    log_epsilon=log_epsilon)

        # create high density space
        dense_log_m_h = np.linspace(log_m_sn, log_m_bh, int(num*0.8))
        # create low density spaces
        sparse_log_m_lower = np.linspace(log_m_sn/2, log_m_sn, int(num*0.1))
        sparse_log_m_upper = np.linspace(log_m_bh, 2*log_m_bh, int(num*0.1))

        log_m_h_table = np.sort(np.concatenate([dense_log_m_h,
                                                sparse_log_m_lower,
                                                sparse_log_m_upper]))
        return(np.unique(log_m_h_table))

class StellarFeedback_free_m_c(StellarBlackholeFeedback_free_m_c):
    def __init__(self, log_m_c, initial_guess, bounds):
        '''
        Model for stellar feedback with free m_c.
        '''
        super().__init__(log_m_c, initial_guess, bounds)
        self.name = 'stellar_free_m_c'
        self.parameter_used = [0,1,2]
        self.fixed_m_c_flag = False

    def calculate_log_quantity(self, log_m_h, log_m_c, log_A, alpha):
        return(StellarBlackholeFeedback_free_m_c.
               calculate_log_quantity(self, log_m_h,
                                      log_m_c,
                                      log_A,
                                      alpha,
                                      beta=0))

    def calculate_log_halo_mass(self, log_quantity, log_m_c, log_A, alpha,
                                num=500):
        return(StellarBlackholeFeedback_free_m_c.
               calculate_log_halo_mass(self, log_quantity,
                                       log_m_c,
                                       log_A,
                                       alpha,
                                       beta=0,
                                       num=num))

    def calculate_dlogquantity_dlogmh(self, log_m_h, log_m_c, log_A, alpha):
        return(StellarBlackholeFeedback_free_m_c.
               calculate_dlogquantity_dlogmh(self, log_m_h,
                                             log_m_c,
                                             log_A,
                                             alpha,
                                             beta=0))

    def calculate_d2logquantity_dlogmh2(self, log_m_h, log_m_c, log_A, alpha):
        return(StellarBlackholeFeedback_free_m_c.
               calculate_d2logquantity_dlogmh2(self, log_m_h,
                                               log_m_c,
                                               log_A,
                                               alpha,
                                               beta=0))

    def _calculate_feedback_regimes(self, log_m_c, log_A, alpha,
                                   log_epsilon=-2, output='halo_mass'):
        return(StellarBlackholeFeedback_free_m_c.
               _calculate_feedback_regimes(self, log_m_c,
                                                log_A,
                                                alpha,
                                                beta=0,
                                                log_epsilon=log_epsilon,
                                                output=output))

class QuasarGrowth_free_m_c(StellarBlackholeFeedback_free_m_c):
    def __init__(self, log_m_c, initial_guess, bounds):
        '''
        Black hole growth model with free m_c.

        '''
        super().__init__(log_m_c, initial_guess, bounds)
        self.name = 'quasargrowth_free_m_c'
        self.parameter_used = [0,1,2]
        self.fixed_m_c_flag = False

    def calculate_log_quantity(self, log_m_h, log_m_c, log_B, eta):
        '''
        Calculate (log of) bolometric luminosity from input halo mass,
        model parameter and chosen log_eddington_ratio.
        '''
        log_m_h = make_array(log_m_h)
        return(log_B + eta*(log_m_h - log_m_c))

    def calculate_log_halo_mass(self, log_quantity, log_m_c, log_B,
                                eta):
        '''
        Calculate (log of) halo mass from input bolometric luminosity,
        model parameter and chosen log_eddington_ratio.
        '''
        log_quantity = make_array(log_quantity)
        return((log_quantity-log_B)/eta + log_m_c)

    def calculate_dlogquantity_dlogmh(self, log_m_h, log_m_c,
                                      log_B, eta):
        '''
        Calculate d/d(log m_h) log_quantity.
        '''
        return(np.full_like(log_m_h, eta))

class QuasarLuminosity_free_m_c_free_ERDF(object):
    def __init__(self, log_m_c, initial_guess, bounds):
        '''
        Black hole bolometric luminosity model with free m_c and ERDF.

        '''
        self.name = 'eddington_free_m_c_free_ERDF'
        self.initial_guess = initial_guess   # initial guess for least_squares
        # fit
        self.bounds = bounds                 # parameter bounds

        self.log_m_c = log_m_c

        # parameter used compared to the unconstrained model
        self.parameter_used = [0,1,2,3,4]
        self.fixed_m_c_flag = False

        # latest parameter used
        self.parameter = None
        self.eddington_distribution = None

    def calculate_log_quantity(self, log_m_h, log_eddington_ratio, log_m_c,
                               log_C, eta):
        '''
        Calculate (log of) bolometric luminosity from input halo mass,
        model parameter and chosen log_eddington_ratio.
        '''
        log_m_h = make_array(log_m_h)
        return(log_eddington_ratio + log_C + eta*(log_m_h - log_m_c))

    def calculate_log_halo_mass(self, log_L, log_eddington_ratio, log_m_c,
                                log_C, eta):
        '''
        Calculate (log of) halo mass from input bolometric luminosity,
        model parameter and chosen log_eddington_ratio.
        '''
        log_L = make_array(log_L)
        return((log_L-(log_eddington_ratio + log_C))/eta + log_m_c)

    def calculate_dlogquantity_dlogmh(self, log_m_h, log_eddington_ratio,
                                      log_m_c, log_C, eta):
        '''
        Calculate d/d(log m_h) log_quantity.
        '''
        return(np.full_like(log_m_h, eta))

    def calculate_log_erdf(self, log_eddington_ratio, log_eddington_star,
                           rho):
        '''
        Calculate (log) value of ERDF (probability of given eddington ratio),
        given input log_eddingtion_ratio and parameter (log_eddington_star,
        rho).

        '''
        self._make_distribution(log_eddington_star, rho)
        log_erdf = self.eddington_distribution.log_probability(
            log_eddington_ratio)
        return(log_erdf)

    def calculate_mean_log_eddington_ratio(self, log_eddington_star,
                                           rho):
        '''
        Calculate mean eddington ratio for the given parameter.

        '''
        self._make_distribution(log_eddington_star, rho)
        mean = self.eddington_distribution.mean()
        return(mean)

    def draw_eddington_ratio(self, log_eddington_star, rho, num=1):
        '''
        Draw random sample of Eddingtion ratio from distribution defined by
        parameter (log_eddington_star, rho). Can draw num samples at
        once.

        '''
        self._make_distribution(log_eddington_star, rho)
        log_eddington_ratio = self.eddington_distribution.rvs(size=num)
        return(log_eddington_ratio)

    def _make_distribution(self, log_eddington_star, rho):
        '''
        Check if distribution function with the given parameter was already
        created and stored, otherwise create it.

        '''
        if self.parameter == [log_eddington_star, rho]:
            pass
        else:
            self.parameter = [log_eddington_star, rho]
            self.eddington_distribution = ERDF(log_eddington_star, rho)
        return(self.eddington_distribution)

################ MODEL WITH FIXED CRITICAL MASS ###############################

class StellarBlackholeFeedback(StellarBlackholeFeedback_free_m_c):
    def __init__(self, log_m_c, initial_guess, bounds):
        '''
        Model for stellar + black hole feedback.

        '''
        super().__init__(log_m_c, initial_guess, bounds)
        self.name = 'stellar_blackhole'
        self.parameter_used = [1,2,3]
        self.fixed_m_c_flag = True

    def calculate_log_quantity(self, log_m_h, log_A, alpha, beta):
        return(StellarBlackholeFeedback_free_m_c.
               calculate_log_quantity(self, log_m_h,
                                      log_m_c=self.log_m_c,
                                      log_A=log_A,
                                      alpha=alpha,
                                      beta=beta))

    def calculate_log_halo_mass(self, log_quantity, log_A, alpha, beta,
                                num=500):
        return(StellarBlackholeFeedback_free_m_c.
               calculate_log_halo_mass(self, log_quantity,
                                       log_m_c=self.log_m_c,
                                       log_A=log_A,
                                       alpha=alpha,
                                       beta=beta,
                                       num=num))

    def calculate_dlogquantity_dlogmh(self, log_m_h, log_A, alpha, beta):
        return(StellarBlackholeFeedback_free_m_c.
               calculate_dlogquantity_dlogmh(self, log_m_h,
                                             log_m_c=self.log_m_c,
                                             log_A=log_A,
                                             alpha=alpha,
                                             beta=beta))

    def calculate_d2logquantity_dlogmh2(self, log_m_h, log_A, alpha, beta):
        return(StellarBlackholeFeedback_free_m_c.
               calculate_d2logquantity_dlogmh2(self, log_m_h,
                                               log_m_c=self.log_m_c,
                                               log_A=log_A,
                                               alpha=alpha,
                                               beta=beta))

    def _calculate_feedback_regimes(self, log_A, alpha, beta, log_epsilon=-2,
                                   output='halo_mass'):
        return(StellarBlackholeFeedback_free_m_c.
               _calculate_feedback_regimes(self,
                                          log_m_c=self.log_m_c,
                                          log_A=log_A,
                                          alpha=alpha,
                                          beta=beta,
                                          log_epsilon=log_epsilon,
                                          output=output))

class StellarFeedback(StellarBlackholeFeedback):
    def __init__(self, log_m_c, initial_guess, bounds):
        '''
        Model for stellar feedback.
        '''
        super().__init__(log_m_c, initial_guess, bounds)
        self.name = 'stellar'
        self.parameter_used = [1,2]
        self.fixed_m_c_flag = True

    def calculate_log_quantity(self, log_m_h, log_A, alpha):
        return(StellarBlackholeFeedback_free_m_c.
               calculate_log_quantity(self, log_m_h,
                                      log_m_c=self.log_m_c,
                                      log_A=log_A,
                                      alpha=alpha,
                                      beta=0))

    def calculate_log_halo_mass(self, log_quantity, log_A, alpha, num=500):
        return(StellarBlackholeFeedback_free_m_c.
               calculate_log_halo_mass(self, log_quantity,
                                       log_m_c=self.log_m_c,
                                       log_A=log_A,
                                       alpha=alpha,
                                       beta=0,
                                       num=num))

    def calculate_dlogquantity_dlogmh(self, log_m_h, log_A, alpha):
        return(StellarBlackholeFeedback_free_m_c.
               calculate_dlogquantity_dlogmh(self, log_m_h,
                                             log_m_c=self.log_m_c,
                                             log_A=log_A,
                                             alpha=alpha,
                                             beta=0))

    def calculate_d2logquantity_dlogmh2(self, log_m_h, log_A, alpha):
        return(StellarBlackholeFeedback_free_m_c.
               calculate_d2logquantity_dlogmh2(self, log_m_h,
                                               log_m_c=self.log_m_c,
                                               log_A=log_A,
                                               alpha=alpha,
                                               beta=0))

    def _calculate_feedback_regimes(self, log_A, alpha, log_epsilon=-2,
                                   output='halo_mass'):
        return(StellarBlackholeFeedback_free_m_c.
               _calculate_feedback_regimes(self, log_m_c=self.log_m_c,
                                          log_A=log_A,
                                          alpha=alpha,
                                          beta=0,
                                          log_epsilon=log_epsilon,
                                          output=output))

class NoFeedback(StellarBlackholeFeedback):
    def __init__(self, log_m_c, initial_guess, bounds):
        '''
        Model without feedback.
        '''
        super().__init__(log_m_c, initial_guess, bounds)
        self.name = 'none'
        self.log2 = np.log10(2)  # so it doesn't need to be calc everytime

    def calculate_log_quantity(self, log_m_h, log_A):
        '''
        Calculate observable quantity from input halo mass and model parameter.
        '''
        log_m_h = make_array(log_m_h)
        return(log_A - self.log2 + log_m_h)

    def calculate_log_halo_mass(self, log_quantity, log_A):
        '''
        Calculate halo mass from input observable quantity and model parameter.
        '''
        log_quantity = make_array(log_quantity)
        return(log_quantity - log_A + self.log2)

    def calculate_dlogquantity_dlogmh(self, log_m_h, log_A):
        log_m_h = make_array(log_m_h)
        return(np.full_like(log_m_h, 1))

    def calculate_d2logquantity_dlogmh2(self, log_m_h, log_A):
        log_m_h = make_array(log_m_h)
        return(np.full_like(log_m_h, 0))

class QuasarGrowth(QuasarGrowth_free_m_c):
    def __init__(self, log_m_c, initial_guess, bounds):
        '''
        Black hole growth model.

        '''
        super().__init__(log_m_c, initial_guess, bounds)
        self.name = 'quasargrowth'
        self.parameter_used = [1,2]
        self.fixed_m_c_flag = True

    def calculate_log_quantity(self, log_m_h, log_B, eta):
        return(QuasarGrowth_free_m_c.
               calculate_log_quantity(self, log_m_h,
                                      log_m_c=self.log_m_c,
                                      log_B=log_B,
                                      eta=eta))

    def calculate_log_halo_mass(self, log_quantity, log_B, eta):
        return(QuasarGrowth_free_m_c.
               calculate_log_halo_mass(self, log_quantity,
                                       log_m_c=self.log_m_c,
                                       log_B=log_B,
                                       eta=eta))

    def calculate_dlogquantity_dlogmh(self, log_m_h, log_B, eta):
        return(QuasarGrowth_free_m_c.
               calculate_dlogquantity_dlogmh(self, log_m_h,
                                             log_m_c=self.log_m_c,
                                             log_B=log_B,
                                             eta=eta))

class QuasarLuminosity_free_ERDF(QuasarLuminosity_free_m_c_free_ERDF):
    def __init__(self, log_m_c, initial_guess, bounds):
        '''
        Black hole bolometric luminosity model with free ERDF.

        '''
        self.name = 'eddington_free_ERDF'
        self.initial_guess = initial_guess   # initial guess for least_squares
        # fit
        self.bounds = bounds                 # parameter bounds

        self.log_m_c = log_m_c

        # parameter used compared to the unconstrained model
        self.parameter_used = [1,2,3,4]
        self.fixed_m_c_flag = True

        # latest parameter used
        self.parameter = None
        self.eddington_distribution = None

    def calculate_log_quantity(self, log_m_h, log_eddington_ratio,
                               log_C, eta):
        return(QuasarLuminosity_free_m_c_free_ERDF.calculate_log_quantity(
                                self,
                                log_m_h=log_m_h,
                                log_eddington_ratio=log_eddington_ratio,
                                log_m_c=self.log_m_c,
                                log_C=log_C,
                                eta=eta))

    def calculate_log_halo_mass(self, log_L, log_eddington_ratio,
                                log_C, eta):
        return(QuasarLuminosity_free_m_c_free_ERDF.calculate_log_halo_mass(
                                self,
                                log_L=log_L,
                                log_eddington_ratio=log_eddington_ratio,
                                log_m_c=self.log_m_c,
                                log_C=log_C,
                                eta=eta))

    def calculate_dlogquantity_dlogmh(self, log_m_h, log_eddington_ratio,
                                      log_C, eta):
        return(QuasarLuminosity_free_m_c_free_ERDF.calculate_dlogquantity_dlogmh(
                                self,
                                log_m_h=log_m_h,
                                log_eddington_ratio=log_eddington_ratio,
                                log_m_c=self.log_m_c,
                                log_C=log_C,
                                eta=eta))

class QuasarLuminosity(QuasarLuminosity_free_ERDF):
    def __init__(self, log_m_c, initial_guess, bounds, eddington_erdf_params):
        '''
        Black hole bolometric luminosity model with fixed ERDF.

        '''
        self.name = 'eddington'
        self.initial_guess = initial_guess   # initial guess for least_squares
        # fit
        self.bounds = bounds                 # parameter bounds

        self.log_m_c = log_m_c

        self.parameter_used = [0,1]
        self.fixed_m_c_flag = True

        # parameter used
        self.parameter = eddington_erdf_params
        self.eddington_distribution = ERDF(*self.parameter)

    def calculate_log_erdf(self, log_eddington_ratio):
        log_erdf = self.eddington_distribution.log_probability(
            log_eddington_ratio)
        return(log_erdf)

    def calculate_mean_log_eddington_ratio(self):
        mean = self.eddington_distribution.mean()
        return(mean)

    def draw_eddington_ratio(self, num=1):
        log_eddington_ratio = self.eddington_distribution.rvs(size=num)
        return(log_eddington_ratio)

    def _make_distribution(self, log_eddington_star, rho):
        raise NotImplementedError('Cannot make new distribution for fixed ERDF'
                                  ' model.')


import os
from collections import Counter
from utils import extension_to_folder, move_file, dir_empty
from tqdm import tqdm

if __name__ == '__main__':
    # Folder names
    fol_to_manage = 'D:\\Edge Download'
    esc_border = "=="* 50
    print(f"\n{esc_border}")
    print(f"\nManaging folder - {fol_to_manage}")
    files = [f for f in os.listdir(fol_to_manage) if os.path.isfile(os.path.join(fol_to_manage, f))]
    print(f"\nThe folder has {len(files)} files")
    n = 5
    print(f'\nFirst {n} files are :')
    for i in range(len(files)):
        print(f"  -->{files[i]}")
        if i == n:
            break

    file_exts = [i.split('.')[-1] for i in files]
    file_ext_dict = Counter(file_exts)
    file_ext_dict = dict(sorted(file_ext_dict.items(), key=lambda item: item[1], reverse= True))
    print(f"\nThere are {len(file_ext_dict)} different types of files:")
    print(f"\nThe first {n} are:")

    for ii, key in enumerate(file_ext_dict.keys()):
        print(f"  -->{key}\t\t{file_ext_dict[key]}")
        if ii == 4:
            break

    folders_available = extension_to_folder()
    c = 0
    c_in = 0
    not_counted = {}

    for ii, key in enumerate(file_ext_dict.keys()):
        if str(key).lower() in list(folders_available.keys()):
            c_in += file_ext_dict[key]

        else:
            c += file_ext_dict[key]
            not_counted[key] = file_ext_dict[key]

    print(f"\nTotal {c_in} types of file accounted for and {c} types of file not accounted\n")
    # for ii, key in enumerate(not_counted.keys()):
    #     print(f"  -->{key}\t\t{not_counted[key]}")

    # Create Folders First
    for folder in folders_available.values():
        os.makedirs(os.path.join(fol_to_manage, folder), exist_ok= True)

    for f in tqdm(files, total = len(files)):
        ext = f.split('.')[-1]
        if ext.lower() in list(folders_available.keys()):
            folder = folders_available[ext.lower()]
        else:
            folder = 'Other'
        source = os.path.join(fol_to_manage, f)
        dest   = os.path.join(fol_to_manage, folder, f)
        move_file(source, dest)

    for folder in folders_available.values():
        if os.path.exists(os.path.join(fol_to_manage, folder)):
            if dir_empty(os.path.join(fol_to_manage, folder)):
                os.rmdir(os.path.join(fol_to_manage, folder))


##########################################################
# self (pre) -> a -> b -> b.next
# self (pre) -> b -> a -> b.next
#                    a(pre) -> a -> b -> b.next
#                    a(pre) -> b -> a -> b.next
# ............
# T -> a -> b -> P
# P = b.next
# T -> b -> a -> P
##########################################################
class ListNode(object):
    def __init__(self, x):
            self.val = x
            self.next = None
    def __repr__(self):
            return "{} -> {}".format(self.val, self.next)

class Solution:
    # @param a ListNode
    # @return a ListNode
    def swapPairs(self, head):
        dummy = ListNode(0)
        dummy.next = head
        current = dummy
        while current.next and current.next.next:
            a, b, c = current.next, current.next.next, current.next.next.next
            current.next = b
            b.next = a
            a.next = c
            current = a
        return dummy.next

head = ListNode(1)
head.next = ListNode(2)
head.next.next = ListNode(3)
head.next.next.next = ListNode(4)
s = Solution()
print(s.swapPairs(head))


import json
from typing import TypeVar, Any

from redis.client import StrictRedis

from config.redis import HOST, PORT
from src.helpers.abstract_classes.AbstractMessageAwaiter import AbstractMessageAwaiter

K = TypeVar('K')

class RedisJsonMessageAwaiter(AbstractMessageAwaiter[K, dict[str, Any]]):
    def __init__(self, topic: str, key_field: str):
        super().__init__()
        red = StrictRedis(HOST, PORT, encoding='utf-8', decode_responses=True)
        self.__topic = topic
        self.__sub = red.pubsub(ignore_subscribe_messages=True)
        self.__sub.subscribe(topic)
        self.__key_field = key_field

    def put_awaitable(self, key: K) -> None:
        if self.awaited_messages_count() == 0 and not self.__sub.subscribed:
            self.__sub.subscribe(self.__topic)
        super().put_awaitable(key)

    def start_receiving_responses(self) -> None:
        try:
            for message in self.__sub.listen():
                message = json.loads(message['data'])
                key = message[self.__key_field]
                if not self.contains_key(key):
                    continue  # The message may be for another consumer of the topic
                del message[self.__key_field]
                self._put_result_and_notify(key, message)
                if self.awaited_messages_count() == 0 and self.__sub.subscribed:
                    self.__sub.reset()

        finally:
            self.__sub.close()