from ctypes import *
from ctypes.util import find_library
import os
import logging
import platform
import re
import sys
import tempfile
import filecmp
import shutil
import codecs

def gen_enum_value(value):
    return 'k' + value[0].upper() + value[1:]

class EnumType:
    name = ''
    enumValues = []
    def __init__(self, typeName = 'ElementWiseOpType', \
                 typeString = "{'avg', 'max', 'sum'}"):
        self.name = typeName
        if (typeString[0] == '{'):  # is a enum type
            isEnum = True
            # parse enum
            self.enumValues = typeString[typeString.find('{') + 1:typeString.find('}')].split(',')
            for i in range(0, len(self.enumValues)):
                self.enumValues[i] = self.enumValues[i].strip().strip("'")
        else:
            logging.warn(f"trying to parse none-enum type as enum: {typeString}")
    def GetDefinitionString(self, indent = 0):
        indentStr = ' ' * indent
        ret = indentStr + 'enum class {} {{\n'.format(self.name)
        for i in range(0, len(self.enumValues)):
            ret = ret + indentStr + f'  {gen_enum_value(self.enumValues[i])} = {i}'
            if (i != len(self.enumValues) -1):
                ret = ret + ","
            ret = ret + "\n"
        ret = ret + "};\n"
        return ret
    def GetDefaultValueString(self, value = ''):
        return self.name + "::" + gen_enum_value(value)
    def GetEnumStringArray(self, indent = 0):
        indentStr = ' ' * indent
        ret = indentStr + 'static const char *{}Values[] = {{\n'.format(self.name)
        for i in range(0, len(self.enumValues)):
            ret = ret + indentStr + f'  "{self.enumValues[i]}"'
            if (i != len(self.enumValues) -1):
                ret = ret + ","
            ret = ret + "\n"
        ret = ret + indentStr + "};\n"
        return ret
    def GetConvertEnumVariableToString(self, variable=''):
        return f"{self.name}Values[int({variable})]"

class Arg:
    typeDict = {'boolean':'bool',\
        'boolean or None':'dmlc::optional<bool>',\
        'Shape(tuple)':'Shape',\
        'Symbol':'Symbol',\
        'NDArray':'Symbol',\
        'NDArray-or-Symbol':'Symbol',\
        'Symbol[]':'const std::vector<Symbol>&',\
        'Symbol or Symbol[]':'const std::vector<Symbol>&',\
        'NDArray[]':'const std::vector<Symbol>&',\
        'caffe-layer-parameter':'::caffe::LayerParameter',\
        'NDArray-or-Symbol[]':'const std::vector<Symbol>&',\
        'float':'mx_float',\
        'real_t':'mx_float',\
        'int':'int',\
        'int (non-negative)': 'uint32_t',\
        'long (non-negative)': 'uint64_t',\
        'int or None':'dmlc::optional<int>',\
        'float or None':'dmlc::optional<float>',\
        'long':'int64_t',\
        'double':'double',\
        'double or None':'dmlc::optional<double>',\
        'Shape or None':'dmlc::optional<Shape>',\
        'string':'const std::string&',\
        'tuple of <float>':'nnvm::Tuple<mx_float>',\
        'tuple of <>':'mxnet::cpp::Shape',\
        '':'index_t'}
    name = ''
    type = ''
    description = ''
    isEnum = False
    enum = None
    hasDefault = False
    defaultString = ''
    def __init__(self, opName = '', argName = '', typeString = '', descString = ''):
        self.name = argName
        self.description = descString
        if (typeString[0] == '{'):  # is enum type
            self.isEnum = True
            self.enum = EnumType(self.ConstructEnumTypeName(opName, argName), typeString)
            self.type = self.enum.name
        else:
            try:
                self.type = self.typeDict[typeString.split(',')[0]]
            except:
                print(f'argument "{argName}" of operator "{opName}" has unknown type "{typeString}"')
                pass
        if typeString.find('default=') != -1:
            self.hasDefault = True
            self.defaultString = typeString.split('default=')[1].strip().strip("'")
            if typeString.startswith('string'):
                self.defaultString = self.MakeCString(self.defaultString)
            elif self.isEnum:
                self.defaultString = self.enum.GetDefaultValueString(self.defaultString)
            elif self.defaultString == 'None':
                self.defaultString = self.type + '()'
            elif self.type == "bool":
                if self.defaultString == "1" or self.defaultString == "True":
                    self.defaultString = "true"
                else:
                    self.defaultString = "false"
            elif self.defaultString[0] == '(':
                self.defaultString = 'Shape' + self.defaultString
            elif self.defaultString[0] == '[':
                self.defaultString = 'Shape(' + self.defaultString[1:-1] + ")"
            elif self.type == 'dmlc::optional<int>':
                self.defaultString = self.type + '(' + self.defaultString + ')'
            elif self.type == 'dmlc::optional<bool>':
                self.defaultString = self.type + '(' + self.defaultString + ')'
            elif typeString.startswith('caffe-layer-parameter'):
                self.defaultString = 'textToCaffeLayerParameter(' + self.MakeCString(self.defaultString) + ')'
                hasCaffe = True

    def MakeCString(self, str):
        str = str.replace('\n', "\\n")
        str = str.replace('\t', "\\t")
        return '\"' + str + '\"'

    def ConstructEnumTypeName(self, opName = '', argName = ''):
        a = opName[0].upper()
        # format ArgName so instead of act_type it returns ActType
        argNameWords = argName.split('_')
        argName = ''
        for an in argNameWords:
            argName = argName + an[0].upper() + an[1:]
        typeName = a + opName[1:] + argName
        return typeName

class Op:
    name = ''
    description = ''
    args = []

    def __init__(self, name = '', description = '', args = []):
        self.name = name
        self.description = description
        # add a 'name' argument
        nameArg = Arg(self.name, \
                      'symbol_name', \
                      'string', \
                      'name of the resulting symbol')
        args.insert(0, nameArg)
        # reorder arguments, put those with default value to the end
        orderedArgs = []
        for arg in args:
            if not arg.hasDefault:
                orderedArgs.append(arg)
        for arg in args:
            if arg.hasDefault:
                orderedArgs.append(arg)
        self.args = orderedArgs

    def WrapDescription(self, desc = ''):
        ret = []
        sentences = desc.split('.')
        lines = desc.split('\n')
        for line in lines:
          line = line.strip()
          if len(line) <= 80:
            ret.append(line.strip())
          else:
            while len(line) > 80:
              pos = line.rfind(' ', 0, 80)+1
              if pos <= 0:
                pos = line.find(' ')
              if pos < 0:
                pos = len(line)
              ret.append(line[:pos].strip())
              line = line[pos:]
        return ret

    def GenDescription(self, desc = '', \
                        firstLineHead = ' * \\brief ', \
                        otherLineHead = ' *        '):
        ret = ''
        descs = self.WrapDescription(desc)
        ret = ret + firstLineHead
        if len(descs) == 0:
          return ret.rstrip()
        ret = (ret + descs[0]).rstrip() + '\n'
        for i in range(1, len(descs)):
            ret = ret + (otherLineHead + descs[i]).rstrip() + '\n'
        return ret

    def GetOpDefinitionString(self, use_name, indent=0):
        ret = ''
        indentStr = ' ' * indent
        # define enums if any
        for arg in self.args:
            if arg.isEnum and use_name:
                # comments
                ret = ret + self.GenDescription(arg.description, \
                                        '/*! \\brief ', \
                                        ' *        ')
                ret = ret + " */\n"
                # definition
                ret = ret + arg.enum.GetDefinitionString(indent) + '\n'
        # create function comments
        ret = ret + self.GenDescription(self.description, \
                                        '/*!\n * \\brief ', \
                                        ' *        ')
        for arg in self.args:
            if arg.name != 'symbol_name' or use_name:
                ret = ret + self.GenDescription(arg.name + ' ' + arg.description, \
                                        ' * \\param ', \
                                        ' *        ')
        ret = ret + " * \\return new symbol\n"
        ret = ret + " */\n"
        # create function header
        declFirstLine = indentStr + f'inline Symbol {self.name}('
        ret = ret + declFirstLine
        argIndentStr = ' ' * len(declFirstLine)
        arg_start = 0 if use_name else 1
        if len(self.args) > arg_start:
            ret = ret + self.GetArgString(self.args[arg_start])
        for i in range(arg_start+1, len(self.args)):
            ret = ret + ',\n'
            ret = ret + argIndentStr + self.GetArgString(self.args[i])
        ret = ret + ') {\n'
        # create function body
        # if there is enum, generate static enum<->string mapping
        for arg in self.args:
            if arg.isEnum:
                ret = ret + arg.enum.GetEnumStringArray(indent + 2)
        # now generate code
        ret = ret + indentStr + f'  return Operator(\"{self.name}\")\n'
        for arg in self.args:   # set params
            if arg.type == 'Symbol' or \
                arg.type == 'const std::string&' or \
                arg.type == 'const std::vector<Symbol>&':
                continue
            v = arg.name
            if arg.isEnum:
                v = arg.enum.GetConvertEnumVariableToString(v)
            ret = ret + indentStr + ' ' * 11 + \
                f'.SetParam(\"{arg.name}\", {v})\n'
        #ret = ret[:-1]  # get rid of the last \n
        symbols = ''
        inputAlreadySet = False
        for arg in self.args:   # set inputs
            if arg.type != 'Symbol':
                continue
            inputAlreadySet = True
            #if symbols != '':
            #    symbols = symbols + ', '
            #symbols = symbols + arg.name
            ret = ret + indentStr + ' ' * 11 + \
                f'.SetInput(\"{arg.name}\", {arg.name})\n'
        for arg in self.args:   # set input arrays vector<Symbol>
            if arg.type != 'const std::vector<Symbol>&':
                continue
            if (inputAlreadySet):
                logging.error(f"op {self.name} has both Symbol[] and Symbol inputs!")
            inputAlreadySet = True
            symbols = arg.name
            ret = ret + f'({symbols})\n'
        ret = ret + indentStr + ' ' * 11
        if use_name:
            ret = ret + '.CreateSymbol(symbol_name);\n'
        else:
            ret = ret + '.CreateSymbol();\n'
        ret = ret + indentStr + '}\n'
        return ret

    def GetArgString(self, arg):
        ret = f'{arg.type} {arg.name}'
        if arg.hasDefault:
            ret = ret + ' = ' + arg.defaultString
        return ret

