from typing import Any, Callable, List, Optional, Tuple

import sqlparse  # type:ignore
from sqlparse.sql import Token  # type:ignore
from sqlparse.tokens import Name, Newline, Punctuation  # type:ignore
from sqlparse.utils import remove_quotes  # type:ignore

class LayerSqlFunction:
    """
    A parsed Layer SQL statement
    """

    SUPPORTED_FUNCTION_TRAIN = "train"
    SUPPORTED_FUNCTION_PREDICT = "predict"
    SUPPORTED_FUNCTION_AUTOML = "automl"
    SUPPORTED_FUNCTION_TYPES = [SUPPORTED_FUNCTION_TRAIN, SUPPORTED_FUNCTION_PREDICT, SUPPORTED_FUNCTION_AUTOML]

    def __init__(self, function_type: str, source_name: str, target_name: str) -> None:
        self.function_type = function_type
        self.source_name = source_name
        self.target_name = target_name

    def __repr__(self) -> str:
        return (
            f"<LayerSQL function_type:{self.function_type}"
            + f" source_name:{self.source_name}"
            + f" target_name:{self.target_name}>"
        )

class LayerPredictFunction(LayerSqlFunction):
    def __init__(
        self,
        source_name: str,
        target_name: str,
        model_name: str,
        select_columns: List[str],
        predict_columns: List[str],
        all_columns: List[str],
        prediction_alias: str,
        sql: str,
    ) -> None:
        super().__init__(
            function_type=self.SUPPORTED_FUNCTION_PREDICT, source_name=source_name, target_name=target_name
        )
        self.model_name = model_name
        self.select_columns = select_columns
        self.predict_columns = predict_columns
        self.all_columns = all_columns
        self.prediction_alias = prediction_alias
        self.sql = sql

class LayerTrainFunction(LayerSqlFunction):
    def __init__(
        self,
        source_name: str,
        target_name: str,
        train_columns: List[str],
    ) -> None:
        super().__init__(function_type=self.SUPPORTED_FUNCTION_TRAIN, source_name=source_name, target_name=target_name)
        self.train_columns = train_columns

class LayerAutoMLFunction(LayerSqlFunction):
    def __init__(
        self,
        source_name: str,
        target_name: str,
        model_type: str,
        feature_columns: List[str],
        target_column: str,
        sql: str,
    ) -> None:
        super().__init__(function_type=self.SUPPORTED_FUNCTION_AUTOML, source_name=source_name, target_name=target_name)
        self.feature_columns = feature_columns
        self.target_column = target_column
        self.model_type = model_type
        self.sql = sql

def find_from_token(tokens: List[Token]) -> List[Token]:
    return expect_tokens(tokens, [keyword("from"), lambda x: isinstance(x, sqlparse.sql.Identifier)])

def get_from_where_clause(select_tokens: List[Token]) -> Tuple[str, str]:
    from_token = find_from_token(select_tokens)
    if not from_token:
        raise ValueError("Invalid sql. Missing 'from' clause.")
    source = from_token[0].value
    where_statement = " ".join((x.value for x in from_token[1:]))
    return source, where_statement

def get_cols_from_container(brace_container_token: Token) -> List[str]:
    _, cols_token, _ = brace_container_token.tokens
    if isinstance(cols_token, sqlparse.sql.IdentifierList):
        return [x.value for x in clean_separators(cols_token.tokens)]
    return [cols_token.value]

def build_sql(cols: List[str], source: str, where_statement: str) -> str:
    col_str = ", ".join(cols)
    if len(where_statement) > 0:
        return f"select {col_str} from {source} {where_statement}"
    else:
        return f"select {col_str} from {source}"

class LayerSQLParser:

    ALL_COLUMNS_WILDCARD = "*"

    def parse(self, sql: str) -> Optional[LayerSqlFunction]:
        """
        returns None if not a layer SQL statement
        returns an instance of LayerSQL if a valid layer SQL statement
        """
        parsed = sqlparse.parse(sql)
        if not parsed:
            return None

        tokens = parsed[0].tokens
        layer_func = next((x for x in find_functions(tokens) if is_layer_function(x)), None)
        if not layer_func:
            return None

        target_name_group = next(
            (x for x in expect_tokens(tokens, [keyword("create or replace"), keyword("table"), group()])), None
        )
        if not target_name_group:
            return None
        target_name = self._get_target_name_from_group(target_name_group)

        if is_automl_function(layer_func):
            return self.parse_automl(layer_func, target_name)
        elif is_predict_function(layer_func):
            return self.parse_predict(layer_func, target_name)
        elif is_train_function(layer_func):
            return self.parse_train(layer_func, target_name)
        else:
            raise ValueError(f"Unsupported function: {layer_func.get_name()}")

    def _get_target_name_from_group(self, token: Token) -> str:
        clean_inner_tokens = remove_tokens(token.flatten(), newline())
        target_name_tokens = slice_between_tokens(clean_inner_tokens, name(), whitespace())
        if target_name_tokens:
            return "".join([x.value for x in target_name_tokens])
        else:
            raise ValueError("Invalid sql syntax for Layer function")

    def parse_predict(self, layer_func_token: Token, target: str) -> LayerPredictFunction:

        # We need the parent `select` statement that contains the function
        # to get access to the selected columns and the source relation
        select_sttmt = find_parent(layer_func_token, lambda x: isinstance(x, sqlparse.sql.Parenthesis))
        clean_select_sttmt = []
        if not select_sttmt:
            raise ValueError("SQL syntax error")
        else:
            clean_select_sttmt = clean_separators(select_sttmt.tokens)

        # sqlparse doesn't seem to parse correctly the inner contents of this parenthesis.
        # Here, we rebuild the sql and use parse it again to get the relevant tokens
        inner_sql_text = " ".join(x.value for x in clean_select_sttmt)
        inner_sql = sqlparse.parse(inner_sql_text)[0]
        clean_inner_sql = remove_tokens(inner_sql.tokens, lambda x: x.is_whitespace)

        source, where_statement = get_from_where_clause(clean_inner_sql)

        # extract the selected columns in the query
        columns_incl_layer = []
        select_columns_list = find_token(clean_inner_sql, lambda x: isinstance(x, sqlparse.sql.IdentifierList))
        if select_columns_list is None:
            raise ValueError("SQL syntax error. Select columns are missing.")
        else:
            columns_incl_layer = find_tokens(
                select_columns_list.tokens, lambda x: isinstance(x, sqlparse.sql.Identifier)
            )
        columns = remove_tokens(columns_incl_layer, lambda t: find_layer_function(t) is not None)
        select_columns = [t.value for t in columns]

        # extract the layer prediction function
        clean_func_tokens = clean_separators(layer_func_token[1].tokens)
        if len(clean_func_tokens) < 3:
            invalid_func = " ".join(t.value for t in clean_func_tokens)
            raise ValueError(f"Invalid predict function syntax {invalid_func}")
        model_name, _, bracket_container = clean_func_tokens
        predict_model = remove_quotes(model_name.value)
        predict_columns = get_cols_from_container(bracket_container)

        all_columns = select_columns + list(set(predict_columns) - set(select_columns))
        sql_text = build_sql(all_columns, source, where_statement)
        sql = sqlparse.format(sql_text, keyword_case="lower", strip_whitespace=True)

        prediction_alias = layer_func_token.parent.get_alias() or "prediction"

        return LayerPredictFunction(
            source,
            target,
            predict_model,
            select_columns,
            predict_columns,
            all_columns,
            prediction_alias,
            sql,
        )

    def parse_automl(self, layer_func_token: Token, target: str) -> LayerAutoMLFunction:

        # We need the parent `select` statement that contains the function
        # to get access to the selected columns and the source relation
        select_sttmt = find_parent(layer_func_token, lambda x: isinstance(x, sqlparse.sql.Parenthesis))
        clean_select_sttmt = []
        if not select_sttmt:
            raise ValueError("SQL syntax error")
        else:
            clean_select_sttmt = clean_separators(select_sttmt.tokens)

        # sqlparse doesn't seem to parse correctly the inner contents of this parenthesis.
        # Here, we rebuild the sql and use parse it again to get the relevant tokens
        inner_sql_text = " ".join(x.value for x in clean_select_sttmt)
        inner_sql = sqlparse.parse(inner_sql_text)[0]
        clean_inner_sql = remove_tokens(inner_sql.tokens, lambda x: x.is_whitespace)

        source, where_statement = get_from_where_clause(clean_inner_sql)

        # extract the layer prediction function
        clean_func_tokens = clean_separators(layer_func_token[1].tokens)
        if len(clean_func_tokens) < 4:
            invalid_func = " ".join(t.value for t in clean_func_tokens)
            raise ValueError(f"Invalid automl function syntax {invalid_func}")
        model_type_token, _, bracket_container, target_column_token = clean_func_tokens
        model_type = remove_quotes(model_type_token.value)
        feature_columns = get_cols_from_container(bracket_container)
        target_column = remove_quotes(target_column_token.value)
        all_columns = list(dict.fromkeys((feature_columns + [target_column])))
        sql_text = build_sql(all_columns, source, where_statement)
        sql = sqlparse.format(sql_text, keyword_case="lower", strip_whitespace=True)

        return LayerAutoMLFunction(source, target, model_type, feature_columns, target_column, sql)

    def parse_train(self, layer_func_token: Token, target: str) -> LayerTrainFunction:
        select = find_parent(layer_func_token, lambda x: isinstance(x, sqlparse.sql.Identifier))
        if not select:
            raise ValueError("SQL syntax error")
        else:
            clean_select = clean_separators(select)
            source, _ = get_from_where_clause(layer_func_token.parent.parent.tokens)
            _, train_func = clean_select
            if len(train_func.tokens) < 2:
                invalid_func = "".join(x.value for x in select.flatten())
                raise ValueError(f"Invalid train function syntax {invalid_func}")
            _, parenthesis_group = train_func.tokens
            cols = self.extract_columns(clean_separators(parenthesis_group.tokens))

            return LayerTrainFunction(source, target, cols)

    def extract_columns(self, content_tokens: List[Token]) -> List[str]:
        if len(content_tokens) < 1:
            return [self.ALL_COLUMNS_WILDCARD]
        if content_tokens[0].value.lower() == "array":
            if len(content_tokens) < 2:
                raise ValueError("Invalid train function syntax")
            identifiers = clean_separators(content_tokens[1].tokens)[0]
            return [x.value for x in identifiers.get_identifiers()]
        else:
            return [content_tokens[0].value]

# Functions to assist the SQL Parsing

TokenPredicate = Callable[[Token], bool]

def get_predict_function_name(layer_column_token: Token) -> Optional[str]:
    if layer_column_token.has_alias():
        return layer_column_token.get_name()
    return "predict"

def find_token(tokens: List[Token], pred: TokenPredicate) -> Optional[Token]:
    for token in tokens:
        if pred(token):
            return token
    return None

def find_tokens(tokens: List[Token], pred: TokenPredicate) -> List[Token]:
    return [x for x in tokens if pred(x)]

def remove_tokens(tokens: List[Token], pred: TokenPredicate) -> List[Token]:
    return [x for x in tokens if not pred(x)]

def clean_separators(tokens: List[Token]) -> List[Token]:
    return remove_tokens(tokens, lambda x: x.is_whitespace or x.ttype == Punctuation)

def expect_tokens(tokens: List[Token], predicates: List[TokenPredicate]) -> List[Token]:
    if not tokens or not predicates:
        return tokens
    pred = predicates[0]
    for i, token in enumerate(tokens):
        if pred(token):
            return expect_tokens(tokens[i:], predicates[1:])
    return []