def ParseAllOps():
    """
    MXNET_DLL int MXSymbolListAtomicSymbolCreators(mx_uint *out_size,
                                                   AtomicSymbolCreator **out_array);

    MXNET_DLL int MXSymbolGetAtomicSymbolInfo(AtomicSymbolCreator creator,
                                              const char **name,
                                              const char **description,
                                              mx_uint *num_args,
                                              const char ***arg_names,
                                              const char ***arg_type_infos,
                                              const char ***arg_descriptions,
                                              const char **key_var_num_args);
    """
    cdll.libmxnet = cdll.LoadLibrary(sys.argv[1])
    ListOP = cdll.libmxnet.MXSymbolListAtomicSymbolCreators
    GetOpInfo = cdll.libmxnet.MXSymbolGetAtomicSymbolInfo
    ListOP.argtypes=[POINTER(c_int), POINTER(POINTER(c_void_p))]
    GetOpInfo.argtypes=[c_void_p, \
        POINTER(c_char_p), \
        POINTER(c_char_p), \
        POINTER(c_int), \
        POINTER(POINTER(c_char_p)), \
        POINTER(POINTER(c_char_p)), \
        POINTER(POINTER(c_char_p)), \
        POINTER(c_char_p), \
        POINTER(c_char_p)
        ]

    nOps = c_int()
    opHandlers = POINTER(c_void_p)()
    r = ListOP(byref(nOps), byref(opHandlers))
    ret = ''
    ret2 = ''
    for i in range(0, nOps.value):
        handler = opHandlers[i]
        name = c_char_p()
        description = c_char_p()
        nArgs = c_int()
        argNames = POINTER(c_char_p)()
        argTypes = POINTER(c_char_p)()
        argDescs = POINTER(c_char_p)()
        varArgName = c_char_p()
        return_type = c_char_p()

        GetOpInfo(handler, byref(name), byref(description), \
            byref(nArgs), byref(argNames), byref(argTypes), \
            byref(argDescs), byref(varArgName), byref(return_type))

        if name.value.decode('utf-8').startswith('_'):     # get rid of functions like __init__
            continue

        args = []

        for i in range(0, nArgs.value):
            arg = Arg(name.value.decode('utf-8'),
                      argNames[i].decode('utf-8'),
                      argTypes[i].decode('utf-8'),
                      argDescs[i].decode('utf-8'))
            args.append(arg)

        op = Op(name.value.decode('utf-8'), description.value.decode('utf-8'), args)

        ret = ret + op.GetOpDefinitionString(True) + "\n"
        ret2 = ret2 + op.GetOpDefinitionString(False) + "\n"
    return ret + ret2

if __name__ == "__main__":
    #et = EnumType(typeName = 'MyET')
    #print(et.GetDefinitionString())
    #print(et.GetEnumStringArray())
    #arg = Arg()
    #print(arg.ConstructEnumTypeName('SoftmaxActivation', 'act_type'))
    #arg = Arg(opName = 'FullConnected', argName='act_type', \
    #    typeString="{'elu', 'leaky', 'prelu', 'rrelu'},optional, default='leaky'", \
    #    descString='Activation function to be applied.')
    #print(arg.isEnum)
    #print(arg.defaultString)
    #arg = Arg("fc", "alpha", "float, optional, default=0.0001", "alpha")
    #decl = "%s %s" % (arg.type, arg.name)
    #if arg.hasDefault:
    #    decl = decl + "=" + arg.defaultString
    #print(decl)

    temp_file_name = ""
    output_file = '../include/mxnet-cpp/op.h'
    try:
        # generate file header
        patternStr = ("/*!\n"
                      "* \\file op.h\n"
                      "* \\brief definition of all the operators\n"
                      "* \\author Chuntao Hong, Xin Li\n"
                      "*/\n"
                      "\n"
                      "#ifndef MXNET_CPP_OP_H_\n"
                      "#define MXNET_CPP_OP_H_\n"
                      "\n"
                      "#include <string>\n"
                      "#include <vector>\n"
                      "#include \"mxnet-cpp/base.h\"\n"
                      "#include \"mxnet-cpp/shape.h\"\n"
                      "#include \"mxnet-cpp/op_util.h\"\n"
                      "#include \"mxnet-cpp/operator.h\"\n"
                      "#include \"dmlc/optional.h\"\n"
                      "#include \"nnvm/tuple.h\"\n"
                      "\n"
                      "namespace mxnet {{\n"
                      "namespace cpp {{\n"
                      "\n"
                      "{}"
                      "}} //namespace cpp\n"
                      "}} //namespace mxnet\n"
                      "#endif  // MXNET_CPP_OP_H_\n")

        # Generate a temporary file name
        tf = tempfile.NamedTemporaryFile()
        temp_file_name = tf.name
        tf.close()
        with codecs.open(temp_file_name, 'w', 'utf-8') as f:
            f.write(patternStr.format(ParseAllOps()))
    except Exception as e:
      if (os.path.exists(output_file)):
        os.remove(output_file)
      if len(temp_file_name) > 0:
        os.remove(temp_file_name)
      raise(e)
    if os.path.exists(output_file):
      if not filecmp.cmp(temp_file_name, output_file):
          os.remove(output_file)
    if not os.path.exists(output_file):
      shutil.move(temp_file_name, output_file)


from django import template
from django.core.urlresolvers import reverse

register = template.Library()

@register.simple_tag
def current(request, name, by_path=False):
    """ Return the string 'current' current request.path is same as name

    This tag is used for top menu in the header.

    Keyword aruguments:
    request  -- Django request object
    name     -- name of the url or the actual path
    by_path  -- True if name contains a url instead of url name
    """

    if by_path:
        path = name
    else:
        path = reverse(name)

    if request.path == path:
        return 'current'

    # Here is a hack for the gallery.
    if request.path[:21] == path:
        return 'current'

    return ''


import requests
import jwt
import datetime
import json
from decouple import config
import urllib3
import logging
import os
from src.utilities import Utilities

class Config:
    def __init__(self, path="../src/config/config.json"):
        self._path = path
        with open(self._path) as file:
            self._data = json.load(file)
            self._file = file

    def reload_file(self):
        with open(self._path) as file:
            self._data = json.load(file)
            self._file = file

    def read_config_file(self, key):
        self.reload_file()
        try:
            return self._data[key]
        except Exception:
            return None

    def set_config_file(self, key, value):
        self.reload_file()
        self._data[key] = value
        with open(self._path, "w") as write_file:
            json.dump(self._data, write_file, ensure_ascii=False, indent=4)
        self.reload_file()

    def get_whole_file(self):
        self.reload_file()
        return self._data

class API:
    def __init__(self):
        self._logger = Utilities.setup_logger()
        urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
        self._session = requests.session()
        self._demo = config("DEMO", False, cast=bool)
        self._id = None
        self._secret = config("SECRET", "admin", cast=str)
        self._access_token = None
        self._refresh_token = None
        self._url = config("IP")
        self.__counter = 0
        self._session.verify = config("SSL_VERIFY", True, cast=bool)
        self._session.trust_env = config("SSL_VERIFY", True, cast=bool)
        self._session.auth = JWTAuth(self)
        self._session.headers['Content-Type'] = 'application/json'

    def login(self):
        # generate new token or use old one
        redo_auth = False
        jwt_options = {
            'verify_signature': False,
            'verify_exp': False,
            'verify_nbf': False,
            'verify_iat': True,
            'verify_aud': False
        }
        # first run of code
        if self._access_token is None:
            redo_auth = True
        # if not first Run
        if redo_auth is False:
            try:
                time = datetime.datetime.utcfromtimestamp(
                    jwt.decode(self._access_token, algorithms=["HS512", "HS256"], options=jwt_options)["exp"])
                # check if code already expired
                if datetime.datetime.utcnow() > (time - datetime.timedelta(seconds=60)):
                    redo_auth = True
            except jwt.exceptions.DecodeError:
                redo_auth = True
        # generate new Code
        if redo_auth:
            use_refresh = True
            if self._refresh_token is not None:
                time = datetime.datetime.utcfromtimestamp(
                    jwt.decode(self._access_token, algorithms=["HS512", "HS256"], options=jwt_options)["exp"])
                if datetime.datetime.utcnow() > (time - datetime.timedelta(seconds=60)):
                    use_refresh = False
            else:
                use_refresh = False
            if use_refresh:
                req = self._session.post(f"{self._url}/api/aggregator-refresh",
                                         auth=JWTAuthRefreshToken(self._refresh_token))
                if req.status_code == requests.codes.ok:
                    output = req.json()
                    self._logger.debug(output)
                    self._access_token = output["access_token"]
                    self._refresh_token = output["refresh_token"]
                    return self._access_token
            else:
                payload = {
                    "token": self._secret
                }
                req = self._session.post(f"{self._url}/api/aggregator-login", json=payload)
                if req.status_code == requests.codes.ok:
                    output = req.json()
                    self._logger.debug(output)
                    self._id = output["aggregator_id"]
                    self._access_token = output["access_token"]
                    self._refresh_token = output["refresh_token"]
                    return self._access_token
                elif req.status_code == requests.codes.unauthorized:
                    self._logger.critical("Invalid client secret!")
                    os._exit(22)
        else:
            return self._access_token

    def send_data(self, data):
        req = self._session.post(f"{self._url}/api/devices/data", auth=JWTAuth(self), json=data)
        # self._logger.error(req.request.body)
        if req.status_code == requests.codes.ok:
            self._logger.info("Data sent successfully.")
            return True
        else:
            self._logger.error(f"Status code: {req.status_code}, response: {req.text}")
            return False

    def get_running_threads(self):
        if not self._demo:
            try:
                req = self._session.get(f"{self._url}/api/aggregator/{self._id}", auth=JWTAuth(self))
                if req.status_code == requests.codes.ok:
                    output = req.json()
                    self._logger.debug(output)
                    return output
                else:
                    self._logger.debug(req.json())
                    return []
            except Exception as ex:
                self._logger.critical(f"get_running_threads: {ex}")
                return []
        else:
            data = []
            modules = [{"name": "unifi_lldp", "config": {}}]
            # data.append(
            #    {"id": "1", "name": "Ubi", "timeout": 1, "type": "Ubiquiti", "ip": "172.31.37.95", "modules": modules})
            data.append({"id": "2", "name": "Zabbi", "timeout": 10, "type": "Ubiquiti", "ip": "172.31.37.77",
                         "modules": modules})
            data.append({"id": "798787", "name": "Zabbix", "timeout": 10, "type": "Zabbix", "ip": "zabbix.htl-vil.local",
                         "modules": [{"name": "events", "config": {}}]})

            # if(self.__conter % 5 == 0):
            #     data.append({"id": "3", "name": "Cisco", "timeout": 10, "type": "Cisco", "ip": "172.31.8.81",
            #                  "modules": modules})
            # else:
            # data.append({"id": "3", "name": "Cisco", "timeout": 10, "type": "Cisco", "ip": "172.31.8.81",
            # "modules": modules})
            self.__counter += 1
            # print(data)
            return data

    def send_version_string(self, version):
        if version == "%VER%":
            version = "v0.0-DEV"
        req = self._session.post(f"{self._url}/api/aggregator/{self._id}/version", json={"version": version}, auth=JWTAuth(self))
        if req.status_code == requests.codes.ok:
            return True
        else:
            self._logger.critical(f"send_version_string: {req.text}")
            return False

    def send_known_modules(self, modules):
        req = self._session.post(f"{self._url}/api/aggregator/{self._id}/modules", json={"modules": modules},
                                 auth=JWTAuth(self))
        if req.status_code == requests.codes.ok:
            self._logger.info(f"send_known_modules: {req.text}")
            return True
        else:
            self._logger.critical(f"send_known_modules: {req.text}")
            return False