def slice_between_tokens(tokens: List[Token], start: TokenPredicate, end: TokenPredicate) -> Optional[List[Token]]:
    if not tokens:
        return None
    start_index = -1
    end_index = -1
    start_seen = False
    for i, token in enumerate(tokens):
        if not start_seen and start(token):
            start_index = i
            start_seen = True
        if start_seen and end(token):
            end_index = i
            break
    if start_index >= 0 and end_index > 0:
        return tokens[start_index:end_index]
    else:
        return None

def find_functions(elem: Any) -> List[Token]:
    from itertools import chain

    funcs: List[List[Token]] = [[]]
    if isinstance(elem, List):
        funcs = [find_functions(x) for x in elem]
    elif isinstance(elem, sqlparse.sql.Statement):
        funcs = [find_functions(x) for x in elem.tokens]
    elif isinstance(elem, sqlparse.sql.IdentifierList):
        funcs = [find_functions(x) for x in elem.get_identifiers()]
    elif isinstance(elem, sqlparse.sql.Identifier):
        funcs = [find_functions(x) for x in elem.tokens]
    elif isinstance(elem, sqlparse.sql.Parenthesis):
        funcs = [find_functions(x) for x in elem.tokens]
    elif isinstance(elem, sqlparse.sql.Function):
        return [elem]
    return list(chain.from_iterable(funcs))

def find_parent(token: Token, pred: TokenPredicate) -> Optional[Token]:
    parent = token.parent
    if not parent:
        return None
    elif pred(parent):
        return parent
    else:
        return find_parent(parent, pred)

def is_layer_function(func: sqlparse.sql.Function) -> bool:
    parent = func.parent
    return isinstance(parent, sqlparse.sql.Identifier) and parent.tokens[0].value == "layer"

def find_layer_function(tokens: List[Token]) -> Optional[Token]:
    return next((x for x in find_functions(tokens) if is_layer_function(x)), None)

def is_predict_function(func: sqlparse.sql.Function) -> bool:
    return func[0].value.lower() == LayerSqlFunction.SUPPORTED_FUNCTION_PREDICT

def is_automl_function(func: sqlparse.sql.Function) -> bool:
    return func[0].value.lower() == LayerSqlFunction.SUPPORTED_FUNCTION_AUTOML

def is_train_function(func: sqlparse.sql.Function) -> bool:
    return func[0].value.lower() == LayerSqlFunction.SUPPORTED_FUNCTION_TRAIN

def keyword(value: str) -> TokenPredicate:
    def check_token(token: Token) -> bool:
        return token.is_keyword and token.value.lower() == value.lower()

    return check_token

def whitespace() -> Callable[[Token], bool]:
    return lambda x: x.is_whitespace

def name() -> Callable[[Token], bool]:
    return lambda x: x.ttype is Name

def newline() -> Callable[[Token], bool]:
    return lambda x: x.ttype is Newline

def group() -> Callable[[Token], bool]:
    return lambda x: x.is_group


from json import load

from dataclasses import dataclass

CONFIG_FILENAME = "config.json"

@dataclass
class Config:
    filters_filename: str
    segment_len: int
    sample_skips: int
    bit_depth: int
    original_audio_dir: str
    filtered_audio_dir: str
    output_filename: str

def load_config(filename: str) -> Config:
    with open(filename) as config_file:
        config_data = load(config_file)
    return Config(
        config_data["filters_filename"],
        config_data["segment_len"],
        config_data["sample_skips"],
        config_data["bit_depth"],
        config_data["original_audio_dir"],
        config_data["filtered_audio_dir"],
        config_data["output_filename"]
    )


import abc
import os
import re
import unicodedata
from typing import List, Union

import numpy as np
import tensorflow as tf
import tensorflow_datasets as tds
from utils.character import en, jp

import frontends.wordpiece as wordpiece

def preprocess_paths(paths: Union[List, str]):
    if isinstance(paths, list):
        return [
            path
            if path.startswith("gs://")
            else os.path.abspath(os.path.expanduser(path))
            for path in paths
        ]
    elif isinstance(paths, str):
        return (
            paths
            if paths.startswith("gs://")
            else os.path.abspath(os.path.expanduser(paths))
        )
    else:
        return None

class TextFeaturizer(metaclass=abc.ABCMeta):
    def __init__(self):
        self.blank = None
        self.tokens2indices = {}
        self.tokens = []
        self.num_classes = None

    def preprocess_text(self, text):
        text = unicodedata.normalize("NFC", text.lower())
        return text.strip("\n")  # remove trailing newline

    def normalize_indices(self, indices: tf.Tensor) -> tf.Tensor:
        """Remove -1 in indices by replacing them with blanks
        Args:
            indices (tf.Tensor): shape any

        Returns:
            (tf.Tensor): normalized indices with shape same as indices
        """
        minus_one = -1 * tf.ones_like(indices, dtype=tf.int32)
        blank_like = self.blank * tf.ones_like(indices, dtype=tf.int32)
        return tf.where(indices == minus_one, blank_like, indices)

    def prepand_blank(self, text: tf.Tensor) -> tf.Tensor:
        """Prepand blank index for transducer models"""
        return tf.concat([[self.blank], text], axis=0)

    @abc.abstractclassmethod
    def init_upoints(self):
        """Create a dictionary of upoints that map to target indicies"""
        raise NotImplementedError()

    @abc.abstractclassmethod
    def string2Indices(self, text):
        """Converts a string to a list of indices

        Args:
            text (string): sequence of characters

        Returns:
            (tf.Tensor): sequence of ints
        """
        raise NotImplementedError()

    @abc.abstractclassmethod
    def indices2String(self, indices):
        """Directly converts a list of indices to string

        Note: this function is intended for only debugging purpose

        Args:
            indices (tf.Tensor): [B, None]

        Returns:
            (tf.Tensor): transcripts of dtype tf.string with dim [B]
        """
        raise NotImplementedError()

    @abc.abstractclassmethod
    def indices2upoints(self, indices):
        """Transforms predicted indices to unicode code points

        Args:
            indices (tf.Tensor): indices in shape [None].

        Returns:
            (tf.Tensor): unicode code points transcript with dtype tf.int32 and shape
            [None]

        """
        raise NotImplementedError()

class CharFeaturizer(TextFeaturizer):
    def __init__(
        self,
        vocab_path: str = None,
        blank_at_zero: bool = True,
        lang: str = "en",
        normalize: bool = False,
    ):
        super().__init__()

        self.vocab_path = vocab_path
        self.blank_at_zero = blank_at_zero
        self.lang = lang
        self.normalize = normalize

        self.init_vocabulary()
        self.init_upoints()

    def init_vocabulary(self):
        lines = []
        if self.vocab_path is not None:
            with open(self.vocab_path, "r", encoding="utf-8") as fin:
                lines.extend(fin.readlines())
        else:
            if self.lang == "en":
                lines = en
            elif self.lang == "jp":
                lines = jp
            else:
                raise ValueError("This language is not supported.")

        self.blank = 0 if self.blank_at_zero else None
        self.tokens2indices = {}
        self.tokens = []

        index = 1 if self.blank == 0 else 0
        for line in lines:
            if self.normalize:
                line = self.preprocess_text(line)

            self.tokens2indices[line[0]] = index
            self.tokens.append(line[0])
            index += 1

        if self.blank is None:
            self.blank = len(self.tokens)  # blank not at zero

        self.vocab_array = self.tokens.copy()
        self.tokens.insert(self.blank, "")  # add blank token to tokens
        self.num_classes = len(self.tokens)
        self.tokens = tf.convert_to_tensor(self.tokens, dtype=tf.string)

    def init_upoints(self):
        self.upoints = tf.strings.unicode_decode(self.tokens, "UTF-8").to_tensor(
            shape=[None, 1]
        )

    def string2Indices(self, text: str) -> tf.Tensor:
        if self.normalize:
            text = self.preprocess_text(text)

        text = list(text.strip())  # remove trailing space
        indices = [self.tokens2indices[token] for token in text]
        return tf.convert_to_tensor(indices, dtype=tf.int32)

    def indices2String(self, indices: tf.Tensor) -> tf.Tensor:
        indices = self.normalize_indices(indices)
        tokens = tf.gather_nd(self.tokens, tf.expand_dims(indices, axis=-1))
        with tf.device("/CPU:0"):  # string data is not supported on GPU
            tokens = tf.strings.reduce_join(tokens, axis=-1)
        return tokens

    @tf.function(input_signature=[tf.TensorSpec([None], dtype=tf.int32)])
    def indices2upoints(self, indices: tf.Tensor) -> tf.Tensor:
        indices = self.normalize_indices(indices)
        upoints = tf.gather_nd(self.upoints, tf.expand_dims(indices, axis=-1))
        return tf.gather_nd(upoints, tf.where(tf.not_equal(upoints, 0)))

    @tf.function(input_signature=[tf.TensorSpec([None, None], dtype=tf.int32)])
    def indices2upointsBatch(self, indices: tf.Tensor) -> tf.Tensor:
        """Transforms predicted indices to unicode code points in batch

        Note: Unlike `indices2upoints`, to keep batch dimension, this function does
        not remove blank tokens.

        L: length of predicted tokens

        Args:
            indices: [B, L]

        """
        indices = self.normalize_indices(indices)
        # [B, L, 1]
        upoints = tf.gather_nd(self.upoints, tf.expand_dims(indices, axis=-1))
        # [B, L]
        upoints = tf.squeeze(upoints, axis=-1)
        return upoints

class WordpieceFeaturizer(TextFeaturizer):
    PAD_TOKEN = "<pad>"  # also used as blank token

    def __init__(
        self,
        subwords_path: str,
        subwords_corpus: str = None,
        target_vocab_size: int = 1024,
        max_token_length: int = 4,
        normalize: bool = False,
    ):
        """Text featurizer based on Wordpiece

        TODO: Since Japanese is the top priority as of 2021/3 and japanese asr performs
        better in character (see refs), wordpiece in production is not supported,
        althought it is just fine in training. This needs to be exportable to TFLite!!
        ref: https://github.com/espnet/espnet/pull/326
        ref: http://www.interspeech2020.org/uploadfile/pdf/Mon-2-2-2.pdf

        """
        super().__init__()

        self.target_vocab_size = target_vocab_size
        self.max_token_length = max_token_length
        self.max_corpus_chars = None
        self.reserved_tokens = [self.PAD_TOKEN]

        self.normalize = normalize

        if subwords_path and os.path.exists(subwords_path):
            self.subwords = self.load_from_file(subwords_path)
        else:
            self.subwords = self.build_from_corpus(subwords_corpus, subwords_path)

        self.blank = 0  # subword treats blank as 0
        self.num_classes = self.subwords.vocab_size

        self.init_upoints()

    def init_upoints(self):
        text = [""]  # first one is reserved for blank token
        # start iteration after blank
        for idx in np.arange(1, self.num_classes, dtype=np.int32):
            token = self.subwords.detokenize(tf.constant([[idx]], tf.int32))
            token = token.values.numpy()[0].decode("utf-8")

            # `##_` in wordpiece means, a beginging space should be trimmed, otherwise
            # prepend a space
            if "#" in token:
                token = re.sub("#", "", token)
            else:
                token = " " + token

            text.append(token)
        self.upoints = tf.strings.unicode_decode(text, "UTF-8")
        self.upoints = self.upoints.to_tensor()  # [num_classes, max_token_length]

    def build_from_corpus(self, corpus_files: list = None, output_file: str = None):
        filename = preprocess_paths(output_file)

        def corpus_generator():
            for file in corpus_files:
                with open(file, "r", encoding="utf-8") as f:
                    lines = f.read().splitlines()
                    lines = lines[1:]
                for line in lines:
                    line = line.split("\t")
                    yield line[-1]

        wordpiece.build_from_corpus(
            corpus_generator(),
            output_file_path=filename,
            target_vocab_size=self.target_vocab_size,
            max_token_length=self.max_token_length,
            max_corpus_chars=self.max_corpus_chars,
            reserved_tokens=self.reserved_tokens,
        )

        subwords = wordpiece.WordpieceTokenizer(filename, token_out_type=tf.int32)
        return subwords

    def load_from_file(self, filename: str = None):
        filename = preprocess_paths(filename)
        subwords = wordpiece.WordpieceTokenizer(filename, token_out_type=tf.int32)
        return subwords

    def string2Indices(self, text: str) -> tf.Tensor:
        if self.normalize:
            text = self.preprocess_text(text)

        # remove trailing space
        text = text.strip()
        splitted_text = text.split()
        indices = self.subwords.tokenize(splitted_text)
        # `tokenize` returns ragged tensor
        indices = indices.merge_dims(0, -1)
        return tf.cast(indices, dtype=tf.int32)

    def indices2String(self, indices: tf.Tensor) -> tf.Tensor:
        with tf.device("/CPU:0"):  # string data is not supported on GPU
            indices = self.normalize_indices(indices)
            # returns ragged tensor with one element corresponds to one word
            # eg) ["i'm", "<pad>", "thirsty", "<pad>", ...]
            text = self.subwords.detokenize(indices)
            # ["i'm", "", "thirsty", ""]
            text = tf.strings.regex_replace(text, self.PAD_TOKEN, "")
            # ["i'm  thirsty "] <- double spaces & trailing space
            text = tf.strings.reduce_join(text, separator=" ", axis=-1)
            # ["i'm thirsty "] <- trailing space
            text = tf.strings.regex_replace(text, " +", " ")
            # ["i'm thirsty"]
            text = tf.strings.strip(text)
            return text

    @tf.function(input_signature=[tf.TensorSpec([None], dtype=tf.int32)])
    def indices2upoints(self, indices: tf.Tensor) -> tf.Tensor:
        indices = self.normalize_indices(indices)
        upoints = tf.gather_nd(self.upoints, tf.expand_dims(indices, axis=-1))
        return tf.gather_nd(upoints, tf.where(tf.not_equal(upoints, 0)))

class DeprecatedWordpieceFeaturizer(TextFeaturizer):
    def __init__(
        self,
        subwords_path: str,
        subwords_corpus: str = None,
        target_vocab_size: int = 1024,
        max_token_length: int = 4,
        normalize: bool = False,
    ):
        super().__init__()

        self.target_vocab_size = target_vocab_size
        self.max_token_length = max_token_length
        self.max_corpus_chars = None
        # unlike undeprecated version, tds.deprecated.text.SubwordTextEncoder
        # automatically register blank token to zero index
        self.reserved_tokens = None

        self.normalize = normalize

        if subwords_path and os.path.exists(subwords_path):
            self.subwords = self.load_from_file(subwords_path)
        else:
            self.subwords = self.build_from_corpus(subwords_corpus)
            self.save_to_file(subwords_path)

        self.blank = 0  # subword treats blank as 0
        self.num_classes = self.subwords.vocab_size

        self.init_upoints()

    def init_upoints(self):
        text = [""]  # first one is reserved for blank token
        for idx in np.arange(1, self.num_classes, dtype=np.int32):
            token = self.subwords.decode([idx])
            text.append(token)
        self.upoints = tf.strings.unicode_decode(text, "UTF-8")
        self.upoints = self.upoints.to_tensor()  # [num_classes, max_token_length]

    def build_from_corpus(self, corpus_files: list = None):
        def corpus_generator():
            for file in corpus_files:
                with open(file, "r", encoding="utf-8") as f:
                    lines = f.read().splitlines()
                    lines = lines[1:]
                for line in lines:
                    line = line.split("\t")
                    yield line[-1]

        subwords = tds.deprecated.text.SubwordTextEncoder.build_from_corpus(
            corpus_generator(),
            self.target_vocab_size,
            self.max_token_length,
            self.max_corpus_chars,
            self.reserved_tokens,
        )
        return subwords

    def load_from_file(self, filename: str = None):
        filename = preprocess_paths(filename)
        filename_prefix = os.path.splitext(filename)[0]
        subwords = tds.deprecated.text.SubwordTextEncoder.load_from_file(
            filename_prefix
        )
        return subwords

    def save_to_file(self, filename: str = None):
        filename = preprocess_paths(filename)
        filename_prefix = os.path.splitext(filename)[0]
        return self.subwords.save_to_file(filename_prefix)

    def string2Indices(self, text: str) -> tf.Tensor:
        if self.normalize:
            text = self.preprocess_text(text)

        text = text.strip()  # remove trailing space
        indices = self.subwords.encode(text)
        return tf.convert_to_tensor(indices, dtype=tf.int32)

    def indices2String(self, indices: tf.Tensor) -> tf.Tensor:
        with tf.device("/CPU:0"):  # string data is not supported on GPU
            total = tf.shape(indices)[0]
            batch = tf.constant(0, dtype=tf.int32)
            transcripts = tf.TensorArray(
                dtype=tf.string,
                size=total,
                dynamic_size=False,
                infer_shape=False,
                clear_after_read=False,
                element_shape=tf.TensorShape([]),
            )

            def cond(batch, total, _):
                return tf.less(batch, total)

            def body(batch, total, transcripts):
                upoints = self.indices2upoints(indices[batch])
                transcripts = transcripts.write(
                    batch, tf.strings.unicode_encode(upoints, "UTF-8")
                )
                return batch + 1, total, transcripts

            _, _, transcripts = tf.while_loop(
                cond, body, loop_vars=[batch, total, transcripts]
            )

            return transcripts.stack()

    @tf.function(input_signature=[tf.TensorSpec([None], dtype=tf.int32)])
    def indices2upoints(self, indices: tf.Tensor) -> tf.Tensor:
        indices = self.normalize_indices(indices)
        upoints = tf.gather_nd(self.upoints, tf.expand_dims(indices, axis=-1))
        return tf.gather_nd(upoints, tf.where(tf.not_equal(upoints, 0)))


import os
from scipy import io
# from glob import glob
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from PIL import Image
from support_functions import relabel
import mediapipe as mp
import cv2

change_background_mp = mp.solutions.selfie_segmentation
change_bg_segment = change_background_mp.SelfieSegmentation()

class DataLoader():
    def __init__(self, img_res=(512,512), test_size=0.1, val_size=0.1):
        self.img_res = img_res[::-1]
        path = os.path.join(os.getcwd(), "dataset/annotations")
        all_files = os.listdir(path)
        self.train_files, self.test_files = train_test_split(all_files, test_size=test_size+val_size)
        self.test_files, self.val_files = train_test_split(self.test_files, test_size=val_size/(test_size+val_size))

    def load_data(self, batch_size=1, is_testing=False):
        # data_type = "train" if not is_testing else "test"
        # path = glob('./datasets/%s/%s/*' % (self.dataset_name, data_type))
        if is_testing:
            path = self.test_files
        else:
            path = self.train_files
        batch = list(np.random.choice(path, size=batch_size))
        imgs_A = []
        imgs_B = []
        for image in batch:
            # Read the images
            img_A = self.imread(image)
            img_A = self.resize_img(img_A)
            # Read the annotations
            img_B = self.annread(image)
            img_B = self.resize_img(img_B)
            # If training => do random flip
            if not is_testing and np.random.random() < 0.5:
                img_A = np.fliplr(img_A)
                img_B = np.fliplr(img_B)
            imgs_A.append(img_A)
            imgs_B.append(img_B)

        imgs_A = np.array(imgs_A)/127.5 - 1.
        imgs_B = np.array(imgs_B)/127.5 - 1.

        return imgs_A, imgs_B

    def load_batch(self, batch_size=1, is_testing=False):
        # data_type = "train" if not is_testing else "val"
        # path = glob('./datasets/%s/%s/*' % (self.dataset_name, data_type))
        if is_testing:
            path = self.val_files
        else:
            path = self.train_files
        self.n_batches = int(len(path) / batch_size)
        for i in range(self.n_batches-1):
            batch = path[i*batch_size:(i+1)*batch_size]
            imgs_A, imgs_B = [], []
            for img in batch:
                img_A = self.imread(img)
                img_B = self.annread(img)

                img_A = self.resize_img(img_A)
                img_B = self.resize_img(img_B)
                # h, w, _ = img.shape
                # half_w = int(w/2)
                # img_A = img[:, :half_w, :]
                # img_B = img[:, half_w:, :]

                # img_A = scipy.misc.imresize(img_A, self.img_res)
                # img_B = scipy.misc.imresize(img_B, self.img_res)

                if not is_testing and np.random.random() > 0.5:
                        img_A = np.fliplr(img_A)
                        img_B = np.fliplr(img_B)

                imgs_A.append(img_A)
                imgs_B.append(img_B)

            imgs_A = np.array(imgs_A)/127.5 - 1.
            imgs_B = np.array(imgs_B)/127.5 - 1.

            yield imgs_A, imgs_B

    def imread(self, path):
        path = os.path.join(os.getcwd(), "dataset/photos", path).replace("mat", "jpg")
        sample_img = cv2.imread(path)
        RGB_sample_img = cv2.cvtColor(sample_img, cv2.COLOR_BGR2RGB)
        result = change_bg_segment.process(RGB_sample_img)
        binary_mask = result.segmentation_mask > 0.9
        binary_mask_3 = np.dstack((binary_mask,binary_mask,binary_mask))
        output_image = np.where(binary_mask_3, sample_img, 255)[:,:,::-1]
        # plt.imshow(output_image);plt.title("Output Image");plt.axis('off');
        # return plt.imread(path).astype(float)
        return output_image.astype(float)

    def annread(self, path):
        path = os.path.join(os.getcwd(), "dataset/annotations", path)
        anno = io.loadmat(path)["groundtruth"].astype(float)
        # anno = np.repeat(anno[:, :, np.newaxis], 3, axis=2)
        anno = np.array(Image.fromarray(anno.astype(np.uint8)).convert('RGB')).astype('int')
        anno = relabel(anno) # Relabelling the segmented image
        return anno

    def resize_img(self, img):
        return np.array(Image.fromarray(img.astype(np.uint8)).resize(self.img_res)).astype('float32')


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

from django.db import models, migrations
from django.conf import settings