class JWTAuth(requests.auth.AuthBase):
    def __init__(self, api: API):
        self.token = api.login()

    def __call__(self, r):
        # print(self.token)
        r.headers["Authorization"] = "Bearer " + self.token
        #print(self.token)
        return r

class JWTAuthRefreshToken(requests.auth.AuthBase):
    def __init__(self, token: str):
        self.token = token

    def __call__(self, r):
        # print(self.token)
        r.headers["Authorization"] = "Bearer " + self.token
        #print(self.token)
        return r


# tiddler2anki by Samuel Vargas # this is a plugin for anki that uses a ghetto regex based parsing system
# to import exported tiddler.json file from tiddlywiki into Anki.
# It checks for the prescence of a <qa time=blah></qa> and matching <an></an>
# tag to determine the question and answer. This is so I can combine my questions and
# flashcards directly in my notes. Feel free to adapt it to your needs.
# TODO: nested lists, nested bullets # __update, __create, __remove need to
# operate on the specific deck only # __remove needs to be created (will
# suspend cards not found anymore and the user can delete them manually) # to
# avoid wiping out important cards

from aqt import mw
from aqt.utils import showInfo
from aqt.qt import *
from PyQt4 import QtGui
from os import getenv
from itertools import izip
from pprint import pformat as Pretty
from BeautifulSoup import BeautifulSoup as Soup
import json, re

def tiddler2html(text):

  struct = {
    'header'  : re.compile("^\!+", re.M),
    'bullets' : re.compile("^\*+", re.M),
    'numlist' : re.compile("^\#+", re.M),
    "table"   : re.compile("((?<=\|)[^\s]*?)(?=\|)", re.M),
    "strings" : re.compile("(?![\!\*\#\|])^.+", re.M),
  }

  fmt = {
    "link_a"  : {'p':re.compile(""), 's':'', 'e': ''}, # [[ABC|Alphabet]] TODO
    "link_b"  : {'p':re.compile(""), 's':'', 'e': ''},  # [[ABC]] TODO
    "latex"   : {'p':re.compile("\$\$"), 's': '[$]', 'e': '[/$]'},
    "bold"    : {'p':re.compile("''"), 's': '', 'e': ''},
    "italics" : {'p':re.compile("\/\/"), 's': '', 'e': ''},
    "under"   : {'p':re.compile("\_\_"), 's': '', 'e': ''},
    "strike"  : {'p':re.compile("\~\~"), 's': '', 'e': ''},
  }

  html = ''
  prev = newrow = None
  for line in text.splitlines():
    for rule in struct:
      newrow = True if prev == "table" and rule == "table" else False
      for m in struct[rule].finditer(line):
        s,e = m.span();
        # Lists
        if (prev != 'bullets' and rule  == "bullets"): html += "

{}
".format(line[e:])
        if (prev == 'bullets' and rule  == "bullets"): html += "
{}
".format(line[e:])
        if (prev == 'bullets' and rule  != "bullets"): html += "

"
        if (prev != 'numlist' and rule  == "numlist"): html += "

{}
".format(line[e:])
        if (prev == 'numlist' and rule  == "numlist"): html += "
{}
".format(line[e:])
        if (prev == 'numlist' and rule  != "numlist"): html += "

"
        # Tables
        if (prev != 'table'   and rule == 'table'): html += "

{}	".format(line[s:e])
        if (prev == 'table'   and rule == "table" and newrow):
          newrow= False;
          html += ("
{}	".format(line[s:e]))
        elif (prev == "table" and rule == "table" and not newrow): html += "{}	".format(line[s:e])
        elif (prev == "table"   and rule != "table"): html += "

";
        # One Liners
        if (rule == "strings"): html += (line[s:e])
        if (rule == "header"): html += "<h{l}>{val}</h{l}>".format(l=len(line[s:e]), val=line[e:])
        prev = rule;

  # Doublecheck Table Ended
  if (prev == "table" and rule != "table"): html += "
\n

\n";

  # Formatting
  for f in fmt:
     count = len(fmt[f]['p'].findall(html))
     for tag in range(count):
       delim = fmt[f]['s'] if (tag+1) % 2 == 1 else fmt[f]['e'] # put the appropriate start / end delimiter
       html = re.sub(fmt[f]['p'].pattern, delim, html, count=1);

  return Soup(html + "").prettify();

class Logic:
  def __init__(self):
    self.__json_db = {}
    self.__anki_db = {}

  def run(self, json_path, deck_str, log_fn):
    self.__log = log_fn
    self.__build_json_db(json_path)
    self.__build_anki_db(deck_str)
    self.__log(">>> Database Built")

    self.__selected_deck = deck_str
    self.__set_card_type('Basic')

    self.__log(">>> Card Type set to Basic")

    self.__create();
    self.__log(">>> Created new entries from " + json_path)

    self.__update();
    self.__log(">>> Updated existing entries from " + json_path)
    self.__log(">>> Peachy Keen")
    #self.__remove();

  #
  # Database Initialisation
  #

  def __json_to_cards(self, text):
    s = Soup(text);
    assert(s != None)
    questions = s.findAll('qa')
    answers = s.findAll('an')
    if (questions == None or answers == None):
      return None
    assert len(questions) == len(answers) # mismatch no corresponding <an> for <qa> somewhere in json
    out = {}
    for q, a in izip(questions, answers):
      assert q['time'] != None # user forgot to put in a TIME attribute
      out[q['time']] = {'Front': q, 'Back': a}
    return out

  # tiddlers.json --> self.__json_db
  def __build_json_db(self, jpath):
    with open(jpath, 'r') as f:
      for node in json.loads(f.read()):
        cards = self.__json_to_cards(node['text'])
        if cards != None:
          for key in cards:
            self.__json_db[key] = cards[key]

  # anki.sqlite --> self.__anki_db
  def __build_anki_db(self, deck):
    nids = mw.col.findNotes('')
    for i in nids:
      c = mw.col.getNote(i)
      if (c['md'] != '' or c['md'] != None):
        self.__anki_db[c['md']] = {'Front': c['Front'], 'Back': c['Back']}

  #
  # Database Operations
  #

  def __set_card_type(self, ctype):
    # TODO: add support for cloze and basic and reverse'd
    did = mw.col.decks.id(self.__selected_deck)
    mw.col.decks.select(did)
    m = mw.col.models.byName(ctype)
    deck = mw.col.decks.get(did)
    deck['mid'] = m['id']
    mw.col.decks.save(deck)

  def __create(self):
    jkeys = self.__json_db.keys()
    akeys = self.__anki_db.keys()
    for k in jkeys:
      if k not in akeys:
        n = mw.col.newNote()
        did = mw.col.decks.id(self.__selected_deck)
        n['Front'] = tiddler2html((self.__json_db[k]['Front'].string))
        n['Back'] = tiddler2html((self.__json_db[k]['Back'].string))
        n['md'] = k
        n.model()['did'] = did
        mw.col.addNote(n)
    mw.reset()

  def __update(self):
    jkeys = self.__json_db.keys()
    akeys = self.__anki_db.keys()
    for k in jkeys:
      if k in akeys:
        nid = mw.col.findNotes(k)
        assert len(nid) == 1 # C O L L I S I O N
        n = mw.col.getNote(nid[0])
        n['Front'] = tiddler2html((self.__json_db[k]['Front'].string))
        n['Back']  = tiddler2html((self.__json_db[k]['Back'].string))
        self.__anki_db['Front'] = self.__json_db[k]['Front']
        self.__anki_db['Back']  = self.__json_db[k]['Back']
        n.flush();
    mw.reset()

  def __delete(self):
    jkeys = self.__json_db.keys()
    akeys = self.__anki_db.keys()
    for k in jkeys:
      if k not in akeys:
        pass

class UI(QtGui.QWidget):
  def __init__(self, decks):
    super(UI, self).__init__()
    self.__data = {'deck': None, 'path': None}
    self.__decks = decks
    self.__callback = None
    self.__json_btn = QtGui.QPushButton("Open JSON")
    self.__json_path = QtGui.QLineEdit(self)
    self.__deck_label = QtGui.QLabel("Deck")
    self.__deck_process = QtGui.QPushButton("Process")
    self.__deck_cancel = QtGui.QPushButton("Cancel")
    self.__deck_list = QComboBox(self)
    self.__log_stats = QtGui.QTextBrowser(self)
    self.__init_actions()
    self.__init_layout()

  def __init_actions(self):
    def set_path():
      fname = QtGui.QFileDialog.getOpenFileName(self, 'json location', getenv('HOME'))
      if fname:
        self.__json_path.setText(fname)
    self.__json_btn.clicked.connect(set_path)

    def process():
      self.__data['path'] = self.__json_path.text()
      self.__data['deck'] = self.__deck_list.currentText()
      self.__callback(self.__data['path'], self.__data['deck'], self.logger)

    self.__deck_process.clicked.connect(process)

  def __init_layout(self):
    frame = QtGui.QVBoxLayout()
    row_a = QtGui.QHBoxLayout()
    row_a.addWidget(self.__json_btn)
    row_a.addWidget(self.__json_path)
    row_b = QtGui.QHBoxLayout()
    row_b.addWidget(self.__deck_label)
    for d in self.__decks:
      self.__deck_list.addItem(d)
    row_b.addWidget(self.__deck_list)
    row_b.addWidget(self.__deck_process)
    row_b.addWidget(self.__deck_cancel)
    frame.addLayout(row_a)
    frame.addLayout(row_b)
    frame.addWidget(self.__log_stats)
    self.setLayout(frame)
    self.show()

  # Public
  def setCallback(self, fn):
    self.__callback = fn

  def logger(self, msg):
    self.__log_stats.append(msg)

def get_decks():
  out = []
  decks = mw.col.decks.all()
  for d in decks:
    out.append(d['name'])
  return out

def main():
  decks = get_decks()
  logic = Logic()
  mw.ui = UI(decks)
  mw.ui.setCallback(logic.run)

action = QAction("tiddler2anki", mw)
mw.connect(action, SIGNAL("triggered()"), main)
mw.form.menuTools.addAction(action)


import sqlite3, os, sys, platform
import PySimpleGUI as sg
from ppadb.client import Client as AdbClient
from datetime import datetime as dt

def pull_from_db(db, command, facebook_name=False):
    '''Send queries to a database and return the results'''
    try:
        conn = sqlite3.connect(db)
        c = conn.cursor()
        c.execute(command)
        return c.fetchall()
    except Exception as e:
        if facebook_name:
            return [("Name Unavailable %s" % e,)]
        else:
            sg.Popup("Error reading the database: %s" % e)

def dump_logcat_by_line(connect):
    file_obj = connect.socket.makefile()
    for index in range(0, 10):
        print("Line {}: {}".format(index, file_obj.readline().strip()))

    file_obj.close()
    connect.close()


def test_get_hint_repr(hints_pep_meta) -> None:
    '''
    Test the :func:`beartype._util.hint.utilhintget.get_hint_repr` getter.

    Parameters
    ----------
    hints_pep_meta : List[beartype_test.a00_unit.data.hint.util.data_hintmetacls.HintPepMetadata]
        List of type hint metadata describing sample type hints exercising edge
        cases in the :mod:`beartype` codebase.
    '''

    # Defer test-specific imports.
    from beartype._util.hint.utilhintget import get_hint_repr
    from beartype_test.a00_unit.data.hint.data_hint import (
        NOT_HINTS, HINTS_NONPEP)

    # Assert this getter returns the expected representations of objects *NOT*
    # supported as either PEP-noncompliant or -compliant type hints.
    for non_hint in NOT_HINTS:
        assert get_hint_repr(non_hint) == repr(non_hint)

    # Assert this getter returns the expected representations of
    # PEP-noncompliant type hints.
    for nonhint_pep in HINTS_NONPEP:
        assert get_hint_repr(nonhint_pep) == repr(nonhint_pep)

    # Assert this getter returns the expected representations of PEP-compliant
    # type hints.
    for hint_pep_meta in hints_pep_meta:
        # PEP-compliant type hint to be tested.
        hint_pep = hint_pep_meta.hint

        # Representation of this hint.
        hint_pep_repr = repr(hint_pep)

        # If this representation embeds the representation of a type union,
        # silently ignore this hint and continue to the next. Why? Because
        # Python intentionally hashes both PEP 484-compliant type unions (e.g.,
        # "typing.Union[int, str]") *AND* PEP 604-compliant type unions (e.g.,
        # "int | str") to the same integers. However, these two kinds of unions
        # have fundamentally different representations.
        #
        # Specifically, ignore this hint if this representation embeds the
        # representation of either...
        if (
            # A PEP 484-compliant type union.
            'typing.Union[' in hint_pep_repr or
            # A PEP 604-compliant type union.
            ' | ' in hint_pep_repr
        ):
            continue
        # Else, this representation contains *NO* type union.

        # Assert this getter returns the expected representation.
        assert get_hint_repr(hint_pep) == repr(hint_pep)


#!/usr/bin/env python
# coding: utf-8

# # Model

# In[2]:

#%load_ext autotime

# ## Directory Settings

# In[3]:

# ====================================================
# Directory settings
# ====================================================
import os

OUTPUT_DIR = './'
if not os.path.exists(OUTPUT_DIR):
    os.makedirs(OUTPUT_DIR)

PREPROCESS_DIR = 'input/inchi-preprocess'
if not os.path.exists(PREPROCESS_DIR):
    os.makedirs(PREPROCESS_DIR)

#if os.path.isfile('train.log'):
#    os.remove('train.log')

# ## Data Loading

# In[4]:

import numpy as np
import pandas as pd
import dill
import torch

train = pd.read_json(f'{PREPROCESS_DIR}/train.json')

def get_train_file_path(image_id):
    return "input/train/{}/{}/{}/{}.png".format(
        image_id[0], image_id[1], image_id[2], image_id
    )

train['file_path'] = train['image_id'].apply(get_train_file_path)

print(f'train.shape: {train.shape}')
print(train.head())

# In[5]:

class Tokenizer(object):

    def __init__(self):
        self.stoi = {}
        self.itos = {}

    def __len__(self):
        return len(self.stoi)

    def fit_on_texts(self, texts):
        vocab = set()
        for text in texts:
            vocab.update(text.split(' '))
        vocab = sorted(vocab)
        vocab.append('<sos>')
        vocab.append('<eos>')
        vocab.append('<pad>')
        for i, s in enumerate(vocab):
            self.stoi[s] = i
        self.itos = {item[1]: item[0] for item in self.stoi.items()}

    def text_to_sequence(self, text):
        sequence = []
        sequence.append(self.stoi['<sos>'])
        for s in text.split(' '):
            sequence.append(self.stoi[s])
        sequence.append(self.stoi['<eos>'])
        return sequence

    def texts_to_sequences(self, texts):
        sequences = []
        for text in texts:
            sequence = self.text_to_sequence(text)
            sequences.append(sequence)
        return sequences

    def sequence_to_text(self, sequence):
        return ''.join(list(map(lambda i: self.itos[i], sequence)))

    def sequences_to_texts(self, sequences):
        texts = []
        for sequence in sequences:
            text = self.sequence_to_text(sequence)
            texts.append(text)
        return texts

    def predict_caption(self, sequence):
        caption = ''
        for i in sequence:
            if i == self.stoi['<eos>'] or i == self.stoi['<pad>']:
                break
            caption += self.itos[i]
        return caption

    def predict_captions(self, sequences):
        captions = []
        for sequence in sequences:
            caption = self.predict_caption(sequence)
            captions.append(caption)
        return captions

tokenizer = torch.load(f'{PREPROCESS_DIR}/tokenizer.pth', pickle_module=dill)
print(f"tokenizer.stoi: {tokenizer.stoi}")

# ## CFG

# In[6]:

# ====================================================
# CFG
# ====================================================
class CFG:
    debug=False
    max_len=275
    print_freq=1000
    num_workers=16 if os.cpu_count() >= 16 else os.cpu_count()
    model_name='resnet34'
    size=224
    scheduler='CosineAnnealingLR' # ['ReduceLROnPlateau', 'CosineAnnealingLR', 'CosineAnnealingWarmRestarts']
    epochs=8 # not to exceed 9h
    #factor=0.2 # ReduceLROnPlateau
    #patience=4 # ReduceLROnPlateau
    #eps=1e-6 # ReduceLROnPlateau
    T_max=4 # CosineAnnealingLR
    #T_0=4 # CosineAnnealingWarmRestarts
    encoder_lr=1e-4
    decoder_lr=4e-4
    min_lr=1e-6
    batch_size=64
    weight_decay=1e-6
    gradient_accumulation_steps=1
    max_grad_norm=5
    attention_dim=256
    embed_dim=256
    decoder_dim=512
    dropout=0.5
    seed=42
    n_fold=2
    trn_fold=[0] # [0, 1, 2, 3, 4]
    train=True

print(f'Using {CFG.num_workers} workers.')

# In[7]:

if CFG.debug:
    CFG.epochs = 1
    train = train[:10000]

# ## Library

# In[8]:

# ====================================================
# Library
# ====================================================
import sys
sys.path.append('input/pytorch-image-models/pytorch-image-models-master')

import os
import gc
import re
import math
import time
import random
import shutil
import pickle
import argparse
from pathlib import Path
from contextlib import contextmanager
from collections import defaultdict, Counter

import scipy as sp
import numpy as np
import pandas as pd
from tqdm.auto import tqdm

import Levenshtein
from sklearn import preprocessing
from sklearn.model_selection import StratifiedKFold, GroupKFold, KFold

from functools import partial

import cv2
from PIL import Image

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Adam, SGD
import torchvision.models as models
from torch.nn.parameter import Parameter
from torch.utils.data import DataLoader, Dataset
from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from adabound import AdaBound

from albumentations import (
    Compose, OneOf, Normalize, Resize, RandomResizedCrop, RandomCrop, HorizontalFlip, VerticalFlip,
    RandomBrightness, RandomContrast, RandomBrightnessContrast, Rotate, ShiftScaleRotate, Cutout,
    IAAAdditiveGaussianNoise, Transpose, Blur
    )
from albumentations.pytorch import ToTensorV2
from albumentations import ImageOnlyTransform

import timm

import warnings
warnings.filterwarnings('ignore')

cuda_available = torch.cuda.is_available()

device = torch.device('cuda' if cuda_available else 'cpu')

if (cuda_available):
    torch.cuda.init()
    print(f'CUDA available: {cuda_available}')
    print(f'Number of available devices: {torch.cuda.device_count()}')
    print(f'Using device: {torch.cuda.get_device_name(torch.cuda.current_device())}')

# ## Utils

# In[9]:

# ====================================================
# Utils
# ====================================================

def get_score(y_true, y_pred):
    scores = []
    for true, pred in zip(y_true, y_pred):
        score = Levenshtein.distance(true, pred)
        scores.append(score)
    avg_score = np.mean(scores)
    return avg_score, scores

def init_logger(log_file=OUTPUT_DIR+'train.log'):
    from logging import getLogger, INFO, FileHandler,  Formatter,  StreamHandler
    logger = getLogger(__name__)
    logger.setLevel(INFO)
    handler1 = StreamHandler()
    handler1.setFormatter(Formatter("%(message)s"))
    #handler2 = FileHandler(filename=log_file)
    #handler2.setFormatter(Formatter("%(message)s"))
    logger.addHandler(handler1)
    #logger.addHandler(handler2)
    return logger

LOGGER = init_logger()

def seed_torch(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

seed_torch(seed=CFG.seed)

# ## CV Split

# In[10]:

folds = train.copy()
Fold = StratifiedKFold(n_splits=CFG.n_fold, shuffle=True, random_state=CFG.seed)
for n, (train_index, val_index) in enumerate(Fold.split(folds, folds['InChI_length'])):
    folds.loc[val_index, 'fold'] = int(n)
folds['fold'] = folds['fold'].astype(int)
print(folds.groupby(['fold']).size())

# ## Dataset

# In[11]:

# ====================================================
# Dataset
# ====================================================
class TrainDataset(Dataset):
    def __init__(self, df, tokenizer, transform=None):
        super().__init__()
        self.df = df
        self.tokenizer = tokenizer
        self.file_paths = df['file_path'].values
        self.labels = df['InChI_text'].values
        self.transform = transform

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        file_path = self.file_paths[idx]
        image = cv2.imread(file_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
        if self.transform:
            augmented = self.transform(image=image)
            image = augmented['image']
        label = self.labels[idx]
        label = self.tokenizer.text_to_sequence(label)
        label_length = len(label)
        label_length = torch.LongTensor([label_length])
        return image, torch.LongTensor(label), label_length

class TestDataset(Dataset):
    def __init__(self, df, transform=None):
        super().__init__()
        self.df = df
        self.file_paths = df['file_path'].values
        self.transform = transform

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        file_path = self.file_paths[idx]
        image = cv2.imread(file_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
        if self.transform:
            augmented = self.transform(image=image)
            image = augmented['image']
        return image

# In[12]:

def bms_collate(batch):
    imgs, labels, label_lengths = [], [], []
    for data_point in batch:
        imgs.append(data_point[0])
        labels.append(data_point[1])
        label_lengths.append(data_point[2])
    labels = pad_sequence(labels, batch_first=True, padding_value=tokenizer.stoi["<pad>"])
    return torch.stack(imgs), labels, torch.stack(label_lengths).reshape(-1, 1)

# ## Transforms

# In[13]:

def get_transforms(*, data):

    if data == 'train':
        return Compose([
            Resize(CFG.size, CFG.size),
            Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225],
            ),
            ToTensorV2(),
        ])

    elif data == 'valid':
        return Compose([
            Resize(CFG.size, CFG.size),
            Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225],
            ),
            ToTensorV2(),
        ])

# In[14]:

from matplotlib import pyplot as plt

train_dataset = TrainDataset(train, tokenizer, transform=get_transforms(data='train'))

for i in range(1):
    image, label, label_length = train_dataset[i]
    text = tokenizer.sequence_to_text(label.numpy())
    #plt.imshow(image.transpose(0, 1).transpose(1, 2))
    plt.title(f'label: {label}  text: {text}  label_length: {label_length}')
    plt.show()