class Migration(migrations.Migration):

    dependencies = [
        ('programs', '0001_initial'),
        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
        ('videos', '0001_initial'),
        ('people', '0001_initial'),
        ('geographies', '0001_initial'),
    ]

    operations = [
        migrations.CreateModel(
            name='PersonAdoptPractice',
            fields=[
                ('time_created', models.DateTimeField(auto_now_add=True, null=True)),
                ('time_modified', models.DateTimeField(auto_now=True, null=True)),
                ('id', models.AutoField(serialize=False, primary_key=True)),
                ('old_coco_id', models.BigIntegerField(null=True, editable=False)),
                ('date_of_adoption', models.DateField()),
                ('group', models.ForeignKey(to='people.PersonGroup')),
                ('partner', models.ForeignKey(verbose_name=b'Supply Partner', to='programs.Partner')),
                ('topic', models.ForeignKey(to='videos.Topic')),
                ('user_created', models.ForeignKey(related_name='activities_personadoptpractice_created', blank=True, editable=False, to=settings.AUTH_USER_MODEL, null=True)),
                ('user_modified', models.ForeignKey(related_name='activities_personadoptpractice_related_modified', blank=True, editable=False, to=settings.AUTH_USER_MODEL, null=True)),
            ],
            options={
            },
            bases=(models.Model,),
        ),
        migrations.CreateModel(
            name='PersonMeetingAttendance',
            fields=[
                ('time_created', models.DateTimeField(auto_now_add=True, null=True)),
                ('time_modified', models.DateTimeField(auto_now=True, null=True)),
                ('id', models.AutoField(serialize=False, primary_key=True)),
                ('old_coco_id', models.BigIntegerField(null=True, editable=False)),
                ('interested', models.BooleanField(db_index=True)),
                ('expressed_question', models.CharField(max_length=500, blank=True)),
                ('expressed_adoption_video', models.ForeignKey(related_name='expressed_adoption_video', blank=True, to='videos.Video', null=True)),
                ('person', models.ForeignKey(to='people.Person')),
            ],
            options={
                'abstract': False,
            },
            bases=(models.Model,),
        ),
        migrations.CreateModel(
            name='Screening',
            fields=[
                ('time_created', models.DateTimeField(auto_now_add=True, null=True)),
                ('time_modified', models.DateTimeField(auto_now=True, null=True)),
                ('id', models.AutoField(serialize=False, primary_key=True)),
                ('old_coco_id', models.BigIntegerField(null=True, editable=False)),
                ('date', models.DateField()),
                ('start_time', models.TimeField()),
                ('end_time', models.TimeField()),
                ('location', models.CharField(max_length=200, blank=True)),
                ('animator', models.ForeignKey(verbose_name=b'Field Officer', to='people.Animator')),
                ('farmer_groups_targeted', models.ManyToManyField(to='people.PersonGroup', verbose_name=b'Farmer Families')),
                ('farmers_attendance', models.ManyToManyField(to='people.Person', null=b'False', through='activities.PersonMeetingAttendance', blank=b'False')),
                ('partner', models.ForeignKey(verbose_name=b'Supply Partner', to='programs.Partner')),
                ('user_created', models.ForeignKey(related_name='activities_screening_created', blank=True, editable=False, to=settings.AUTH_USER_MODEL, null=True)),
                ('user_modified', models.ForeignKey(related_name='activities_screening_related_modified', blank=True, editable=False, to=settings.AUTH_USER_MODEL, null=True)),
                ('videoes_screened', models.ManyToManyField(to='videos.Video', verbose_name=b'Videos Screened')),
                ('village', models.ForeignKey(to='geographies.Village')),
            ],
            options={
            },
            bases=(models.Model,),
        ),
        migrations.AlterUniqueTogether(
            name='screening',
            unique_together=set([('date', 'start_time', 'end_time', 'animator', 'village')]),
        ),
        migrations.AddField(
            model_name='personmeetingattendance',
            name='screening',
            field=models.ForeignKey(to='activities.Screening'),
            preserve_default=True,
        ),
        migrations.AddField(
            model_name='personmeetingattendance',
            name='user_created',
            field=models.ForeignKey(related_name='activities_personmeetingattendance_created', blank=True, editable=False, to=settings.AUTH_USER_MODEL, null=True),
            preserve_default=True,
        ),
        migrations.AddField(
            model_name='personmeetingattendance',
            name='user_modified',
            field=models.ForeignKey(related_name='activities_personmeetingattendance_related_modified', blank=True, editable=False, to=settings.AUTH_USER_MODEL, null=True),
            preserve_default=True,
        ),
        migrations.AlterUniqueTogether(
            name='personadoptpractice',
            unique_together=set([('group', 'topic', 'date_of_adoption')]),
        ),
    ]


import pandas as pd
from FlokAlgorithmLocal import FlokDataFrame, FlokAlgorithmLocal
import json
import sys, os
import numpy as np
import librosa
import cv2
from pandas import Series, DataFrame

class Batch_AudioMFCC(FlokAlgorithmLocal):
    def array2str(self,arr):
        return " ".join([str(x) for x in arr])
    def run(self, inputDataSets, params):
        audio_dict = inputDataSets.get(0)
        text_dict={}
        for audio_name, audio in audio_dict.items():
            y = audio[0]
            sr = audio[1]
            mfccs = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=int(params.get("vector_size")))
            text_name=audio_name.split('.')[0]+"_MFCC.txt"
            #将向量变为行
            mfccs = mfccs.T
            ans_list=[self.array2str(vec) for vec in mfccs]
            text_dict[text_name]="\n".join(ans_list)
        result = FlokDataFrame()
        result.addDF(text_dict)
        return result

if __name__ == "__main__":
    all_info = json.loads(sys.argv[1])
    # f = open("test.json", encoding = 'utf-8')
    # all_info = json.loads(f)
    # all_info = {
    #         "input": ["in.mp3"],
    #         "inputFormat": ["mp3"],
    #         "inputLocation":["local_fs"],
    #         "output": ["test_mfcc.csv"],
    #         "outputFormat": ["csv"],
    #         "outputLocation": ['local_fs'],
    #         "parameters": {"mfcc_num": 20
    #                        }
    #     }

    params = all_info["parameters"]
    inputPaths = all_info["input"]
    inputTypes = all_info["inputFormat"]
    inputLocation = all_info["inputLocation"]
    outputPaths = all_info["output"]
    outputTypes = all_info["outputFormat"]
    outputLocation = all_info["outputLocation"]
    algorithm = Batch_AudioMFCC()
    audioList = algorithm.read(inputPaths, inputTypes, inputLocation, outputPaths, outputTypes)
    result = algorithm.run(audioList, params)
    algorithm.write(outputPaths, result, outputTypes, outputLocation)


from KalmanNet_nn import KalmanNetNN
import torch

if torch.cuda.is_available():
    dev = torch.device("cuda:0")
    torch.set_default_tensor_type("torch.cuda.FloatTensor")
else:
    dev = torch.device("cpu")

def NNBuild(SysModel):

    Model = KalmanNetNN()

    Model.InitSystemDynamics(SysModel.f, SysModel.h, SysModel.m, SysModel.n, SysModel.R, infoString = "partialInfo")
    Model.InitSequence(SysModel.m1x_0, SysModel.m2x_0, SysModel.T)

    Model.InitKGainNet(SysModel.prior_Q, SysModel.prior_Sigma, SysModel.prior_S)

    return Model


import sys
from datetime import datetime
import json
import utilities
import numpy as np

#---------------------------------
#python3 main.py  ./crimes-new-york-city/NYPD_Complaint_Data_Historic.csv

#Uncomment this if you want to output a json file dump of dateString => noOfCrimes for that date

# pathToNYPDCrimeDataset = sys.argv[1]
# f = open(pathToNYPDCrimeDataset, "r")
# firstLine = True
# headerLine = None
# currentDate = None
# dateToNoOfCrimesDict = {} #maps dateString => noOfCrimes
# for line in f:
#    if firstLine == True:
#       headerLine = line
#       firstLine = False
#    else:
#       tokens = line.split(",")
#       dateString = tokens[1].strip()
#       if "2015" not in dateString and "2014" not in dateString:
#          continue
#       try:
#          currentLineDate = datetime.strptime(dateString, "%m/%d/%Y").date()
#          utilities.addCountInDictionaryForKey(dateToNoOfCrimesDict, dateString)
#       except:
#          continue

# with open("./dateToNoOfCrimesDict.json", "w") as dateToNoOfCrimesDictFileHandler:
#    json.dump(dateToNoOfCrimesDict, dateToNoOfCrimesDictFileHandler)
# f.close()
#---------------------------------

#1048576 lines
# with open("./dateToNoOfCrimesDict.json") as dateToNoOfCrimesDictFileHandler:
#    dateToNoOfCrimesDict = json.load(dateToNoOfCrimesDictFileHandler)
#    summation = 0
#    for key, value in dateToNoOfCrimesDict.items():
#       print(key + " " + str(value))
#       summation += value
#    print(summation)

#---------------------------------
#/Users/joewijoyo/anaconda3/bin/python3.7 main.py ./historical-hourly-weather-data/temperature.csv ./historical-hourly-weather-data/wind_speed.csv ./historical-hourly-weather-data/humidity.csv ./historical-hourly-weather-data/weather_description.csv historical-hourly-weather-data/pressure.csv historical-hourly-weather-data/wind_direction.csv

#Uncomment this to get dateToWeatherDataDict outputted to the current directory as a json dump file

# WEATHER_DATASET_DATE_COLUMN_INDEX = 0
# WEATHER_DATASET_NEW_YORK_COLUMN_INDEX = 28

# def getWeatherDataFromWeatherCSVFile(pathToWeatherDataset, dateToWeatherDataDict, keyString):
#    firstLine = True
#    f = open(pathToWeatherDataset, "r")
#    for line in f:
#       if firstLine == True:
#          headerLine = line
#          firstLine = False
#          #print(headerLine.split(",")[WEATHER_DATASET_NEW_YORK_COLUMN_INDEX])
#       else:
#          tokens = line.split(",")
#          dateString = tokens[WEATHER_DATASET_DATE_COLUMN_INDEX].split(" ")[0]
#          if "2015" not in dateString and "2014" not in dateString:
#             continue
#          try:
#             currentLineDate = datetime.strptime(dateString, "%Y-%m-%d").date()
#             currentLineDate = datetime.strftime(currentLineDate, "%m/%d/%Y")
#             if currentLineDate not in dateToWeatherDataDict:
#                dateToWeatherDataDict[currentLineDate] = {
#                   "humidityValuesList" : [],
#                   "pressureValuesList" : [],
#                   "temperatureValuesList" : [],
#                   "weatherDescriptionValuesList" : [],
#                   "windDirectionValuesList" : [],
#                   "windSpeedValuesList" : []
#                }
#             appendValue = None
#             try:
#                appendValue = float(tokens[WEATHER_DATASET_NEW_YORK_COLUMN_INDEX])
#             except:
#                appendValue = tokens[WEATHER_DATASET_NEW_YORK_COLUMN_INDEX].strip()
#             if appendValue != "":
#                dateToWeatherDataDict[currentLineDate][keyString].append(appendValue)

#          except:
#             continue

#    f.close()

# dateToWeatherDataDict = {}
# """
#    dateString => {
#       humidity:
#       pressure:
#       temperature:
#       weatherDescription:
#       windDirection:
#       windSpeed:
#    }
# """
# for pathToWeatherDataset in sys.argv[1:]:
#    keyString = None
#    if "temperature" in pathToWeatherDataset:
#       keyString = "temperatureValuesList"
#    elif "wind_speed" in pathToWeatherDataset:
#       keyString = "windSpeedValuesList"
#    elif "humidity" in pathToWeatherDataset:
#       keyString = "humidityValuesList"
#    elif "weather_description" in pathToWeatherDataset:
#       keyString = "weatherDescriptionValuesList"
#    elif "pressure" in pathToWeatherDataset:
#       keyString = "pressureValuesList"
#    else:
#       keyString = "windDirectionValuesList"
#    getWeatherDataFromWeatherCSVFile(pathToWeatherDataset, dateToWeatherDataDict, keyString)

# for dateString, weatherDataDict in dateToWeatherDataDict.items():
#    dateToWeatherDataDict[dateString]["humidity"] = np.mean(weatherDataDict["humidityValuesList"])
#    dateToWeatherDataDict[dateString]["temperature"] = np.mean(weatherDataDict["temperatureValuesList"])
#    dateToWeatherDataDict[dateString]["windSpeed"] = np.mean(weatherDataDict["windSpeedValuesList"])
#    dateToWeatherDataDict[dateString]["weatherDescription"] = utilities.getMostAppearingValueInList(weatherDataDict["weatherDescriptionValuesList"])
#    dateToWeatherDataDict[dateString]["pressure"] = np.mean(weatherDataDict["pressureValuesList"])
#    dateToWeatherDataDict[dateString]["windDirection"] = np.mean(weatherDataDict["windDirectionValuesList"])
#    del(weatherDataDict["humidityValuesList"])
#    del(weatherDataDict["temperatureValuesList"])
#    del(weatherDataDict["windSpeedValuesList"])
#    del(weatherDataDict["weatherDescriptionValuesList"])
#    del(weatherDataDict["pressureValuesList"])
#    del(weatherDataDict["windDirectionValuesList"])
#    print(weatherDataDict)

# with open("./dateToWeatherDataDict.json", "w") as dateToWeatherDataDictFileHandler:
#     json.dump(dateToWeatherDataDict, dateToWeatherDataDictFileHandler)

#need to make a csv file where weather attributes are on the line and
#---------------------------------

import xml.etree.ElementTree
import xml.etree.cElementTree as ET

dateToWeatherDataDict = {}
with open("./dateToWeatherDataDict.json") as dateToWeatherDataDictFileHandler:
   dateToWeatherDataDict = json.load(dateToWeatherDataDictFileHandler)

dateToNoOfCrimesDict = {}
with open("./dateToNoOfCrimesDict.json") as dateToNoOfCrimesDictFileHandler:
   dateToNoOfCrimesDict = json.load(dateToNoOfCrimesDictFileHandler)

allUniqueValuesDict = {
   "humidity" : [],
   "temperature" : [],
   "windSpeed" : [],
   "weatherDescription" : [],
   "pressure" : [],
   "windDirection" : [],
   "crimeRate" : [],
} #maps column => [unique values of column]
count = 0
noOfLessThanOrEqualTo800 = 0
noOf801To1000 = 0
noOf1001To1200 = 0
noOf1201To1400 = 0
noOfGreaterThan1400 = 0
for dateString, noOfCrimes in dateToNoOfCrimesDict.items():
   if dateString in dateToWeatherDataDict:
      weatherDataDict = dateToWeatherDataDict[dateString]
      crimeRateString = None
      if noOfCrimes <= 800:
         crimeRateString = "<=800"
         noOfLessThanOrEqualTo800 += 1
      elif noOfCrimes <= 1000:
         crimeRateString = "801-1000"
         noOf801To1000 += 1
      elif noOfCrimes <= 1200:
         crimeRateString = "1001-1200"
         noOf1001To1200 += 1
      elif noOfCrimes <= 1400:
         crimeRateString = "1201-1400"
         noOf1201To1400 += 1
      else:
         crimeRateString = ">1400"
         noOfGreaterThan1400 += 1
      if str(weatherDataDict["weatherDescription"]) not in allUniqueValuesDict["weatherDescription"]:
         allUniqueValuesDict["weatherDescription"].append(str(weatherDataDict["weatherDescription"]))
      if crimeRateString not in allUniqueValuesDict["crimeRate"]:
         allUniqueValuesDict["crimeRate"].append(crimeRateString)
      count += 1
print(count)
print("No of <=800: " + str(noOfLessThanOrEqualTo800))
print("No of 801-1000: " + str(noOf801To1000))
print("No of 1001-1200: " + str(noOf1001To1200))
print("No of 1201-1400: " + str(noOf1201To1400))
print("No of >1400: " + str(noOfGreaterThan1400))

# with open("./weatherDataAndNoOfCrimesData.csv", "w") as outputCSVFileHandler:
#    outputCSVFileHandler.write("humidity,temperature,windSpeed,weatherDescription,pressure,windDirection,crimeRate\n")
#    outputCSVFileHandler.write("1,1,1,1,1,1,1\n")
#    outputCSVFileHandler.write("crimeRate\n")
#    for dateString, noOfCrimes in dateToNoOfCrimesDict.items():
#       if dateString in dateToWeatherDataDict:
#          crimeRateString = None
#          if noOfCrimes <= 800:
#             crimeRateString = "<=800"
#          elif noOfCrimes <= 1000:
#             crimeRateString = "801-1000"
#          elif noOfCrimes <= 1200:
#             crimeRateString = "1001-1200"
#          elif noOfCrimes <= 1400:
#             crimeRateString = "1201-1400"
#          else:
#             crimeRateString = ">1400"
#          weatherDataDict = dateToWeatherDataDict[dateString]
#          writeString = str(weatherDataDict["humidity"]) + "," + str(weatherDataDict["temperature"]) + "," \
#                           + str(weatherDataDict["windSpeed"]) + "," + str(allUniqueValuesDict["weatherDescription"].index(weatherDataDict["weatherDescription"]) + 1) + "," \
#                           + str(weatherDataDict["pressure"]) + "," + str(weatherDataDict["windDirection"]) + "," \
#                           + str(allUniqueValuesDict["crimeRate"].index(crimeRateString) + 1) + "\n"
#          outputCSVFileHandler.write(writeString)
#          print(noOfCrimes)

#          if str(weatherDataDict["humidity"]) not in allUniqueValuesDict["humidity"]:
#             allUniqueValuesDict["humidity"].append(str(weatherDataDict["humidity"]))
#          if str(weatherDataDict["temperature"]) not in allUniqueValuesDict["temperature"]:
#             allUniqueValuesDict["temperature"].append(str(weatherDataDict["temperature"]))
#          if str(weatherDataDict["windSpeed"]) not in allUniqueValuesDict["windSpeed"]:
#             allUniqueValuesDict["windSpeed"].append(str(weatherDataDict["windSpeed"]))
#          # if str(weatherDataDict["weatherDescription"]) not in allUniqueValuesDict["weatherDescription"]:
#          #    allUniqueValuesDict["weatherDescription"].append(str(weatherDataDict["weatherDescription"]))
#          if str(weatherDataDict["pressure"]) not in allUniqueValuesDict["pressure"]:
#             allUniqueValuesDict["pressure"].append(str(weatherDataDict["pressure"]))
#          if str(weatherDataDict["windDirection"]) not in allUniqueValuesDict["windDirection"]:
#             allUniqueValuesDict["windDirection"].append(str(weatherDataDict["windDirection"]))
#          # if crimeRateString not in allUniqueValuesDict["crimeRate"]:
#          #    allUniqueValuesDict["crimeRate"].append(crimeRateString)

#          #outputCSVFileHandler.write("")

# #CODE BELOW TO CREATE XML DOMAIN FILE
# domainRootElement = ET.Element("domain")

# humidityElement = ET.SubElement(domainRootElement, "variable", name="humidity", continuous="true")
# for uniqueElement in allUniqueValuesDict["humidity"]:
#    ET.SubElement(humidityElement, "group", name=uniqueElement, p=".00")

# temperatureElement = ET.SubElement(domainRootElement, "variable", name="temperature", continuous="true")
# for uniqueElement in allUniqueValuesDict["temperature"]:
#    ET.SubElement(temperatureElement, "group", name=uniqueElement, p=".00")

# windSpeedElement = ET.SubElement(domainRootElement, "variable", name="windSpeed", continuous="true")
# for uniqueElement in allUniqueValuesDict["windSpeed"]:
#    ET.SubElement(windSpeedElement, "group", name=uniqueElement, p=".00")

# weatherDescriptionElement = ET.SubElement(domainRootElement, "variable", name="weatherDescription")
# for uniqueElement in allUniqueValuesDict["weatherDescription"]:
#    ET.SubElement(weatherDescriptionElement, "group", name=uniqueElement, p=".00")

# pressureElement = ET.SubElement(domainRootElement, "variable", name="pressure", continuous="true")
# for uniqueElement in allUniqueValuesDict["pressure"]:
#    ET.SubElement(pressureElement, "group", name=uniqueElement, p=".00")

# windDirectionElement = ET.SubElement(domainRootElement, "variable", name="windDirection", continuous="true")
# for uniqueElement in allUniqueValuesDict["windDirection"]:
#    ET.SubElement(windDirectionElement, "group", name=uniqueElement, p=".00")

# crimeRateElement = ET.SubElement(domainRootElement, "Category", name="crimeRate")
# for index, uniqueElement in enumerate(allUniqueValuesDict["crimeRate"]):
#    ET.SubElement(crimeRateElement, "choice", name=uniqueElement, type=str(index))

# ET.ElementTree(domainRootElement).write("weatherDataAndNoOfCrimesDataDomain.xml")


import typing
import uuid
from dataclasses import dataclass, field
from enum import IntEnum, auto

from .message import Message, Trace

@dataclass(frozen=True, kw_only=True)
class CommandReply(Message):
    """
    A command reply message.

    Every command needs to receive a reply in the form of a ``CommandReply`` message. The reply contains
    information about its ``success``, as well as a text message which is usually used to describe reasons for
    failures.

    Attributes:
        success: Whether the command succeeded.
        message: Arbitrary text, usually used to describe reasons for failures.
        unique: The unique identifier of its corresponding command.
    """
    class FailType(IntEnum):
        """
        Used when a command failed.
        """
        NONE = 0
        TIMEOUT = auto()
        EXCEPTION = auto()
        UNKNOWN = auto()

    success: bool = True
    message: str = ""

    unique: Trace = field(default_factory=uuid.uuid4)

CommandReplyType = typing.TypeVar("CommandReplyType", bound=CommandReply)  # pylint: disable=invalid-name

CommandDoneCallback = typing.Callable[[CommandReplyType, bool, str], None]
CommandFailCallback = typing.Callable[[CommandReply.FailType, str], None]


#Package Imports
import json
import multiprocessing as mp
import os
import random
import glob
import copy
import math
from os.path import abspath, dirname, exists, join
import sys
sys.path.append("./")
from pathlib import Path
from sklearn.model_selection import train_test_split

import ase
from ase import Atoms, Atom
from ase.visualize import view
from ase.io import read, write
from ase.io.trajectory import Trajectory
from ase.visualize.plot import plot_atoms
from ase.constraints import dict2constraint
from ase.calculators.singlepoint import SinglePointCalculator
from ase.db import connect
from ase.neighborlist import NeighborList, natural_cutoffs
from ase.io import extxyz

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

import seaborn as sns
import networkx as nx
import mendeleev
import itertools
import pickle
import shutil
from datetime import datetime

import torch
from torch_geometric.loader import DataLoader
from torch_geometric.utils import from_networkx
from torch_geometric.data import Data
from torch_geometric.utils import to_networkx

from pymatgen.io.ase import AseAtomsAdaptor
from scipy.spatial import distance_matrix