# ## MODEL

# In[15]:

class Encoder(nn.Module):
    def __init__(self, model_name='resnet18', pretrained=False):
        super().__init__()
        self.cnn = timm.create_model(model_name, pretrained=pretrained)
        self.n_features = self.cnn.fc.in_features
        self.cnn.global_pool = nn.Identity()
        self.cnn.fc = nn.Identity()

    def forward(self, x):
        bs = x.size(0)
        features = self.cnn(x)
        features = features.permute(0, 2, 3, 1)
        return features

# In[16]:

class Attention(nn.Module):
    """
    Attention network for calculate attention value
    """
    def __init__(self, encoder_dim, decoder_dim, attention_dim):
        """
        :param encoder_dim: input size of encoder network
        :param decoder_dim: input size of decoder network
        :param attention_dim: input size of attention network
        """
        super(Attention, self).__init__()
        self.encoder_att = nn.Linear(encoder_dim, attention_dim)  # linear layer to transform encoded image
        self.decoder_att = nn.Linear(decoder_dim, attention_dim)  # linear layer to transform decoder's output
        self.full_att = nn.Linear(attention_dim, 1)  # linear layer to calculate values to be softmax-ed
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)  # softmax layer to calculate weights

    def forward(self, encoder_out, decoder_hidden):
        att1 = self.encoder_att(encoder_out)  # (batch_size, num_pixels, attention_dim)
        att2 = self.decoder_att(decoder_hidden)  # (batch_size, attention_dim)
        att = self.full_att(self.relu(att1 + att2.unsqueeze(1))).squeeze(2)  # (batch_size, num_pixels)
        alpha = self.softmax(att)  # (batch_size, num_pixels)
        attention_weighted_encoding = (encoder_out * alpha.unsqueeze(2)).sum(dim=1)  # (batch_size, encoder_dim)
        return attention_weighted_encoding, alpha

class DecoderWithAttention(nn.Module):
    """
    Decoder network with attention network used for training
    """

    def __init__(self, attention_dim, embed_dim, decoder_dim, vocab_size, device, encoder_dim=512, dropout=0.5):
        """
        :param attention_dim: input size of attention network
        :param embed_dim: input size of embedding network
        :param decoder_dim: input size of decoder network
        :param vocab_size: total number of characters used in training
        :param encoder_dim: input size of encoder network
        :param dropout: dropout rate
        """
        super(DecoderWithAttention, self).__init__()
        self.encoder_dim = encoder_dim
        self.attention_dim = attention_dim
        self.embed_dim = embed_dim
        self.decoder_dim = decoder_dim
        self.vocab_size = vocab_size
        self.dropout = dropout
        self.device = device
        self.attention = Attention(encoder_dim, decoder_dim, attention_dim)  # attention network
        self.embedding = nn.Embedding(vocab_size, embed_dim)  # embedding layer
        self.dropout = nn.Dropout(p=self.dropout)
        self.decode_step = nn.LSTMCell(embed_dim + encoder_dim, decoder_dim, bias=True)  # decoding LSTMCell
        self.init_h = nn.Linear(encoder_dim, decoder_dim)  # linear layer to find initial hidden state of LSTMCell
        self.init_c = nn.Linear(encoder_dim, decoder_dim)  # linear layer to find initial cell state of LSTMCell
        self.f_beta = nn.Linear(decoder_dim, encoder_dim)  # linear layer to create a sigmoid-activated gate
        self.sigmoid = nn.Sigmoid()
        self.fc = nn.Linear(decoder_dim, vocab_size)  # linear layer to find scores over vocabulary
        self.init_weights()  # initialize some layers with the uniform distribution

    def init_weights(self):
        self.embedding.weight.data.uniform_(-0.1, 0.1)
        self.fc.bias.data.fill_(0)
        self.fc.weight.data.uniform_(-0.1, 0.1)

    def load_pretrained_embeddings(self, embeddings):
        self.embedding.weight = nn.Parameter(embeddings)

    def fine_tune_embeddings(self, fine_tune=True):
        for p in self.embedding.parameters():
            p.requires_grad = fine_tune

    def init_hidden_state(self, encoder_out):
        mean_encoder_out = encoder_out.mean(dim=1)
        h = self.init_h(mean_encoder_out)  # (batch_size, decoder_dim)
        c = self.init_c(mean_encoder_out)
        return h, c

    def forward(self, encoder_out, encoded_captions, caption_lengths):
        """
        :param encoder_out: output of encoder network
        :param encoded_captions: transformed sequence from character to integer
        :param caption_lengths: length of transformed sequence
        """
        batch_size = encoder_out.size(0)
        encoder_dim = encoder_out.size(-1)
        vocab_size = self.vocab_size
        encoder_out = encoder_out.view(batch_size, -1, encoder_dim)  # (batch_size, num_pixels, encoder_dim)
        num_pixels = encoder_out.size(1)
        caption_lengths, sort_ind = caption_lengths.squeeze(1).sort(dim=0, descending=True)
        encoder_out = encoder_out[sort_ind]
        encoded_captions = encoded_captions[sort_ind]
        # embedding transformed sequence for vector
        embeddings = self.embedding(encoded_captions)  # (batch_size, max_caption_length, embed_dim)
        # initialize hidden state and cell state of LSTM cell
        h, c = self.init_hidden_state(encoder_out)  # (batch_size, decoder_dim)
        # set decode length by caption length - 1 because of omitting start token
        decode_lengths = (caption_lengths - 1).tolist()
        predictions = torch.zeros(batch_size, max(decode_lengths), vocab_size).to(self.device)
        alphas = torch.zeros(batch_size, max(decode_lengths), num_pixels).to(self.device)
        # predict sequence
        for t in range(max(decode_lengths)):
            batch_size_t = sum([l > t for l in decode_lengths])
            attention_weighted_encoding, alpha = self.attention(encoder_out[:batch_size_t], h[:batch_size_t])
            gate = self.sigmoid(self.f_beta(h[:batch_size_t]))  # gating scalar, (batch_size_t, encoder_dim)
            attention_weighted_encoding = gate * attention_weighted_encoding
            h, c = self.decode_step(
                torch.cat([embeddings[:batch_size_t, t, :], attention_weighted_encoding], dim=1),
                (h[:batch_size_t], c[:batch_size_t]))  # (batch_size_t, decoder_dim)
            preds = self.fc(self.dropout(h))  # (batch_size_t, vocab_size)
            predictions[:batch_size_t, t, :] = preds
            alphas[:batch_size_t, t, :] = alpha
        return predictions, encoded_captions, decode_lengths, alphas, sort_ind

    def predict(self, encoder_out, decode_lengths, tokenizer):
        batch_size = encoder_out.size(0)
        encoder_dim = encoder_out.size(-1)
        vocab_size = self.vocab_size
        encoder_out = encoder_out.view(batch_size, -1, encoder_dim)  # (batch_size, num_pixels, encoder_dim)
        num_pixels = encoder_out.size(1)
        # embed start tocken for LSTM input
        start_tockens = torch.ones(batch_size, dtype=torch.long).to(self.device) * tokenizer.stoi["<sos>"]
        embeddings = self.embedding(start_tockens)
        # initialize hidden state and cell state of LSTM cell
        h, c = self.init_hidden_state(encoder_out)  # (batch_size, decoder_dim)
        predictions = torch.zeros(batch_size, decode_lengths, vocab_size).to(self.device)
        # predict sequence
        for t in range(decode_lengths):
            attention_weighted_encoding, alpha = self.attention(encoder_out, h)
            gate = self.sigmoid(self.f_beta(h))  # gating scalar, (batch_size_t, encoder_dim)
            attention_weighted_encoding = gate * attention_weighted_encoding
            h, c = self.decode_step(
                torch.cat([embeddings, attention_weighted_encoding], dim=1),
                (h, c))  # (batch_size_t, decoder_dim)
            preds = self.fc(self.dropout(h))  # (batch_size_t, vocab_size)
            predictions[:, t, :] = preds
            if np.argmax(preds.detach().cpu().numpy()) == tokenizer.stoi["<eos>"]:
                break
            embeddings = self.embedding(torch.argmax(preds, -1))
        return predictions

# ## Helper functions

# In[17]:

# ====================================================
# Helper functions
# ====================================================
class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def asMinutes(s):
    m = math.floor(s / 60)
    s -= m * 60
    return '%dm %ds' % (m, s)

def timeSince(since, percent):
    now = time.time()
    s = now - since
    es = s / (percent)
    rs = es - s
    return '%s (remain %s)' % (asMinutes(s), asMinutes(rs))

def train_fn(train_loader, encoder, decoder, criterion,
             encoder_optimizer, decoder_optimizer, epoch,
             encoder_scheduler, decoder_scheduler, device):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    # switch to train mode
    LOGGER.info('Switching to train mode ..')
    encoder.train()
    decoder.train()
    start = end = time.time()
    global_step = 0
    #LOGGER.info('Enumerating train loader ..')
    for step, (images, labels, label_lengths) in enumerate(tqdm(train_loader)):
        LOGGER.info(f'Step: {step}')
        #LOGGER.info('Training - Measuring data loading time ..')
        # measure data loading time
        data_time.update(time.time() - end)
        #LOGGER.info('Training - loading data ..')
        images = images.to(device)
        labels = labels.to(device)
        label_lengths = label_lengths.to(device)
        batch_size = images.size(0)
        #LOGGER.info('Training - encoding ..')
        features = encoder(images)
        #LOGGER.info('Training - decoding ..')
        predictions, caps_sorted, decode_lengths, alphas, sort_ind = decoder(features, labels, label_lengths)
        targets = caps_sorted[:, 1:]
        predictions = pack_padded_sequence(predictions, decode_lengths, batch_first=True).data
        targets = pack_padded_sequence(targets, decode_lengths, batch_first=True).data
        loss = criterion(predictions, targets)
        # record loss
        #LOGGER.info('Training - recording ..')
        losses.update(loss.item(), batch_size)
        if CFG.gradient_accumulation_steps > 1:
            loss = loss / CFG.gradient_accumulation_steps
        loss.backward()
        encoder_grad_norm = torch.nn.utils.clip_grad_norm_(encoder.parameters(), CFG.max_grad_norm)
        decoder_grad_norm = torch.nn.utils.clip_grad_norm_(decoder.parameters(), CFG.max_grad_norm)
        if (step + 1) % CFG.gradient_accumulation_steps == 0:
            encoder_optimizer.step()
            decoder_optimizer.step()
            encoder_optimizer.zero_grad()
            decoder_optimizer.zero_grad()
            global_step += 1
        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()
        if step % CFG.print_freq == 0 or step == (len(train_loader)-1):
            print('Epoch: [{0}][{1}/{2}] '
                  'Data {data_time.val:.3f} ({data_time.avg:.3f}) '
                  'Elapsed {remain:s} '
                  'Loss: {loss.val:.4f}({loss.avg:.4f}) '
                  'Encoder Grad: {encoder_grad_norm:.4f}  '
                  'Decoder Grad: {decoder_grad_norm:.4f}  '
                  #'Encoder LR: {encoder_lr:.6f}  '
                  #'Decoder LR: {decoder_lr:.6f}  '
                  .format(
                   epoch+1, step, len(train_loader), batch_time=batch_time,
                   data_time=data_time, loss=losses,
                   remain=timeSince(start, float(step+1)/len(train_loader)),
                   encoder_grad_norm=encoder_grad_norm,
                   decoder_grad_norm=decoder_grad_norm,
                   #encoder_lr=encoder_scheduler.get_lr()[0],
                   #decoder_lr=decoder_scheduler.get_lr()[0],
                   ))
    return losses.avg

def valid_fn(valid_loader, encoder, decoder, tokenizer, criterion, device):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    # switch to evaluation mode
    LOGGER.info('Switching to evaluation mode ..')
    encoder.eval()
    decoder.eval()
    text_preds = []
    start = end = time.time()
    for step, (images) in enumerate(tqdm(valid_loader)):
        LOGGER.info(f'Step: {step}')
        # measure data loading time
        #LOGGER.info('Evaluation - loading data ..')
        data_time.update(time.time() - end)
        images = images.to(device)
        batch_size = images.size(0)
        with torch.no_grad():
            #LOGGER.info('Evaluation - encoding ..')
            features = encoder(images)
            #LOGGER.info('Evaluation - decoding ..')
            predictions = decoder.predict(features, CFG.max_len, tokenizer)
        predicted_sequence = torch.argmax(predictions.detach().cpu(), -1).numpy()
        _text_preds = tokenizer.predict_captions(predicted_sequence)
        text_preds.append(_text_preds)
        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()
        if step % CFG.print_freq == 0 or step == (len(valid_loader)-1):
            print('EVAL: [{0}/{1}] '
                  'Data {data_time.val:.3f} ({data_time.avg:.3f}) '
                  'Elapsed {remain:s} '
                  .format(
                   step, len(valid_loader), batch_time=batch_time,
                   data_time=data_time,
                   remain=timeSince(start, float(step+1)/len(valid_loader)),
                   ))
    text_preds = np.concatenate(text_preds)
    return text_preds

# ## Train loop

# In[18]:

# ====================================================
# Train loop
# ====================================================
def train_loop(gpu, folds, fold, args):

    LOGGER.info(f"========== fold: {fold} training ==========")

    LOGGER.info('Initializing process group ..')

    rank = args.nr * args.gpus + gpu
    dist.init_process_group(
        backend='nccl',
        init_method='env://',
        world_size=args.world_size,
        rank=rank
    )

    # ====================================================
    # loader
    # ====================================================

    LOGGER.info('Loading ..')

    trn_idx = folds[folds['fold'] != fold].index
    val_idx = folds[folds['fold'] == fold].index

    train_folds = folds.loc[trn_idx].reset_index(drop=True)
    valid_folds = folds.loc[val_idx].reset_index(drop=True)
    valid_labels = valid_folds['InChI'].values

    train_dataset = TrainDataset(train_folds, tokenizer, transform=get_transforms(data='train'))
    valid_dataset = TestDataset(valid_folds, transform=get_transforms(data='valid'))

    train_sampler = torch.utils.data.distributed.DistributedSampler(
        train_dataset,
        num_replicas=args.world_size,
        rank=rank
    )

    train_loader = DataLoader(train_dataset,
                              batch_size=CFG.batch_size,
                              #shuffle=True,
                              num_workers=0,
                              pin_memory=True,
                              drop_last=True,
                              collate_fn=bms_collate,
                              sampler=train_sampler)
    valid_loader = DataLoader(valid_dataset,
                              batch_size=CFG.batch_size,
                              shuffle=False,
                              num_workers=0,
                              pin_memory=True,
                              drop_last=False)

    # ====================================================
    # scheduler
    # ====================================================
    def get_scheduler(optimizer):
        if CFG.scheduler=='ReduceLROnPlateau':
            scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=CFG.factor, patience=CFG.patience, verbose=True, eps=CFG.eps)
        elif CFG.scheduler=='CosineAnnealingLR':
            scheduler = CosineAnnealingLR(optimizer, T_max=CFG.T_max, eta_min=CFG.min_lr, last_epoch=-1)
        elif CFG.scheduler=='CosineAnnealingWarmRestarts':
            scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=CFG.T_0, T_mult=1, eta_min=CFG.min_lr, last_epoch=-1)
        return scheduler

    # ====================================================
    # model & optimizer
    # ====================================================

    torch.cuda.set_device(gpu)

    LOGGER.info('Creating encoder ..')

    encoder = Encoder(CFG.model_name, pretrained=True)
    encoder.to(device)
    encoder_optimizer = Adam(encoder.parameters(), lr=CFG.encoder_lr, weight_decay=CFG.weight_decay, amsgrad=False)
    encoder_scheduler = get_scheduler(encoder_optimizer)

    encoder = DDP(encoder, device_ids=[gpu])

    LOGGER.info('Creating decoder ..')

    decoder = DecoderWithAttention(attention_dim=CFG.attention_dim,
                                   embed_dim=CFG.embed_dim,
                                   decoder_dim=CFG.decoder_dim,
                                   vocab_size=len(tokenizer),
                                   dropout=CFG.dropout,
                                   device=device)
    decoder.to(device)
    decoder_optimizer = Adam(decoder.parameters(), lr=CFG.decoder_lr, weight_decay=CFG.weight_decay, amsgrad=False)
    decoder_scheduler = get_scheduler(decoder_optimizer)

    # ====================================================
    # loop
    # ====================================================
    criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.stoi["<pad>"]).cuda(gpu)

    best_score = np.inf
    best_loss = np.inf

    LOGGER.info('Starting loop ..')

    for epoch in range(CFG.epochs):
        LOGGER.info(f'Epoch: {epoch}')

        start_time = time.time()

        LOGGER.info('Starting training ..')

        # train
        avg_loss = train_fn(train_loader, encoder, decoder, criterion,
                            encoder_optimizer, decoder_optimizer, epoch,
                            encoder_scheduler, decoder_scheduler, device)

        LOGGER.info('Evaluating ..')

        # eval
        text_preds = valid_fn(valid_loader, encoder, decoder, tokenizer, criterion, device)
        text_preds = [f"InChI=1S/{text}" for text in text_preds]
        LOGGER.info(f"labels: {valid_labels[:5]}")
        LOGGER.info(f"preds: {text_preds[:5]}")

        LOGGER.info('Scoring ..')

        # scoring
        score, scores = get_score(valid_labels, text_preds)
        LOGGER.info(f'Average levenshtein distance: {score}')

        image_keys = pd.read_csv('train_labels.csv')

        LOGGER.info('Writing output csv')

        LOGGER.info(f'Num scores: {len(scores)}')
        LOGGER.info(f'Num preds: {len(text_preds)}')

        scores_computed = pd.DataFrame({
            'predicted_inchi': text_preds,
            'score': scores
        })

        LOGGER.info(f"Num image ids: {len(valid_folds['image_id'])}")
        LOGGER.info(f"Num actual inchis: {len(valid_folds['InChI'])}")

        scores_actual = pd.DataFrame({
            'image_id': valid_folds['image_id'],
            'actual_inchi': valid_folds['InChI']
        })

        scores_out = pd.concat([scores_actual, scores_computed], axis=1)

        scores_out.sort_values('score', ascending=False).to_csv('output.csv', index=False)

        if isinstance(encoder_scheduler, ReduceLROnPlateau):
            encoder_scheduler.step(score)
        elif isinstance(encoder_scheduler, CosineAnnealingLR):
            encoder_scheduler.step()
        elif isinstance(encoder_scheduler, CosineAnnealingWarmRestarts):
            encoder_scheduler.step()

        if isinstance(decoder_scheduler, ReduceLROnPlateau):
            decoder_scheduler.step(score)
        elif isinstance(decoder_scheduler, CosineAnnealingLR):
            decoder_scheduler.step()
        elif isinstance(decoder_scheduler, CosineAnnealingWarmRestarts):
            decoder_scheduler.step()

        elapsed = time.time() - start_time

        LOGGER.info(f'Epoch {epoch+1} - avg_train_loss: {avg_loss:.4f}  time: {elapsed:.0f}s')
        LOGGER.info(f'Epoch {epoch+1} - Score: {score:.4f}')

        if score < best_score:
            best_score = score
            LOGGER.info(f'Epoch {epoch+1} - Save Best Score: {best_score:.4f} Model')
            torch.save({'encoder': encoder.state_dict(),
                        'encoder_optimizer': encoder_optimizer.state_dict(),
                        'encoder_scheduler': encoder_scheduler.state_dict(),
                        'decoder': decoder.state_dict(),
                        'decoder_optimizer': decoder_optimizer.state_dict(),
                        'decoder_scheduler': decoder_scheduler.state_dict(),
                        'text_preds': text_preds,
                       },
                        OUTPUT_DIR+f'{CFG.model_name}_fold{fold}_best.pth')

# ## Main

# In[19]:

# ====================================================
# main
# ====================================================
def main():

    """
    Prepare: 1.train  2.folds
    """

    LOGGER.info(f"Using CUDA: {torch.cuda.is_available()}")
    LOGGER.info(f"Using {CFG.num_workers} workers")

    if CFG.train:
        # train
        oof_df = pd.DataFrame()
        for fold in range(CFG.n_fold):
            if fold in CFG.trn_fold:
                parser = argparse.ArgumentParser()
                parser.add_argument('-n', '--nodes', default=1, type=int, metavar='N')
                parser.add_argument('-g', '--gpus', default=1, type=int, help='number of gpus per node')
                parser.add_argument('-nr', '--nr', default=0, type=int, help='ranking within the nodes')
                parser.add_argument('--epochs', default=2, type=int, metavar='N', help='number of total epochs to run')
                args = parser.parse_args()

                args.world_size = args.gpus * args.nodes
                os.environ['MASTER_ADDR'] = '127.0.0.1'
                os.environ['MASTER_PORT'] = '8888'
                mp.spawn(train_loop, nprocs=args.gpus, args=(folds, fold, args,))

# In[20]:

if __name__ == '__main__':
    main()


# -*- encoding: utf-8 -*-
"""
@File    : akshare_tryout.py
@Time    : 2022/11/13 16:17
@Author  : ZhengKai
@Email   : 156252108@qq.com
@Software: PyCharm
"""

import akshare as ak
import mplfinance as mpf  # Please install mplfinance as follows: pip install mplfinance
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
'''
stock_daily_df = ak.stock_us_daily(symbol="AAPL", adjust="qfq")
stock_daily_df.index = stock_daily_df['date']
print(stock_daily_df.columns)
# print(stock_daily_df.index)
stock_daily_df = stock_daily_df[["open", "high", "low", "close", "volume"]]
stock_daily_df.columns = ["Open", "High", "Low", "Close", "Volume"]
stock_daily_df.index.name = "Date"
stock_daily_df = stock_daily_df.loc["2020-11-01": "2020-11-13"]
mpf.plot(stock_daily_df, type='candle', mav=(3, 6, 9), volume=True, show_nontrading=False)
'''