class GaussianDistance():
    """
    Expands the distance by Gaussian basis.
    Unit: angstrom
    """

    def __init__(self, dmin=0.0, dmax=6, step=0.2, var=None):
        """
        Args:
            dmin (float): Minimum interatomic distance
            dmax (float): Maximum interatomic distance
            step (float): Step size for the Gaussian filter
            var (float, optional): Variance of Gaussian basis. Defaults to step if not given
        """
        assert dmin < dmax
        assert dmax - dmin > step

        self.filter = np.arange(dmin, dmax + step, step)
        self.embedding_size = len(self.filter)

        if var is None:
            var = step

        self.var = var

    def expand(self, distances):
        """Apply Gaussian distance filter to a numpy distance array
        Args:
            distances (ArrayLike): A distance matrix of any shape
        Returns:
            Expanded distance matrix with the last dimension of length
            len(self.filter)
        """
        distances = np.array(distances)

        return np.exp(-((distances[..., np.newaxis] - self.filter) ** 2) / self.var ** 2)

    def get_layer(self,slb, thresh = 1):
        """
        * slb (atoms object): The system, WITHOUT adsorbate attached
        * thresh (float): The threshold in Angstroms of the comparison between z heights
        * layer_remapped (list of floats): An ordered list of the layer heights identified in the system
        * num_layers (int): # of layers in system
        * num_atoms (int): # of atoms in system
        * layer_ID_Indices (list of lists): A list of lists of atom ID #s corresponding to the atoms placed in each layer classification.
          NOTE: layer_ID_Indices[0] contains the ID for atoms in the ASE system that were identified as belonging to the height in layer_remapped[0]

        """

        layer_ID_Indices = []
        similar_idx = []
        num_atoms = len(slb)
        slab_z_list= slb.positions[:,2].tolist()
        unique_z = np.unique(slab_z_list).tolist()

        # For the unique heights in the system...
        # loop, and see if any other entries are close
        # If they are, save the indices of these redundant heights so we can eliminate them
        for h1 in unique_z:
            for h2 in unique_z:
                if h1 != h2 and unique_z.index(h2) > unique_z.index(h1) and math.isclose(h1, h2, abs_tol = thresh):
                    similar_idx.append(h2)

        # This gets all the distinct z values together by removing the entries in the similar list
        layer_z = [unique_z[i] for i in range(len(unique_z)) if unique_z[i] not in similar_idx]
        # This is how many layers we think we have
        num_layers = len(layer_z)
        # This is the set of layer heights we found
        layer_heights = layer_z

        # Find the indices present in each layer we have found
        slotted_indices = []
        for i in range(num_layers):
            layer_ID_Indices.append([])
            for j in range(num_atoms):
                if math.isclose(slab_z_list[j], layer_heights[i], abs_tol = 1) == True and j not in slotted_indices:
                    layer_ID_Indices[i].append(j)
                    slotted_indices.append(j)
        # This next section organizes the indices so that we have each list of indices for a layer ordered in terms of layer ascending/descending instead of randomly
        layer_remapped = copy.deepcopy(layer_heights)
        layer_remapped.sort()
        ordering = [layer_heights.index(layers) for layers in layer_remapped]
        #print("Ordering", ordering)
        layer_ID_Indices = [layer_ID_Indices[i] for i in ordering]

        return layer_remapped, ordering, layer_ID_Indices

#This class handle model/data imports
class ASE_To_Graph():
    def __init__(self,db_file,embedding):
        '''
        inputs:
        db_file (str) - filename/path of the ASE db file
        embedding (str) - name of the embedding you want to implement
        '''

        self.db_filename = db_file
        self.embedding = embedding

    def import_db(self):
        '''
        self explanatory function
        '''

        with connect(self.db_filename) as db:
            ase_db = db
        self.db = ase_db
        return ase_db

    #Not used in algorithm, just helps with data inspection
    def bond_symbol(self,atoms, a1, a2):
        '''
        self explanatory function
        '''

        return "{}{}".format(*sorted((atoms[a1].symbol, atoms[a2].symbol)))

    #Get electronegativity (EN)
    def get_EN(self,atoms, a1, a2):
        '''
        self explanatory function
        '''

        try:
            avg = (mendeleev.element(atoms[a1].symbol).en_pauling + mendeleev.element(atoms[a2].symbol).en_pauling)/2
        except:
            avg = 0.0
        try:
            dif = abs(mendeleev.element(atoms[a1].symbol).en_pauling - mendeleev.element(atoms[a2].symbol).en_pauling)
        except:
            dif = 0
        return np.array([avg, dif])

    def node_symbol(self,atom):
        '''
        self explanatory function
        '''

        return "".format(atom.symbol)

    def add_atoms_node(self,graph, atoms, a1):
        '''
        self explanatory function
        '''

        graph.add_node(self.node_symbol(atoms[a1]), index=a1)

    def add_atoms_edge(self,graph, atoms, a1, a2):
        '''
        self explanatory function
        '''

        EN = self.get_EN(atoms, a1, a2)
        distance = atoms.get_distances(a1,a2)
        graph.add_edge(a1,
                       a2,
                       EN= EN,
                       distance = distance)

    #Main surface analysis automation function
    def analyse_surface(self,atoms, adsorbate,radii_multiplier=1.1, skin = 0.25, **kwargs):
        '''
        inputs:
        atoms (ASE atoms) - atom to analyze
        adsorbate (str) - adsorbate to analyze
        radii_multiplier (float) - atomic radius factor
        skin (float) - skin factor

        outputs:
        binding_atoms (list) - binding atom index
        adsorbate_atoms (list) - adsorbate atom index
        bonds (list) - connections
        '''

        #Extract nearest neighbor indexes
        nl = NeighborList(natural_cutoffs(atoms, radii_multiplier), self_interaction=False,
                              bothways=True, skin=skin)

        nl.update(atoms)
        #Extract index of adsorbate atoms
        adsorbate_atoms = [atom.index for atom in atoms if atom.symbol in adsorbate]

        # Add all atoms to graph
        distances = atoms.get_all_distances(mic=True)

        binding_atoms = []
        bonds = []
        for index, atom in enumerate(atoms):
            neighbors, offsets = nl.get_neighbors(index)
            for neighbor, offset in zip(neighbors, offsets):
                if np.all(offset==np.zeros(3)):
                    if distances[index][neighbor] > 2.5 and (bool(index in adsorbate_atoms) ^ bool(neighbor in adsorbate_atoms)):

                        continue
                    if (bool(index not in adsorbate_atoms) ^ bool(neighbor not in adsorbate_atoms)) and (index not in adsorbate_atoms):
                        binding_atoms.append(index)

                bonds.append((index, neighbor))

        return binding_atoms, adsorbate_atoms, bonds

    def extract_ads_geom(self, atms, adsorbate, elem_features, plot=False):
        '''
        inputs:
        atms (ASE Atoms) - slab to analyze
        adsorbate (str) - adsorbate to analyze
        elem_features (str) - embedding features initialized with the class
        plot (bool) - plot the exactracted structure

        outputs:
        surf_n (ASE atoms) - extracted surface
        full (nx graph) - resulting graph
        '''

        #Extract neighbors
        binding_atoms, adsorbate_atoms, bonds = self.analyse_surface(atms, adsorbate=adsorbate)
        nl = NeighborList(natural_cutoffs(atms, 1.1), self_interaction=False,
                                  bothways=True, skin=0.25)
        nl.update(atms)
        poss = [atms.positions[i] for i in binding_atoms]
        symbols = [atms[i].symbol for i in binding_atoms]
        distances =atms.get_all_distances(mic=True)
        #extract coordiantes of atoms
        for atm in  binding_atoms:
            neighbors, offsets = nl.get_neighbors(atm)
            for i, offset in zip(neighbors, offsets):
                poss.append(atms.positions[i] + np.dot(offset, atms.get_cell()))
                symbols.append(atms[i].symbol)
        #extract ads atoms coords
        for atm in  adsorbate_atoms:
            poss.append(atms.positions[atm])
            symbols.append(atms[atm].symbol)

        coords = np.vstack(poss)

        #generate extracted neighbors structure
        surf_n = Atoms(symbols,positions=coords )
        ase.geometry.get_duplicate_atoms(surf_n, cutoff=0.1, delete=True)
        if plot:
            fig, ax = plt.subplots(1,2)
            plot_atoms(surf_n,ax[0],radii=0.9, rotation=('0x,0y,0z'))
            plot_atoms(surf_n,ax[1],radii=0.9, rotation=('-90x,45y,0z'))

        params = {"binding_atoms":binding_atoms,
                "adsorbate_atoms":adsorbate_atoms,
                "bonds":bonds}

        adsorbate_atms = [atom.index for atom in surf_n if atom.symbol in adsorbate]

        #Generate Graph
        full = nx.Graph()
        full.add_nodes_from(range(len(surf_n)))
        distances = surf_n.get_all_distances(mic=True)

        numbers = surf_n.get_atomic_numbers().tolist()
        symbols = surf_n.get_chemical_symbols()

        for i, feat_dict in full.nodes(data=True):
            feat_dict.update({'symbol':  symbols[i]})
            feat_dict.update({'atomic_number':  numbers[i]})
            feat_dict.update({'node_attr':  np.array(elem_features[symbols[i]]).astype(np.float32)})
        comb = list(itertools.combinations([i for i in range(len(surf_n))],2))

        for atm1, atm2 in comb:
            if atm1==atm2:
                continue
            if distances[atm1][atm2] > 2.8 :
                continue

            if distances[atm1][atm2] > 2.3 and (bool(atm1 in adsorbate_atms) ^ bool(atm2 in adsorbate_atms)):
                continue
            if distances[atm1][atm2] > 2.3 and (bool(atm1 in adsorbate_atms) and bool(atm2 in adsorbate_atms)):
                continue
            self.add_atoms_edge(full, surf_n,  atm1, atm2)
        return surf_n,  full #params

    #convert a ASE atoms object from an ASE db into a graph
    def get_data(self,row, atom_features,map_warning=False):
        '''
        inputs:
        row (ASE rows) - row to convert into a torch graph
        atom_features (str) - embedding features
        map_warning (bool) - print warnings about failed mappings
        '''

        gdf = GaussianDistance()
        adsorbate=row.Adsorbate
        atm = row.toatoms()
        slabId = row.Hea_ID
        ads_energy = row.Energy
        #embed graph object with features
        try:
            atms, G = self.extract_ads_geom(atms=atm, adsorbate=list(set(list(adsorbate))), elem_features=atom_features)
            data =from_networkx(G)

            data.x = data.node_attr
            edge_attr = data.EN[:,0]*data.distance[:,0]
            edge_attr.unsqueeze(-1)
            edge_attr = gdf.expand(edge_attr)
            edge_attr = torch.tensor(edge_attr)
            data.edge_attr = edge_attr
            data.__setattr__('idx',slabId)
            data.__setattr__('adsorbate',adsorbate)
            data.__setattr__('postions',torch.from_numpy(atms.positions))
            data.y = torch.tensor([ads_energy])

            try:
                atm1 = copy.deepcopy(atms)
                del(atm1[-len(list(adsorbate)):])
                atm1.center()
                _,_, layers = gdf.get_layer( atm1, thresh = 1)
                surface = layers[-1]
                subsurface = layers[-2]
                data.__setattr__('surface',[ atm1.symbols[i] for i in surface])
                data.__setattr__('subsurface',[ atm1.symbols[i] for i in subsurface])
            #Throw error
            except Exception as er1:
                if map_warning:
                    print(f" Warnn ing!!! surface and subsurface atoms not found in ID[{slabId}]")
                data.__setattr__('surface',[])
                data.__setattr__('subsurface',[])
            return data
        except Exception as err:
            print(f"{err}")
            return None

    def get_data_from_asedb(self, db,  adsorbate, emb="cgcnn92", ncpus=1):
        '''
        inputs:
        db (ASE db) - db to fully extract
        adsorbate (str) - adsorbate to analyze/extract
        emb (str) - embedding features
        ncpus (int) - core cpus
        '''

        gdf = GaussianDistance()
        rows = [row for row in db.select('Energy', Adsorbate=adsorbate) ]

        elem_emb = join(os.path.dirname(os.path.realpath("__file__")), f"element/{emb}.json")
        with open(elem_emb) as f:
                atom_features = json.load(f)

        if len(rows) > 100:
            print(f'Warning, {len(rows)} ASE atoms will likely take {len(rows)/10}+ min to complete')

        results = []
        for i in range(len(rows)):
            if i==0:
                results.append(self.get_data( rows[i],atom_features))
            elif i%20==0:
                print(f'{i}/{len(rows)} Graphs generated')
            else:
                results.append(self.get_data( rows[i],atom_features))

        datas = results

        with open(f"data_{adsorbate}_{emb}.pickle", 'wb') as handle:
            pickle.dump(datas, handle)

        self.adsorbate = adsorbate
        self.emb = emb
        self.pickle_name = f"data_{adsorbate}_{emb}.pickle"

    #Provide a sample visual of AGRA extraction
    def sample_visual(self):
        with open('element/cgcnn92.json') as f:
            atom_features = json.load(f)

        slab_sample_visual = self.db.get_atoms(1)
        extraction_sample_visual, g = self.extract_ads_geom(slab_sample_visual,
                                                           adsorbate=['O','C'],
                                                           elem_features=atom_features)

        ocp_a2g = OCP_AtomsToGraphs(
            max_neigh=50,
            radius=6,
            r_energy=False,
            r_forces=False,
            r_distances=True,
            r_edges=True,
            r_fixed=True,)

        ocp_slab = ocp_a2g.convert(slab_sample_visual)
        ocp_graph = to_networkx(ocp_slab)

        fig, ax = plt.subplots(2,2,figsize=(8,6),subplot_kw={'aspect': 'equal'})

        gs = gridspec.GridSpec(2, 2, width_ratios=[1, 1], height_ratios=[1, 1])
        ax[0, 0] = plt.subplot(gs[0, 0])
        ax[0, 1] = plt.subplot(gs[0, 1])
        ax[1, 0] = plt.subplot(gs[1, 0])
        ax[1, 1] = plt.subplot(gs[1, 1])

        plot_atoms(slab_sample_visual,ax[0,0],radii=0.9, rotation=('0x,0y,0z'))
        ax[0,0].set_title("ASE db Extracted Structure")
        ax[0,0].set_xticks([])
        ax[0,0].set_yticks([])

        plot_atoms(extraction_sample_visual,ax[0,1],radii=0.9, rotation=('0x,0y,0z'))
        ax[0,1].set_title("AGRA Refined Structure")
        ax[0,1].set_xticks([])
        ax[0,1].set_yticks([])

        nx.draw_networkx(g,with_labels=True,ax=ax[1,1])
        ax[1,1].set_title("Resulting AGRA Graph")

        nx.draw_networkx(ocp_graph,with_labels=True,ax=ax[1,0])
        ax[1,0].set_title("Resulting OCP Graph")

        plt.tight_layout()
        return None

    #split generated graph data into train/test data
    def split_data(self,split=0.1):
        '''
        inputs:
        split (float) - train/test split

        outputs:
        train/test (np.array) - train test datasets
        '''

        with open(self.pickle_name, 'rb') as handle:
            dataset= pickle.load(handle)

        train,test,ytrain,ytest = train_test_split(dataset,dataset,test_size=split)

        with open(f"data_{self.adsorbate}_Train_{self.emb}.pickle", 'wb') as handle:
            pickle.dump(train, handle)
        with open(f"data_{self.adsorbate}_Test_{self.emb}.pickle", 'wb') as handle:
            pickle.dump(test, handle)
        return train,test

#This entire class is copied over from the OCP github solely for the purpose of graph generation comparison (https://github.com/Open-Catalyst-Project/ocp)
class OCP_AtomsToGraphs:
    def __init__(
        self,
        max_neigh=200,
        radius=6,
        r_energy=False,
        r_forces=False,
        r_distances=False,
        r_edges=True,
        r_fixed=True,
        r_pbc=False,
    ):
        self.max_neigh = max_neigh
        self.radius = radius
        self.r_energy = r_energy
        self.r_forces = r_forces
        self.r_distances = r_distances
        self.r_fixed = r_fixed
        self.r_edges = r_edges
        self.r_pbc = r_pbc

    def _get_neighbors_pymatgen(self, atoms):
        """Preforms nearest neighbor search and returns edge index, distances,
        and cell offsets"""
        struct = AseAtomsAdaptor.get_structure(atoms)
        _c_index, _n_index, _offsets, n_distance = struct.get_neighbor_list(
            r=self.radius, numerical_tol=0, exclude_self=True
        )

        _nonmax_idx = []
        for i in range(len(atoms)):
            idx_i = (_c_index == i).nonzero()[0]
            # sort neighbors by distance, remove edges larger than max_neighbors
            idx_sorted = np.argsort(n_distance[idx_i])[: self.max_neigh]
            _nonmax_idx.append(idx_i[idx_sorted])
        _nonmax_idx = np.concatenate(_nonmax_idx)

        _c_index = _c_index[_nonmax_idx]
        _n_index = _n_index[_nonmax_idx]
        n_distance = n_distance[_nonmax_idx]
        _offsets = _offsets[_nonmax_idx]

        return _c_index, _n_index, n_distance, _offsets

    def _reshape_features(self, c_index, n_index, n_distance, offsets):
        """Stack center and neighbor index and reshapes distances,
        takes in np.arrays and returns torch tensors"""
        edge_index = torch.LongTensor(np.vstack((n_index, c_index)))
        edge_distances = torch.FloatTensor(n_distance)
        cell_offsets = torch.LongTensor(offsets)

        # remove distances smaller than a tolerance ~ 0. The small tolerance is
        # needed to correct for pymatgen's neighbor_list returning self atoms
        # in a few edge cases.
        nonzero = torch.where(edge_distances >= 1e-8)[0]
        edge_index = edge_index[:, nonzero]
        edge_distances = edge_distances[nonzero]
        cell_offsets = cell_offsets[nonzero]

        return edge_index, edge_distances, cell_offsets

    def convert(
        self,
        atoms,
    ):
        """Convert a single atomic stucture to a graph.
        Args:
            atoms (ase.atoms.Atoms): An ASE atoms object.
        Returns:
            data (torch_geometric.data.Data): A torch geometic data object with positions, atomic_numbers, tags,
            and optionally, energy, forces, distances, edges, and periodic boundary conditions.
            Optional properties can included by setting r_property=True when constructing the class.
        """

        # set the atomic numbers, positions, and cell
        atomic_numbers = torch.Tensor(atoms.get_atomic_numbers())
        positions = torch.Tensor(atoms.get_positions())
        cell = torch.Tensor(atoms.get_cell()).view(1, 3, 3)
        natoms = positions.shape[0]
        # initialized to torch.zeros(natoms) if tags missing.
        # https://wiki.fysik.dtu.dk/ase/_modules/ase/atoms.html#Atoms.get_tags
        tags = torch.Tensor(atoms.get_tags())

        # put the minimum data in torch geometric data object
        data = Data(
            cell=cell,
            pos=positions,
            atomic_numbers=atomic_numbers,
            natoms=natoms,
            tags=tags,
        )

        # optionally include other properties
        if self.r_edges:
            # run internal functions to get padded indices and distances
            split_idx_dist = self._get_neighbors_pymatgen(atoms)
            edge_index, edge_distances, cell_offsets = self._reshape_features(
                *split_idx_dist
            )

            data.edge_index = edge_index
            data.cell_offsets = cell_offsets
        if self.r_energy:
            energy = atoms.get_potential_energy(apply_constraint=False)
            data.y = energy
        if self.r_forces:
            forces = torch.Tensor(atoms.get_forces(apply_constraint=False))
            data.force = forces
        if self.r_distances and self.r_edges:
            data.distances = edge_distances
        if self.r_fixed:
            fixed_idx = torch.zeros(natoms)
            if hasattr(atoms, "constraints"):
                from ase.constraints import FixAtoms

                for constraint in atoms.constraints:
                    if isinstance(constraint, FixAtoms):
                        fixed_idx[constraint.index] = 1
            data.fixed = fixed_idx
        if self.r_pbc:
            data.pbc = torch.tensor(atoms.pbc)

        return data

    def convert_all(
        self,
        atoms_collection,
        processed_file_path=None,
        collate_and_save=False,
        disable_tqdm=False,
    ):
        """Convert all atoms objects in a list or in an ase.db to graphs.
        Args:
            atoms_collection (list of ase.atoms.Atoms or ase.db.sqlite.SQLite3Database):
            Either a list of ASE atoms objects or an ASE database.
            processed_file_path (str):
            A string of the path to where the processed file will be written. Default is None.
            collate_and_save (bool): A boolean to collate and save or not. Default is False, so will not write a file.
        Returns:
            data_list (list of torch_geometric.data.Data):
            A list of torch geometric data objects containing molecular graph info and properties.
        """

        # list for all data
        data_list = []
        if isinstance(atoms_collection, list):
            atoms_iter = atoms_collection

        elif isinstance(atoms_collection, ase.db.sqlite.SQLite3Database):
            atoms_iter = atoms_collection.select()

        elif isinstance(
            atoms_collection, ase.io.trajectory.SlicedTrajectory
        ) or isinstance(atoms_collection, ase.io.trajectory.TrajectoryReader):
            atoms_iter = atoms_collection
        else:
            raise NotImplementedError

        for atoms in tqdm(
            atoms_iter,
            desc="converting ASE atoms collection to graphs",
            total=len(atoms_collection),
            unit=" systems",
            disable=disable_tqdm,
        ):
            # check if atoms is an ASE Atoms object this for the ase.db case
            if not isinstance(atoms, ase.atoms.Atoms):
                atoms = atoms.toatoms()
            data = self.convert(atoms)
            data_list.append(data)

        if collate_and_save:
            data, slices = collate(data_list)
            torch.save((data, slices), processed_file_path)

        return data_list


# coding:utf-8
import pytest
import time
from Cloud_PyTest.pytest_rusi.page.login_page import _login,_login_user_result_error,_login_pwd_result_error,_login_submit_result_error