stock_daily_df = ak.stock_zh_a_hist(symbol="002594", period="daily",  adjust="")
print(stock_daily_df.columns)
stock_daily_df.index = pd.to_datetime(stock_daily_df['日期'])
# print(stock_daily_df.index)
stock_daily_df = stock_daily_df[['开盘',  '最高', '最低','收盘', '成交量']]
stock_daily_df.columns = ["Open", "High", "Low", "Close", "Volume"]
stock_daily_df.index.name = "Date"
stock_daily_df = stock_daily_df.loc["2020-10-01": "2020-11-14"]
# mpf.plot(stock_daily_df, type='candle', mav=(3, 6, 9), volume=True, show_nontrading=True)

# data是测试数据，可以直接下载后读取，在下例中只显示其中100个交易日的数据
plot_data = stock_daily_df
# 读取显示区间最后一个交易日的数据
last_data = plot_data.iloc[-1]
# 设置mplfinance的蜡烛颜色，up为阳线颜色，down为阴线颜色
my_color = mpf.make_marketcolors(up='r',
                                 down='g',
                                 edge='inherit',
                                 wick='inherit',
                                 volume='inherit')
# 设置图表的背景色
my_style = mpf.make_mpf_style(marketcolors=my_color,
                              figcolor='(0.82, 0.83, 0.85)',
                              gridcolor='(0.82, 0.83, 0.85)')
# 使用mpf.figure()函数可以返回一个figure对象，从而进入External Axes Mode，从而实现对Axes对象和figure对象的自由控制
fig = mpf.figure(style=my_style, figsize=(12, 8), facecolor=(0.82, 0.83, 0.85))
# 添加三个图表，四个数字分别代表图表左下角在figure中的坐标，以及图表的宽（0.88）、高（0.60）
ax1 = fig.add_axes([0.06, 0.25, 0.88, 0.60])
# 添加第二、三张图表时，使用sharex关键字指明与ax1在x轴上对齐，且共用x轴
ax2 = fig.add_axes([0.06, 0.15, 0.88, 0.10], sharex=ax1)
ax3 = fig.add_axes([0.06, 0.05, 0.88, 0.10], sharex=ax1)

# 设置三张图表的Y轴标签
ax1.set_ylabel('price')
ax2.set_ylabel('volume')
ax3.set_ylabel('macd')

# 调用mpf.plot()函数，注意调用的方式跟上一节不同，这里需要指定ax=ax1，volume=ax2，将K线图显示在ax1中，交易量显示在ax2中
mpf.plot(plot_data,
		 ax=ax1,
		 volume=ax2,
		 type='candle',
		 style=my_style)
fig.show()


from rest_framework import serializers

from apps.authentication.models import OnlineUser as User
from apps.authentication.serializers import (
    PositionReadOnlySerializer,
    SpecialPositionSerializer,
)
from apps.profiles.models import Privacy

class ProfileSerializer(serializers.ModelSerializer):
    year = serializers.IntegerField()
    positions = PositionReadOnlySerializer(many=True)
    special_positions = SpecialPositionSerializer(many=True)

    class Meta:
        model = User
        fields = (
            "id",
            "first_name",
            "last_name",
            "username",
            "nickname",
            "ntnu_username",
            "year",
            "email",
            "online_mail",
            "phone_number",
            "address",
            "website",
            "github",
            "linkedin",
            "positions",
            "special_positions",
            "rfid",
            "field_of_study",
            "started_date",
            "compiled",
            "infomail",
            "jobmail",
            "zip_code",
            "allergies",
            "mark_rules_accepted",
            "gender",
            "bio",
            "saldo",
            "is_committee",
            "is_member",
            "image",
            "has_expiring_membership",
        )
        read_only_fields = (
            "username",
            "ntnu_username",
            "online_mail",
            "field_of_study",
            "started_date",
            "compiled",
            "saldo",
        )

class ExposableUserField(serializers.Field):
    """
    Serializes a field for a user based on the users Privacy settings.

    Defaults to a defined variable naming:
    For a field 'foo' on OnlineUser, the field 'expose_foo' from Privacy will be used.
    This behavior can be overwritten with the 'privacy_field' kwarg.
    """

    def __init__(self, **kwargs):
        self.field_name = kwargs.pop("field_name", None)
        defualt_privacy_field = "expose_{}".format(self.field_name)
        self.privacy_field = kwargs.pop("privacy_field", defualt_privacy_field)
        super().__init__(**kwargs)

    def get_attribute(self, obj):
        return obj

    def to_representation(self, obj):
        privacy = Privacy.objects.get(user=obj)
        expose = getattr(privacy, self.privacy_field)
        if expose:
            return getattr(obj, self.field_name)
        return None

class PublicProfileSerializer(serializers.ModelSerializer):
    address = ExposableUserField(field_name="address")
    nickname = ExposableUserField(field_name="nickname")
    email = ExposableUserField(field_name="email")
    phone_number = ExposableUserField(field_name="phone_number")
    zip_code = ExposableUserField(field_name="zip_code", privacy_field="expose_address")

    positions = PositionReadOnlySerializer(many=True)
    special_positions = SpecialPositionSerializer(many=True)

    class Meta:
        model = User
        fields = (
            "id",
            "username",
            "nickname",
            "first_name",
            "last_name",
            "phone_number",
            "online_mail",
            "address",
            "zip_code",
            "email",
            "website",
            "github",
            "linkedin",
            "ntnu_username",
            "field_of_study",
            "year",
            "bio",
            "positions",
            "special_positions",
            "image",
            "started_date",
        )

class PrivacySerializer(serializers.ModelSerializer):
    class Meta:
        model = Privacy
        fields = (
            "visible_for_other_users",
            "expose_nickname",
            "expose_email",
            "expose_phone_number",
            "expose_address",
            "visible_as_attending_events",
            "allow_pictures",
        )


"""
Here we compute the state transition matrices from the ALlen Atlas,
for several mappings (Allen, Cosmos, Beryl)
"""
from pathlib import Path
import numpy as np
import scipy.sparse as sp
import seaborn as sns
import matplotlib.pyplot as plt

from iblatlas.atlas import AllenAtlas
from iblutil.numerical import ismember

# from iblapps.atlasview import atlasview
# av = atlasview.view()
LOCAL_DATA_PATH = Path("/datadisk/Data/paper-ephys-atlas/ephys-atlas-decoding")

# instantiate the Allen atlas
ba = AllenAtlas(res_um=25)
ba.compute_surface()

str_mapping = 'Cosmos'
str_mapping = 'Beryl'
# str_mapping = 'Allen'
volume = ba.regions.mappings[str_mapping][ba.label]  # ap, ml, dv
mask = ba.mask()
volume[~mask] = -1

# getting the unique set of regions for the given mapping
aids_unique = np.unique(ba.regions.id[ba.regions.mappings[str_mapping]])
_, ir_unique = ismember(aids_unique, ba.regions.id)

up = volume[:, :, :-1].flatten()
lo = volume[:, :, 1:].flatten()
iok = np.logical_and(up >= 0, lo >= 0)
_, icc_up = ismember(up[iok], ir_unique)
_, icc_lo = ismember(lo[iok], ir_unique)

# here we count the number of voxel from each reagion
state_transitions = sp.coo_matrix(  # (data, (i, j))
    (np.ones_like(icc_lo), (icc_up, icc_lo)), shape=(ir_unique.size, ir_unique.size)).todense()
states_count = sp.coo_matrix(
(np.ones_like(icc_lo), (icc_up, icc_up*0 )), shape=(ir_unique.size,1)).todense()

up2down = state_transitions / state_transitions.sum(axis=1)

np.savez(LOCAL_DATA_PATH.joinpath(f'region_transition_{str_mapping.lower()}.npz'),
         region_transitions=state_transitions, region_counts=states_count.squeeze(), region_aids=ba.regions.id[ir_unique]
         )

## %%

fig, ax = plt.subplots(figsize=(12, 10))
if str_mapping == 'Cosmos':
    sns.heatmap(up2down * 100, vmin=0, vmax=1.2, annot=ir_unique.size < 50, fmt='.1f', ax=ax, cmap='Blues')
    ax.set(
        xticklabels=ba.regions.acronym[ir_unique],
        yticklabels=ba.regions.acronym[ir_unique],
    )
else:
    sns.heatmap(up2down * 100, vmin=0, vmax=1.2, annot=ir_unique.size < 50, fmt='.1f', ax=ax, cmap='Blues')

ax.set(
    title=f'Probability of next channel switching {str_mapping} region (%)',
    ylabel='Current channel',
    xlabel='Next channel down',
)

## %%


import json, re
from datetime import datetime

with open("Google Chat/Groups/DM uZ8_QQAAAAE/messages.json", "r") as d:
    dict = json.load(d)

    messages = dict['messages']

    dates = []

    for message in messages:
        try:
            sender = message['creator']['name']
            date = message['created_date']
            text = message['text']

            # parse created_date string to datetime
            date_time_obj = datetime.strptime(date, '%A, %B %d, %Y at %X %p UTC')

            dates.append(date_time_obj.date())

            print(sender, ' - ' , date, '\n', text, '\n\n --- \n')
        except KeyError:
            continue

# firstDate = dates[0].__str__()
# lastDate = dates[-1].__str__()
# chat_date_duration = firstDate + ' - ' + lastDate

# print(chat_date_duration)


from pathlib import Path
import os
import sys
home_dir = Path.home()
modules_dir = os.path.join(home_dir, 'services/milestones_chains/modules/')
if modules_dir not in sys.path: sys.path.append(modules_dir)
from config import *

# Directories
if 'run_dir' in os.listdir(working_dir):
	shutil.rmtree(run_dir_path)
os.mkdir(run_dir_path)
run_dirs = os.listdir(run_dir_path)
if 'chunks' not in run_dirs:
	os.mkdir(chunks_path)
if 'scaffolds' not in run_dirs:
	os.mkdir(scaffolds_path)
if 'sub_graphs' not in run_dirs:
	os.mkdir(sub_graphs_path)


'''
    example book api call
'''
import urllib.request
import json

ISBN = '9780007550258'
NAME = "I%20Robot"
NAME = "V%20for%20Vendetta"

while True:

    BASE_API_LINK = "https://www.googleapis.com/books/v1/volumes?q=ISBN:"
    # title search
    BASE_API_LINK = "https://www.googleapis.com/books/v1/volumes?q="
    # USER_INPUT = input("Enter ISBN: ").strip()
    # USER_INPUT = ISBN
    USER_INPUT = NAME
    # this is a test

    with urllib.request.urlopen(BASE_API_LINK + USER_INPUT) as f:
        text = f.read()

    decoded_text = text.decode('utf-8-sig')
    obj = json.loads(decoded_text)
    # deserializes decoded_text to a Python object
    volume_info = obj["items"][0]
    authors = obj["items"][0]["volumeInfo"]["authors"]

    # displays title, summary, author, domain, page count and language
    print("\nTitle:", volume_info["volumeInfo"]["title"])
    print("\n")
    # print("\nISBN: ", volume_info["volumeInfo"]["industryIdentifiers"]
    #                    ["ISBN_13"])
    # print("\nSummary:\n")
    # print(textwrap.fill(volume_info["searchInfo"]["textSnippet"], width=65))
    print("\nAuthor(s):", ",".join(authors))
    # print("\nPublic Domain:", volume_info["accessInfo"]["publicDomain"])
    # print("\nPage count:", volume_info["volumeInfo"]["pageCount"])
    # print("\nLanguage:", volume_info["volumeInfo"]["language"])
    print("\n***")

    status_update = input("\nEnter another ISBN? y or n: ").lower().strip()

    if status_update == "n":
        print("\nThank you! Have a nice day.")
        break