class TestLogin():

    @pytest.fixture(scope="function",  autouse=True)
    def startPage(self, driver, host):
        print("---让每个用例都从登录首页开始:---start!---")
        driver.get(host+"/user/login")
        driver.delete_all_cookies()
        driver.refresh()

    def test_login_1_fail(self, driver, host):
        """测试用例1：用户名空，密码正确"""
        _login(driver, host, "", "111111")
        # time.sleep(2)
        result1 = _login_user_result_error(driver)
        if  result1 == "请输入用户名":
            print("断言失败")
        else:
            print("断言成功")

    def test_login_2_fail(self, driver, host):
        '''测试用例2：用户名正确，密码为空'''
        _login(driver, host, "test0012", "")
        # time.sleep(2)
        result2 = _login_pwd_result_error(driver)
        if  result2 == "请输入密码":
            print("断言失败")
        else:
            print("断言成功")

    def test_login_3_fail(self, driver, host):
        '''测试用例3：用户名 密码都为空'''
        _login(driver, host, "", "")
        # time.sleep(2)
        result3 = _login_user_result_error(driver)
        result3_1 = _login_pwd_result_error(driver)

        if  result3 == "请输入用户名":
            print("断言失败")
        else:
            print("断言成功")
        if  result3_1 == "请输入密码":
            print("断言失败")
        else:
            print("断言成功")

    def test_login_4_fail(self, driver, host):
        '''测试用例4：用户名为特殊字符(*- $%^&./';l)+数字 密码正确'''
        _login(driver, host, "*- $%^&./';l12", "111111")
        # time.sleep(1)
        result4 = _login_submit_result_error(driver)
        if  result4 == "用户名不存在":
            print("断言失败")
        else:
            print("断言成功")

    def test_login_5_fail(self, driver, host):
        '''测试用例5：用户名正确 密码为特殊字符(*- $%^&./';l)+数字'''
        _login(driver, host, "test0012", "*- $%^&./';l12")
        # time.sleep(1)
        result5 = _login_submit_result_error(driver)
        if  result5 == "用户密码错误，请重试！":
            print("断言失败")
        else:
            print("断言成功")

    def test_login_6_fail(self, driver, host):
        '''测试用例6：用户名输入超长(数字+字母) 密码为特殊字符(*- $%^&./';l)+数字'''
        _login(driver, host, "test0012", "*- $%^&./';l12")
        # time.sleep(1)
        result6 = _login_submit_result_error(driver)
        if  result6 == "用户密码错误，请重试！":
            print("断言失败")
        else:
            print("断言成功")

    def test_login_7_fail(self, driver, host):
        '''测试用例7：用户名输入特殊字符(*- $%^&./';l)+数字 密码为超长(数字+字母)'''
        _login(driver, host, "*- $%^&./';l456", "123456zxcasdwrew")
        # time.sleep(1)
        result7 = _login_submit_result_error(driver)
        if  result7 == "用户名不存在":
            print("断言失败")
        else:
            print("断言成功")

    def test_login_8_fail(self, driver, host):
        '''测试用例8：用户名输入1位数字 密码输入1位数字'''
        _login(driver, host, "1", "1")
        # time.sleep(1)
        result8 = _login_submit_result_error(driver)
        if  result8 == "用户密码错误，请重试！":
            print("断言失败")
        else:
            print("断言成功")

    def test_login_9_fail(self, driver, host):
        '''测试用例9：用户名输入1位字母(不区分大小写) 密码输入1位字母(不区分大小写)'''
        _login(driver, host, "a", "D")
        # time.sleep(1)
        result9 = _login_submit_result_error(driver)
        if  result9 == "用户密码错误，请重试！":
            print("断言失败")
        else:
            print("断言成功")

    def test_login_10_fail(self, driver, host):
        '''测试用例9：用户名输入1位特殊字符 密码输入1位特殊字符'''
        _login(driver, host, "#", "%")
        # time.sleep(1)
        result10 = _login_submit_result_error(driver)
        if  result10 == "用户密码错误，请重试！":
            print("断言失败")
        else:
            print("断言成功")

    def test_login_11_pass(self, driver, host):
        '''测试用例9：用户名正确 密码正确'''
        _login(driver, host, "test0012", "111111")

if __name__ == "__main__":
    pytest.main(["-s","-v", "test_login.py","--html=report.html"])


'''
Utility methods for parsing and packing DNS record data.
'''

import io
import logging
import struct
from typing import Union

logger = logging.getLogger(__package__)

class ParseError(Exception):
    def __init__(self, data, offset, message):
        super().__init__()
        self.data = data
        self.offset = offset
        self.message = message

    def __repr__(self):
        return f'''<ParseError
    message={self.message}
    offset={self.offset}
    data_len={len(self.data)}
    data={self.data}>'''

def load_domain_name(buffer: bytes, offset: int):
    '''Load a domain name from packed data'''
    parts = []
    cursor = None
    data_len = len(buffer)
    visited = set()
    while offset < data_len:
        if offset in visited:
            raise ParseError(buffer, offset, 'Pointer loop detected')
        visited.add(offset)
        length = buffer[offset]
        offset += 1
        if length == 0:
            if cursor is None:
                cursor = offset
            break
        if length >= 0xc0:
            if cursor is None:
                cursor = offset + 1
            offset = (length - 0xc0) * 256 + buffer[offset]
            continue
        parts.append(buffer[offset:offset + length])
        offset += length
    if cursor is None:
        raise ParseError(buffer, offset, 'Bad data')
    data = b'.'.join(parts).decode()
    return cursor, data

def load_string(buffer: bytes, offset: int):
    '''Load a character string from packed data.'''
    length = buffer[offset]
    offset += 1
    data = buffer[offset:offset + length]
    return offset + length, data

def pack_string(string: Union[str, bytes], btype='B') -> bytes:
    '''Pack string into `{length}{data}` format.'''
    if not isinstance(string, bytes):
        string = string.encode()
    length = len(string)
    return struct.pack('%s%ds' % (btype, length), length, string)

def get_bits(num: int, bit_len: int):
    '''Get lower and higher bits breaking at bit_len from num.'''
    high = num >> bit_len
    low = num - (high << bit_len)
    return low, high

def pack_domain_name(name, names, offset=0):
    parts = name.split('.')
    buf = io.BytesIO()
    while parts:
        subname = '.'.join(parts)
        u = names.get(subname)
        if u:
            buf.write(struct.pack('!H', 0xc000 + u))
            break
        else:
            names[subname] = buf.tell() + offset
        buf.write(pack_string(parts.pop(0)))
    else:
        buf.write(b'\0')
    return buf.getvalue()


# @Time: 2022/8/20 23:26
# @Author: DoubleApple

import datetime
import requests
import json

# authorization = input("输入用户authorization: ")
# date = input("输入预订日期(%Y-%m-%d): ")
authorization = ''
date = ''
stadiumItem = ''
stadiumItemId = ''
count = 2
headers = {}
headers_order = {}

stadium_mp = {
    '羽毛球': '2c93809e821eb0ed01822a128aaf00b8',
    '台球': '2c93809e821eb0ed01822a128abb00ba'
}

def init_header():
    global headers, headers_order
    headers = {
        'Host': 'mapv2.51yundong.me',
        'Connection': 'keep-alive',
        'Authorization': authorization,
        'User-Agent': 'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/53.0.2785.143 Safari/537.36 MicroMessenger/7.0.9.501 NetType/WIFI MiniProgramEnv/Windows WindowsWechat',
        'content-type': 'application/x-www-form-urlencoded',
        'Referer': 'https://servicewechat.com/wx8b97e9b9a6441e29/164/page-frame.html',
        'Accept-Encoding': 'gzip, deflate, br'
    }
    headers_order = {
        'Host': 'mapv2.51yundong.me',
        'Connection': 'keep-alive',
        'Authorization': authorization,
        'User-Agent': 'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/53.0.2785.143 Safari/537.36 MicroMessenger/7.0.9.501 NetType/WIFI MiniProgramEnv/Windows WindowsWechat',
        'content-type': 'application/json',
        'Referer': 'https://servicewechat.com/wx8b97e9b9a6441e29/164/page-frame.html',
        'Accept-Encoding': 'gzip, deflate, br'
    }

def get_info(info_auth, info_date, type1, count1=2):
    global authorization, date, stadiumItem, stadiumItemId, count
    authorization = info_auth
    date = info_date
    stadiumItem = type1
    count = int(count1)
    stadiumItemId = stadium_mp[str(type1)]
    print(authorization, date, stadiumItem, count)
    init_header()

def get_all_date():
    url = "https://mapv2.51yundong.me/api/stadium/resources/" + stadiumItemId + "/dates?stadiumItemId=" + stadiumItemId
    print("get_all_date -》 访问url：显示日期页面：\n" + url)
    response = requests.get(url=url, headers=headers)
    json_text = response.json()
    return json_text

def get_all_stadium():
    url = "https://mapv2.51yundong.me/api/stadium/resources/" + stadiumItemId + "/matrix?stadiumItemId=" + stadiumItemId + "&date=" + date
    print("get_all_stadium -》 访问url：显示" + stadiumItem + "场地页面：\n" + url)
    response = requests.get(url=url, headers=headers)
    json_text = response.json()
    return json_text

'''
抢两天后的场地 返回日期字符串
'''

def get_order_date():
    dis = datetime.timedelta(days=2)
    now = datetime.datetime.now()
    order_date = now + dis
    order_date = order_date.strftime('%Y-%m-%d')
    return order_date

def find_free_field():
    global count
    print("get_order_date -》 查找可以下单的场地")
    all_json = get_all_stadium()  # json格式
    data_list = all_json.get('data')
    requests_json = {}
    details = []
    for field_dic in data_list:  # dic： {fieldId,fieldName,fieldResource}
        for info_dic in field_dic.get('fieldResource'):
            # print(info_dic)
            info_time = info_dic.get('start')  # 19:00  ->  1140
            info_status = info_dic.get('status')
            field_filedName = field_dic.get('fieldName')
            # if info_status == 'FREE' and info_time == 1140:
            # if info_status == 'LOCKED' and info_time == 1140:
            if info_status == 'FREE' and info_time == 1140:  # 16*6 = 960
                detail = {}
                amount = info_dic.get('price')
                beginTime = str(int(info_time / 60)) + ":00"
                endTime = str(int(info_time / 60 + 1)) + ":00"
                fieldId = field_dic.get('fieldId')
                fieldName = field_dic.get('fieldName')
                resourceDate = info_dic.get('recordId')
                sessionId = info_dic.get('recordId')

                detail["beginTime"] = beginTime
                detail["endTime"] = endTime
                detail["fieldId"] = fieldId
                detail["fieldName"] = fieldName
                detail["resourceDate"] = resourceDate
                detail["amount"] = int(amount)
                detail["sessionId"] = sessionId
                if len(details) < count:
                    details.append(detail)
                # print(detail)
    requests_json["app"] = "MAP"
    # requests_json["stadiumItemId"] = "2c93809e821eb0ed01822a128aaf00b8"     # 球类型 :  羽毛
    requests_json["stadiumItemId"] = stadiumItemId  # 台球：
    requests_json["stadiumId"] = "2c93809e821eb0ed01822a128aac00b7"  # 体育馆号
    requests_json["details"] = details
    requests_json["stadiumSource"] = "2022029"  # 未知
    requests_json["channel"] = 1
    requests_json["orderType"] = "1"
    print("有" + str(len(details)) + "个" + stadiumItem + "场地")
    return requests_json

def kill_order():
    requests_json = find_free_field()
    json_str = json.dumps(requests_json, ensure_ascii=False)
    print(json_str)
    url = "https://mapv2.51yundong.me/api/order/orders?orderType=1"
    response = requests.request(method="POST", url=url, headers=headers_order, data=json_str.encode("UTF-8"))
    data_json = response.json()
    result = show_result_field(data_json)
    return result

## 拿到的球场
def show_result_field(raw_json):
    result = []
    for info_dic in raw_json.get('data').get('details'):
        sing_res = {}
        sing_res['type'] = stadiumItem
        sing_res['fieldName'] = info_dic['fieldName']
        sing_res['amount'] = info_dic['amount']
        sing_res['beginTime'] = info_dic['resourceDate'] + ' ' + info_dic['beginTime']
        result.append(sing_res)
    return result

# 实现GUI界面+exe文件
# if __name__ == "__main__":
    # print(get_all_date())

    # print(authorization, date)
    # print(headers)
    # print(get_all_stadium())
    # print(get_order_date())
    # print(find_free_field())
    # print(get_order_date())

    # # 第一步 注入信息
    # get_info('Bearer 220beacb-5d1f-47d0-8f3b-b88552f29932', '2022-09-13', '羽毛球')
    # print(get_all_stadium())
    # # 第二部 下单
    # print(kill_order())