from argparse import ArgumentParser
import os
import math
import librosa
import yaml
import numpy as np
import pickle
from preprocess.tacotron.hyperparams import Hyperparams as hy
from preprocess.tacotron import utils

# file_path=r'D:\document\pycharmproject\mouth_voice\preprocess\phoneme_video_model_file.txt'
# audio_file_prefix=r'D:\document\paper\personpaper\audio-visual_consistance\data\timit_audio2'
# mean_std_path=r'D:\document\pycharmproject\mouth_voice\output'
test_set=['mwbt0','msjs1','mrgg0','mpgl0',
          'fram1','fjwb0','fjem0','felc0']

def get_audio_mel(data_dir,phoneme_info_path,out_dir):
    audio_signal=[]
    loaded_audio_list=[]
    audio_feature_map = {}
    train_audio_feature_list = []
    with open(phoneme_info_path,'r') as phoneme_info_file:
        #fadg0_sa1 SH 21 24 860.544 980.27 119.72
        for i, line in enumerate(phoneme_info_file):
            print(line)
            phoneme_info = line.strip().split()
            figure_id=phoneme_info[0].split("_")[0]
            word_id=phoneme_info[0].split("_")[1]
            phoneme_label=phoneme_info[1]
            # fadg0_sa1_SH_i
            phoneme_unit_label=phoneme_info[0]+"_"+phoneme_label+'_'+str(i)
            # The corresponding video frame
            start_frame = int(phoneme_info[2])
            end_frame = int(phoneme_info[3])
            # interval start and end time
            start_time = float(phoneme_info[4])
            end_time = float(phoneme_info[5])
            # audio tag
            audio_label=figure_id+'_'+word_id+'.wav'
            if audio_label not in loaded_audio_list:
                audio_file=os.path.join(data_dir,figure_id,word_id+'.wav')
                audio_signal, _ = librosa.load(audio_file, hy.sr)
                loaded_audio_list.append(audio_label)

            # clip segment with matching strategy
            phoneme_unit_mel_features=[]
            for index in range(start_frame, end_frame + 1):
                if index == start_frame:
                    start_point = math.floor(start_time * hy.sr * 1.0/1000 )
                    end_point = start_point + int(hy.frame_length * hy.sr * 1.0 )
                elif index == end_frame:
                    end_point = math.floor(end_time * hy.sr * 1.0 / 1000.0)
                    start_point = end_point - int(hy.frame_length * hy.sr * hy.sr * 1.0 )
                else:
                    start_point = math.floor(index*hy.frame_length * hy.sr * 1.0 )
                    end_point=start_point+int(hy.frame_length * hy.sr * 1.0/1000 )
                clip_signal = audio_signal[start_point:end_point]
                mel, _ = utils.get_spectrogramsfromsignal(clip_signal)
                # n_mels 512
                assert len(mel[0]) == hy.n_mels
                phoneme_unit_mel_features.append(mel[0])
                # calcute for mean and std
                if figure_id not in test_set:
                    train_audio_feature_list.append(mel[0])
            audio_feature_map[phoneme_unit_label]=phoneme_unit_mel_features

    train_audio_features = np.concatenate(train_audio_feature_list)
    print('len(train_audio_feature_list)', len(train_audio_feature_list))
    mean = np.mean(train_audio_features, axis=0)
    std = np.std(train_audio_features, axis=0)
    attr = {'mean': mean, 'std': std}
    with open(os.path.join(out_dir, 'audio_attr.pkl'), 'wb') as f:
        pickle.dump(attr, f)

    normalized_audio_feature_map = {}
    for key, val in audio_feature_map.items():
        processed_value_list=[(i - mean) / std for i in val]
        normalized_audio_feature_map[key]=processed_value_list
    # audio_feature.pkl
    audio_feature = {'audio_feature': normalized_audio_feature_map}
    with open(os.path.join(out_dir, 'audio_feature.pkl'), 'wb') as f:
        pickle.dump(audio_feature, f)

if  __name__  == '__main__':
    parser = ArgumentParser()
    parser.add_argument('-config', '-c', default='config.yaml')
    args = parser.parse_args()

    args.config = r'D:\document\pycharmproject\AVCDetection\preprocess\preprocess_config.yaml'

    # load config file
    with open(args.config) as f:
        config = yaml.load(f, Loader=yaml.FullLoader)

    data_dir = config['audio']['data_dir']
    phoneme_info_path = config['audio']['phoneme_info_path']
    out_dir = config['audio']['out_dir']

    get_audio_mel(data_dir,phoneme_info_path,out_dir)


import json
import os
import allure
import pytest
from common.handle_test import Handle
from config import Conf
from config.Conf import ConfigYaml
from utils.Assert import Assertions
from utils.YamlUtil import YamlReader

getDataIndexValueTreeList = os.path.join(Conf.get_data_path(), "TITAN_12getDataIndexValueTreeList.yml")
data_index_value_tree_list = YamlReader(getDataIndexValueTreeList).read_data_all()

def get_indicator_data(response_data):
    conf_indicator_list = ConfigYaml().get_titan_info()
    for indicator_name in conf_indicator_list:
        test_indicator_list = list()
        file = Conf.get_titan_indicator_path() + os.sep + indicator_name + ".txt"
        with open(file, "r", encoding="utf-8") as f:
            standard_indicator_list = eval(f.read())
            for first_data in response_data:
                disease_title = first_data["title"]
                if disease_title == indicator_name:
                    if "children" in first_data and first_data["children"]:
                        for second_data in first_data["children"]:
                            if second_data["title"] is None:
                                second_data["title"] = " "
                            primary_classification_title = second_data["title"]
                            if "children" in second_data and second_data["children"]:
                                for third_data in second_data["children"]:
                                    if third_data["title"] is None or "title" not in third_data.keys():
                                        secondary_classification_title = " "
                                    else:
                                        secondary_classification_title = third_data["title"]
                                    if "children" in third_data and third_data["children"]:
                                        for fourth_data in third_data["children"]:
                                            if "title" not in fourth_data.keys():
                                                fourth_data["title"] = " "
                                            elif fourth_data["title"] is None:
                                                fourth_data["title"] = " "
                                            tertiary_classification_title = fourth_data["title"]
                                            if "children" in fourth_data and fourth_data["children"]:
                                                for fifth_data in fourth_data["children"]:
                                                    if "title" not in fifth_data.keys():
                                                        fifth_data["title"] = " "
                                                    elif fifth_data["title"] is None:
                                                        fifth_data["title"] = " "
                                                    fourth_classification_title = fifth_data["title"]
                                                    if "children" in fifth_data and fifth_data["children"]:
                                                        for sixth_data in fifth_data["children"]:
                                                            if "title" not in sixth_data.keys():
                                                                sixth_data["title"] = " "
                                                            elif sixth_data["title"] is None:
                                                                sixth_data["title"] = " "
                                                            fifth_classification_title = sixth_data["title"]
                                                            if "children" in sixth_data and sixth_data["children"]:
                                                                for seventh_data in sixth_data["children"]:
                                                                    if seventh_data["title"] is None:
                                                                        seventh_data["title"] = " "
                                                                    seventh_classification_title = seventh_data["title"]
                                                                    title = disease_title + " -- " + primary_classification_title + " -- " + secondary_classification_title + " -- " + tertiary_classification_title + " -- " + fourth_classification_title + " -- " + fifth_classification_title + " -- " + seventh_classification_title
                                                                    test_indicator_list.append(title)
                                                            else:
                                                                title = disease_title + " -- " + primary_classification_title + " -- " + secondary_classification_title + " -- " + tertiary_classification_title + " -- " + fourth_classification_title + " -- " + fifth_classification_title
                                                                test_indicator_list.append(title)
                                                    else:
                                                        title = disease_title + " -- " + primary_classification_title + " -- " + secondary_classification_title + " -- " + tertiary_classification_title + " -- " + fourth_classification_title
                                                        test_indicator_list.append(title)
                                            else:
                                                title = disease_title + " -- " + primary_classification_title + " -- " + secondary_classification_title + " -- " + tertiary_classification_title
                                                test_indicator_list.append(title)
                                    else:
                                        title = disease_title + " -- " + primary_classification_title + " -- " + secondary_classification_title
                                        test_indicator_list.append(title)
            Assertions().assert_in_titan_data(standard_indicator_list, test_indicator_list, indicator_name)
    # return test_indicator_list, standard_indicator_list

# title_list = list()
#
# def get_data(first_data, data_list, title3):
#     if first_data["title"] is None or "title" not in first_data.keys():
#         title1 = ""
#     else:
#         title1 = first_data["title"]
#     if "children" in first_data and first_data["children"]:
#         for second_data in first_data["children"]:
#             # if second_data["title"] is None or "title" not in second_data.keys():
#             #     title2 = ""
#             # else:
#             #     title2 = second_data["title"]
#             title3 = title3 + "--" + title1
#             return get_data(second_data, data_list, title3)
#     else:
#         data_list.append(title3)
#         # b = "--".join(data_list)
#         # data_list.append(b)
#         return data_list

# def test(data):
#     test_list = list()
#     title = ""
#     for i in data:
#         p = i
#         # while "child" in p.keys():
#         while "children" in p and p["children"]:
#             title = p["title"] + "--" + title
#             p = p["children"]
#         test_list.append(title)
#     return test_list

# def get_indicator_data(response_data):
#     conf_indicator_list = ConfigYaml().get_titan_info()
#     for indicator_name in conf_indicator_list:
#         test_indicator_list = list()
#         file = Conf.get_titan_indicator_path() + os.sep + indicator_name + ".txt"
#         with open(file, "r", encoding="utf-8") as f:
#             standard_indicator_list = eval(f.read())
#             for first_data in response_data:
#                 disease_title = first_data["title"]
#                 if disease_title == indicator_name:
#                     test_indicator_list = test(first_data)
#                     print(test_indicator_list)
#             Assertions().assert_in_titan_data(standard_indicator_list, test_indicator_list, indicator_name)

@allure.feature("Titan筛选")
class TestTitan:
    @pytest.mark.parametrize("index_value_tree_list", data_index_value_tree_list)
    def test_data_index_value_tree_list(self, index_value_tree_list):
        body = Handle().handle_res_body(index_value_tree_list)
        response_data = body["responseData"]
        # test_indicator_list, standard_indicator_list = get_indicator_data(response_data)
        get_indicator_data(response_data)
        # print(test_indicator_list)
        # write_titan_file(test_indicator_list)
        # standard_indicator_list = read_titan_file()
        # Assertions().assert_in_titan_data(test_indicator_list, standard_indicator_list)